diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
index c58cfeb50..679e57ff2 100644
--- a/src/accelerate/utils/modeling.py
+++ b/src/accelerate/utils/modeling.py
@@ -666,7 +666,7 @@ def load_checkpoint_in_model(
         elif len(potential_index) == 1:
             index_filename = os.path.join(checkpoint, potential_index[0])
         else:
-            raise ValueError(f"{checkpoint} containing mote than one `.index.json` file, delete the irrelevant ones.")
+            raise ValueError(f"{checkpoint} containing more than one `.index.json` file, delete the irrelevant ones.")
     else:
         raise ValueError(
             "`checkpoint` should be the path to a file containing a whole state dict, or the index of a sharded "

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/955
diff --git a/src/accelerate/big_modeling.py b/src/accelerate/big_modeling.py
index d3247a464..75a2c06f6 100644
--- a/src/accelerate/big_modeling.py
+++ b/src/accelerate/big_modeling.py
@@ -404,6 +404,7 @@ def load_checkpoint_and_dispatch(
         offload_folder=offload_folder,
         dtype=dtype,
         offload_state_dict=offload_state_dict,
+        offload_buffers=offload_buffers,
     )
     if device_map is None:
         return model
diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
index 2774b18ce..c58cfeb50 100644
--- a/src/accelerate/utils/modeling.py
+++ b/src/accelerate/utils/modeling.py
@@ -607,6 +607,7 @@ def load_checkpoint_in_model(
     offload_folder: Optional[Union[str, os.PathLike]] = None,
     dtype: Optional[Union[str, torch.dtype]] = None,
     offload_state_dict: bool = False,
+    offload_buffers: bool = False,
 ):
     """
     Loads a (potentially sharded) checkpoint inside a model, potentially sending weights to a given device as they are
@@ -636,6 +637,8 @@ def load_checkpoint_in_model(
         offload_state_dict (`bool`, *optional*, defaults to `False`):
             If `True`, will temporarily offload the CPU state dict on the hard drive to avoid getting out of CPU RAM if
             the weight of the CPU state dict + the biggest shard does not fit.
+        offload_buffers (`bool`, *optional*, defaults to `False):
+            Whether or not to include the buffers in the weights offloaded to disk.
     """
     if offload_folder is None and device_map is not None and "disk" in device_map.values():
         raise ValueError(
@@ -687,6 +690,8 @@ def load_checkpoint_in_model(
         state_dict_folder = tempfile.mkdtemp()
         state_dict_index = {}
 
+    buffer_names = [name for name, _ in model.named_buffers()]
+
     for checkpoint_file in checkpoint_files:
         checkpoint = torch.load(checkpoint_file)
         if device_map is None:
@@ -703,7 +708,8 @@ def load_checkpoint_in_model(
                 param_device = device_map[module_name]
 
                 if param_device == "disk":
-                    set_module_tensor_to_device(model, param_name, "meta")
+                    if offload_buffers or param_name not in buffer_names:
+                        set_module_tensor_to_device(model, param_name, "meta")
                     offload_weight(param, param_name, offload_folder, index=offload_index)
                 elif param_device == "cpu" and offload_state_dict:
                     set_module_tensor_to_device(model, param_name, "meta")
diff --git a/tests/test_modeling_utils.py b/tests/test_modeling_utils.py
index 16243d5e1..644d297b2 100644
--- a/tests/test_modeling_utils.py
+++ b/tests/test_modeling_utils.py
@@ -275,6 +275,31 @@ def test_load_checkpoint_in_model_one_gpu(self):
         self.assertEqual(model.batchnorm.weight.device, torch.device("cpu"))
         self.assertEqual(model.linear2.weight.device, torch.device("cpu"))
 
+    @require_cuda
+    def test_load_checkpoint_in_model_disk_offload(self):
+        device_map = {"linear1": "cpu", "batchnorm": "disk", "linear2": "cpu"}
+
+        model = ModelForTest()
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            fname = os.path.join(tmp_dir, "pt_model.bin")
+            torch.save(model.state_dict(), fname)
+            load_checkpoint_in_model(model, fname, device_map=device_map, offload_folder=tmp_dir)
+        self.assertEqual(model.linear1.weight.device, torch.device("cpu"))
+        self.assertEqual(model.batchnorm.weight.device, torch.device("meta"))
+        # Buffers are not offloaded by default
+        self.assertEqual(model.batchnorm.running_mean.device, torch.device("cpu"))
+        self.assertEqual(model.linear2.weight.device, torch.device("cpu"))
+
+        model = ModelForTest()
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            fname = os.path.join(tmp_dir, "pt_model.bin")
+            torch.save(model.state_dict(), fname)
+            load_checkpoint_in_model(model, fname, device_map=device_map, offload_folder=tmp_dir, offload_buffers=True)
+        self.assertEqual(model.linear1.weight.device, torch.device("cpu"))
+        self.assertEqual(model.batchnorm.weight.device, torch.device("meta"))
+        self.assertEqual(model.batchnorm.running_mean.device, torch.device("meta"))
+        self.assertEqual(model.linear2.weight.device, torch.device("cpu"))
+
     @require_multi_gpu
     def test_load_checkpoint_in_model_two_gpu(self):
         device_map = {"linear1": 0, "batchnorm": "cpu", "linear2": 1}

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/951
diff --git a/tests/deepspeed/test_deepspeed.py b/tests/deepspeed/test_deepspeed.py
index fceb0bd23..03a4e346c 100644
--- a/tests/deepspeed/test_deepspeed.py
+++ b/tests/deepspeed/test_deepspeed.py
@@ -243,10 +243,7 @@ def test_deepspeed_plugin(self, stage):
 
     @parameterized.expand([FP16, BF16], name_func=parameterized_custom_name_func)
     def test_accelerate_state_deepspeed(self, dtype):
-        state = AcceleratorState(_from_accelerator=True)
-        if state.initialized:
-            state.initialized = False
-
+        AcceleratorState._reset_state()
         deepspeed_plugin = DeepSpeedPlugin(
             gradient_accumulation_steps=1,
             gradient_clipping=1.0,
@@ -259,7 +256,6 @@ def test_accelerate_state_deepspeed(self, dtype):
         with mockenv_context(**self.dist_env):
             state = Accelerator(mixed_precision=dtype, deepspeed_plugin=deepspeed_plugin).state
             self.assertTrue(state.deepspeed_plugin.deepspeed_config[dtype]["enabled"])
-            state.initialized = False
 
     def test_init_zero3(self):
         deepspeed_plugin = DeepSpeedPlugin(

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/950
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index 31a6d64a2..d97327de4 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -35,6 +35,14 @@
     import torch_xla.core.xla_model as xm
 
 
+def is_initialized() -> bool:
+    """
+    Checks if the `AcceleratorState` has been initialized from `Accelerator`. Same as `AcceleratorState.initialized`,
+    but works as a module method.
+    """
+    return AcceleratorState._shared_state != {}
+
+
 # Inspired by Alex Martelli's 'Borg'.
 class AcceleratorState:
     """
@@ -45,6 +53,7 @@ class AcceleratorState:
         - **device** (`torch.device`) -- The device to use.
         - **distributed_type** ([`~accelerate.state.DistributedType`]) -- The type of distributed environment currently
           in use.
+        - **initialized** (`bool`) -- Whether or not the `AcceleratorState` has been initialized from `Accelerator`.
         - **local_process_index** (`int`) -- The index of the current process on the current server.
         - **mixed_precision** (`str`) -- Whether or not the current script will use mixed precision, and if so the type
           of mixed precision being performed.
@@ -69,8 +78,7 @@ def __init__(
         if parse_flag_from_env("ACCELERATE_USE_CPU"):
             cpu = True
         self._check_initialized(mixed_precision, cpu)
-        self.fork_launched = parse_flag_from_env("FORK_LAUNCHED", 0)
-        if not getattr(self, "initialized", False):
+        if not self.initialized:
             self.backend = None
             self.deepspeed_plugin = None
             mixed_precision = (
@@ -245,18 +253,17 @@ def __init__(
                 and self.device.type == "cuda"
             ):
                 torch.backends.cuda.matmul.allow_tf32 = True
-            self.initialized = True
 
-    def __repr__(self):
-        mixed_precision = self.mixed_precision
+        self.fork_launched = parse_flag_from_env("FORK_LAUNCHED", 0)
 
+    def __repr__(self):
         repr = (
             f"Distributed environment: {self.distributed_type}{('  Backend: ' + self.backend) if self.backend else ''}\n"
             f"Num processes: {self.num_processes}\n"
             f"Process index: {self.process_index}\n"
             f"Local process index: {self.local_process_index}\n"
             f"Device: {self.device}\n"
-            f"Mixed precision type: {mixed_precision}\n"
+            f"Mixed precision type: {self.mixed_precision}\n"
         )
         if self.distributed_type == DistributedType.DEEPSPEED:
             repr += f"ds_config: {self.deepspeed_plugin.deepspeed_config}\n"
@@ -286,9 +293,14 @@ def _reset_state():
         "Resets `_shared_state`, is used internally and should not be called"
         AcceleratorState._shared_state = {}
 
+    @property
+    def initialized(self) -> bool:
+        "Returns whether the `AcceleratorState` has been initialized"
+        return self._shared_state != {}
+
     def _check_initialized(self, mixed_precision=None, cpu=None):
         "Checks if a modification is trying to be made and the `AcceleratorState` has already been initialized"
-        if getattr(self, "initialized", False):
+        if self.initialized:
             err = "AcceleratorState has already been initialized and cannot be changed, restart your runtime completely and pass `{flag}` to `Accelerate()`."
             if cpu and self.device.type != "cpu":
                 raise ValueError(err.format(flag="cpu=True"))
@@ -311,11 +323,15 @@ class GradientState:
 
     def __init__(self):
         self.__dict__ = self._shared_state
-        if not getattr(self, "initialized", False):
+        if not self.initialized:
             self.sync_gradients = True
             self.end_of_dataloader = False
             self.remainder = -1
-        self.initialized = True
+
+    @property
+    def initialized(self) -> bool:
+        "Returns whether the `GradientState` has been initialized"
+        return GradientState._shared_state != {}
 
     def __repr__(self):
         return (

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/949
diff --git a/src/accelerate/utils/offload.py b/src/accelerate/utils/offload.py
index 22f51cce1..84c115fd3 100644
--- a/src/accelerate/utils/offload.py
+++ b/src/accelerate/utils/offload.py
@@ -34,7 +34,7 @@ def offload_weight(weight, weight_name, offload_folder, index=None):
         # Need to reinterpret the underlined data as int16 since NumPy does not handle bfloat16s.
         weight = weight.view(torch.int16)
         dtype = "bfloat16"
-    array = weight.numpy()
+    array = weight.cpu().numpy()
     tensor_file = os.path.join(offload_folder, f"{weight_name}.dat")
     if index is not None:
         if dtype is None:

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/945
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index ec5041bfa..7c07e9105 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -363,7 +363,7 @@ def __init__(
         if (
             self.state.mixed_precision == "fp16"
             and self.device.type != "cpu"
-            and self.distributed_type != DistributedType.MEGATRON_LM
+            and self.distributed_type not in (DistributedType.DEEPSPEED, DistributedType.MEGATRON_LM)
         ):
             self.native_amp = True
             if not torch.cuda.is_available() and not parse_flag_from_env("ACCELERATE_USE_MPS_DEVICE"):
@@ -375,10 +375,10 @@ def __init__(
                 self.scaler = ShardedGradScaler(**kwargs)
             else:
                 self.scaler = torch.cuda.amp.GradScaler(**kwargs)
-        elif (
-            self.state.mixed_precision == "bf16"
-            and self.distributed_type != DistributedType.FSDP
-            and self.distributed_type != DistributedType.MEGATRON_LM
+        elif self.state.mixed_precision == "bf16" and self.distributed_type not in (
+            DistributedType.DEEPSPEED,
+            DistributedType.FSDP,
+            DistributedType.MEGATRON_LM,
         ):
             if self.device.type == "cpu":
                 self.native_amp = is_torch_version(">=", "1.10")

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/943
diff --git a/src/accelerate/tracking.py b/src/accelerate/tracking.py
index 7afebea48..379d6454e 100644
--- a/src/accelerate/tracking.py
+++ b/src/accelerate/tracking.py
@@ -37,7 +37,10 @@
 _available_trackers = []
 
 if is_tensorboard_available():
-    from torch.utils import tensorboard
+    try:
+        from torch.utils import tensorboard
+    except ModuleNotFoundError:
+        import tensorboardX as tensorboard
 
     _available_trackers.append(LoggerType.TENSORBOARD)
 

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/942
diff --git a/docs/source/usage_guides/deepspeed.mdx b/docs/source/usage_guides/deepspeed.mdx
index 29561c77b..0377296c1 100644
--- a/docs/source/usage_guides/deepspeed.mdx
+++ b/docs/source/usage_guides/deepspeed.mdx
@@ -395,6 +395,196 @@ We will look at the changes needed in the code when using these.
 based on model, dataloaders, dummy optimizer and dummy schedulers provided to `prepare` method. 
 Only the `auto` fields specified in above examples are handled by `prepare` method and the rest have to be explicitly specified by the user.
 
+**Things to note when using DeepSpeed Config File**
+
+Below is a sample script using `deepspeed_config_file` in different scenarios.
+
+Code `test.py`:
+
+```python
+from accelerate import Accelerator
+from accelerate.state import AcceleratorState
+
+
+def main():
+    accelerator = Accelerator()
+    accelerator.print(f"{AcceleratorState()}")
+
+
+if __name__ == "__main__":
+    main()
+```
+
+**Scenario 1**: Manually tampered accelerate config file having `deepspeed_config_file` along with other entries.
+
+1. Content of the `accelerate` config:
+
+```yaml
+command_file: null
+commands: null
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+  gradient_accumulation_steps: 1
+  gradient_clipping: 1.0
+  offload_optimizer_device: 'cpu'
+  offload_param_device: 'cpu'
+  zero3_init_flag: true
+  zero3_save_16bit_model: true
+  zero_stage: 3
+  deepspeed_config_file: 'ds_config.json'
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+dynamo_backend: 'NO'
+fsdp_config: {}
+gpu_ids: null
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+megatron_lm_config: {}
+num_machines: 1
+num_processes: 2
+rdzv_backend: static
+same_network: true
+tpu_name: null
+tpu_zone: null
+use_cpu: false
+```
+
+2. `ds_config.json`:
+
+```json
+{
+    "bf16": {
+        "enabled": true
+    },
+    "zero_optimization": {
+        "stage": 3,
+        "stage3_gather_16bit_weights_on_model_save": false,
+        "offload_optimizer": {
+            "device": "none"
+        },
+        "offload_param": {
+            "device": "none"
+        }
+    },
+    "gradient_clipping": 1.0,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "gradient_accumulation_steps": 10,
+    "steps_per_print": 2000000
+}
+```
+
+3. Output of `accelerate launch test.py`:
+
+```bash
+ValueError: When using `deepspeed_config_file`, the following accelerate config variables will be ignored: 
+['gradient_accumulation_steps', 'gradient_clipping', 'zero_stage', 'offload_optimizer_device', 'offload_param_device', 
+'zero3_save_16bit_model', 'mixed_precision'].
+Please specify them appropriately in the DeepSpeed config file.
+If you are using an accelerate config file, remove others config variables mentioned in the above specified list.
+The easiest method is to create a new config following the questionnaire via `accelerate config`.
+It will only ask for the necessary config variables when using `deepspeed_config_file`.
+```
+
+**Scenario 2**: Use the solution of the error to create new accelerate config and check that no ambiguity error is now thrown.
+
+1. Run `accelerate config`:
+
+```bash
+$ accelerate config
+-------------------------------------------------------------------------------------------------------------------------------
+In which compute environment are you running?
+This machine                                                                                                                   
+-------------------------------------------------------------------------------------------------------------------------------
+Which type of machine are you using?                                                                                           
+multi-GPU                                                                                                                      
+How many different machines will you use (use more than 1 for multi-node training)? [1]:                                       
+Do you wish to optimize your script with torch dynamo?[yes/NO]:                                                                
+Do you want to use DeepSpeed? [yes/NO]: yes                                                                                    
+Do you want to specify a json file to a DeepSpeed config? [yes/NO]: yes                                                        
+Please enter the path to the json DeepSpeed config file: ds_config.json                                                        
+Do you want to enable `deepspeed.zero.Init` when using ZeRO Stage-3 for constructing massive models? [yes/NO]: yes
+How many GPU(s) should be used for distributed training? [1]:4
+accelerate configuration saved at ds_config_sample.yaml
+```
+
+2. Content of the `accelerate` config:
+
+```yaml
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+  deepspeed_config_file: ds_config.json
+  zero3_init_flag: true
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+dynamo_backend: 'NO'
+fsdp_config: {}
+machine_rank: 0
+main_training_function: main
+megatron_lm_config: {}
+num_machines: 1
+num_processes: 4
+rdzv_backend: static
+same_network: true
+use_cpu: false
+```
+
+3. Output of `accelerate launch test.py`:
+
+```bash
+Distributed environment: DEEPSPEED  Backend: nccl
+Num processes: 4
+Process index: 0
+Local process index: 0
+Device: cuda:0
+Mixed precision type: bf16
+ds_config: {'bf16': {'enabled': True}, 'zero_optimization': {'stage': 3, 'stage3_gather_16bit_weights_on_model_save': False, 'offload_optimizer': {'device': 'none'}, 'offload_param': {'device': 'none'}}, 'gradient_clipping': 1.0, 'train_batch_size': 'auto', 'train_micro_batch_size_per_gpu': 'auto', 'gradient_accumulation_steps': 10, 'steps_per_print': inf, 'fp16': {'enabled': False}}
+```
+
+**Scenario 3**: Setting the `accelerate launch` command arguments related to DeepSpeed as `"auto"` in the DeepSpeed` configuration file and check that things work as expected.
+
+1. New `ds_config.json` with `"auto"` for the `accelerate launch` DeepSpeed command arguments:
+
+```json
+{
+    "bf16": {
+        "enabled": "auto"
+    },
+    "zero_optimization": {
+        "stage": "auto",
+        "stage3_gather_16bit_weights_on_model_save": "auto",
+        "offload_optimizer": {
+            "device": "auto"
+        },
+        "offload_param": {
+            "device": "auto"
+        }
+    },
+    "gradient_clipping": "auto",
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "gradient_accumulation_steps": "auto",
+    "steps_per_print": 2000000
+}
+```
+
+2. Output of `accelerate launch --mixed_precision="fp16" --zero_stage=3 --gradient_accumulation_steps=5 --gradient_clipping=1.0 --offload_param_device="cpu" --offload_optimizer_device="nvme" --zero3_save_16bit_model="true" test.py`:
+
+```bash
+Distributed environment: DEEPSPEED  Backend: nccl
+Num processes: 4
+Process index: 0
+Local process index: 0
+Device: cuda:0
+Mixed precision type: fp16
+ds_config: {'bf16': {'enabled': False}, 'zero_optimization': {'stage': 3, 'stage3_gather_16bit_weights_on_model_save': True, 'offload_optimizer': {'device': 'nvme'}, 'offload_param': {'device': 'cpu'}}, 'gradient_clipping': 1.0, 'train_batch_size': 'auto', 'train_micro_batch_size_per_gpu': 'auto', 'gradient_accumulation_steps': 5, 'steps_per_print': inf, 'fp16': {'enabled': True, 'auto_cast': True}}
+```
+
+**Note**: Remaining `"auto"` values are handled in `accelerator.prepare()` call as explained in point 2 of 
+`Important code changes when using DeepSpeed Config File`.
+
 ## Saving and loading
 
 1. Saving and loading of models is unchanged for ZeRO Stage-1 and Stage-2.
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index a3dcd2dcb..50acf55a1 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -460,7 +460,7 @@ def get_cluster_input():
 
     if distributed_type != DistributedType.TPU:
         if distributed_type == DistributedType.DEEPSPEED and use_deepspeed_config:
-            mixed_precision = "no"
+            mixed_precision = None
         else:
             mixed_precision = _ask_options(
                 "Do you wish to use FP16 or BF16 (mixed precision)?",
diff --git a/src/accelerate/commands/config/config_args.py b/src/accelerate/commands/config/config_args.py
index ba492802e..8b4a28292 100644
--- a/src/accelerate/commands/config/config_args.py
+++ b/src/accelerate/commands/config/config_args.py
@@ -78,6 +78,7 @@ def to_dict(self):
         for key, value in result.items():
             if isinstance(value, Enum):
                 result[key] = value.value
+        result = {k: v for k, v in result.items() if v is not None}
         return result
 
     @classmethod
@@ -88,7 +89,7 @@ def from_json_file(cls, json_file=None):
         if "compute_environment" not in config_dict:
             config_dict["compute_environment"] = ComputeEnvironment.LOCAL_MACHINE
         if "mixed_precision" not in config_dict:
-            config_dict["mixed_precision"] = "fp16" if ("fp16" in config_dict and config_dict["fp16"]) else "no"
+            config_dict["mixed_precision"] = "fp16" if ("fp16" in config_dict and config_dict["fp16"]) else None
         if "fp16" in config_dict:  # Convert the config to the new format.
             del config_dict["fp16"]
         if "use_cpu" not in config_dict:
@@ -111,7 +112,7 @@ def from_yaml_file(cls, yaml_file=None):
             config_dict["compute_environment"] = ComputeEnvironment.LOCAL_MACHINE
 
         if "mixed_precision" not in config_dict:
-            config_dict["mixed_precision"] = "fp16" if ("fp16" in config_dict and config_dict["fp16"]) else "no"
+            config_dict["mixed_precision"] = "fp16" if ("fp16" in config_dict and config_dict["fp16"]) else None
         if "fp16" in config_dict:  # Convert the config to the new format.
             del config_dict["fp16"]
         if "use_cpu" not in config_dict:
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 91d4427ac..b5f831b47 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -301,45 +301,50 @@ def launch_command_parser(subparsers=None):
         "--zero_stage",
         default=None,
         type=int,
-        help="DeepSpeed's ZeRO optimization stage (useful only when `use_deepspeed` flag is passed).",
+        help="DeepSpeed's ZeRO optimization stage (useful only when `use_deepspeed` flag is passed). "
+        "If unspecified, will default to `2`.",
     )
     deepspeed_args.add_argument(
         "--offload_optimizer_device",
         default=None,
         type=str,
-        help="Decides where (none|cpu|nvme) to offload optimizer states (useful only when `use_deepspeed` flag is passed).",
+        help="Decides where (none|cpu|nvme) to offload optimizer states (useful only when `use_deepspeed` flag is passed). "
+        "If unspecified, will default to 'none'.",
     )
     deepspeed_args.add_argument(
         "--offload_param_device",
         default=None,
         type=str,
-        help="Decides where (none|cpu|nvme) to offload parameters (useful only when `use_deepspeed` flag is passed).",
+        help="Decides where (none|cpu|nvme) to offload parameters (useful only when `use_deepspeed` flag is passed). "
+        "If unspecified, will default to 'none'.",
     )
     deepspeed_args.add_argument(
         "--gradient_accumulation_steps",
         default=None,
         type=int,
-        help="No of gradient_accumulation_steps used in your training script (useful only when `use_deepspeed` flag is passed).",
+        help="No of gradient_accumulation_steps used in your training script (useful only when `use_deepspeed` flag is passed). "
+        "If unspecified, will default to `1`.",
     )
     deepspeed_args.add_argument(
         "--gradient_clipping",
         default=None,
         type=float,
-        help="gradient clipping value used in your training script (useful only when `use_deepspeed` flag is passed).",
+        help="gradient clipping value used in your training script (useful only when `use_deepspeed` flag is passed). "
+        "If unspecified, will default to `1.0`.",
     )
     deepspeed_args.add_argument(
         "--zero3_init_flag",
-        default="true",
+        default=None,
         type=str,
         help="Decides Whether (true|false) to enable `deepspeed.zero.Init` for constructing massive models. "
-        "Only applicable with DeepSpeed ZeRO Stage-3.",
+        "Only applicable with DeepSpeed ZeRO Stage-3. If unspecified, will default to `true`.",
     )
     deepspeed_args.add_argument(
         "--zero3_save_16bit_model",
-        default="false",
+        default=None,
         type=str,
         help="Decides Whether (true|false) to save 16-bit model weights when using ZeRO Stage-3. "
-        "Only applicable with DeepSpeed ZeRO Stage-3.",
+        "Only applicable with DeepSpeed ZeRO Stage-3. If unspecified, will default to `false`.",
     )
     deepspeed_args.add_argument(
         "--deepspeed_hostfile",
@@ -363,7 +368,7 @@ def launch_command_parser(subparsers=None):
         "--deepspeed_multinode_launcher",
         default=None,
         type=str,
-        help="DeepSpeed multi-node launcher to use.",
+        help="DeepSpeed multi-node launcher to use. If unspecified, will default to `pdsh`.",
     )
 
     # fsdp arguments
@@ -717,14 +722,22 @@ def deepspeed_launcher(args):
 
     current_env["PYTHONPATH"] = env_var_path_add("PYTHONPATH", os.path.abspath("."))
     current_env["ACCELERATE_MIXED_PRECISION"] = str(mixed_precision)
+    current_env["ACCELERATE_CONFIG_DS_FIELDS"] = str(args.deepspeed_fields_from_accelerate_config).lower()
     current_env["ACCELERATE_USE_DEEPSPEED"] = "true"
-    current_env["ACCELERATE_DEEPSPEED_ZERO_STAGE"] = str(args.zero_stage)
-    current_env["ACCELERATE_GRADIENT_ACCUMULATION_STEPS"] = str(args.gradient_accumulation_steps)
-    current_env["ACCELERATE_GRADIENT_CLIPPING"] = str(args.gradient_clipping).lower()
-    current_env["ACCELERATE_DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE"] = str(args.offload_optimizer_device).lower()
-    current_env["ACCELERATE_DEEPSPEED_OFFLOAD_PARAM_DEVICE"] = str(args.offload_param_device).lower()
-    current_env["ACCELERATE_DEEPSPEED_ZERO3_INIT"] = str(args.zero3_init_flag).lower()
-    current_env["ACCELERATE_DEEPSPEED_ZERO3_SAVE_16BIT_MODEL"] = str(args.zero3_save_16bit_model).lower()
+    if args.zero_stage is not None:
+        current_env["ACCELERATE_DEEPSPEED_ZERO_STAGE"] = str(args.zero_stage)
+    if args.gradient_accumulation_steps is not None:
+        current_env["ACCELERATE_GRADIENT_ACCUMULATION_STEPS"] = str(args.gradient_accumulation_steps)
+    if args.gradient_clipping is not None:
+        current_env["ACCELERATE_GRADIENT_CLIPPING"] = str(args.gradient_clipping).lower()
+    if args.offload_optimizer_device is not None:
+        current_env["ACCELERATE_DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE"] = str(args.offload_optimizer_device).lower()
+    if args.offload_param_device is not None:
+        current_env["ACCELERATE_DEEPSPEED_OFFLOAD_PARAM_DEVICE"] = str(args.offload_param_device).lower()
+    if args.zero3_init_flag is not None:
+        current_env["ACCELERATE_DEEPSPEED_ZERO3_INIT"] = str(args.zero3_init_flag).lower()
+    if args.zero3_save_16bit_model is not None:
+        current_env["ACCELERATE_DEEPSPEED_ZERO3_SAVE_16BIT_MODEL"] = str(args.zero3_save_16bit_model).lower()
     if args.deepspeed_config_file is not None:
         current_env["ACCELERATE_DEEPSPEED_CONFIG_FILE"] = str(args.deepspeed_config_file)
 
@@ -966,6 +979,7 @@ def launch_command(args):
 
     defaults = None
     warned = []
+    mp_from_config_flag = False
     # Get the default from the config file.
     if args.config_file is not None or os.path.isfile(default_config_file) and not args.cpu:
         defaults = load_config_from_file(args.config_file)
@@ -1013,7 +1027,12 @@ def launch_command(args):
                 ):
                     setattr(args, name, attr)
         if not args.mixed_precision:
-            args.mixed_precision = defaults.mixed_precision
+            if defaults.mixed_precision is None:
+                args.mixed_precision = "no"
+            else:
+                args.mixed_precision = defaults.mixed_precision
+                mp_from_config_flag = True
+
         if args.dynamo_backend is None:
             warned.append("\t`--dynamo_backend` was set to a value of `'no'`")
             args.dynamo_backend = "no"
@@ -1056,6 +1075,10 @@ def launch_command(args):
 
     # Use the proper launcher
     if args.use_deepspeed and not args.cpu:
+        args.deepspeed_fields_from_accelerate_config = list(defaults.deepspeed_config.keys()) if defaults else []
+        if mp_from_config_flag:
+            args.deepspeed_fields_from_accelerate_config.append("mixed_precision")
+        args.deepspeed_fields_from_accelerate_config = ",".join(args.deepspeed_fields_from_accelerate_config)
         deepspeed_launcher(args)
     elif args.use_fsdp and not args.cpu:
         multi_gpu_launcher(args)
diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py
index 5a87d7860..01f174a10 100644
--- a/src/accelerate/utils/dataclasses.py
+++ b/src/accelerate/utils/dataclasses.py
@@ -394,6 +394,28 @@ class DeepSpeedPlugin:
     def __post_init__(self):
         from .deepspeed import HfDeepSpeedConfig
 
+        if self.gradient_accumulation_steps is None:
+            self.gradient_accumulation_steps = int(os.environ.get("ACCELERATE_GRADIENT_ACCUMULATION_STEPS", 1))
+
+        if self.gradient_clipping is None:
+            gradient_clipping = os.environ.get("ACCELERATE_GRADIENT_CLIPPING", "none")
+            if gradient_clipping != "none":
+                self.gradient_clipping = float(gradient_clipping)
+
+        if self.zero_stage is None:
+            self.zero_stage = int(os.environ.get("ACCELERATE_DEEPSPEED_ZERO_STAGE", 2))
+
+        if self.offload_optimizer_device is None:
+            self.offload_optimizer_device = os.environ.get("ACCELERATE_DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE", "none")
+
+        if self.offload_param_device is None:
+            self.offload_param_device = os.environ.get("ACCELERATE_DEEPSPEED_OFFLOAD_PARAM_DEVICE", "none")
+
+        if self.zero3_save_16bit_model is None:
+            self.zero3_save_16bit_model = (
+                os.environ.get("ACCELERATE_DEEPSPEED_ZERO3_SAVE_16BIT_MODEL", "false") == "true"
+            )
+
         if self.hf_ds_config is None:
             self.hf_ds_config = os.environ.get("ACCELERATE_DEEPSPEED_CONFIG_FILE", "none")
         if (
@@ -405,33 +427,22 @@ def __post_init__(self):
                 self.hf_ds_config = HfDeepSpeedConfig(self.hf_ds_config)
             if "gradient_accumulation_steps" not in self.hf_ds_config.config:
                 self.hf_ds_config.config["gradient_accumulation_steps"] = 1
-            elif self.hf_ds_config.config["gradient_accumulation_steps"] == "auto":
-                raise ValueError("gradient_accumulation_steps cannot be set to 'auto' in the DeepSpeed config.")
             if "zero_optimization" not in self.hf_ds_config.config:
                 raise ValueError("Please specify the ZeRO optimization config in the DeepSpeed config.")
-        else:
-            if self.gradient_accumulation_steps is None:
-                self.gradient_accumulation_steps = int(os.environ.get("ACCELERATE_GRADIENT_ACCUMULATION_STEPS", 1))
-
-            if self.gradient_clipping is None:
-                gradient_clipping = os.environ.get("ACCELERATE_GRADIENT_CLIPPING", "none")
-                if gradient_clipping != "none":
-                    self.gradient_clipping = float(gradient_clipping)
-
-            if self.zero_stage is None:
-                self.zero_stage = int(os.environ.get("ACCELERATE_DEEPSPEED_ZERO_STAGE", 2))
-
-            if self.offload_optimizer_device is None:
-                self.offload_optimizer_device = os.environ.get("ACCELERATE_DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE", "none")
-
-            if self.offload_param_device is None:
-                self.offload_param_device = os.environ.get("ACCELERATE_DEEPSPEED_OFFLOAD_PARAM_DEVICE", "none")
-
-            if self.zero3_save_16bit_model is None:
-                self.zero3_save_16bit_model = (
-                    os.environ.get("ACCELERATE_DEEPSPEED_ZERO3_SAVE_16BIT_MODEL", "false") == "true"
-                )
 
+            self._deepspeed_config_checks()
+            kwargs = {
+                "gradient_accumulation_steps": self.gradient_accumulation_steps,
+                "gradient_clipping": self.gradient_clipping if self.gradient_clipping else 1.0,
+                "zero_optimization.stage": self.zero_stage,
+                "zero_optimization.offload_optimizer.device": self.offload_optimizer_device,
+                "zero_optimization.offload_param.device": self.offload_param_device,
+                "zero_optimization.stage3_gather_16bit_weights_on_model_save": self.zero3_save_16bit_model,
+            }
+            for key in kwargs.keys():
+                self.fill_match(key, **kwargs, must_match=False)
+            self.hf_ds_config.set_stage_and_offload()
+        else:
             config = {
                 "train_batch_size": "auto",
                 "train_micro_batch_size_per_gpu": "auto",
@@ -450,15 +461,19 @@ def __post_init__(self):
             if self.gradient_clipping:
                 config["gradient_clipping"] = self.gradient_clipping
             self.hf_ds_config = HfDeepSpeedConfig(config)
+
         self.deepspeed_config = self.hf_ds_config.config
         self.deepspeed_config["steps_per_print"] = float("inf")  # this will stop deepspeed from logging @ stdout
         if self.zero3_init_flag is None:
-            self.zero3_init_flag = os.environ.get("ACCELERATE_DEEPSPEED_ZERO3_INIT", "false") == "true"
+            self.zero3_init_flag = (
+                strtobool(os.environ.get("ACCELERATE_DEEPSPEED_ZERO3_INIT", str(self.hf_ds_config.is_zero3()))) == 1
+            )
         if self.zero3_init_flag and not self.hf_ds_config.is_zero3():
             warnings.warn("DeepSpeed Zero3 Init flag is only applicable for ZeRO Stage 3. Setting it to False.")
             self.zero3_init_flag = False
 
-    def fill_match(self, ds_key_long, mismatches, must_match=True, **kwargs):
+    def fill_match(self, ds_key_long, mismatches=None, must_match=True, **kwargs):
+        mismatches = [] if mismatches is None else mismatches
         config, ds_key = self.hf_ds_config.find_config_node(ds_key_long)
         if config is None:
             return
@@ -503,10 +518,28 @@ def deepspeed_config_process(self, prefix="", mismatches=None, config=None, must
 
     def set_mixed_precision(self, mixed_precision):
         ds_config = self.deepspeed_config
-        if mixed_precision == "fp16" and "fp16" not in ds_config and "bf16" not in ds_config:
-            ds_config.update({"fp16": {"enabled": True, "auto_cast": True}})
-        elif mixed_precision == "bf16" and "fp16" not in ds_config and "bf16" not in ds_config:
-            ds_config.update({"bf16": {"enabled": True}})
+        kwargs = {
+            "fp16.enabled": mixed_precision == "fp16",
+            "bf16.enabled": mixed_precision == "bf16",
+        }
+        if mixed_precision == "fp16":
+            if "fp16" not in ds_config:
+                ds_config["fp16"] = {"enabled": True, "auto_cast": True}
+        elif mixed_precision == "bf16":
+            if "bf16" not in ds_config:
+                ds_config["bf16"] = {"enabled": True}
+
+        if mixed_precision != "no":
+            diff_dtype = "bf16" if mixed_precision == "fp16" else "fp16"
+            if str(ds_config.get(diff_dtype, {}).get("enabled", "False")).lower() == "true":
+                raise ValueError(
+                    f"`--mixed_precision` arg cannot be set to `{mixed_precision}` when `{diff_dtype}` is set in the DeepSpeed config file."
+                )
+        for dtype in ["fp16", "bf16"]:
+            if dtype not in ds_config:
+                ds_config[dtype] = {"enabled": False}
+        self.fill_match("fp16.enabled", must_match=False, **kwargs)
+        self.fill_match("bf16.enabled", must_match=False, **kwargs)
 
     def set_deepspeed_weakref(self):
         from .imports import is_transformers_available
@@ -549,6 +582,31 @@ def zero3_init_context_manager(self, enable=False):
             self.dschf = None
             self.set_deepspeed_weakref()
 
+    def _deepspeed_config_checks(self):
+        env_variable_names_to_ignore = [
+            "ACCELERATE_GRADIENT_ACCUMULATION_STEPS",
+            "ACCELERATE_GRADIENT_CLIPPING",
+            "ACCELERATE_DEEPSPEED_ZERO_STAGE",
+            "ACCELERATE_DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE",
+            "ACCELERATE_DEEPSPEED_OFFLOAD_PARAM_DEVICE",
+            "ACCELERATE_DEEPSPEED_ZERO3_SAVE_16BIT_MODEL",
+            "ACCELERATE_MIXED_PRECISION",
+        ]
+        env_variable_names_to_ignore = [
+            name.replace("ACCELERATE_", "").replace("DEEPSPEED_", "").lower() for name in env_variable_names_to_ignore
+        ]
+
+        deepspeed_fields_from_accelerate_config = os.environ.get("ACCELERATE_CONFIG_DS_FIELDS", "").split(",")
+
+        if any(name in env_variable_names_to_ignore for name in deepspeed_fields_from_accelerate_config):
+            raise ValueError(
+                f"When using `deepspeed_config_file`, the following accelerate config variables will be ignored: {env_variable_names_to_ignore}.\n"
+                "Please specify them appropriately in the DeepSpeed config file.\n"
+                "If you are using an accelerate config file, remove others config variables mentioned in the above specified list.\n"
+                "The easiest method is to create a new config following the questionnaire via `accelerate config`.\n"
+                "It will only ask for the necessary config variables when using `deepspeed_config_file`."
+            )
+
 
 @dataclass
 class FullyShardedDataParallelPlugin:
diff --git a/src/accelerate/utils/deepspeed.py b/src/accelerate/utils/deepspeed.py
index 02d1ab8bc..69dc5c7f8 100644
--- a/src/accelerate/utils/deepspeed.py
+++ b/src/accelerate/utils/deepspeed.py
@@ -50,6 +50,9 @@ def __init__(self, config_file_or_dict):
             raise ValueError("expecting either a path to a DeepSpeed config file or a pre-populated dict")
         self.config = config
 
+        self.set_stage_and_offload()
+
+    def set_stage_and_offload(self):
         # zero stage - this is done as early as possible, before model is created, to allow
         # ``is_deepspeed_zero3_enabled`` query and getting to the early deepspeed config object
         # during ``zero.Init()`` which needs to know the dtype, and some other hparams.
diff --git a/tests/deepspeed/test_deepspeed.py b/tests/deepspeed/test_deepspeed.py
index eca75c1ed..fceb0bd23 100644
--- a/tests/deepspeed/test_deepspeed.py
+++ b/tests/deepspeed/test_deepspeed.py
@@ -285,8 +285,6 @@ def test_prepare_deepspeed(self, optim_type, scheduler_type):
         from deepspeed.runtime.engine import DeepSpeedEngine
 
         kwargs = {
-            "fp16.enabled": True,
-            "bf16.enabled": False,
             "optimizer.params.lr": 5e-5,
             "optimizer.params.weight_decay": 0.0,
             "scheduler.params.warmup_min_lr": 0.0,
@@ -370,7 +368,7 @@ def test_prepare_deepspeed(self, optim_type, scheduler_type):
             # Test DeepSpeed optimizer + DeepSpeed scheduler
             deepspeed_plugin = DeepSpeedPlugin(hf_ds_config=self.ds_config_file[ZERO2])
             with mockenv_context(**self.dist_env):
-                accelerator = Accelerator(deepspeed_plugin=deepspeed_plugin)
+                accelerator = Accelerator(deepspeed_plugin=deepspeed_plugin, mixed_precision="fp16")
                 train_set = RegressionDataset(length=80)
                 eval_set = RegressionDataset(length=20)
                 train_dataloader = DataLoader(train_set, batch_size=10, shuffle=True)
@@ -430,7 +428,7 @@ def test_prepare_deepspeed(self, optim_type, scheduler_type):
             # Test custom optimizer + DeepSpeed scheduler
             deepspeed_plugin = DeepSpeedPlugin(hf_ds_config=self.ds_config_file[ZERO2])
             with mockenv_context(**self.dist_env):
-                accelerator = Accelerator(deepspeed_plugin=deepspeed_plugin)
+                accelerator = Accelerator(deepspeed_plugin=deepspeed_plugin, mixed_precision="fp16")
                 train_set = RegressionDataset(length=80)
                 eval_set = RegressionDataset(length=20)
                 train_dataloader = DataLoader(train_set, batch_size=10, shuffle=True)
@@ -463,7 +461,7 @@ def test_prepare_deepspeed(self, optim_type, scheduler_type):
             # Test deepspeed optimizer + custom scheduler
             deepspeed_plugin = DeepSpeedPlugin(hf_ds_config=self.ds_config_file[ZERO2])
             with mockenv_context(**self.dist_env):
-                accelerator = Accelerator(deepspeed_plugin=deepspeed_plugin)
+                accelerator = Accelerator(deepspeed_plugin=deepspeed_plugin, mixed_precision="fp16")
                 train_set = RegressionDataset(length=80)
                 eval_set = RegressionDataset(length=20)
                 train_dataloader = DataLoader(train_set, batch_size=10, shuffle=True)
@@ -501,8 +499,6 @@ def test_save_checkpoints(self):
         )
         del deepspeed_plugin.deepspeed_config["bf16"]
         kwargs = {
-            "fp16.enabled": True,
-            "bf16.enabled": False,
             "optimizer.params.lr": 5e-5,
             "optimizer.params.weight_decay": 0.0,
             "scheduler.params.warmup_min_lr": 0.0,
@@ -518,7 +514,7 @@ def test_save_checkpoints(self):
         }
 
         with mockenv_context(**self.dist_env):
-            accelerator = Accelerator(deepspeed_plugin=deepspeed_plugin)
+            accelerator = Accelerator(deepspeed_plugin=deepspeed_plugin, mixed_precision="fp16")
             kwargs["train_batch_size"] = (
                 kwargs["train_micro_batch_size_per_gpu"]
                 * deepspeed_plugin.deepspeed_config["gradient_accumulation_steps"]
@@ -594,6 +590,81 @@ def test_autofill_dsconfig(self):
                 accelerator.deepspeed_config["zero_optimization"]["stage3_gather_16bit_weights_on_model_save"]
             )
 
+    @parameterized.expand([FP16, BF16], name_func=parameterized_custom_name_func)
+    def test_autofill_dsconfig_from_ds_plugin(self, dtype):
+        ds_config = self.ds_config_dict["zero3"]
+        if dtype == BF16:
+            del ds_config["fp16"]
+        else:
+            del ds_config["bf16"]
+        ds_config[dtype]["enabled"] = "auto"
+        ds_config["zero_optimization"]["stage"] = "auto"
+        ds_config["zero_optimization"]["stage3_gather_16bit_weights_on_model_save"] = "auto"
+        ds_config["zero_optimization"]["offload_optimizer"]["device"] = "auto"
+        ds_config["zero_optimization"]["offload_param"]["device"] = "auto"
+        ds_config["gradient_accumulation_steps"] = "auto"
+        ds_config["gradient_clipping"] = "auto"
+
+        deepspeed_plugin = DeepSpeedPlugin(
+            hf_ds_config=ds_config,
+            zero3_init_flag=True,
+            gradient_accumulation_steps=1,
+            gradient_clipping=1.0,
+            zero_stage=2,
+            offload_optimizer_device="cpu",
+            offload_param_device="cpu",
+            zero3_save_16bit_model=True,
+        )
+
+        with mockenv_context(**self.dist_env):
+            accelerator = Accelerator(deepspeed_plugin=deepspeed_plugin, mixed_precision=dtype)
+            deepspeed_plugin = accelerator.state.deepspeed_plugin
+            self.assertEqual(deepspeed_plugin.deepspeed_config["gradient_clipping"], 1.0)
+            self.assertEqual(deepspeed_plugin.deepspeed_config["gradient_accumulation_steps"], 1)
+            self.assertEqual(deepspeed_plugin.deepspeed_config["zero_optimization"]["stage"], 2)
+            self.assertEqual(
+                deepspeed_plugin.deepspeed_config["zero_optimization"]["offload_optimizer"]["device"], "cpu"
+            )
+            self.assertEqual(deepspeed_plugin.deepspeed_config["zero_optimization"]["offload_param"]["device"], "cpu")
+            self.assertTrue(
+                deepspeed_plugin.deepspeed_config["zero_optimization"]["stage3_gather_16bit_weights_on_model_save"]
+            )
+            self.assertTrue(deepspeed_plugin.deepspeed_config[dtype]["enabled"])
+
+        AcceleratorState._reset_state()
+        diff_dtype = "bf16" if dtype == "fp16" else "fp16"
+        with mockenv_context(**self.dist_env):
+            with self.assertRaises(ValueError) as cm:
+                accelerator = Accelerator(deepspeed_plugin=deepspeed_plugin, mixed_precision=diff_dtype)
+            self.assertTrue(
+                f"`--mixed_precision` arg cannot be set to `{diff_dtype}` when `{dtype}` is set in the DeepSpeed config file."
+                in str(cm.exception)
+            )
+
+    def test_ds_config_assertions(self):
+        ambiguous_env = self.dist_env.copy()
+        ambiguous_env[
+            "ACCELERATE_CONFIG_DS_FIELDS"
+        ] = "gradient_accumulation_steps,gradient_clipping,zero_stage,offload_optimizer_device,offload_param_device,zero3_save_16bit_model,mixed_precision"
+
+        with mockenv_context(**ambiguous_env):
+            with self.assertRaises(ValueError) as cm:
+                deepspeed_plugin = DeepSpeedPlugin(
+                    hf_ds_config=self.ds_config_file[ZERO3],
+                    zero3_init_flag=True,
+                    gradient_accumulation_steps=1,
+                    gradient_clipping=1.0,
+                    zero_stage=ZERO2,
+                    offload_optimizer_device="cpu",
+                    offload_param_device="cpu",
+                    zero3_save_16bit_model=True,
+                )
+                _ = Accelerator(deepspeed_plugin=deepspeed_plugin, mixed_precision=FP16)
+            self.assertTrue(
+                "If you are using an accelerate config file, remove others config variables mentioned in the above specified list."
+                in str(cm.exception)
+            )
+
     def test_basic_run(self):
         mod_file = inspect.getfile(accelerate.test_utils)
         test_file_path = os.path.sep.join(

Nit: all `none`->`None`
@@ -301,45 +301,50 @@ def launch_command_parser(subparsers=None):
         "--zero_stage",
         default=None,
         type=int,
-        help="DeepSpeed's ZeRO optimization stage (useful only when `use_deepspeed` flag is passed).",
+        help="DeepSpeed's ZeRO optimization stage (useful only when `use_deepspeed` flag is passed). "
+        "If unspecified, will default to `2`.",
     )
     deepspeed_args.add_argument(
         "--offload_optimizer_device",
         default=None,
         type=str,
-        help="Decides where (none|cpu|nvme) to offload optimizer states (useful only when `use_deepspeed` flag is passed).",
+        help="Decides where (none|cpu|nvme) to offload optimizer states (useful only when `use_deepspeed` flag is passed). "
+        "If unspecified, will default to `none`.",
```suggestion
                "If you are using an accelerate config file, set `mixed_precision=no` "
```
@@ -549,6 +551,40 @@ def zero3_init_context_manager(self, enable=False):
             self.dschf = None
             self.set_deepspeed_weakref()
 
+    def _deepspeed_config_checks(self):
+        env_variable_names_to_ignore = [
+            "ACCELERATE_GRADIENT_ACCUMULATION_STEPS",
+            "ACCELERATE_GRADIENT_CLIPPING",
+            "ACCELERATE_DEEPSPEED_ZERO_STAGE",
+            "ACCELERATE_DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE",
+            "ACCELERATE_DEEPSPEED_OFFLOAD_PARAM_DEVICE",
+            "ACCELERATE_DEEPSPEED_ZERO3_SAVE_16BIT_MODEL",
+            "ACCELERATE_MIXED_PRECISION",
+        ]
+        duplicate_values_flag = False
+        for name in env_variable_names_to_ignore:
+            if name != "ACCELERATE_MIXED_PRECISION" and os.environ.get(name, None) is not None:
+                duplicate_values_flag = True
+            elif name == "ACCELERATE_MIXED_PRECISION" and os.environ.get(name, "no") != "no":
+                duplicate_values_flag = True
+            if duplicate_values_flag:
+                break
+
+        if duplicate_values_flag:
+            env_variable_names_to_ignore = [
+                name.replace("ACCELERATE_", "").replace("DEEPSPEED_", "").lower()
+                for name in env_variable_names_to_ignore
+            ]
+            raise ValueError(
+                f"When using `deepspeed_config_file`, the following accelerate config variables will be ignored: {env_variable_names_to_ignore}.\n"
+                "Please specify them appropriately in the DeepSpeed config file.\n"
+                "If you are using accelerate config file, set `mixed_precision=no` "
```suggestion
                "and make sure to not specify these config variables in `accelerate launch` command. \n"
```
@@ -549,6 +551,40 @@ def zero3_init_context_manager(self, enable=False):
             self.dschf = None
             self.set_deepspeed_weakref()
 
+    def _deepspeed_config_checks(self):
+        env_variable_names_to_ignore = [
+            "ACCELERATE_GRADIENT_ACCUMULATION_STEPS",
+            "ACCELERATE_GRADIENT_CLIPPING",
+            "ACCELERATE_DEEPSPEED_ZERO_STAGE",
+            "ACCELERATE_DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE",
+            "ACCELERATE_DEEPSPEED_OFFLOAD_PARAM_DEVICE",
+            "ACCELERATE_DEEPSPEED_ZERO3_SAVE_16BIT_MODEL",
+            "ACCELERATE_MIXED_PRECISION",
+        ]
+        duplicate_values_flag = False
+        for name in env_variable_names_to_ignore:
+            if name != "ACCELERATE_MIXED_PRECISION" and os.environ.get(name, None) is not None:
+                duplicate_values_flag = True
+            elif name == "ACCELERATE_MIXED_PRECISION" and os.environ.get(name, "no") != "no":
+                duplicate_values_flag = True
+            if duplicate_values_flag:
+                break
+
+        if duplicate_values_flag:
+            env_variable_names_to_ignore = [
+                name.replace("ACCELERATE_", "").replace("DEEPSPEED_", "").lower()
+                for name in env_variable_names_to_ignore
+            ]
+            raise ValueError(
+                f"When using `deepspeed_config_file`, the following accelerate config variables will be ignored: {env_variable_names_to_ignore}.\n"
+                "Please specify them appropriately in the DeepSpeed config file.\n"
+                "If you are using accelerate config file, set `mixed_precision=no` "
+                "and remove others config variables mentioned in the above specified list; "
+                "else don't specify these config variables in `accelerate launch` command. \n"
Hello, `none` is the string option as possible values are `cpu`|`nvme`l`none`
@@ -301,45 +301,50 @@ def launch_command_parser(subparsers=None):
         "--zero_stage",
         default=None,
         type=int,
-        help="DeepSpeed's ZeRO optimization stage (useful only when `use_deepspeed` flag is passed).",
+        help="DeepSpeed's ZeRO optimization stage (useful only when `use_deepspeed` flag is passed). "
+        "If unspecified, will default to `2`.",
     )
     deepspeed_args.add_argument(
         "--offload_optimizer_device",
         default=None,
         type=str,
-        help="Decides where (none|cpu|nvme) to offload optimizer states (useful only when `use_deepspeed` flag is passed).",
+        help="Decides where (none|cpu|nvme) to offload optimizer states (useful only when `use_deepspeed` flag is passed). "
+        "If unspecified, will default to `none`.",
Oh then add quotes?
@@ -301,45 +301,50 @@ def launch_command_parser(subparsers=None):
         "--zero_stage",
         default=None,
         type=int,
-        help="DeepSpeed's ZeRO optimization stage (useful only when `use_deepspeed` flag is passed).",
+        help="DeepSpeed's ZeRO optimization stage (useful only when `use_deepspeed` flag is passed). "
+        "If unspecified, will default to `2`.",
     )
     deepspeed_args.add_argument(
         "--offload_optimizer_device",
         default=None,
         type=str,
-        help="Decides where (none|cpu|nvme) to offload optimizer states (useful only when `use_deepspeed` flag is passed).",
+        help="Decides where (none|cpu|nvme) to offload optimizer states (useful only when `use_deepspeed` flag is passed). "
+        "If unspecified, will default to `none`.",
Done.
@@ -301,45 +301,50 @@ def launch_command_parser(subparsers=None):
         "--zero_stage",
         default=None,
         type=int,
-        help="DeepSpeed's ZeRO optimization stage (useful only when `use_deepspeed` flag is passed).",
+        help="DeepSpeed's ZeRO optimization stage (useful only when `use_deepspeed` flag is passed). "
+        "If unspecified, will default to `2`.",
     )
     deepspeed_args.add_argument(
         "--offload_optimizer_device",
         default=None,
         type=str,
-        help="Decides where (none|cpu|nvme) to offload optimizer states (useful only when `use_deepspeed` flag is passed).",
+        help="Decides where (none|cpu|nvme) to offload optimizer states (useful only when `use_deepspeed` flag is passed). "
+        "If unspecified, will default to `none`.",
Did I understand it right that here you suggest to the user to set `mixed_precision=no` even if they use mixed precision in ds_config.json? If so this now is worse than before, as now you're proposing to force 2 opposite values in the same config entry. As in `mixed_precision=no` and `fp16 { enabled: true}`.

Unless I'm misunderstanding the wording that is. 

I think it should be either `mixed_precision=xyz` or `ds_config.json` (with fp16 or bf16 blocks or no block at all as fp16 is the default), but not both.
@@ -549,6 +551,40 @@ def zero3_init_context_manager(self, enable=False):
             self.dschf = None
             self.set_deepspeed_weakref()
 
+    def _deepspeed_config_checks(self):
+        env_variable_names_to_ignore = [
+            "ACCELERATE_GRADIENT_ACCUMULATION_STEPS",
+            "ACCELERATE_GRADIENT_CLIPPING",
+            "ACCELERATE_DEEPSPEED_ZERO_STAGE",
+            "ACCELERATE_DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE",
+            "ACCELERATE_DEEPSPEED_OFFLOAD_PARAM_DEVICE",
+            "ACCELERATE_DEEPSPEED_ZERO3_SAVE_16BIT_MODEL",
+            "ACCELERATE_MIXED_PRECISION",
+        ]
+        duplicate_values_flag = False
+        for name in env_variable_names_to_ignore:
+            if name != "ACCELERATE_MIXED_PRECISION" and os.environ.get(name, None) is not None:
+                duplicate_values_flag = True
+            elif name == "ACCELERATE_MIXED_PRECISION" and os.environ.get(name, "no") != "no":
+                duplicate_values_flag = True
+            if duplicate_values_flag:
+                break
+
+        if duplicate_values_flag:
+            env_variable_names_to_ignore = [
+                name.replace("ACCELERATE_", "").replace("DEEPSPEED_", "").lower()
+                for name in env_variable_names_to_ignore
+            ]
+            raise ValueError(
+                f"When using `deepspeed_config_file`, the following accelerate config variables will be ignored: {env_variable_names_to_ignore}.\n"
+                "Please specify them appropriately in the DeepSpeed config file.\n"
+                "If you are using an accelerate config file, set `mixed_precision=no` "
```suggestion
                "The easiest method is to create a new config following the questionnaire via `accelerate config`.\n"
```
@@ -578,10 +578,9 @@ def _deepspeed_config_checks(self):
             raise ValueError(
                 f"When using `deepspeed_config_file`, the following accelerate config variables will be ignored: {env_variable_names_to_ignore}.\n"
                 "Please specify them appropriately in the DeepSpeed config file.\n"
-                "If you are using an accelerate config file, set `mixed_precision=no` "
-                "and remove others config variables mentioned in the above specified list; "
+                "If you are using an accelerate config file, remove others config variables mentioned in the above specified list; "
                 "and make sure to not specify these config variables in `accelerate launch` command. \n"
-                "The easiest method is to create new config following the questionnaire via  `accelerate config`.\n"
+                "The easiest method is to create new config following the questionnaire via `accelerate config`.\n"
Done. I've simplified the error message wherein user can just remove all the ambiguous entries from  accelerate config file and not specify them in launch command
@@ -549,6 +551,40 @@ def zero3_init_context_manager(self, enable=False):
             self.dschf = None
             self.set_deepspeed_weakref()
 
+    def _deepspeed_config_checks(self):
+        env_variable_names_to_ignore = [
+            "ACCELERATE_GRADIENT_ACCUMULATION_STEPS",
+            "ACCELERATE_GRADIENT_CLIPPING",
+            "ACCELERATE_DEEPSPEED_ZERO_STAGE",
+            "ACCELERATE_DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE",
+            "ACCELERATE_DEEPSPEED_OFFLOAD_PARAM_DEVICE",
+            "ACCELERATE_DEEPSPEED_ZERO3_SAVE_16BIT_MODEL",
+            "ACCELERATE_MIXED_PRECISION",
+        ]
+        duplicate_values_flag = False
+        for name in env_variable_names_to_ignore:
+            if name != "ACCELERATE_MIXED_PRECISION" and os.environ.get(name, None) is not None:
+                duplicate_values_flag = True
+            elif name == "ACCELERATE_MIXED_PRECISION" and os.environ.get(name, "no") != "no":
+                duplicate_values_flag = True
+            if duplicate_values_flag:
+                break
+
+        if duplicate_values_flag:
+            env_variable_names_to_ignore = [
+                name.replace("ACCELERATE_", "").replace("DEEPSPEED_", "").lower()
+                for name in env_variable_names_to_ignore
+            ]
+            raise ValueError(
+                f"When using `deepspeed_config_file`, the following accelerate config variables will be ignored: {env_variable_names_to_ignore}.\n"
+                "Please specify them appropriately in the DeepSpeed config file.\n"
+                "If you are using an accelerate config file, set `mixed_precision=no` "
to focus we are dealing with `mixed_precision` duplicity here:

I've just tried:

```
deepspeed_config:
  deepspeed_multinode_launcher: standard
  deepspeed_config_file: ./configs/vopt-large-z3/ds_config.json
  zero3_init_flag: true
distributed_type: DEEPSPEED
mixed_precision: 'no'
```

```
{
    "fp16": {
        "enabled": true,
        [...]
```
and it doesn't assert.

with the latest commit of `dd2d57b6a1`


@@ -549,6 +551,40 @@ def zero3_init_context_manager(self, enable=False):
             self.dschf = None
             self.set_deepspeed_weakref()
 
+    def _deepspeed_config_checks(self):
+        env_variable_names_to_ignore = [
+            "ACCELERATE_GRADIENT_ACCUMULATION_STEPS",
+            "ACCELERATE_GRADIENT_CLIPPING",
+            "ACCELERATE_DEEPSPEED_ZERO_STAGE",
+            "ACCELERATE_DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE",
+            "ACCELERATE_DEEPSPEED_OFFLOAD_PARAM_DEVICE",
+            "ACCELERATE_DEEPSPEED_ZERO3_SAVE_16BIT_MODEL",
+            "ACCELERATE_MIXED_PRECISION",
+        ]
+        duplicate_values_flag = False
+        for name in env_variable_names_to_ignore:
+            if name != "ACCELERATE_MIXED_PRECISION" and os.environ.get(name, None) is not None:
+                duplicate_values_flag = True
+            elif name == "ACCELERATE_MIXED_PRECISION" and os.environ.get(name, "no") != "no":
+                duplicate_values_flag = True
+            if duplicate_values_flag:
+                break
+
+        if duplicate_values_flag:
+            env_variable_names_to_ignore = [
+                name.replace("ACCELERATE_", "").replace("DEEPSPEED_", "").lower()
+                for name in env_variable_names_to_ignore
+            ]
+            raise ValueError(
+                f"When using `deepspeed_config_file`, the following accelerate config variables will be ignored: {env_variable_names_to_ignore}.\n"
+                "Please specify them appropriately in the DeepSpeed config file.\n"
+                "If you are using an accelerate config file, set `mixed_precision=no` "
(incidentally `mixed_precision: fp16|bf16|'no'` is odd - why do different values don't follow the same style - i.e 'no' in quotes but not the other values)
@@ -549,6 +551,40 @@ def zero3_init_context_manager(self, enable=False):
             self.dschf = None
             self.set_deepspeed_weakref()
 
+    def _deepspeed_config_checks(self):
+        env_variable_names_to_ignore = [
+            "ACCELERATE_GRADIENT_ACCUMULATION_STEPS",
+            "ACCELERATE_GRADIENT_CLIPPING",
+            "ACCELERATE_DEEPSPEED_ZERO_STAGE",
+            "ACCELERATE_DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE",
+            "ACCELERATE_DEEPSPEED_OFFLOAD_PARAM_DEVICE",
+            "ACCELERATE_DEEPSPEED_ZERO3_SAVE_16BIT_MODEL",
+            "ACCELERATE_MIXED_PRECISION",
+        ]
+        duplicate_values_flag = False
+        for name in env_variable_names_to_ignore:
+            if name != "ACCELERATE_MIXED_PRECISION" and os.environ.get(name, None) is not None:
+                duplicate_values_flag = True
+            elif name == "ACCELERATE_MIXED_PRECISION" and os.environ.get(name, "no") != "no":
+                duplicate_values_flag = True
+            if duplicate_values_flag:
+                break
+
+        if duplicate_values_flag:
+            env_variable_names_to_ignore = [
+                name.replace("ACCELERATE_", "").replace("DEEPSPEED_", "").lower()
+                for name in env_variable_names_to_ignore
+            ]
+            raise ValueError(
+                f"When using `deepspeed_config_file`, the following accelerate config variables will be ignored: {env_variable_names_to_ignore}.\n"
+                "Please specify them appropriately in the DeepSpeed config file.\n"
+                "If you are using an accelerate config file, set `mixed_precision=no` "
Hello @stas, this case can't be helped as 'no' is the default value and if one doesn't specify the entry in accelerate config file, it will get default value of 'no' and thereby we have no way to check if user has given the default in config file or was it set to default when reading the config file. It would mean a lot of code rewriting as `mixed_precision` is used in various other non DeepSpeed parts.

The default value of `mixed_precision` would be overriden by that in `deepspeed_config_file`. Now, if you answer the questionnaire via `accelerate config`, the entry of `mixed_precision` won't be there when using `deepspeed_config_file`. 


@sgugger, what are your thoughts around this. 
@@ -549,6 +551,40 @@ def zero3_init_context_manager(self, enable=False):
             self.dschf = None
             self.set_deepspeed_weakref()
 
+    def _deepspeed_config_checks(self):
+        env_variable_names_to_ignore = [
+            "ACCELERATE_GRADIENT_ACCUMULATION_STEPS",
+            "ACCELERATE_GRADIENT_CLIPPING",
+            "ACCELERATE_DEEPSPEED_ZERO_STAGE",
+            "ACCELERATE_DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE",
+            "ACCELERATE_DEEPSPEED_OFFLOAD_PARAM_DEVICE",
+            "ACCELERATE_DEEPSPEED_ZERO3_SAVE_16BIT_MODEL",
+            "ACCELERATE_MIXED_PRECISION",
+        ]
+        duplicate_values_flag = False
+        for name in env_variable_names_to_ignore:
+            if name != "ACCELERATE_MIXED_PRECISION" and os.environ.get(name, None) is not None:
+                duplicate_values_flag = True
+            elif name == "ACCELERATE_MIXED_PRECISION" and os.environ.get(name, "no") != "no":
+                duplicate_values_flag = True
+            if duplicate_values_flag:
+                break
+
+        if duplicate_values_flag:
+            env_variable_names_to_ignore = [
+                name.replace("ACCELERATE_", "").replace("DEEPSPEED_", "").lower()
+                for name in env_variable_names_to_ignore
+            ]
+            raise ValueError(
+                f"When using `deepspeed_config_file`, the following accelerate config variables will be ignored: {env_variable_names_to_ignore}.\n"
+                "Please specify them appropriately in the DeepSpeed config file.\n"
+                "If you are using an accelerate config file, set `mixed_precision=no` "
understood! in which case I think this should work:

1. if the `mixed_precision` config isn't in the accelerate config file all is good
2. if the `mixed_precision` config is in the accelerate config file and ds_config file is used, the value of the former must match the value of the latter or assert.
@@ -549,6 +551,40 @@ def zero3_init_context_manager(self, enable=False):
             self.dschf = None
             self.set_deepspeed_weakref()
 
+    def _deepspeed_config_checks(self):
+        env_variable_names_to_ignore = [
+            "ACCELERATE_GRADIENT_ACCUMULATION_STEPS",
+            "ACCELERATE_GRADIENT_CLIPPING",
+            "ACCELERATE_DEEPSPEED_ZERO_STAGE",
+            "ACCELERATE_DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE",
+            "ACCELERATE_DEEPSPEED_OFFLOAD_PARAM_DEVICE",
+            "ACCELERATE_DEEPSPEED_ZERO3_SAVE_16BIT_MODEL",
+            "ACCELERATE_MIXED_PRECISION",
+        ]
+        duplicate_values_flag = False
+        for name in env_variable_names_to_ignore:
+            if name != "ACCELERATE_MIXED_PRECISION" and os.environ.get(name, None) is not None:
+                duplicate_values_flag = True
+            elif name == "ACCELERATE_MIXED_PRECISION" and os.environ.get(name, "no") != "no":
+                duplicate_values_flag = True
+            if duplicate_values_flag:
+                break
+
+        if duplicate_values_flag:
+            env_variable_names_to_ignore = [
+                name.replace("ACCELERATE_", "").replace("DEEPSPEED_", "").lower()
+                for name in env_variable_names_to_ignore
+            ]
+            raise ValueError(
+                f"When using `deepspeed_config_file`, the following accelerate config variables will be ignored: {env_variable_names_to_ignore}.\n"
+                "Please specify them appropriately in the DeepSpeed config file.\n"
+                "If you are using an accelerate config file, set `mixed_precision=no` "
```suggestion
Below is a sample script using `deepspeed_config_file` in different scenarios.
```
@@ -395,6 +395,194 @@ We will look at the changes needed in the code when using these.
 based on model, dataloaders, dummy optimizer and dummy schedulers provided to `prepare` method. 
 Only the `auto` fields specified in above examples are handled by `prepare` method and the rest have to be explicitly specified by the user.
 
+**Things to note when using DeepSpeed Config File**
+
+Below is a smaple script using `deepspeed_config_file` in different scenarios.
Are the lines of `-` intended? Should there be a new line before the text?
@@ -395,6 +395,194 @@ We will look at the changes needed in the code when using these.
 based on model, dataloaders, dummy optimizer and dummy schedulers provided to `prepare` method. 
 Only the `auto` fields specified in above examples are handled by `prepare` method and the rest have to be explicitly specified by the user.
 
+**Things to note when using DeepSpeed Config File**
+
+Below is a smaple script using `deepspeed_config_file` in different scenarios.
+
+Code `test.py`:
+
+```python
+from accelerate import Accelerator
+from accelerate.state import AcceleratorState
+
+
+def main():
+    accelerator = Accelerator()
+    accelerator.print(f"{AcceleratorState()}")
+
+
+if __name__ == "__main__":
+    main()
+```
+
+**Scenario 1**: Manually tampered accelerate config file having `deepspeed_config_file` along with other entries.
+
+1. `accelerate config`:
+
+```yaml
+command_file: null
+commands: null
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+  gradient_accumulation_steps: 1
+  gradient_clipping: 1.0
+  offload_optimizer_device: 'cpu'
+  offload_param_device: 'cpu'
+  zero3_init_flag: true
+  zero3_save_16bit_model: true
+  zero_stage: 3
+  deepspeed_config_file: 'ds_config.json'
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+dynamo_backend: 'NO'
+fsdp_config: {}
+gpu_ids: null
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+megatron_lm_config: {}
+num_machines: 1
+num_processes: 2
+rdzv_backend: static
+same_network: true
+tpu_name: null
+tpu_zone: null
+use_cpu: false
+```
+
+2. `ds_config.json`:
+
+```json
+{
+    "bf16": {
+        "enabled": true
+    },
+    "zero_optimization": {
+        "stage": 3,
+        "stage3_gather_16bit_weights_on_model_save": false,
+        "offload_optimizer": {
+            "device": "none"
+        },
+        "offload_param": {
+            "device": "none"
+        }
+    },
+    "gradient_clipping": 1.0,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "gradient_accumulation_steps": 10,
+    "steps_per_print": 2000000
+}
+```
+
+3. Output of `accelerate launch test.py`:
+
+```bash
+ValueError: When using `deepspeed_config_file`, the following accelerate config variables will be ignored: 
+['gradient_accumulation_steps', 'gradient_clipping', 'zero_stage', 'offload_optimizer_device', 'offload_param_device', 
+'zero3_save_16bit_model', 'mixed_precision'].
+Please specify them appropriately in the DeepSpeed config file.
+If you are using an accelerate config file, remove others config variables mentioned in the above specified list.
+The easiest method is to create a new config following the questionnaire via `accelerate config`.
+It will only ask for the necessary config variables when using `deepspeed_config_file`.
+```
+
+**Scenario 2**: Use the solution of the error to create new accelerate config and check that no ambiguity error is now thrown.
+
+1. Run `accelerate config`:
+
+```bash
+$ accelerate config
+-------------------------------------------------------------------------------------------------------------------------------In which compute environment are you running?
+This machine                                                                                                                   
+-------------------------------------------------------------------------------------------------------------------------------Which type of machine are you using?                                                                                           
```suggestion
1. Content of the `accelerate` config:
```
(to not mistake this with the command `accelerate config`
@@ -395,6 +395,194 @@ We will look at the changes needed in the code when using these.
 based on model, dataloaders, dummy optimizer and dummy schedulers provided to `prepare` method. 
 Only the `auto` fields specified in above examples are handled by `prepare` method and the rest have to be explicitly specified by the user.
 
+**Things to note when using DeepSpeed Config File**
+
+Below is a smaple script using `deepspeed_config_file` in different scenarios.
+
+Code `test.py`:
+
+```python
+from accelerate import Accelerator
+from accelerate.state import AcceleratorState
+
+
+def main():
+    accelerator = Accelerator()
+    accelerator.print(f"{AcceleratorState()}")
+
+
+if __name__ == "__main__":
+    main()
+```
+
+**Scenario 1**: Manually tampered accelerate config file having `deepspeed_config_file` along with other entries.
+
+1. `accelerate config`:
```suggestion
2. Content of the `accelerate` config:
```
@@ -395,6 +395,194 @@ We will look at the changes needed in the code when using these.
 based on model, dataloaders, dummy optimizer and dummy schedulers provided to `prepare` method. 
 Only the `auto` fields specified in above examples are handled by `prepare` method and the rest have to be explicitly specified by the user.
 
+**Things to note when using DeepSpeed Config File**
+
+Below is a smaple script using `deepspeed_config_file` in different scenarios.
+
+Code `test.py`:
+
+```python
+from accelerate import Accelerator
+from accelerate.state import AcceleratorState
+
+
+def main():
+    accelerator = Accelerator()
+    accelerator.print(f"{AcceleratorState()}")
+
+
+if __name__ == "__main__":
+    main()
+```
+
+**Scenario 1**: Manually tampered accelerate config file having `deepspeed_config_file` along with other entries.
+
+1. `accelerate config`:
+
+```yaml
+command_file: null
+commands: null
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+  gradient_accumulation_steps: 1
+  gradient_clipping: 1.0
+  offload_optimizer_device: 'cpu'
+  offload_param_device: 'cpu'
+  zero3_init_flag: true
+  zero3_save_16bit_model: true
+  zero_stage: 3
+  deepspeed_config_file: 'ds_config.json'
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+dynamo_backend: 'NO'
+fsdp_config: {}
+gpu_ids: null
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+megatron_lm_config: {}
+num_machines: 1
+num_processes: 2
+rdzv_backend: static
+same_network: true
+tpu_name: null
+tpu_zone: null
+use_cpu: false
+```
+
+2. `ds_config.json`:
+
+```json
+{
+    "bf16": {
+        "enabled": true
+    },
+    "zero_optimization": {
+        "stage": 3,
+        "stage3_gather_16bit_weights_on_model_save": false,
+        "offload_optimizer": {
+            "device": "none"
+        },
+        "offload_param": {
+            "device": "none"
+        }
+    },
+    "gradient_clipping": 1.0,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "gradient_accumulation_steps": 10,
+    "steps_per_print": 2000000
+}
+```
+
+3. Output of `accelerate launch test.py`:
+
+```bash
+ValueError: When using `deepspeed_config_file`, the following accelerate config variables will be ignored: 
+['gradient_accumulation_steps', 'gradient_clipping', 'zero_stage', 'offload_optimizer_device', 'offload_param_device', 
+'zero3_save_16bit_model', 'mixed_precision'].
+Please specify them appropriately in the DeepSpeed config file.
+If you are using an accelerate config file, remove others config variables mentioned in the above specified list.
+The easiest method is to create a new config following the questionnaire via `accelerate config`.
+It will only ask for the necessary config variables when using `deepspeed_config_file`.
+```
+
+**Scenario 2**: Use the solution of the error to create new accelerate config and check that no ambiguity error is now thrown.
+
+1. Run `accelerate config`:
+
+```bash
+$ accelerate config
+-------------------------------------------------------------------------------------------------------------------------------In which compute environment are you running?
+This machine                                                                                                                   
+-------------------------------------------------------------------------------------------------------------------------------Which type of machine are you using?                                                                                           
+multi-GPU                                                                                                                      
+How many different machines will you use (use more than 1 for multi-node training)? [1]:                                       
+Do you wish to optimize your script with torch dynamo?[yes/NO]:                                                                
+Do you want to use DeepSpeed? [yes/NO]: yes                                                                                    
+Do you want to specify a json file to a DeepSpeed config? [yes/NO]: yes                                                        
+Please enter the path to the json DeepSpeed config file: ds_config.json                                                        
+Do you want to enable `deepspeed.zero.Init` when using ZeRO Stage-3 for constructing massive models? [yes/NO]: yes
+How many GPU(s) should be used for distributed training? [1]:4
+accelerate configuration saved at ds_config_sample.yaml
+```
+
+2. `accelerate config`:
```suggestion
**Scenario 3**: Setting the `accelerate launch` command arguments related to DeepSpeed as `"auto"` in the DeepSpeed` configuration file and check that things work as expected.
```
@@ -395,6 +395,194 @@ We will look at the changes needed in the code when using these.
 based on model, dataloaders, dummy optimizer and dummy schedulers provided to `prepare` method. 
 Only the `auto` fields specified in above examples are handled by `prepare` method and the rest have to be explicitly specified by the user.
 
+**Things to note when using DeepSpeed Config File**
+
+Below is a smaple script using `deepspeed_config_file` in different scenarios.
+
+Code `test.py`:
+
+```python
+from accelerate import Accelerator
+from accelerate.state import AcceleratorState
+
+
+def main():
+    accelerator = Accelerator()
+    accelerator.print(f"{AcceleratorState()}")
+
+
+if __name__ == "__main__":
+    main()
+```
+
+**Scenario 1**: Manually tampered accelerate config file having `deepspeed_config_file` along with other entries.
+
+1. `accelerate config`:
+
+```yaml
+command_file: null
+commands: null
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+  gradient_accumulation_steps: 1
+  gradient_clipping: 1.0
+  offload_optimizer_device: 'cpu'
+  offload_param_device: 'cpu'
+  zero3_init_flag: true
+  zero3_save_16bit_model: true
+  zero_stage: 3
+  deepspeed_config_file: 'ds_config.json'
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+dynamo_backend: 'NO'
+fsdp_config: {}
+gpu_ids: null
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+megatron_lm_config: {}
+num_machines: 1
+num_processes: 2
+rdzv_backend: static
+same_network: true
+tpu_name: null
+tpu_zone: null
+use_cpu: false
+```
+
+2. `ds_config.json`:
+
+```json
+{
+    "bf16": {
+        "enabled": true
+    },
+    "zero_optimization": {
+        "stage": 3,
+        "stage3_gather_16bit_weights_on_model_save": false,
+        "offload_optimizer": {
+            "device": "none"
+        },
+        "offload_param": {
+            "device": "none"
+        }
+    },
+    "gradient_clipping": 1.0,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "gradient_accumulation_steps": 10,
+    "steps_per_print": 2000000
+}
+```
+
+3. Output of `accelerate launch test.py`:
+
+```bash
+ValueError: When using `deepspeed_config_file`, the following accelerate config variables will be ignored: 
+['gradient_accumulation_steps', 'gradient_clipping', 'zero_stage', 'offload_optimizer_device', 'offload_param_device', 
+'zero3_save_16bit_model', 'mixed_precision'].
+Please specify them appropriately in the DeepSpeed config file.
+If you are using an accelerate config file, remove others config variables mentioned in the above specified list.
+The easiest method is to create a new config following the questionnaire via `accelerate config`.
+It will only ask for the necessary config variables when using `deepspeed_config_file`.
+```
+
+**Scenario 2**: Use the solution of the error to create new accelerate config and check that no ambiguity error is now thrown.
+
+1. Run `accelerate config`:
+
+```bash
+$ accelerate config
+-------------------------------------------------------------------------------------------------------------------------------In which compute environment are you running?
+This machine                                                                                                                   
+-------------------------------------------------------------------------------------------------------------------------------Which type of machine are you using?                                                                                           
+multi-GPU                                                                                                                      
+How many different machines will you use (use more than 1 for multi-node training)? [1]:                                       
+Do you wish to optimize your script with torch dynamo?[yes/NO]:                                                                
+Do you want to use DeepSpeed? [yes/NO]: yes                                                                                    
+Do you want to specify a json file to a DeepSpeed config? [yes/NO]: yes                                                        
+Please enter the path to the json DeepSpeed config file: ds_config.json                                                        
+Do you want to enable `deepspeed.zero.Init` when using ZeRO Stage-3 for constructing massive models? [yes/NO]: yes
+How many GPU(s) should be used for distributed training? [1]:4
+accelerate configuration saved at ds_config_sample.yaml
+```
+
+2. `accelerate config`:
+
+```yaml
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+  deepspeed_config_file: ds_config.json
+  zero3_init_flag: true
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+dynamo_backend: 'NO'
+fsdp_config: {}
+machine_rank: 0
+main_training_function: main
+megatron_lm_config: {}
+num_machines: 1
+num_processes: 4
+rdzv_backend: static
+same_network: true
+use_cpu: false
+```
+
+3. Output of `accelerate launch test.py`:
+
+```bash
+Distributed environment: DEEPSPEED  Backend: nccl
+Num processes: 4
+Process index: 0
+Local process index: 0
+Device: cuda:0
+Mixed precision type: bf16
+ds_config: {'bf16': {'enabled': True}, 'zero_optimization': {'stage': 3, 'stage3_gather_16bit_weights_on_model_save': False, 'offload_optimizer': {'device': 'none'}, 'offload_param': {'device': 'none'}}, 'gradient_clipping': 1.0, 'train_batch_size': 'auto', 'train_micro_batch_size_per_gpu': 'auto', 'gradient_accumulation_steps': 10, 'steps_per_print': inf, 'fp16': {'enabled': False}}
+```
+
+**Scenario 3**: Setting the `accelerate launch` cmd args related to deepspeed as `auto` in  `deepspeed_config_file` and check that things work as expected.
```suggestion
1. New `ds_config.json` with `"auto"` for the `accelerate launch` DeepSpeed command arguments:
```
@@ -395,6 +395,194 @@ We will look at the changes needed in the code when using these.
 based on model, dataloaders, dummy optimizer and dummy schedulers provided to `prepare` method. 
 Only the `auto` fields specified in above examples are handled by `prepare` method and the rest have to be explicitly specified by the user.
 
+**Things to note when using DeepSpeed Config File**
+
+Below is a smaple script using `deepspeed_config_file` in different scenarios.
+
+Code `test.py`:
+
+```python
+from accelerate import Accelerator
+from accelerate.state import AcceleratorState
+
+
+def main():
+    accelerator = Accelerator()
+    accelerator.print(f"{AcceleratorState()}")
+
+
+if __name__ == "__main__":
+    main()
+```
+
+**Scenario 1**: Manually tampered accelerate config file having `deepspeed_config_file` along with other entries.
+
+1. `accelerate config`:
+
+```yaml
+command_file: null
+commands: null
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+  gradient_accumulation_steps: 1
+  gradient_clipping: 1.0
+  offload_optimizer_device: 'cpu'
+  offload_param_device: 'cpu'
+  zero3_init_flag: true
+  zero3_save_16bit_model: true
+  zero_stage: 3
+  deepspeed_config_file: 'ds_config.json'
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+dynamo_backend: 'NO'
+fsdp_config: {}
+gpu_ids: null
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+megatron_lm_config: {}
+num_machines: 1
+num_processes: 2
+rdzv_backend: static
+same_network: true
+tpu_name: null
+tpu_zone: null
+use_cpu: false
+```
+
+2. `ds_config.json`:
+
+```json
+{
+    "bf16": {
+        "enabled": true
+    },
+    "zero_optimization": {
+        "stage": 3,
+        "stage3_gather_16bit_weights_on_model_save": false,
+        "offload_optimizer": {
+            "device": "none"
+        },
+        "offload_param": {
+            "device": "none"
+        }
+    },
+    "gradient_clipping": 1.0,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "gradient_accumulation_steps": 10,
+    "steps_per_print": 2000000
+}
+```
+
+3. Output of `accelerate launch test.py`:
+
+```bash
+ValueError: When using `deepspeed_config_file`, the following accelerate config variables will be ignored: 
+['gradient_accumulation_steps', 'gradient_clipping', 'zero_stage', 'offload_optimizer_device', 'offload_param_device', 
+'zero3_save_16bit_model', 'mixed_precision'].
+Please specify them appropriately in the DeepSpeed config file.
+If you are using an accelerate config file, remove others config variables mentioned in the above specified list.
+The easiest method is to create a new config following the questionnaire via `accelerate config`.
+It will only ask for the necessary config variables when using `deepspeed_config_file`.
+```
+
+**Scenario 2**: Use the solution of the error to create new accelerate config and check that no ambiguity error is now thrown.
+
+1. Run `accelerate config`:
+
+```bash
+$ accelerate config
+-------------------------------------------------------------------------------------------------------------------------------In which compute environment are you running?
+This machine                                                                                                                   
+-------------------------------------------------------------------------------------------------------------------------------Which type of machine are you using?                                                                                           
+multi-GPU                                                                                                                      
+How many different machines will you use (use more than 1 for multi-node training)? [1]:                                       
+Do you wish to optimize your script with torch dynamo?[yes/NO]:                                                                
+Do you want to use DeepSpeed? [yes/NO]: yes                                                                                    
+Do you want to specify a json file to a DeepSpeed config? [yes/NO]: yes                                                        
+Please enter the path to the json DeepSpeed config file: ds_config.json                                                        
+Do you want to enable `deepspeed.zero.Init` when using ZeRO Stage-3 for constructing massive models? [yes/NO]: yes
+How many GPU(s) should be used for distributed training? [1]:4
+accelerate configuration saved at ds_config_sample.yaml
+```
+
+2. `accelerate config`:
+
+```yaml
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+  deepspeed_config_file: ds_config.json
+  zero3_init_flag: true
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+dynamo_backend: 'NO'
+fsdp_config: {}
+machine_rank: 0
+main_training_function: main
+megatron_lm_config: {}
+num_machines: 1
+num_processes: 4
+rdzv_backend: static
+same_network: true
+use_cpu: false
+```
+
+3. Output of `accelerate launch test.py`:
+
+```bash
+Distributed environment: DEEPSPEED  Backend: nccl
+Num processes: 4
+Process index: 0
+Local process index: 0
+Device: cuda:0
+Mixed precision type: bf16
+ds_config: {'bf16': {'enabled': True}, 'zero_optimization': {'stage': 3, 'stage3_gather_16bit_weights_on_model_save': False, 'offload_optimizer': {'device': 'none'}, 'offload_param': {'device': 'none'}}, 'gradient_clipping': 1.0, 'train_batch_size': 'auto', 'train_micro_batch_size_per_gpu': 'auto', 'gradient_accumulation_steps': 10, 'steps_per_print': inf, 'fp16': {'enabled': False}}
+```
+
+**Scenario 3**: Setting the `accelerate launch` cmd args related to deepspeed as `auto` in  `deepspeed_config_file` and check that things work as expected.
+
+1. New `ds_config.json` with `auto` for the `accelerate launch` deepspeed cmd args:
```suggestion
**Note**: Remaining `"auto"` values are handled in `accelerator.prepare()` call as explained in point 2 of 
```
@@ -395,6 +395,194 @@ We will look at the changes needed in the code when using these.
 based on model, dataloaders, dummy optimizer and dummy schedulers provided to `prepare` method. 
 Only the `auto` fields specified in above examples are handled by `prepare` method and the rest have to be explicitly specified by the user.
 
+**Things to note when using DeepSpeed Config File**
+
+Below is a smaple script using `deepspeed_config_file` in different scenarios.
+
+Code `test.py`:
+
+```python
+from accelerate import Accelerator
+from accelerate.state import AcceleratorState
+
+
+def main():
+    accelerator = Accelerator()
+    accelerator.print(f"{AcceleratorState()}")
+
+
+if __name__ == "__main__":
+    main()
+```
+
+**Scenario 1**: Manually tampered accelerate config file having `deepspeed_config_file` along with other entries.
+
+1. `accelerate config`:
+
+```yaml
+command_file: null
+commands: null
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+  gradient_accumulation_steps: 1
+  gradient_clipping: 1.0
+  offload_optimizer_device: 'cpu'
+  offload_param_device: 'cpu'
+  zero3_init_flag: true
+  zero3_save_16bit_model: true
+  zero_stage: 3
+  deepspeed_config_file: 'ds_config.json'
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+dynamo_backend: 'NO'
+fsdp_config: {}
+gpu_ids: null
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+megatron_lm_config: {}
+num_machines: 1
+num_processes: 2
+rdzv_backend: static
+same_network: true
+tpu_name: null
+tpu_zone: null
+use_cpu: false
+```
+
+2. `ds_config.json`:
+
+```json
+{
+    "bf16": {
+        "enabled": true
+    },
+    "zero_optimization": {
+        "stage": 3,
+        "stage3_gather_16bit_weights_on_model_save": false,
+        "offload_optimizer": {
+            "device": "none"
+        },
+        "offload_param": {
+            "device": "none"
+        }
+    },
+    "gradient_clipping": 1.0,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "gradient_accumulation_steps": 10,
+    "steps_per_print": 2000000
+}
+```
+
+3. Output of `accelerate launch test.py`:
+
+```bash
+ValueError: When using `deepspeed_config_file`, the following accelerate config variables will be ignored: 
+['gradient_accumulation_steps', 'gradient_clipping', 'zero_stage', 'offload_optimizer_device', 'offload_param_device', 
+'zero3_save_16bit_model', 'mixed_precision'].
+Please specify them appropriately in the DeepSpeed config file.
+If you are using an accelerate config file, remove others config variables mentioned in the above specified list.
+The easiest method is to create a new config following the questionnaire via `accelerate config`.
+It will only ask for the necessary config variables when using `deepspeed_config_file`.
+```
+
+**Scenario 2**: Use the solution of the error to create new accelerate config and check that no ambiguity error is now thrown.
+
+1. Run `accelerate config`:
+
+```bash
+$ accelerate config
+-------------------------------------------------------------------------------------------------------------------------------In which compute environment are you running?
+This machine                                                                                                                   
+-------------------------------------------------------------------------------------------------------------------------------Which type of machine are you using?                                                                                           
+multi-GPU                                                                                                                      
+How many different machines will you use (use more than 1 for multi-node training)? [1]:                                       
+Do you wish to optimize your script with torch dynamo?[yes/NO]:                                                                
+Do you want to use DeepSpeed? [yes/NO]: yes                                                                                    
+Do you want to specify a json file to a DeepSpeed config? [yes/NO]: yes                                                        
+Please enter the path to the json DeepSpeed config file: ds_config.json                                                        
+Do you want to enable `deepspeed.zero.Init` when using ZeRO Stage-3 for constructing massive models? [yes/NO]: yes
+How many GPU(s) should be used for distributed training? [1]:4
+accelerate configuration saved at ds_config_sample.yaml
+```
+
+2. `accelerate config`:
+
+```yaml
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+  deepspeed_config_file: ds_config.json
+  zero3_init_flag: true
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+dynamo_backend: 'NO'
+fsdp_config: {}
+machine_rank: 0
+main_training_function: main
+megatron_lm_config: {}
+num_machines: 1
+num_processes: 4
+rdzv_backend: static
+same_network: true
+use_cpu: false
+```
+
+3. Output of `accelerate launch test.py`:
+
+```bash
+Distributed environment: DEEPSPEED  Backend: nccl
+Num processes: 4
+Process index: 0
+Local process index: 0
+Device: cuda:0
+Mixed precision type: bf16
+ds_config: {'bf16': {'enabled': True}, 'zero_optimization': {'stage': 3, 'stage3_gather_16bit_weights_on_model_save': False, 'offload_optimizer': {'device': 'none'}, 'offload_param': {'device': 'none'}}, 'gradient_clipping': 1.0, 'train_batch_size': 'auto', 'train_micro_batch_size_per_gpu': 'auto', 'gradient_accumulation_steps': 10, 'steps_per_print': inf, 'fp16': {'enabled': False}}
+```
+
+**Scenario 3**: Setting the `accelerate launch` cmd args related to deepspeed as `auto` in  `deepspeed_config_file` and check that things work as expected.
+
+1. New `ds_config.json` with `auto` for the `accelerate launch` deepspeed cmd args:
+
+```json
+{
+    "bf16": {
+        "enabled": "auto"
+    },
+    "zero_optimization": {
+        "stage": "auto",
+        "stage3_gather_16bit_weights_on_model_save": "auto",
+        "offload_optimizer": {
+            "device": "auto"
+        },
+        "offload_param": {
+            "device": "auto"
+        }
+    },
+    "gradient_clipping": "auto",
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "gradient_accumulation_steps": "auto",
+    "steps_per_print": 2000000
+}
+```
+
+2. Output of `accelerate launch --mixed_precision="fp16" --zero_stage=3 --gradient_accumulation_steps=5 --gradient_clipping=1.0 --offload_param_device="cpu" --offload_optimizer_device="nvme" --zero3_save_16bit_model="true" test.py`:
+
+```bash
+Distributed environment: DEEPSPEED  Backend: nccl
+Num processes: 4
+Process index: 0
+Local process index: 0
+Device: cuda:0
+Mixed precision type: fp16
+ds_config: {'bf16': {'enabled': False}, 'zero_optimization': {'stage': 3, 'stage3_gather_16bit_weights_on_model_save': True, 'offload_optimizer': {'device': 'nvme'}, 'offload_param': {'device': 'cpu'}}, 'gradient_clipping': 1.0, 'train_batch_size': 'auto', 'train_micro_batch_size_per_gpu': 'auto', 'gradient_accumulation_steps': 5, 'steps_per_print': inf, 'fp16': {'enabled': True, 'auto_cast': True}}
+```
+
+**Note**: Remaining `auto` values are handled in `accelerator.prepare()` call as explained in point 2 of 
Why is this set to True by default now?
@@ -450,15 +461,17 @@ def __post_init__(self):
             if self.gradient_clipping:
                 config["gradient_clipping"] = self.gradient_clipping
             self.hf_ds_config = HfDeepSpeedConfig(config)
+
         self.deepspeed_config = self.hf_ds_config.config
         self.deepspeed_config["steps_per_print"] = float("inf")  # this will stop deepspeed from logging @ stdout
         if self.zero3_init_flag is None:
-            self.zero3_init_flag = os.environ.get("ACCELERATE_DEEPSPEED_ZERO3_INIT", "false") == "true"
+            self.zero3_init_flag = os.environ.get("ACCELERATE_DEEPSPEED_ZERO3_INIT", "true") == "true"
these are the exact outputs from the new config #830. Added new lines.  
@@ -395,6 +395,194 @@ We will look at the changes needed in the code when using these.
 based on model, dataloaders, dummy optimizer and dummy schedulers provided to `prepare` method. 
 Only the `auto` fields specified in above examples are handled by `prepare` method and the rest have to be explicitly specified by the user.
 
+**Things to note when using DeepSpeed Config File**
+
+Below is a smaple script using `deepspeed_config_file` in different scenarios.
+
+Code `test.py`:
+
+```python
+from accelerate import Accelerator
+from accelerate.state import AcceleratorState
+
+
+def main():
+    accelerator = Accelerator()
+    accelerator.print(f"{AcceleratorState()}")
+
+
+if __name__ == "__main__":
+    main()
+```
+
+**Scenario 1**: Manually tampered accelerate config file having `deepspeed_config_file` along with other entries.
+
+1. `accelerate config`:
+
+```yaml
+command_file: null
+commands: null
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+  gradient_accumulation_steps: 1
+  gradient_clipping: 1.0
+  offload_optimizer_device: 'cpu'
+  offload_param_device: 'cpu'
+  zero3_init_flag: true
+  zero3_save_16bit_model: true
+  zero_stage: 3
+  deepspeed_config_file: 'ds_config.json'
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+dynamo_backend: 'NO'
+fsdp_config: {}
+gpu_ids: null
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+megatron_lm_config: {}
+num_machines: 1
+num_processes: 2
+rdzv_backend: static
+same_network: true
+tpu_name: null
+tpu_zone: null
+use_cpu: false
+```
+
+2. `ds_config.json`:
+
+```json
+{
+    "bf16": {
+        "enabled": true
+    },
+    "zero_optimization": {
+        "stage": 3,
+        "stage3_gather_16bit_weights_on_model_save": false,
+        "offload_optimizer": {
+            "device": "none"
+        },
+        "offload_param": {
+            "device": "none"
+        }
+    },
+    "gradient_clipping": 1.0,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "gradient_accumulation_steps": 10,
+    "steps_per_print": 2000000
+}
+```
+
+3. Output of `accelerate launch test.py`:
+
+```bash
+ValueError: When using `deepspeed_config_file`, the following accelerate config variables will be ignored: 
+['gradient_accumulation_steps', 'gradient_clipping', 'zero_stage', 'offload_optimizer_device', 'offload_param_device', 
+'zero3_save_16bit_model', 'mixed_precision'].
+Please specify them appropriately in the DeepSpeed config file.
+If you are using an accelerate config file, remove others config variables mentioned in the above specified list.
+The easiest method is to create a new config following the questionnaire via `accelerate config`.
+It will only ask for the necessary config variables when using `deepspeed_config_file`.
+```
+
+**Scenario 2**: Use the solution of the error to create new accelerate config and check that no ambiguity error is now thrown.
+
+1. Run `accelerate config`:
+
+```bash
+$ accelerate config
+-------------------------------------------------------------------------------------------------------------------------------In which compute environment are you running?
+This machine                                                                                                                   
+-------------------------------------------------------------------------------------------------------------------------------Which type of machine are you using?                                                                                           
Hello, as stas mentioned in Issue #922 

> this flag should be True by default as zero stage 3 is for large models, it's very unlikely the user will be able to load those models w/o zero.Init
@@ -450,15 +461,17 @@ def __post_init__(self):
             if self.gradient_clipping:
                 config["gradient_clipping"] = self.gradient_clipping
             self.hf_ds_config = HfDeepSpeedConfig(config)
+
         self.deepspeed_config = self.hf_ds_config.config
         self.deepspeed_config["steps_per_print"] = float("inf")  # this will stop deepspeed from logging @ stdout
         if self.zero3_init_flag is None:
-            self.zero3_init_flag = os.environ.get("ACCELERATE_DEEPSPEED_ZERO3_INIT", "false") == "true"
+            self.zero3_init_flag = os.environ.get("ACCELERATE_DEEPSPEED_ZERO3_INIT", "true") == "true"
Is this only used if zero3 is enabled in the rest of the code? It should at least only default to True when Zero-3 is enabled otherwise.
@@ -450,15 +461,17 @@ def __post_init__(self):
             if self.gradient_clipping:
                 config["gradient_clipping"] = self.gradient_clipping
             self.hf_ds_config = HfDeepSpeedConfig(config)
+
         self.deepspeed_config = self.hf_ds_config.config
         self.deepspeed_config["steps_per_print"] = float("inf")  # this will stop deepspeed from logging @ stdout
         if self.zero3_init_flag is None:
-            self.zero3_init_flag = os.environ.get("ACCELERATE_DEEPSPEED_ZERO3_INIT", "false") == "true"
+            self.zero3_init_flag = os.environ.get("ACCELERATE_DEEPSPEED_ZERO3_INIT", "true") == "true"
Yes, only used when zero3 is enabled else will give a warning and gets set to False as mentioned in on of the above threads. Done, defaults to True only for stage 3.
@@ -450,15 +461,17 @@ def __post_init__(self):
             if self.gradient_clipping:
                 config["gradient_clipping"] = self.gradient_clipping
             self.hf_ds_config = HfDeepSpeedConfig(config)
+
         self.deepspeed_config = self.hf_ds_config.config
         self.deepspeed_config["steps_per_print"] = float("inf")  # this will stop deepspeed from logging @ stdout
         if self.zero3_init_flag is None:
-            self.zero3_init_flag = os.environ.get("ACCELERATE_DEEPSPEED_ZERO3_INIT", "false") == "true"
+            self.zero3_init_flag = os.environ.get("ACCELERATE_DEEPSPEED_ZERO3_INIT", "true") == "true"
_The documentation is not available anymore as the PR was closed or merged._
I tried this branch, getting:

```
Traceback (most recent call last):
  File "m4/training/main.py", line 47, in <module>
    accelerator = Accelerator(
  File "/mnt/nvme0/code/huggingface/accelerate-master/src/accelerate/accelerator.py", line 246, in __init__
    DeepSpeedPlugin() if os.environ.get("ACCELERATE_USE_DEEPSPEED", "false") == "true" else None
  File "<string>", line 12, in __init__
  File "/mnt/nvme0/code/huggingface/accelerate-master/src/accelerate/utils/dataclasses.py", line 412, in __post_init__
    self._deepspeed_config_checks()
  File "/mnt/nvme0/code/huggingface/accelerate-master/src/accelerate/utils/dataclasses.py", line 560, in _deepspeed_config_checks
    if ds_gradient_accumulation_steps != int(accelerate_gradient_accumulation_steps):
ValueError: invalid literal for int() with base 10: 'None'
```

configs:

```
{
    "fp16": {
        "enabled": true,
        "auto_cast": true,
        "loss_scale": 0.0,
        "initial_scale_power": 10,
        "loss_scale_window": 1000,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "zero_optimization": {
        "stage": 3,
        "allgather_partitions": true,
        "allgather_bucket_size": 5e8,
        "overlap_comm": false,
        "reduce_scatter": true,
        "reduce_bucket_size": "auto",
        "contiguous_gradients": true,
        "stage3_gather_16bit_weights_on_model_save": false,
        "offload_optimizer": {
            "device": "none"
        },
        "offload_param": {
            "device": "none"
        }
    },
    "gradient_clipping": 1.0,
    "gradient_accumulation_steps": 2,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "steps_per_print": 2000000
}

compute_environment: LOCAL_MACHINE
deepspeed_config:
  deepspeed_multinode_launcher: standard
  deepspeed_config_file: ./configs/vopt-large-z3/ds_config.json
  zero3_init_flag: true
distributed_type: DEEPSPEED
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
num_machines: 1
num_processes: 1
#num_machines: 20
#num_processes: 80
use_cpu: false


```
@pacman100 default needs to be 1 instead of `None` :)
Also clarifying defaults for args in this PR. Now, `accelerate launch --use_deepspeed --help` shows below output wherein default value info is specified `If unspecified, will default to *`:
```
...

DeepSpeed Arguments:
  Arguments related to DeepSpeed.

  --deepspeed_config_file DEEPSPEED_CONFIG_FILE
                        DeepSpeed config file.
  --zero_stage ZERO_STAGE
                        DeepSpeed's ZeRO optimization stage (useful only when `use_deepspeed` flag
                        is passed). If unspecified, will default to `2`.
  --offload_optimizer_device OFFLOAD_OPTIMIZER_DEVICE
                        Decides where (none|cpu|nvme) to offload optimizer states (useful only
                        when `use_deepspeed` flag is passed). If unspecified, will default to
                        `none`.
  --offload_param_device OFFLOAD_PARAM_DEVICE
                        Decides where (none|cpu|nvme) to offload parameters (useful only when
                        `use_deepspeed` flag is passed). If unspecified, will default to `none`.
  --gradient_accumulation_steps GRADIENT_ACCUMULATION_STEPS
                        No of gradient_accumulation_steps used in your training script (useful
                        only when `use_deepspeed` flag is passed). If unspecified, will default to
                        `1`.
  --gradient_clipping GRADIENT_CLIPPING
                        gradient clipping value used in your training script (useful only when
                        `use_deepspeed` flag is passed). If unspecified, will default to `1.0`.
  --zero3_init_flag ZERO3_INIT_FLAG
                        Decides Whether (true|false) to enable `deepspeed.zero.Init` for
                        constructing massive models. Only applicable with DeepSpeed ZeRO Stage-3.
                        If unspecified, will default to `true`.
  --zero3_save_16bit_model ZERO3_SAVE_16BIT_MODEL
                        Decides Whether (true|false) to save 16-bit model weights when using ZeRO
                        Stage-3. Only applicable with DeepSpeed ZeRO Stage-3. If unspecified, will
                        default to `false`.
  --deepspeed_hostfile DEEPSPEED_HOSTFILE
                        DeepSpeed hostfile for configuring multi-node compute resources.
  --deepspeed_exclusion_filter DEEPSPEED_EXCLUSION_FILTER
                        DeepSpeed exclusion filter string when using mutli-node setup.
  --deepspeed_inclusion_filter DEEPSPEED_INCLUSION_FILTER
                        DeepSpeed inclusion filter string when using mutli-node setup.
  --deepspeed_multinode_launcher DEEPSPEED_MULTINODE_LAUNCHER
                        DeepSpeed multi-node launcher to use. If unspecified, will default to
                        `pdsh`.

```
> ```shell
> ValueError: When using `deepspeed_config_file`, the following accelerate config variables will be 
> ignored: ['gradient_accumulation_steps', 'gradient_clipping', 'zero_stage', 
> 'offload_optimizer_device', 'offload_param_device', 'zero3_save_16bit_model', 'mixed_precision'].
> Please specify them appropriately in the DeepSpeed config file.
> If you are using accelerate config file, set `mixed_precision=no` and remove others config variables
> mentioned in the above specified list; else don't specify these config variables in `accelerate 
> launch` command. 
> The easiest method is to create new config following the questionnaire via  `accelerate config`.
> It will only ask for the necessary config variables when using `deepspeed_config_file`.
> ```

Now, the error will be this
> but really what we want is one definitive source of information.

Hello @stas00, let us know if this addresses the issue. When `deepspeed_config_file` is specified, it is the single definitive source of information and error is raised when duplicates are found either in accelerate config file or through the arguments of `accelerate launch` command
> Will this still work with `auto` entries in `ds_config.json`

Hello Stas, they will work as before because nothing is changing other than the fact that an error is raised if their is any duplicate entry in `accelerate config` with things that are in `deepspeed_config_file`. Here, these fields which can be duplicates are never `auto` in either case and hence fields that had `auto` support won't be impacted.
Thank you for confirming that, @pacman100 - let me try it out as it's easier to see things in action.
But this shows otherwise:
```
  File "/mnt/nvme0/code/huggingface/accelerate-master/src/accelerate/utils/dataclasses.py", line 409, in __post_init__
    raise ValueError("gradient_accumulation_steps cannot be set to 'auto' in the DeepSpeed config.")
ValueError: gradient_accumulation_steps cannot be set to 'auto' in the DeepSpeed config.
ERROR:torch.distributed.elastic.multiprocessing.api:failed (exitcode: 1) local_rank: 0 (pid
```

why can't it be `auto`? I'd expect it to be set either (1) to the default value or (2) the value passed via the `--gradient_accumulation_steps` cmd line arg.

I wonder if we are having a miscommunication here. I brought up the issue of duplicity of the 2 styles of settings the ds config, since accelerate used its own config files from the beginning, but I have never suggested that setting values via cmd line args support should be dropped.
Regarding `gradient_accumulation_steps` and `auto`, that piece of code wasn't changed at all and behaviour is the same as before. If one doesn't specify that entry in the config file, it is set to default of `1`, if it is `auto` error is raised else thale given value is used. When not using `deepspeed_config_file`, it is asked when using `accelerate config` cmd.
> Regarding `gradient_accumulation_steps` and `auto`, that piece of code wasn't changed at all and behaviour is the same as before. If one doesn't specify that entry in the config file, it is set to default of `1`, if it is `auto` error is raised else thale given value is used. When not using `deepspeed_config_file`, it is asked when using `accelerate config` cmd.

OK, so your logic is different from HF Trainer then.

The HF Trainer was directing cmd args into the config file's `auto` values so that the user could override them via cmd line args.

I'm not saying the 2 logics have to match. If I am not mistaken the accelerate logic is less flexible, but it's ok if you prefer it that way. 

In HF Trainer the `auto` feature was designed to be used:
1. when the value can't be known before running - "boot"-time calculated configs
2. values to be set via cmd line args and defaults
As accelerate is meant to work with all models apart from Transformers and user being in control of the training loop, they are in charge of all the arguments and the naming convention of arguments will be different across different users. On the other hand, in Trainer, users are restricted to a given args set and as such those can be used to fill the DeepSpeed config due to clear mapping between args and DS config params. The idea is that artifacts sent to `accelerator.prepare` have the params required by DS config and we exactly know the mapping between them and makes filling of the params independent of the user's training loop and their arguments naming convention.

In accelerate, the `auto` values are those that can be filled via artifacts being sent to `accelerator.prepare` as all the other places,  the  user has complete control over the training loop, argument naming and filling. The user still has flexibility to fill in all the `auto` values themselves as mentioned here https://github.com/huggingface/accelerate/pull/676#issuecomment-1318139353
> As accelerate is meant to work with all models apart from Transformers and user being in control of the training loop, they are in charge of all the arguments and the naming convention of arguments will be different across different users. On the other hand, in Trainer, users are restricted to a given args set and as such those can be used to fill the DeepSpeed config due to clear mapping between args and DS config params. The idea is that artifacts sent to `accelerator.prepare` have the params required by DS config and we exactly know the mapping between them and makes filling of the params independent of the user's training loop and their arguments naming convention.

Thank you for explaining this to me, Sourab, but I'm having a hard time following how Accelerate is any different from HF Trainer wrt sending cmd line arg values to the unfilled out config values in ds_config. e.g.  the Accelerate launcher provides an explicit list of cmd line args for the deepspeed use. There is a 1:1 mapping here as well. Could you please explain how is this different from the HF Trainer?

But as I said above it's totally fine if you prefer to do it this way, Sourab. This doesn't prevent users from doing what they need.

> In accelerate, the auto values are those that can be filled via artifacts being sent to accelerator.prepare as all the other places, the user has complete control over the training loop, argument naming and filling. The user still has flexibility to fill in all the auto values themselves as mentioned here https://github.com/huggingface/accelerate/pull/676#issuecomment-1318139353

Understood. more work, but doable. Thank you for the explanations.


> Could you please explain how is this different from the HF Trainer?

Users can have `bs` or `batch_size` as cmd arguments in their code and as such we can't fill ds config's `train_micro_batch_size_per_gpu` whereas Trainer always maps `args.per_device_train_batch_size` to it. The same reason can go for other configs. 

Please note that `accelerate launch` cmd args are primarily used for setting accelerate config's deespseed fields rather than setting `auto` values of `deepspeed_config_file`. Now I understood that you meant using `accelerate launch` cmd args for filling in `auto` values of `deepspeed_config_file` and I've made respective changes.

Latest changes:

Code `test.py`:
```python
from accelerate import Accelerator
from accelerate.state import AcceleratorState

def main():
    accelerator = Accelerator()
    accelerator.print(f"{AcceleratorState()}")
if __name__ == "__main__":
    main()
```

Scenario 1: manually tampered accelerate config file having `deepspeed_config_file` along with other entries.

1. `accelerate config`:
```yaml
command_file: null
commands: null
compute_environment: LOCAL_MACHINE
deepspeed_config:
  gradient_accumulation_steps: 1
  gradient_clipping: 1.0
  offload_optimizer_device: 'cpu'
  offload_param_device: 'cpu'
  zero3_init_flag: true
  zero3_save_16bit_model: true
  zero_stage: 3
  deepspeed_config_file: 'ds_config.json'
distributed_type: DEEPSPEED
downcast_bf16: 'no'
dynamo_backend: 'NO'
fsdp_config: {}
gpu_ids: null
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
megatron_lm_config: {}
num_machines: 1
num_processes: 2
rdzv_backend: static
same_network: true
tpu_name: null
tpu_zone: null
use_cpu: false
```

2. `ds_config.json`:
```json
{
    "bf16": {
        "enabled": true
    },
    "zero_optimization": {
        "stage": 3,
        "stage3_gather_16bit_weights_on_model_save": false,
        "offload_optimizer": {
            "device": "none"
        },
        "offload_param": {
            "device": "none"
        }
    },
    "gradient_clipping": 1.0,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "gradient_accumulation_steps": 10,
    "steps_per_print": 2000000
}
```

3. Output of `accelerate launch test.py`:
```
ValueError: When using `deepspeed_config_file`, the following accelerate config variables will be ignored: 
['gradient_accumulation_steps', 'gradient_clipping', 'zero_stage', 'offload_optimizer_device', 'offload_param_device', 
'zero3_save_16bit_model', 'mixed_precision'].
Please specify them appropriately in the DeepSpeed config file.
If you are using an accelerate config file, remove others config variables mentioned in the above specified list.
The easiest method is to create a new config following the questionnaire via `accelerate config`.
It will only ask for the necessary config variables when using `deepspeed_config_file`.
```

Scenario 2: Use the solution of the error to create new accelerate config and check that no ambiguity error is now thrown.
1. Run `accelerate config`:
```
$ accelerate config
-------------------------------------------------------------------------------------------------------------------------------In which compute environment are you running?
This machine                                                                                                                   
-------------------------------------------------------------------------------------------------------------------------------Which type of machine are you using?                                                                                           
multi-GPU                                                                                                                      
How many different machines will you use (use more than 1 for multi-node training)? [1]:                                       
Do you wish to optimize your script with torch dynamo?[yes/NO]:                                                                
Do you want to use DeepSpeed? [yes/NO]: yes                                                                                    
Do you want to specify a json file to a DeepSpeed config? [yes/NO]: yes                                                        
Please enter the path to the json DeepSpeed config file: ds_config.json                                                        
Do you want to enable `deepspeed.zero.Init` when using ZeRO Stage-3 for constructing massive models? [yes/NO]: yes
How many GPU(s) should be used for distributed training? [1]:4
accelerate configuration saved at ds_config_sample.yaml
```

2. `accelerate config`:
```yaml
compute_environment: LOCAL_MACHINE
deepspeed_config:
  deepspeed_config_file: ds_config.json
  zero3_init_flag: true
distributed_type: DEEPSPEED
downcast_bf16: 'no'
dynamo_backend: 'NO'
fsdp_config: {}
machine_rank: 0
main_training_function: main
megatron_lm_config: {}
num_machines: 1
num_processes: 4
rdzv_backend: static
same_network: true
use_cpu: false
```

3. Output of `accelerate launch test.py`:
```
Distributed environment: DEEPSPEED  Backend: nccl
Num processes: 4
Process index: 0
Local process index: 0
Device: cuda:0
Mixed precision type: bf16
ds_config: {'bf16': {'enabled': True}, 'zero_optimization': {'stage': 3, 'stage3_gather_16bit_weights_on_model_save': False, 'offload_optimizer': {'device': 'none'}, 'offload_param': {'device': 'none'}}, 'gradient_clipping': 1.0, 'train_batch_size': 'auto', 'train_micro_batch_size_per_gpu': 'auto', 'gradient_accumulation_steps': 10, 'steps_per_print': inf, 'fp16': {'enabled': False}}
```

Scenario 3: Setting the `accelerate launch` cmd args related to deepspeed as `auto` in  `deepspeed_config_file` and check that things work as expected.
1. new `ds_config.json` with `auto` for the `accelerate launch` deepspeed cmd args:
```json
{
    "bf16": {
        "enabled": "auto"
    },
    "zero_optimization": {
        "stage": "auto",
        "stage3_gather_16bit_weights_on_model_save": "auto",
        "offload_optimizer": {
            "device": "auto"
        },
        "offload_param": {
            "device": "auto"
        }
    },
    "gradient_clipping": "auto",
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "gradient_accumulation_steps": "auto",
    "steps_per_print": 2000000
}
```

2. Output of `accelerate launch --mixed_precision="fp16" --zero_stage=3 --gradient_accumulation_steps=5 --gradient_clipping=1.0 --offload_param_device="cpu" --offload_optimizer_device="nvme" --zero3_save_16bit_model="true" test.py`:
```bash
Distributed environment: DEEPSPEED  Backend: nccl
Num processes: 4
Process index: 0
Local process index: 0
Device: cuda:0
Mixed precision type: fp16
ds_config: {'bf16': {'enabled': False}, 'zero_optimization': {'stage': 3, 'stage3_gather_16bit_weights_on_model_save': True, 'offload_optimizer': {'device': 'nvme'}, 'offload_param': {'device': 'cpu'}}, 'gradient_clipping': 1.0, 'train_batch_size': 'auto', 'train_micro_batch_size_per_gpu': 'auto', 'gradient_accumulation_steps': 5, 'steps_per_print': inf, 'fp16': {'enabled': True, 'auto_cast': True}}
```

**Note**: Remaining `auto` values are handled in `accelerator.prepare()` call.



Looks fantastic, Sourab! Thank you for the improvements and taking the time to layout out the different scenarios - if I'm not mistaken those would make for perfect additions to the documentation if it resonates. (at the very least the last one to demo how `auto` values work and why would one want to use those).

BTW, the config generates things like:
```
fsdp_config: {}
megatron_lm_config: {}
```
why not just skip parts that the user hasn't asked for? It just makes the config scarier than it is, no? I'm asking since when I first looked at it I wasn't a all sure which of the empty placeholders were safe to remove and which aren't. My personal preference is for active config - that is to only ever list config entries that I work with and any defaults should be just that defaults and not be listed at all. Which I suppose isn't the case with typical configs where everything is listed out whether it's being used or not.

And I can of course remove all those, so definitely it's not an issue, I'm just asking if my thinking resonates with you.
Sourab, I found one more ambiguous combo in one of our tests:

```
  zero3_init_flag: true
  zero_stage: 2
```

This combo is quietly getting accepted. I'm concerned that a developer may see `zero3_init_flag: true` and think it's zero3?

Do you think accelerate should assert when `zero3_init_flag==True` is used with `stage < 3`?
Hello @stas00, with current setup below warning is given which I think is fine:
```bash
UserWarning: DeepSpeed Zero3 Init flag is only applicable for ZeRO Stage 3. Setting it to False.
```
oh boy. I didn't see it. :(

If a tree falls in a forest and no one is around to hear it, does it make a sound?

I guess I need to start using this pragma to turn warnings into errors, but then some warnings can't be acted upon :(

```
import warnings
warnings.filterwarnings("error")
```

This is for example an even larger issue for tests, where distributed setup hides most warnings or again there are too many of warnings to see anything.

https://api.github.com/repos/huggingface/accelerate/pulls/941
diff --git a/docs/source/quicktour.mdx b/docs/source/quicktour.mdx
index 6ea99f9bd..13a8d7afa 100644
--- a/docs/source/quicktour.mdx
+++ b/docs/source/quicktour.mdx
@@ -206,7 +206,7 @@ Now that this is done, you can run your script with the following command:
 accelerate launch path_to_script.py --args_for_the_script
 ```
 
-If you stored the config file in a non-default location, you can indicate it to the launcher like his:
+If you stored the config file in a non-default location, you can indicate it to the launcher like this:
 
 ```bash
 accelerate launch --config_file path_to_config.yaml path_to_script.py --args_for_the_script

_The documentation is not available anymore as the PR was closed or merged._
Thanks for the fix! 
https://api.github.com/repos/huggingface/accelerate/pulls/939
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index 8a45856cd..7dafed0e5 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -711,7 +711,9 @@ def prepare_data_loader(
     # Need to provide batch_size as batch_sampler is None for Iterable dataset
     if new_batch_sampler is None:
         kwargs["drop_last"] = dataloader.drop_last
-        kwargs["batch_size"] = dataloader.batch_size // num_processes if split_batches else dataloader.batch_size
+        kwargs["batch_size"] = (
+            dataloader.batch_size // num_processes if split_batches and not dispatch_batches else dataloader.batch_size
+        )
 
     if dispatch_batches:
         kwargs.pop("generator")

_The documentation is not available anymore as the PR was closed or merged._
Looks like the right fix to me, thanks!
https://api.github.com/repos/huggingface/accelerate/pulls/937
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index ff880cffd..ec5041bfa 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -473,17 +473,7 @@ def is_last_process(self):
 
     @property
     def mixed_precision(self):
-        if self.distributed_type == DistributedType.DEEPSPEED:
-            config = self.state.deepspeed_plugin.deepspeed_config
-            if config.get("fp16", {}).get("enabled", False):
-                mixed_precision = "fp16"
-            elif config.get("bf16", {}).get("enabled", False):
-                mixed_precision = "bf16"
-            else:
-                mixed_precision = "no"
-        else:
-            mixed_precision = self.state.mixed_precision
-        return mixed_precision
+        return self.state.mixed_precision
 
     def on_main_process(func):
         """
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index a8a5e7960..31a6d64a2 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -104,7 +104,7 @@ def __init__(
                     self.local_process_index = int(os.environ.get("LOCAL_RANK", -1))
                     self.device = torch.device("cuda", self.local_process_index)
                     torch.cuda.set_device(self.device)
-                    self.mixed_precision = mixed_precision
+                    self._mixed_precision = mixed_precision
             elif is_tpu_available() and not cpu:
                 self.distributed_type = DistributedType.TPU
                 self.num_processes = xm.xrt_world_size()
@@ -120,7 +120,7 @@ def __init__(
                         os.environ["XLA_USE_BF16"] = str(1)
                         os.environ["XLA_DOWNCAST_BF16"] = str(0)
                         self.downcast_bfloat = False
-                self.mixed_precision = mixed_precision
+                self._mixed_precision = mixed_precision
             elif os.environ.get("ACCELERATE_USE_DEEPSPEED", "false") == "true" and not cpu:
                 assert (
                     is_deepspeed_available()
@@ -142,7 +142,7 @@ def __init__(
                 self.local_process_index = int(os.environ.get("LOCAL_RANK", -1))
                 self.device = torch.device("cuda", self.local_process_index)
                 torch.cuda.set_device(self.device)
-                self.mixed_precision = "no"  # deepspeed handles mixed_precision using deepspeed_config
+                self._mixed_precision = "no"  # deepspeed handles mixed_precision using deepspeed_config
                 self.deepspeed_plugin = deepspeed_plugin
             elif int(os.environ.get("LOCAL_RANK", -1)) != -1 and not cpu:
                 self.distributed_type = DistributedType.MULTI_GPU
@@ -154,15 +154,15 @@ def __init__(
                 self.local_process_index = int(os.environ.get("LOCAL_RANK", -1))
                 self.device = torch.device("cuda", self.local_process_index)
                 torch.cuda.set_device(self.device)
-                self.mixed_precision = mixed_precision
+                self._mixed_precision = mixed_precision
                 if os.environ.get("ACCELERATE_USE_FSDP", "false") == "true":
                     self.distributed_type = DistributedType.FSDP
-                    if self.mixed_precision != "no":
-                        fsdp_plugin.set_mixed_precision(self.mixed_precision)
+                    if self._mixed_precision != "no":
+                        fsdp_plugin.set_mixed_precision(self._mixed_precision)
                     self.fsdp_plugin = fsdp_plugin
                 if os.environ.get("ACCELERATE_USE_MEGATRON_LM", "false") == "true":
                     self.distributed_type = DistributedType.MEGATRON_LM
-                    megatron_lm_plugin.set_mixed_precision(self.mixed_precision)
+                    megatron_lm_plugin.set_mixed_precision(self._mixed_precision)
                     self.megatron_lm_plugin = megatron_lm_plugin
             elif get_int_from_env(["PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE", "WORLD_SIZE"], 1) > 1:
                 self.distributed_type = DistributedType.MULTI_CPU
@@ -204,7 +204,7 @@ def __init__(
                 self.process_index = torch.distributed.get_rank()
                 self.local_process_index = local_rank
                 self.device = torch.device("cpu")
-                self.mixed_precision = mixed_precision
+                self._mixed_precision = mixed_precision
             else:
                 self.distributed_type = DistributedType.NO
                 self.num_processes = 1
@@ -237,9 +237,13 @@ def __init__(
                     self.device = torch.device("cpu")
                 else:
                     self.device = torch.device("cuda")
-                self.mixed_precision = mixed_precision
+                self._mixed_precision = mixed_precision
 
-            if self.dynamo_backend != DynamoBackend.NO and self.mixed_precision == "no" and self.device.type == "cuda":
+            if (
+                self.dynamo_backend != DynamoBackend.NO
+                and self._mixed_precision == "no"
+                and self.device.type == "cuda"
+            ):
                 torch.backends.cuda.matmul.allow_tf32 = True
             self.initialized = True
 
@@ -252,17 +256,30 @@ def __repr__(self):
             f"Process index: {self.process_index}\n"
             f"Local process index: {self.local_process_index}\n"
             f"Device: {self.device}\n"
+            f"Mixed precision type: {mixed_precision}\n"
         )
         if self.distributed_type == DistributedType.DEEPSPEED:
             repr += f"ds_config: {self.deepspeed_plugin.deepspeed_config}\n"
-        else:
-            repr += f"Mixed precision type: {mixed_precision}\n"
         return repr
 
     # For backward compatibility
     @property
     def use_fp16(self):
-        return self.mixed_precision != "no"
+        return self._mixed_precision != "no"
+
+    @property
+    def mixed_precision(self):
+        if self.distributed_type == DistributedType.DEEPSPEED:
+            config = self.deepspeed_plugin.deepspeed_config
+            if config.get("fp16", {}).get("enabled", False):
+                mixed_precision = "fp16"
+            elif config.get("bf16", {}).get("enabled", False):
+                mixed_precision = "bf16"
+            else:
+                mixed_precision = "no"
+        else:
+            mixed_precision = self._mixed_precision
+        return mixed_precision
 
     @staticmethod
     def _reset_state():
@@ -275,7 +292,7 @@ def _check_initialized(self, mixed_precision=None, cpu=None):
             err = "AcceleratorState has already been initialized and cannot be changed, restart your runtime completely and pass `{flag}` to `Accelerate()`."
             if cpu and self.device.type != "cpu":
                 raise ValueError(err.format(flag="cpu=True"))
-            if mixed_precision is not None and mixed_precision != self.mixed_precision:
+            if mixed_precision is not None and mixed_precision != self._mixed_precision:
                 raise ValueError(err.format(flag=f"mixed_precision='{mixed_precision}'"))
 
 

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_935). All of your documentation changes will be reflected on that endpoint.
> This is not ideal in terms of naming. I'm wondering if we shouldn't rename the current attribute `_mixed_precision` and have the property named `mixed_precision` which would be clearer. There might be a bit of code to adapt, but I think it's worth the effort, no?

Yes, I thought of that but changed my mind in order to be least invasive in terms of the changes required. Addressed this in latest commit. Thank you.
Example with recent changes:

Code:
```
from accelerate import Accelerator
from accelerate.state import AcceleratorState

def main():
    accelerator = Accelerator()
    accelerator.print(f"{AcceleratorState()}")
    accelerator.print(f"{AcceleratorState().mixed_precision=}")
    accelerator.print(f"{accelerator.mixed_precision=}")

if __name__ == "__main__":
    main()
```

config:
```yaml
command_file: null
commands: null
compute_environment: LOCAL_MACHINE
deepspeed_config:
  gradient_accumulation_steps: 1
  gradient_clipping: 1.0
  offload_optimizer_device: 'cpu'
  offload_param_device: 'cpu'
  zero3_init_flag: true
  zero3_save_16bit_model: true
  zero_stage: 3
distributed_type: DEEPSPEED
downcast_bf16: 'no'
dynamo_backend: 'NO'
fsdp_config: {}
gpu_ids: null
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
megatron_lm_config: {}
mixed_precision: 'bf16'
num_machines: 1
num_processes: 2
rdzv_backend: static
same_network: true
tpu_name: null
tpu_zone: null
use_cpu: false
```

Output
```bash
Distributed environment: DEEPSPEED  Backend: nccl
Num processes: 2
Process index: 0
Local process index: 0
Device: cuda:0
Mixed precision type: bf16
ds_config: {'train_batch_size': 'auto', 'train_micro_batch_size_per_gpu': 'auto', 'gradient_accumulation_steps': 1, 'zero_optimization': {'stage': 3, 'offload_optimizer': {'device': 'cpu'}, 'offload_param': {'device': 'cpu'}, 'stage3_gather_16bit_weights_on_model_save': True}, 'gradient_clipping': 1.0, 'steps_per_print': inf, 'bf16': {'enabled': True}}

AcceleratorState().mixed_precision='bf16'
accelerator.mixed_precision='bf16'
```
https://api.github.com/repos/huggingface/accelerate/pulls/935
diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py
index dc8883425..4f9617711 100644
--- a/src/accelerate/utils/dataclasses.py
+++ b/src/accelerate/utils/dataclasses.py
@@ -23,6 +23,7 @@
 import os
 import typing
 import warnings
+from contextlib import contextmanager
 from dataclasses import dataclass, field
 from datetime import timedelta
 from distutils.util import strtobool
@@ -529,6 +530,23 @@ def set_deepspeed_weakref(self):
 
             self.dschf = HfDeepSpeedConfig(ds_config)  # keep this object alive # noqa
 
+    def is_zero3_init_enabled(self):
+        return self.zero3_init_flag
+
+    @contextmanager
+    def zero3_init_context_manager(self, enable=False):
+        old = self.zero3_init_flag
+        if old == enable:
+            yield
+        else:
+            self.zero3_init_flag = enable
+            self.dschf = None
+            self.set_deepspeed_weakref()
+            yield
+            self.zero3_init_flag = old
+            self.dschf = None
+            self.set_deepspeed_weakref()
+
 
 @dataclass
 class FullyShardedDataParallelPlugin:

Might be clearer to have an `else` block here instead of the early `return`.
@@ -529,6 +530,23 @@ def set_deepspeed_weakref(self):
 
             self.dschf = HfDeepSpeedConfig(ds_config)  # keep this object alive # noqa
 
+    def is_zero3_init_enabled(self):
+        return self.zero3_init_flag
+
+    @contextmanager
+    def set_zero3_init(self, enable=False):
+        old = self.zero3_init_flag
+        if old == enable:
+            yield
+            return
+        self.zero3_init_flag = enable
+        self.dschf = None
+        self.set_deepspeed_weakref()
+        yield
+        self.zero3_init_flag = old
+        self.dschf = None
+        self.set_deepspeed_weakref()
Done. 
@@ -529,6 +530,23 @@ def set_deepspeed_weakref(self):
 
             self.dschf = HfDeepSpeedConfig(ds_config)  # keep this object alive # noqa
 
+    def is_zero3_init_enabled(self):
+        return self.zero3_init_flag
+
+    @contextmanager
+    def set_zero3_init(self, enable=False):
+        old = self.zero3_init_flag
+        if old == enable:
+            yield
+            return
+        self.zero3_init_flag = enable
+        self.dschf = None
+        self.set_deepspeed_weakref()
+        yield
+        self.zero3_init_flag = old
+        self.dschf = None
+        self.set_deepspeed_weakref()
_The documentation is not available anymore as the PR was closed or merged._
here is how I used your new method:

```
        from transformer.utils import ContextManagers
        deepspeed_plugin = get_deepspeed_plugin()
        if deepspeed_plugin is not None:
            zero_init_disabled_context = [deepspeed_plugin.set_zero3_init(enable=False)]
        else:
            zero_init_disabled_context = []

        with ContextManagers(zero_init_disabled_context):
            config = AutoConfig.from_pretrained(vision_model_name, **vision_model_params)
            model = AutoModel.from_config(config, torch_dtype=torch_dtype)
```
> Would it be more intuitive to call it `zero_init_context_manager`?

renamed it to `zero3_init_context_manager`
https://api.github.com/repos/huggingface/accelerate/pulls/932
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index fd5ad7d60..04ff21127 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -343,14 +343,14 @@ def launch_command_parser(subparsers=None):
     )
     deepspeed_args.add_argument(
         "--zero3_init_flag",
-        default=None,
+        default="true",
         type=str,
         help="Decides Whether (true|false) to enable `deepspeed.zero.Init` for constructing massive models. "
         "Only applicable with DeepSpeed ZeRO Stage-3.",
     )
     deepspeed_args.add_argument(
         "--zero3_save_16bit_model",
-        default=None,
+        default="false",
         type=str,
         help="Decides Whether (true|false) to save 16-bit model weights when using ZeRO Stage-3. "
         "Only applicable with DeepSpeed ZeRO Stage-3.",
@@ -759,15 +759,15 @@ def deepspeed_launcher(args):
     current_env["PYTHONPATH"] = env_var_path_add("PYTHONPATH", os.path.abspath("."))
     current_env["ACCELERATE_MIXED_PRECISION"] = str(mixed_precision)
     current_env["ACCELERATE_USE_DEEPSPEED"] = "true"
-    current_env["DEEPSPEED_ZERO_STAGE"] = str(args.zero_stage)
-    current_env["GRADIENT_ACCUMULATION_STEPS"] = str(args.gradient_accumulation_steps)
-    current_env["GRADIENT_CLIPPING"] = str(args.gradient_clipping).lower()
-    current_env["DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE"] = str(args.offload_optimizer_device).lower()
-    current_env["DEEPSPEED_OFFLOAD_PARAM_DEVICE"] = str(args.offload_param_device).lower()
-    current_env["DEEPSPEED_ZERO3_INIT"] = str(args.zero3_init_flag).lower()
-    current_env["DEEPSPEED_ZERO3_SAVE_16BIT_MODEL"] = str(args.zero3_save_16bit_model).lower()
+    current_env["ACCELERATE_DEEPSPEED_ZERO_STAGE"] = str(args.zero_stage)
+    current_env["ACCELERATE_GRADIENT_ACCUMULATION_STEPS"] = str(args.gradient_accumulation_steps)
+    current_env["ACCELERATE_GRADIENT_CLIPPING"] = str(args.gradient_clipping).lower()
+    current_env["ACCELERATE_DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE"] = str(args.offload_optimizer_device).lower()
+    current_env["ACCELERATE_DEEPSPEED_OFFLOAD_PARAM_DEVICE"] = str(args.offload_param_device).lower()
+    current_env["ACCELERATE_DEEPSPEED_ZERO3_INIT"] = str(args.zero3_init_flag).lower()
+    current_env["ACCELERATE_DEEPSPEED_ZERO3_SAVE_16BIT_MODEL"] = str(args.zero3_save_16bit_model).lower()
     if args.deepspeed_config_file is not None:
-        current_env["DEEPSPEED_CONFIG_FILE"] = str(args.deepspeed_config_file)
+        current_env["ACCELERATE_DEEPSPEED_CONFIG_FILE"] = str(args.deepspeed_config_file)
 
     if args.num_machines > 1 and args.deepspeed_multinode_launcher != DEEPSPEED_MULTINODE_LAUNCHERS[1]:
         with open(".deepspeed_env", "a") as f:
@@ -1041,8 +1041,7 @@ def launch_command(args):
             for name, attr in defaults.__dict__.items():
                 if isinstance(attr, dict):
                     for k in defaults.deepspeed_config:
-                        if getattr(args, k) is None:
-                            setattr(args, k, defaults.deepspeed_config[k])
+                        setattr(args, k, defaults.deepspeed_config[k])
                     for k in defaults.fsdp_config:
                         arg_to_set = k
                         if "fsdp" not in arg_to_set:
diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py
index eafad8a09..eea8c3ab3 100644
--- a/src/accelerate/utils/dataclasses.py
+++ b/src/accelerate/utils/dataclasses.py
@@ -361,7 +361,7 @@ def __post_init__(self):
         from .deepspeed import HfDeepSpeedConfig
 
         if self.hf_ds_config is None:
-            self.hf_ds_config = os.environ.get("DEEPSPEED_CONFIG_FILE", "none")
+            self.hf_ds_config = os.environ.get("ACCELERATE_DEEPSPEED_CONFIG_FILE", "none")
         if (
             isinstance(self.hf_ds_config, dict)
             or (isinstance(self.hf_ds_config, str) and self.hf_ds_config != "none")
@@ -377,24 +377,26 @@ def __post_init__(self):
                 raise ValueError("Please specify the ZeRO optimization config in the DeepSpeed config.")
         else:
             if self.gradient_accumulation_steps is None:
-                self.gradient_accumulation_steps = int(os.environ.get("GRADIENT_ACCUMULATION_STEPS", 1))
+                self.gradient_accumulation_steps = int(os.environ.get("ACCELERATE_GRADIENT_ACCUMULATION_STEPS", 1))
 
             if self.gradient_clipping is None:
-                gradient_clipping = os.environ.get("GRADIENT_CLIPPING", "none")
+                gradient_clipping = os.environ.get("ACCELERATE_GRADIENT_CLIPPING", "none")
                 if gradient_clipping != "none":
                     self.gradient_clipping = float(gradient_clipping)
 
             if self.zero_stage is None:
-                self.zero_stage = int(os.environ.get("DEEPSPEED_ZERO_STAGE", 2))
+                self.zero_stage = int(os.environ.get("ACCELERATE_DEEPSPEED_ZERO_STAGE", 2))
 
             if self.offload_optimizer_device is None:
-                self.offload_optimizer_device = os.environ.get("DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE", "none")
+                self.offload_optimizer_device = os.environ.get("ACCELERATE_DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE", "none")
 
             if self.offload_param_device is None:
-                self.offload_param_device = os.environ.get("DEEPSPEED_OFFLOAD_PARAM_DEVICE", "none")
+                self.offload_param_device = os.environ.get("ACCELERATE_DEEPSPEED_OFFLOAD_PARAM_DEVICE", "none")
 
             if self.zero3_save_16bit_model is None:
-                self.zero3_save_16bit_model = os.environ.get("DEEPSPEED_ZERO3_SAVE_16BIT_MODEL", "false") == "true"
+                self.zero3_save_16bit_model = (
+                    os.environ.get("ACCELERATE_DEEPSPEED_ZERO3_SAVE_16BIT_MODEL", "false") == "true"
+                )
 
             config = {
                 "train_batch_size": "auto",
@@ -417,7 +419,7 @@ def __post_init__(self):
         self.deepspeed_config = self.hf_ds_config.config
         self.deepspeed_config["steps_per_print"] = float("inf")  # this will stop deepspeed from logging @ stdout
         if self.zero3_init_flag is None:
-            self.zero3_init_flag = os.environ.get("DEEPSPEED_ZERO3_INIT", "false") == "true"
+            self.zero3_init_flag = os.environ.get("ACCELERATE_DEEPSPEED_ZERO3_INIT", "false") == "true"
         if self.zero3_init_flag and not self.hf_ds_config.is_zero3():
             warnings.warn("DeepSpeed Zero3 Init flag is only applicable for ZeRO Stage 3. Setting it to False.")
             self.zero3_init_flag = False

this logic needs to be recoded since the default is no longer `None` and thus this code will never run.
@@ -377,24 +377,26 @@ def __post_init__(self):
                 raise ValueError("Please specify the ZeRO optimization config in the DeepSpeed config.")
         else:
             if self.gradient_accumulation_steps is None:
-                self.gradient_accumulation_steps = int(os.environ.get("GRADIENT_ACCUMULATION_STEPS", 1))
+                self.gradient_accumulation_steps = int(os.environ.get("ACCELERATE_GRADIENT_ACCUMULATION_STEPS", 1))
 
             if self.gradient_clipping is None:
-                gradient_clipping = os.environ.get("GRADIENT_CLIPPING", "none")
+                gradient_clipping = os.environ.get("ACCELERATE_GRADIENT_CLIPPING", "none")
                 if gradient_clipping != "none":
                     self.gradient_clipping = float(gradient_clipping)
 
             if self.zero_stage is None:
-                self.zero_stage = int(os.environ.get("DEEPSPEED_ZERO_STAGE", 2))
+                self.zero_stage = int(os.environ.get("ACCELERATE_DEEPSPEED_ZERO_STAGE", 2))
 
             if self.offload_optimizer_device is None:
-                self.offload_optimizer_device = os.environ.get("DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE", "none")
+                self.offload_optimizer_device = os.environ.get("ACCELERATE_DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE", "none")
 
             if self.offload_param_device is None:
-                self.offload_param_device = os.environ.get("DEEPSPEED_OFFLOAD_PARAM_DEVICE", "none")
+                self.offload_param_device = os.environ.get("ACCELERATE_DEEPSPEED_OFFLOAD_PARAM_DEVICE", "none")
 
             if self.zero3_save_16bit_model is None:
-                self.zero3_save_16bit_model = os.environ.get("DEEPSPEED_ZERO3_SAVE_16BIT_MODEL", "false") == "true"
+                self.zero3_save_16bit_model = (
+                    os.environ.get("ACCELERATE_DEEPSPEED_ZERO3_SAVE_16BIT_MODEL", "false") == "true"
+                )
this logic needs to be recoded since the default is no longer `None` and thus this code will never run.
@@ -417,7 +419,7 @@ def __post_init__(self):
         self.deepspeed_config = self.hf_ds_config.config
         self.deepspeed_config["steps_per_print"] = float("inf")  # this will stop deepspeed from logging @ stdout
         if self.zero3_init_flag is None:
-            self.zero3_init_flag = os.environ.get("DEEPSPEED_ZERO3_INIT", "false") == "true"
+            self.zero3_init_flag = os.environ.get("ACCELERATE_DEEPSPEED_ZERO3_INIT", "false") == "true"
may be add to `help` that `ACCELERATE_DEEPSPEED_ZERO3_INIT` env var can override it?
@@ -343,14 +343,14 @@ def launch_command_parser(subparsers=None):
     )
     deepspeed_args.add_argument(
         "--zero3_init_flag",
-        default=None,
+        default="true",
         type=str,
         help="Decides Whether (true|false) to enable `deepspeed.zero.Init` for constructing massive models. "
         "Only applicable with DeepSpeed ZeRO Stage-3.",
may be add to `help` that `ACCELERATE_DEEPSPEED_ZERO3_SAVE_16BIT_MODEL` env var can override it?
@@ -343,14 +343,14 @@ def launch_command_parser(subparsers=None):
     )
     deepspeed_args.add_argument(
         "--zero3_init_flag",
-        default=None,
+        default="true",
         type=str,
         help="Decides Whether (true|false) to enable `deepspeed.zero.Init` for constructing massive models. "
         "Only applicable with DeepSpeed ZeRO Stage-3.",
     )
     deepspeed_args.add_argument(
         "--zero3_save_16bit_model",
-        default=None,
+        default="false",
         type=str,
         help="Decides Whether (true|false) to save 16-bit model weights when using ZeRO Stage-3. "
         "Only applicable with DeepSpeed ZeRO Stage-3.",
Hello, `ACCELERATE_DEEPSPEED_ZERO3_INIT` is set using `zero3_init_flag` arg, so I don't think explicit mention is required as it is only used internally.
@@ -343,14 +343,14 @@ def launch_command_parser(subparsers=None):
     )
     deepspeed_args.add_argument(
         "--zero3_init_flag",
-        default=None,
+        default="true",
         type=str,
         help="Decides Whether (true|false) to enable `deepspeed.zero.Init` for constructing massive models. "
         "Only applicable with DeepSpeed ZeRO Stage-3.",
Same explanation as above
@@ -343,14 +343,14 @@ def launch_command_parser(subparsers=None):
     )
     deepspeed_args.add_argument(
         "--zero3_init_flag",
-        default=None,
+        default="true",
         type=str,
         help="Decides Whether (true|false) to enable `deepspeed.zero.Init` for constructing massive models. "
         "Only applicable with DeepSpeed ZeRO Stage-3.",
     )
     deepspeed_args.add_argument(
         "--zero3_save_16bit_model",
-        default=None,
+        default="false",
         type=str,
         help="Decides Whether (true|false) to save 16-bit model weights when using ZeRO Stage-3. "
         "Only applicable with DeepSpeed ZeRO Stage-3.",
ah, understood, thank you!
@@ -343,14 +343,14 @@ def launch_command_parser(subparsers=None):
     )
     deepspeed_args.add_argument(
         "--zero3_init_flag",
-        default=None,
+        default="true",
         type=str,
         help="Decides Whether (true|false) to enable `deepspeed.zero.Init` for constructing massive models. "
         "Only applicable with DeepSpeed ZeRO Stage-3.",
Hello, here dataclass object attributes are still `None` initially if the user hasn't passed them explicitly by directly creating `DeepSpeedPlugin` object. In `__post_init__` call, the env variables which were set via [default] cmd args are used for setting dataclass object attributes.
@@ -377,24 +377,26 @@ def __post_init__(self):
                 raise ValueError("Please specify the ZeRO optimization config in the DeepSpeed config.")
         else:
             if self.gradient_accumulation_steps is None:
-                self.gradient_accumulation_steps = int(os.environ.get("GRADIENT_ACCUMULATION_STEPS", 1))
+                self.gradient_accumulation_steps = int(os.environ.get("ACCELERATE_GRADIENT_ACCUMULATION_STEPS", 1))
 
             if self.gradient_clipping is None:
-                gradient_clipping = os.environ.get("GRADIENT_CLIPPING", "none")
+                gradient_clipping = os.environ.get("ACCELERATE_GRADIENT_CLIPPING", "none")
                 if gradient_clipping != "none":
                     self.gradient_clipping = float(gradient_clipping)
 
             if self.zero_stage is None:
-                self.zero_stage = int(os.environ.get("DEEPSPEED_ZERO_STAGE", 2))
+                self.zero_stage = int(os.environ.get("ACCELERATE_DEEPSPEED_ZERO_STAGE", 2))
 
             if self.offload_optimizer_device is None:
-                self.offload_optimizer_device = os.environ.get("DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE", "none")
+                self.offload_optimizer_device = os.environ.get("ACCELERATE_DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE", "none")
 
             if self.offload_param_device is None:
-                self.offload_param_device = os.environ.get("DEEPSPEED_OFFLOAD_PARAM_DEVICE", "none")
+                self.offload_param_device = os.environ.get("ACCELERATE_DEEPSPEED_OFFLOAD_PARAM_DEVICE", "none")
 
             if self.zero3_save_16bit_model is None:
-                self.zero3_save_16bit_model = os.environ.get("DEEPSPEED_ZERO3_SAVE_16BIT_MODEL", "false") == "true"
+                self.zero3_save_16bit_model = (
+                    os.environ.get("ACCELERATE_DEEPSPEED_ZERO3_SAVE_16BIT_MODEL", "false") == "true"
+                )
Same explanation as above
@@ -417,7 +419,7 @@ def __post_init__(self):
         self.deepspeed_config = self.hf_ds_config.config
         self.deepspeed_config["steps_per_print"] = float("inf")  # this will stop deepspeed from logging @ stdout
         if self.zero3_init_flag is None:
-            self.zero3_init_flag = os.environ.get("DEEPSPEED_ZERO3_INIT", "false") == "true"
+            self.zero3_init_flag = os.environ.get("ACCELERATE_DEEPSPEED_ZERO3_INIT", "false") == "true"
my bad then! my apologies for the false alert - thank you for explaining, @pacman100!
@@ -377,24 +377,26 @@ def __post_init__(self):
                 raise ValueError("Please specify the ZeRO optimization config in the DeepSpeed config.")
         else:
             if self.gradient_accumulation_steps is None:
-                self.gradient_accumulation_steps = int(os.environ.get("GRADIENT_ACCUMULATION_STEPS", 1))
+                self.gradient_accumulation_steps = int(os.environ.get("ACCELERATE_GRADIENT_ACCUMULATION_STEPS", 1))
 
             if self.gradient_clipping is None:
-                gradient_clipping = os.environ.get("GRADIENT_CLIPPING", "none")
+                gradient_clipping = os.environ.get("ACCELERATE_GRADIENT_CLIPPING", "none")
                 if gradient_clipping != "none":
                     self.gradient_clipping = float(gradient_clipping)
 
             if self.zero_stage is None:
-                self.zero_stage = int(os.environ.get("DEEPSPEED_ZERO_STAGE", 2))
+                self.zero_stage = int(os.environ.get("ACCELERATE_DEEPSPEED_ZERO_STAGE", 2))
 
             if self.offload_optimizer_device is None:
-                self.offload_optimizer_device = os.environ.get("DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE", "none")
+                self.offload_optimizer_device = os.environ.get("ACCELERATE_DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE", "none")
 
             if self.offload_param_device is None:
-                self.offload_param_device = os.environ.get("DEEPSPEED_OFFLOAD_PARAM_DEVICE", "none")
+                self.offload_param_device = os.environ.get("ACCELERATE_DEEPSPEED_OFFLOAD_PARAM_DEVICE", "none")
 
             if self.zero3_save_16bit_model is None:
-                self.zero3_save_16bit_model = os.environ.get("DEEPSPEED_ZERO3_SAVE_16BIT_MODEL", "false") == "true"
+                self.zero3_save_16bit_model = (
+                    os.environ.get("ACCELERATE_DEEPSPEED_ZERO3_SAVE_16BIT_MODEL", "false") == "true"
+                )
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/928
diff --git a/src/accelerate/__init__.py b/src/accelerate/__init__.py
index d22006d54..9129fc1a5 100644
--- a/src/accelerate/__init__.py
+++ b/src/accelerate/__init__.py
@@ -5,7 +5,14 @@
 __version__ = "0.15.0.dev0"
 
 from .accelerator import Accelerator
-from .big_modeling import cpu_offload, disk_offload, dispatch_model, init_empty_weights, load_checkpoint_and_dispatch
+from .big_modeling import (
+    cpu_offload,
+    disk_offload,
+    dispatch_model,
+    init_empty_weights,
+    init_on_device,
+    load_checkpoint_and_dispatch,
+)
 from .launchers import debug_launcher, notebook_launcher
 from .utils import (
     DeepSpeedPlugin,
diff --git a/src/accelerate/big_modeling.py b/src/accelerate/big_modeling.py
index 24857ad25..d3247a464 100644
--- a/src/accelerate/big_modeling.py
+++ b/src/accelerate/big_modeling.py
@@ -62,6 +62,31 @@ def init_empty_weights(include_buffers: bool = False):
     """
     if not is_torch_version(">=", "1.9.0"):
         raise NotImplementedError("Initializing empty weights to a meta device requires torch >= 1.9.0")
+    with init_on_device(torch.device("meta"), include_buffers=include_buffers) as f:
+        yield f
+
+
+@contextmanager
+def init_on_device(device: torch.device, include_buffers: bool = False):
+    """
+    A context manager under which models are initialized with all parameters on the specified device.
+
+    Args:
+        device (`torch.device`):
+            Device to initialize all parameters on.
+        include_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to also put all buffers on the meta device while initializing.
+
+    Example:
+
+    ```python
+    import torch.nn as nn
+    from accelerate import init_on_device
+
+    with init_on_device(device=torch.device("cuda")):
+        tst = nn.Liner(100, 100)  # on `cuda` device
+    ```
+    """
     old_register_parameter = nn.Module.register_parameter
     if include_buffers:
         old_register_buffer = nn.Module.register_buffer
@@ -71,12 +96,12 @@ def register_empty_parameter(module, name, param):
         if param is not None:
             param_cls = type(module._parameters[name])
             kwargs = module._parameters[name].__dict__
-            module._parameters[name] = param_cls(module._parameters[name].to(torch.device("meta")), **kwargs)
+            module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs)
 
     def register_empty_buffer(module, name, buffer):
         old_register_buffer(module, name, buffer)
         if buffer is not None:
-            module._buffers[name] = module._buffers[name].to(torch.device("meta"))
+            module._buffers[name] = module._buffers[name].to(device)
 
     # Patch tensor creation
     if include_buffers:
@@ -89,7 +114,7 @@ def register_empty_buffer(module, name, buffer):
 
     def patch_tensor_constructor(fn):
         def wrapper(*args, **kwargs):
-            kwargs["device"] = torch.device("meta")
+            kwargs["device"] = device
             return fn(*args, **kwargs)
 
         return wrapper
diff --git a/src/accelerate/test_utils/__init__.py b/src/accelerate/test_utils/__init__.py
index faf8a7da8..b58b932b1 100644
--- a/src/accelerate/test_utils/__init__.py
+++ b/src/accelerate/test_utils/__init__.py
@@ -8,6 +8,7 @@
     require_cpu,
     require_cuda,
     require_huggingface_suite,
+    require_mps,
     require_multi_gpu,
     require_single_gpu,
     require_torch_min_version,
diff --git a/src/accelerate/test_utils/testing.py b/src/accelerate/test_utils/testing.py
index 94e13412a..3a681bf32 100644
--- a/src/accelerate/test_utils/testing.py
+++ b/src/accelerate/test_utils/testing.py
@@ -87,6 +87,15 @@ def require_cuda(test_case):
     return unittest.skipUnless(torch.cuda.is_available(), "test requires a GPU")(test_case)
 
 
+def require_mps(test_case):
+    """
+    Decorator marking a test that requires MPS backend. These tests are skipped when torch doesn't support `mps`
+    backend.
+    """
+    is_mps_supported = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
+    return unittest.skipUnless(is_mps_supported, "test requires a `mps` backend support in `torch`")(test_case)
+
+
 def require_huggingface_suite(test_case):
     """
     Decorator marking a test that requires transformers and datasets. These tests are skipped when they are not.
diff --git a/tests/test_big_modeling.py b/tests/test_big_modeling.py
index 4f738f4d8..467a91424 100644
--- a/tests/test_big_modeling.py
+++ b/tests/test_big_modeling.py
@@ -24,10 +24,11 @@
     disk_offload,
     dispatch_model,
     init_empty_weights,
+    init_on_device,
     load_checkpoint_and_dispatch,
 )
 from accelerate.hooks import remove_hook_from_submodules
-from accelerate.test_utils import require_cuda, require_multi_gpu, require_torch_min_version, slow
+from accelerate.test_utils import require_cuda, require_mps, require_multi_gpu, require_torch_min_version, slow
 from accelerate.utils import offload_state_dict
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
@@ -109,6 +110,22 @@ def test_init_empty_weights_very_large_model(self):
         with init_empty_weights():
             _ = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])
 
+    @require_cuda
+    def test_init_on_device_cuda(self):
+        device = torch.device("cuda:0")
+        with init_on_device(device):
+            model = nn.Linear(10, 10)
+        self.assertEqual(model.weight.device, device)
+        self.assertEqual(model.weight.device, device)
+
+    @require_mps
+    def test_init_on_device_mps(self):
+        device = torch.device("mps:0")
+        with init_on_device(device):
+            model = nn.Linear(10, 10)
+        self.assertEqual(model.weight.device, device)
+        self.assertEqual(model.weight.device, device)
+
     def test_cpu_offload(self):
         model = ModelForTest()
         x = torch.randn(2, 3)

Ah ah, that example needs adapting ;-p

@@ -62,6 +62,32 @@ def init_empty_weights(include_buffers: bool = False):
     """
     if not is_torch_version(">=", "1.9.0"):
         raise NotImplementedError("Initializing empty weights to a meta device requires torch >= 1.9.0")
+    with init_on_device(torch.device("meta"), include_buffers=include_buffers) as f:
+        yield f
+
+
+@contextmanager
+def init_on_device(device: torch.device, include_buffers: bool = False):
+    """
+    A context manager under which models are initialized with all parameters on the specified device.
+
+    Args:
+        device (`torch.device`):
+            Device to initialize all parameters on.
+        include_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to also put all buffers on the meta device while initializing.
+
+    Example:
+
+    ```python
+    import torch.nn as nn
+    from accelerate import init_empty_weights
+
+    # Initialize a model with 100 billions parameters in no time and without using any RAM.
+    with init_on_device(device=torch.device("cuda")):
+        tst = nn.Liner(100, 100)  # on `cuda` device
_The documentation is not available anymore as the PR was closed or merged._
I guess we can also extend to specific dtypes as well now.
Which dtype did you have in mind? For floating types there is a default you can set in PyTorch directly.
You mean https://pytorch.org/docs/stable/generated/torch.set_default_dtype.html ? Does fp16/bf16 also work? If so my bad.
Yes it works pretty well. We use it in Transformers [here](https://github.com/huggingface/transformers/blob/17292440c069118fbdb992b9a17da2098fab5b87/src/transformers/modeling_utils.py#L1109) for instance.
https://api.github.com/repos/huggingface/accelerate/pulls/926
diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml
index 6b9825523..ccaafbc8e 100644
--- a/docs/source/_toctree.yml
+++ b/docs/source/_toctree.yml
@@ -28,7 +28,7 @@
   - local: usage_guides/tracking
     title: Using experiment trackers
   - local: usage_guides/big_modeling
-    title: How to use large models with small resources
+    title: How perform inference on large models with small resources
   - local: usage_guides/memory
     title: How to avoid CUDA Out-of-Memory
   - local: usage_guides/sagemaker
diff --git a/docs/source/usage_guides/big_modeling.mdx b/docs/source/usage_guides/big_modeling.mdx
index 1e13849c6..d7523a03b 100644
--- a/docs/source/usage_guides/big_modeling.mdx
+++ b/docs/source/usage_guides/big_modeling.mdx
@@ -10,7 +10,7 @@ an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express o
 specific language governing permissions and limitations under the License.
 -->
 
-# Handling big models
+# Handling big models for inference
 
 When loading a pretrained model in PyTorch, the usual workflow looks like this:
 

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/921
diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
index ff3e95ead..2774b18ce 100644
--- a/src/accelerate/utils/modeling.py
+++ b/src/accelerate/utils/modeling.py
@@ -85,7 +85,11 @@ def dtype_byte_size(dtype: torch.dtype):
 
 
 def set_module_tensor_to_device(
-    module: nn.Module, tensor_name: str, device: Union[int, str, torch.device], value: Optional[torch.Tensor] = None
+    module: nn.Module,
+    tensor_name: str,
+    device: Union[int, str, torch.device],
+    value: Optional[torch.Tensor] = None,
+    dtype: Optional[Union[str, torch.dtype]] = None,
 ):
     """
     A helper function to set a given tensor (parameter of buffer) of a module on a specific device (note that doing
@@ -97,6 +101,9 @@ def set_module_tensor_to_device(
         device (`int`, `str` or `torch.device`): The device on which to set the tensor.
         value (`torch.Tensor`, *optional*): The value of the tensor (useful when going from the meta device to any
             other device).
+        dtype (`torch.dtype`, *optional*):
+            If passed along the value of the parameter will be cast to this `dtype`. Otherwise, `value` will be cast to
+            the dtype of the existing parameter in the model.
     """
     # Recurse if needed
     if "." in tensor_name:
@@ -116,6 +123,13 @@ def set_module_tensor_to_device(
     if old_value.device == torch.device("meta") and device not in ["meta", torch.device("meta")] and value is None:
         raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {device}.")
 
+    if value is not None:
+        if dtype is None:
+            # For compatibility with PyTorch load_state_dict which converts state dict dtype to existing dtype in model
+            value = value.to(old_value.dtype)
+        elif str(value.dtype).startswith(("torch.uint", "torch.int", "torch.bool")):
+            value = value.to(dtype)
+
     with torch.no_grad():
         if value is None:
             new_value = old_value.to(device)
@@ -680,8 +694,7 @@ def load_checkpoint_in_model(
         else:
             for param_name, param in checkpoint.items():
                 module_name = param_name
-                if dtype is not None and not str(param.dtype).startswith(("torch.uint", "torch.int", "torch.bool")):
-                    param = param.to(dtype)
+
                 while len(module_name) > 0 and module_name not in device_map:
                     module_name = ".".join(module_name.split(".")[:-1])
                 if module_name == "" and "" not in device_map:
@@ -696,7 +709,7 @@ def load_checkpoint_in_model(
                     set_module_tensor_to_device(model, param_name, "meta")
                     offload_weight(param, param_name, state_dict_folder, index=state_dict_index)
                 else:
-                    set_module_tensor_to_device(model, param_name, param_device, value=param)
+                    set_module_tensor_to_device(model, param_name, param_device, value=param, dtype=dtype)
 
         # Force Python to clean up.
         del checkpoint

should this not be better done in ` set_module_tensor_to_device` ? Or maybe additionally add a `torch_dtype` arg to `set_module_tensor_to_device` that handles the param correctly if `value=param` is used? 
@@ -680,8 +680,23 @@ def load_checkpoint_in_model(
         else:
             for param_name, param in checkpoint.items():
                 module_name = param_name
+
                 if dtype is not None and not str(param.dtype).startswith(("torch.uint", "torch.int", "torch.bool")):
                     param = param.to(dtype)
+
+                # For compatibility with PyTorch load_state_dict which converts state dict dtype to existing dtype in
+                # model
+                if dtype is None:
+                    old_param = model
+                    splits = param_name.split(".")
+                    for split in splits:
+                        old_param = getattr(old_param, split)
+                        if old_param is None:
+                            break
+
+                    if old_param is not None:
+                        param = param.to(old_param.dtype)
This is moved to `set_module_tensor_to_device`.
@@ -680,8 +694,7 @@ def load_checkpoint_in_model(
         else:
             for param_name, param in checkpoint.items():
                 module_name = param_name
-                if dtype is not None and not str(param.dtype).startswith(("torch.uint", "torch.int", "torch.bool")):
_The documentation is not available anymore as the PR was closed or merged._
Actually before merging, could it maybe be better to handle this in `set_module_tensor_to_device` ? E.g. add a `dtype` argument to the function there? This would be easier for `diffusers` to by in line with `accelerate` I think - see: https://github.com/huggingface/diffusers/blob/727434c206f6c22b746e460293035a1324f0bc13/src/diffusers/modeling_utils.py#L491
https://api.github.com/repos/huggingface/accelerate/pulls/920
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 1c3782067..7aaee2de7 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -196,7 +196,6 @@ def __init__(
         self,
         device_placement: bool = True,
         split_batches: bool = False,
-        fp16: bool = None,
         mixed_precision: Union[PrecisionType, str] = None,
         gradient_accumulation_steps: int = 1,
         cpu: bool = False,
@@ -220,13 +219,6 @@ def __init__(
                     f"Unknown mixed_precision mode: {mixed_precision}. Choose between {PrecisionType.list()}"
                 )
 
-        if fp16:
-            warnings.warn(
-                "`fp16=True` is deprecated and will be removed in version 0.15.0 of 🤗 Accelerate. Use `mixed_precision='fp16'` instead.",
-                FutureWarning,
-            )
-            mixed_precision = "fp16"
-
         if dynamo_backend is not None:
             dynamo_backend = DynamoBackend(dynamo_backend.upper())
 
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index fd5ad7d60..afb6498f3 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -20,7 +20,6 @@
 import os
 import subprocess
 import sys
-import warnings
 from ast import literal_eval
 from pathlib import Path
 from typing import Dict, List
@@ -64,7 +63,6 @@
     "--multi-gpu": "Distributed GPUs",
     "--tpu": "TPU",
     "--mps": "MPS",
-    "--use_mps_device": "MPS",
     "--use_deepspeed": "DeepSpeed Arguments",
     "--use_fsdp": "FSDP Arguments",
     "--use_megatron_lm": "Megatron-LM Arguments",
@@ -165,12 +163,6 @@ def launch_command_parser(subparsers=None):
     hardware_args.add_argument(
         "--tpu", default=False, action="store_true", help="Whether or not this should launch a TPU training."
     )
-    hardware_args.add_argument(
-        "--use_mps_device",
-        default=False,
-        action="store_true",
-        help="This argument is deprecated, use `--mps` instead.",
-    )
 
     # Resource selection arguments
     resource_args = parser.add_argument_group(
@@ -191,12 +183,6 @@ def launch_command_parser(subparsers=None):
         "Choose between FP16 and BF16 (bfloat16) training. "
         "BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.",
     )
-    resource_args.add_argument(
-        "--fp16",
-        default=False,
-        action="store_true",
-        help="This argument is deprecated, use `--mixed_precision fp16` instead.",
-    )
     resource_args.add_argument(
         "--num_processes", type=int, default=None, help="The total number of processes to be launched in parallel."
     )
@@ -525,12 +511,6 @@ def simple_launcher(args):
 
     current_env = os.environ.copy()
     current_env["ACCELERATE_USE_CPU"] = str(args.cpu or args.use_cpu)
-    if args.use_mps_device:
-        warnings.warn(
-            '`use_mps_device` flag is deprecated and will be removed in version 0.15.0 of 🤗 Accelerate. Use "--mps" instead.',
-            FutureWarning,
-        )
-        args.mps = True
     current_env["ACCELERATE_USE_MPS_DEVICE"] = str(args.mps)
     if args.mps:
         current_env["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
@@ -550,13 +530,6 @@ def simple_launcher(args):
             f"Unknown mixed_precision mode: {args.mixed_precision.lower()}. Choose between {PrecisionType.list()}."
         )
 
-    if args.fp16:
-        warnings.warn(
-            "`fp16` is deprecated and will be removed in version 0.15.0 of 🤗 Accelerate. Use `mixed_precision fp16` instead.",
-            FutureWarning,
-        )
-        mixed_precision = "fp16"
-
     current_env["ACCELERATE_MIXED_PRECISION"] = str(mixed_precision)
 
     try:
@@ -614,13 +587,6 @@ def multi_gpu_launcher(args):
     except ValueError:
         raise ValueError(f"Unknown mixed_precision mode: {mixed_precision}. Choose between {PrecisionType.list()}.")
 
-    if args.fp16:
-        warnings.warn(
-            "`fp16` is deprecated and will be removed in version 0.15.0 of 🤗 Accelerate. Use `mixed_precision fp16` instead.",
-            FutureWarning,
-        )
-        mixed_precision = "fp16"
-
     current_env["ACCELERATE_MIXED_PRECISION"] = str(mixed_precision)
 
     try:
@@ -749,13 +715,6 @@ def deepspeed_launcher(args):
             f"Unknown mixed_precision mode: {args.mixed_precision.lower()}. Choose between {PrecisionType.list()}."
         )
 
-    if args.fp16:
-        warnings.warn(
-            '--fp16 flag is deprecated and will be removed in version 0.15.0 of 🤗 Accelerate. Use "--mixed_precision fp16" instead.',
-            FutureWarning,
-        )
-        mixed_precision = "fp16"
-
     current_env["PYTHONPATH"] = env_var_path_add("PYTHONPATH", os.path.abspath("."))
     current_env["ACCELERATE_MIXED_PRECISION"] = str(mixed_precision)
     current_env["ACCELERATE_USE_DEEPSPEED"] = "true"
@@ -925,10 +884,6 @@ def sagemaker_launcher(sagemaker_config: SageMakerConfig, args):
             f"Unknown mixed_precision mode: {args.mixed_precision.lower()}. Choose between {PrecisionType.list()}."
         )
 
-    if args.fp16:
-        warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', FutureWarning)
-        mixed_precision = "fp16"
-
     try:
         dynamo_backend = DynamoBackend(args.dynamo_backend.upper())
     except ValueError:
@@ -1054,15 +1009,12 @@ def launch_command(args):
 
                 # Those args are handled separately
                 if (
-                    name not in ["compute_environment", "fp16", "mixed_precision", "distributed_type"]
+                    name not in ["compute_environment", "mixed_precision", "distributed_type"]
                     and getattr(args, name, None) is None
                 ):
                     setattr(args, name, attr)
         if not args.mixed_precision:
-            if args.fp16:
-                args.mixed_precision = "fp16"
-            else:
-                args.mixed_precision = defaults.mixed_precision
+            args.mixed_precision = defaults.mixed_precision
         if args.dynamo_backend is None:
             warned.append("\t`--dynamo_backend` was set to a value of `'no'`")
             args.dynamo_backend = "no"
diff --git a/src/accelerate/launchers.py b/src/accelerate/launchers.py
index 5768ec06b..2c9db6a75 100644
--- a/src/accelerate/launchers.py
+++ b/src/accelerate/launchers.py
@@ -15,7 +15,6 @@
 import os
 import sys
 import tempfile
-import warnings
 
 import torch
 
@@ -23,7 +22,7 @@
 from .utils import PrecisionType, PrepareForLaunch, patch_environment
 
 
-def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, mixed_precision="no", use_port="29500"):
+def notebook_launcher(function, args=(), num_processes=None, mixed_precision="no", use_port="29500"):
     """
     Launches a training function, using several processes if it's possible in the current environment (TPU with
     multiple cores for instance).
@@ -104,13 +103,6 @@ def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, mix
                     "function."
                 )
 
-            if use_fp16:
-                warnings.warn(
-                    "`fp16=True` is deprecated and will be removed in version 0.15.0 of 🤗 Accelerate. Use `mixed_precision='fp16'` instead.",
-                    FutureWarning,
-                )
-                mixed_precision = "fp16"
-
             # torch.distributed will expect a few environment variable to be here. We set the ones common to each
             # process here (the other ones will be set be the launcher).
             with patch_environment(
diff --git a/src/accelerate/test_utils/scripts/test_script.py b/src/accelerate/test_utils/scripts/test_script.py
index 6897d9084..ddb053a8a 100644
--- a/src/accelerate/test_utils/scripts/test_script.py
+++ b/src/accelerate/test_utils/scripts/test_script.py
@@ -275,29 +275,6 @@ def training_check():
         assert torch.allclose(old_model.a, model.a), "Did not obtain the same model on CPU or distributed training."
         assert torch.allclose(old_model.b, model.b), "Did not obtain the same model on CPU or distributed training."
 
-        # TEST that previous fp16 flag still works
-        print("Legacy FP16 training check.")
-        AcceleratorState._reset_state()
-        accelerator = Accelerator(fp16=True)
-        train_dl = DataLoader(train_set, batch_size=batch_size, shuffle=True, generator=generator)
-        model = RegressionModel()
-        optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
-
-        train_dl, model, optimizer = accelerator.prepare(train_dl, model, optimizer)
-        set_seed(42)
-        generator.manual_seed(42)
-        for _ in range(3):
-            for batch in train_dl:
-                model.zero_grad()
-                output = model(batch["x"])
-                loss = torch.nn.functional.mse_loss(output, batch["y"])
-                accelerator.backward(loss)
-                optimizer.step()
-
-        model = accelerator.unwrap_model(model).cpu()
-        assert torch.allclose(old_model.a, model.a), "Did not obtain the same model on CPU or distributed training."
-        assert torch.allclose(old_model.b, model.b), "Did not obtain the same model on CPU or distributed training."
-
     # BF16 support is only for CPU + TPU, and some GPU
     if is_bf16_available():
         # Mostly a test that BF16 doesn't crash as the operation inside the model is not converted to BF16

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/917
diff --git a/docs/source/package_reference/utilities.mdx b/docs/source/package_reference/utilities.mdx
index 3c1a1065e..9ebb60d34 100644
--- a/docs/source/package_reference/utilities.mdx
+++ b/docs/source/package_reference/utilities.mdx
@@ -24,6 +24,8 @@ These are basic dataclasses used throughout 🤗 Accelerate and they can be pass
 
 [[autodoc]] utils.PrecisionType
 
+[[autodoc]] utils.ProjectConfiguration
+
 ## Data Manipulation and Operations
 
 These include data operations that mimic the same `torch` ops but can be used on distributed processes.
diff --git a/docs/source/quicktour.mdx b/docs/source/quicktour.mdx
index 6d92e7958..6ea99f9bd 100644
--- a/docs/source/quicktour.mdx
+++ b/docs/source/quicktour.mdx
@@ -370,7 +370,11 @@ Note that since all the model parameters are references to tensors, this will lo
 ## Saving/loading entire states
 
 When training your model, you may want to save the current state of the model, optimizer, random generators, and potentially LR schedulers to be restored in the _same script_.
-You can use [`~Accelerator.save_state`] and [`~Accelerator.load_state`] respectively to do so, just by simply passing in a save location. 
+You can use [`~Accelerator.save_state`] and [`~Accelerator.load_state`] respectively to do so.
+
+To further customize where and how states saved through [`~Accelerator.save_state`] the [`~utils.ProjectConfiguration`] class can be used. For example 
+if `automatic_checkpoint_naming` is enabled each saved checkpoint will be located then at `Accelerator.project_dir/checkpoints/checkpoint_{checkpoint_number}`.
+
 If you have registered any other stateful items to be stored through [`~Accelerator.register_for_checkpointing`] they will also be saved and/or loaded.
 
 <Tip>
diff --git a/docs/source/usage_guides/checkpoint.mdx b/docs/source/usage_guides/checkpoint.mdx
index f2684971b..354549205 100644
--- a/docs/source/usage_guides/checkpoint.mdx
+++ b/docs/source/usage_guides/checkpoint.mdx
@@ -17,6 +17,9 @@ saving and loading the model, optimizer, RNG generators, and the GradScaler. Ins
 - Use [`~Accelerator.save_state`] for saving everything mentioned above to a folder location
 - Use [`~Accelerator.load_state`] for loading everything stored from an earlier `save_state`
 
+To further customize where and how states saved through [`~Accelerator.save_state`] the [`~utils.ProjectConfiguration`] class can be used. For example 
+if `automatic_checkpoint_naming` is enabled each saved checkpoint will be located then at `Accelerator.project_dir/checkpoints/checkpoint_{checkpoint_number}`.
+
 It should be noted that the expectation is that those states come from the same training script, they should not be from two separate scripts.
 
 - By using [`~Accelerator.register_for_checkpointing`], you can register custom objects to be automatically stored or loaded from the two prior functions,
@@ -28,7 +31,7 @@ Below is a brief example using checkpointing to save and reload a state during t
 from accelerate import Accelerator
 import torch
 
-accelerator = Accelerator()
+accelerator = Accelerator(project_dir="my/save/path")
 
 my_scheduler = torch.optim.lr_scheduler.StepLR(my_optimizer, step_size=1, gamma=0.99)
 my_model, my_optimizer, my_training_dataloader = accelerator.prepare(my_model, my_optimizer, my_training_dataloader)
@@ -37,7 +40,7 @@ my_model, my_optimizer, my_training_dataloader = accelerator.prepare(my_model, m
 accelerator.register_for_checkpointing(my_scheduler)
 
 # Save the starting state
-accelerator.save_state("my/save/path")
+accelerator.save_state()
 
 device = accelerator.device
 my_model.to(device)
@@ -56,5 +59,5 @@ for epoch in range(num_epochs):
     my_scheduler.step()
 
 # Restore previous state
-accelerator.load_state("my/save/path")
+accelerator.load_state("my/save/path/checkpointing/checkpoint_0")
 ```
diff --git a/docs/source/usage_guides/tracking.mdx b/docs/source/usage_guides/tracking.mdx
index b22a4dd14..3ff40a4a4 100644
--- a/docs/source/usage_guides/tracking.mdx
+++ b/docs/source/usage_guides/tracking.mdx
@@ -83,6 +83,12 @@ for iteration in config["num_iterations"]:
 accelerator.end_training()
 ```
 
+If a tracker requires a directory to save data to such as `TensorBoard` then a `logging_dir` or `project_dir` can be passed in. `project_dir` is useful 
+if there are other further configurations such as those which can be combined with the [`~utils.ProjectConfiguration`] dataclass.
+
+```python
+accelerator = Accelerator(log_with="tensorboard", logging_dir=".")
+```
 
 ## Implementing Custom Trackers
 
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 1c3782067..151d1db35 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -16,6 +16,7 @@
 import gc
 import math
 import os
+import shutil
 import sys
 import warnings
 from contextlib import contextmanager
@@ -44,6 +45,7 @@
     LoggerType,
     MegatronLMPlugin,
     PrecisionType,
+    ProjectConfiguration,
     RNGType,
     compare_versions,
     convert_outputs_to_fp32,
@@ -157,8 +159,11 @@ class Accelerator:
             - `"comet_ml"`
             If `"all"` is selected, will pick up all available trackers in the environment and initialize them. Can
             also accept implementations of `GeneralTracker` for custom trackers, and can be combined with `"all"`.
-        logging_dir (`str`, `os.PathLike`, *optional*):
-            A path to a directory for storing logs of locally-compatible loggers.
+        project_config (`ProjectConfiguration`, *optional*):
+            A configuration for how saving the state can be handled.
+        project_dir (`str`, `os.PathLike`, *optional*):
+            A path to a directory for storing data such as logs of locally-compatible loggers and potentially saved
+            checkpoints.
         dispatch_batches (`bool`, *optional*):
             If set to `True`, the dataloader prepared by the Accelerator is only iterated through on the main process
             and then the batches are split and broadcast to each process. Will default to `True` for `DataLoader` whose
@@ -205,6 +210,8 @@ def __init__(
         megatron_lm_plugin: MegatronLMPlugin = None,
         rng_types: Optional[List[Union[str, RNGType]]] = None,
         log_with: Optional[List[Union[str, LoggerType, GeneralTracker]]] = None,
+        project_dir: Optional[Union[str, os.PathLike]] = None,
+        project_config: Optional[ProjectConfiguration] = None,
         logging_dir: Optional[Union[str, os.PathLike]] = None,
         dispatch_batches: Optional[bool] = None,
         even_batches: bool = True,
@@ -212,7 +219,19 @@ def __init__(
         kwargs_handlers: Optional[List[KwargsHandler]] = None,
         dynamo_backend: Union[DynamoBackend, str] = None,
     ):
-        self.logging_dir = logging_dir
+        if project_config is not None:
+            self.project_configuration = project_config
+        else:
+            self.project_configuration = ProjectConfiguration(project_dir=project_dir)
+
+        if logging_dir is not None:
+            warnings.warn(
+                "`logging_dir` is deprecated and will be removed in version 0.18.0 of 🤗 Accelerate. Use `project_dir` instead.",
+                FutureWarning,
+            )
+            self.project_configuration.logging_dir = logging_dir
+        if project_dir is not None and self.project_dir is None:
+            self.project_configuration.project_dir = project_dir
         if mixed_precision is not None:
             mixed_precision = str(mixed_precision)
             if mixed_precision not in PrecisionType:
@@ -420,6 +439,18 @@ def local_process_index(self):
     def device(self):
         return self.state.device
 
+    @property
+    def project_dir(self):
+        return self.project_configuration.project_dir
+
+    @property
+    def logging_dir(self):
+        return self.project_configuration.logging_dir
+
+    @property
+    def save_iteration(self):
+        return self.project_configuration.iteration
+
     @property
     def is_main_process(self):
         """True for one process only."""
@@ -1596,9 +1627,16 @@ def save(self, obj, f):
         """
         save(obj, f)
 
-    def save_state(self, output_dir: str):
+    def save_state(self, output_dir: str = None):
         """
-        Saves the current states of the model, optimizer, scaler, RNG generators, and registered objects.
+        Saves the current states of the model, optimizer, scaler, RNG generators, and registered objects to a folder.
+
+        If a `ProjectConfiguration` was passed to the `Accelerator` object with `automatic_checkpoint_naming` enabled
+        then checkpoints will be saved to `self.project_dir/checkpoints`. If the number of current saves is greater
+        than `total_limit` then the oldest save is deleted. Each checkpoint is saved in seperate folders named
+        `checkpoint_<iteration>`.
+
+        Otherwise they are just saved to `output_dir`.
 
         <Tip>
 
@@ -1611,8 +1649,25 @@ def save_state(self, output_dir: str):
             output_dir (`str` or `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)
+        if self.project_configuration.automatic_checkpoint_naming:
+            output_dir = os.path.join(self.project_dir, "checkpoints")
+        os.makedirs(output_dir, exist_ok=True)
+        if self.project_configuration.automatic_checkpoint_naming:
+            folders = [os.path.join(output_dir, folder) for folder in os.listdir(output_dir)]
+            if self.project_configuration.total_limit is not None and (
+                len(folders) + 1 > self.project_configuration.total_limit
+            ):
+                folders.sort()
+                logger.warning(
+                    f"Deleting {len(folders) + 1 - self.project_configuration.total_limit} checkpoints to make room for new checkpoint."
+                )
+                for folder in folders[: len(folders) + 1 - self.project_configuration.total_limit]:
+                    shutil.rmtree(folder)
+            output_dir = os.path.join(output_dir, f"checkpoint_{self.save_iteration}")
+            if os.path.exists(output_dir):
+                raise ValueError(
+                    f"Checkpoint directory {output_dir} ({self.save_iteration}) already exists. Please manually override `self.save_iteration` with what iteration to start with."
+                )
         os.makedirs(output_dir, exist_ok=True)
         logger.info(f"Saving current state to {output_dir}")
 
@@ -1660,6 +1715,7 @@ def save_state(self, output_dir: str):
         )
         for i, obj in enumerate(self._custom_objects):
             save_custom_state(obj, output_dir, i)
+        self.project_configuration.iteration += 1
         return save_location
 
     def load_state(self, input_dir: str):
diff --git a/src/accelerate/tracking.py b/src/accelerate/tracking.py
index 75694ba72..7afebea48 100644
--- a/src/accelerate/tracking.py
+++ b/src/accelerate/tracking.py
@@ -147,7 +147,7 @@ class TensorBoardTracker(GeneralTracker):
     name = "tensorboard"
     requires_logging_directory = True
 
-    def __init__(self, run_name: str, logging_dir: Optional[Union[str, os.PathLike]], **kwargs):
+    def __init__(self, run_name: str, logging_dir: Optional[Union[str, os.PathLike]] = None, **kwargs):
         self.run_name = run_name
         self.logging_dir = os.path.join(logging_dir, run_name)
         self.writer = tensorboard.SummaryWriter(self.logging_dir, **kwargs)
@@ -451,7 +451,6 @@ def __init__(
         run_name: Optional[str] = None,
         description: Optional[str] = None,
     ):
-
         experiment_name = os.getenv("MLFLOW_EXPERIMENT_NAME", experiment_name)
         run_id = os.getenv("MLFLOW_RUN_ID", run_id)
         tags = os.getenv("MLFLOW_TAGS", tags)
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
index be78d0089..d3fb916e8 100644
--- a/src/accelerate/utils/__init__.py
+++ b/src/accelerate/utils/__init__.py
@@ -16,6 +16,7 @@
     LoggerType,
     MegatronLMPlugin,
     PrecisionType,
+    ProjectConfiguration,
     RNGType,
     SageMakerDistributedType,
     TensorInformation,
diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py
index eafad8a09..dc8883425 100644
--- a/src/accelerate/utils/dataclasses.py
+++ b/src/accelerate/utils/dataclasses.py
@@ -313,6 +313,39 @@ class TensorInformation:
     dtype: torch.dtype
 
 
+@dataclass
+class ProjectConfiguration:
+    """
+    Configuration for the Accelerator object based on inner-project needs.
+    """
+
+    project_dir: str = field(default=None, metadata={"help": "A path to a directory for storing data."})
+    logging_dir: str = field(
+        default=None,
+        metadata={
+            "help": "A path to a directory for storing logs of locally-compatible loggers. If None, defaults to `project_dir`."
+        },
+    )
+    automatic_checkpoint_naming: bool = field(
+        default=False,
+        metadata={"help": "Whether saved states should be automatically iteratively named."},
+    )
+
+    total_limit: int = field(
+        default=None,
+        metadata={"help": "The maximum number of total saved states to keep."},
+    )
+
+    iteration: int = field(
+        default=0,
+        metadata={"help": "The current save iteration."},
+    )
+
+    def __post_init__(self):
+        if self.logging_dir is None:
+            self.logging_dir = self.project_dir
+
+
 @dataclass
 class DeepSpeedPlugin:
     """
diff --git a/tests/test_state_checkpointing.py b/tests/test_state_checkpointing.py
index 87b2d3b06..dd2f6a98b 100644
--- a/tests/test_state_checkpointing.py
+++ b/tests/test_state_checkpointing.py
@@ -23,7 +23,7 @@
 from torch.utils.data import DataLoader, TensorDataset
 
 from accelerate import Accelerator
-from accelerate.utils import set_seed
+from accelerate.utils import ProjectConfiguration, set_seed
 
 
 logger = logging.getLogger(__name__)
@@ -75,7 +75,26 @@ def forward(self, x):
 
 
 class CheckpointTest(unittest.TestCase):
-    def test_can_resume_training(self):
+    def test_with_save_limit(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()
+            project_config = ProjectConfiguration(total_limit=1, project_dir=tmpdir, automatic_checkpoint_naming=True)
+            # Train baseline
+            accelerator = Accelerator(project_config=project_config)
+            model, optimizer, train_dataloader, valid_dataloader = accelerator.prepare(
+                model, optimizer, train_dataloader, valid_dataloader
+            )
+            # Save initial
+            accelerator.save_state()
+
+            # Save second state
+            accelerator.save_state()
+            self.assertEqual(len(os.listdir(accelerator.project_dir)), 1)
+
+    def test_can_resume_training_with_folder(self):
         with tempfile.TemporaryDirectory() as tmpdir:
             set_seed(42)
             model = DummyModel()
@@ -126,6 +145,58 @@ def test_can_resume_training(self):
             self.assertEqual(opt_state1, opt_state3)
             self.assertEqual(ground_truth_rands, test_rands)
 
+    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()
+            project_config = ProjectConfiguration(automatic_checkpoint_naming=True)
+
+            # Train baseline
+            accelerator = Accelerator(project_dir=tmpdir, project_config=project_config)
+            model, optimizer, train_dataloader, valid_dataloader = accelerator.prepare(
+                model, optimizer, train_dataloader, valid_dataloader
+            )
+            # Save initial
+            accelerator.save_state()
+            (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()
+            project_config = ProjectConfiguration(iteration=1, automatic_checkpoint_naming=True)
+            accelerator = Accelerator(project_dir=tmpdir, project_config=project_config)
+            model, optimizer, train_dataloader, valid_dataloader = accelerator.prepare(
+                model, optimizer, train_dataloader, valid_dataloader
+            )
+            accelerator.load_state(os.path.join(tmpdir, "checkpoints", "checkpoint_0"))
+            (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
+            accelerator.save_state()
+
+            # Load everything back in and make sure all states work
+            accelerator.load_state(os.path.join(tmpdir, "checkpoints", "checkpoint_1"))
+            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)
+
     def test_invalid_registration(self):
         t = torch.tensor([1, 2, 3])
         t1 = torch.tensor([2, 3, 4])
@@ -147,19 +218,18 @@ def test_with_scheduler(self):
             optimizer = torch.optim.Adam(params=model.parameters(), lr=1e-3)
             scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.99)
             train_dataloader, valid_dataloader = dummy_dataloaders()
+            project_config = ProjectConfiguration(automatic_checkpoint_naming=True)
             # Train baseline
-            accelerator = Accelerator()
-            model, optimizer, train_dataloader, valid_dataloader = accelerator.prepare(
-                model, optimizer, train_dataloader, valid_dataloader
+            accelerator = Accelerator(project_dir=tmpdir, project_config=project_config)
+            model, optimizer, train_dataloader, valid_dataloader, scheduler = accelerator.prepare(
+                model, optimizer, train_dataloader, valid_dataloader, scheduler
             )
-            accelerator.register_for_checkpointing(scheduler)
             # Save initial
-            initial = os.path.join(tmpdir, "initial")
-            accelerator.save_state(initial)
+            accelerator.save_state()
             scheduler_state = scheduler.state_dict()
             train(3, model, train_dataloader, optimizer, accelerator, scheduler)
             self.assertNotEqual(scheduler_state, scheduler.state_dict())
 
             # Load everything back in and make sure all states work
-            accelerator.load_state(initial)
+            accelerator.load_state(os.path.join(tmpdir, "checkpoints", "checkpoint_0"))
             self.assertEqual(scheduler_state, scheduler.state_dict())
diff --git a/tests/test_tracking.py b/tests/test_tracking.py
index d6ff4908a..1aa341f9b 100644
--- a/tests/test_tracking.py
+++ b/tests/test_tracking.py
@@ -70,9 +70,11 @@ def test_log(self):
             log = list(filter(lambda x: x.is_file(), Path(f"{dirpath}/{project_name}").iterdir()))[0]
             self.assertNotEqual(str(log), "")
 
-    def test_logging_dir(self):
+    def test_project_dir(self):
         with self.assertRaisesRegex(ValueError, "Logging with `tensorboard` requires a `logging_dir`"):
             _ = Accelerator(log_with="tensorboard")
+        with tempfile.TemporaryDirectory() as dirpath:
+            _ = Accelerator(log_with="tensorboard", project_dir=dirpath)
         with tempfile.TemporaryDirectory() as dirpath:
             _ = Accelerator(log_with="tensorboard", logging_dir=dirpath)
 

This is starting to make too many arguments.
@@ -158,7 +159,16 @@ class Accelerator:
             If `"all"` is selected, will pick up all available trackers in the environment and initialize them. Can
             also accept implementations of `GeneralTracker` for custom trackers, and can be combined with `"all"`.
         logging_dir (`str`, `os.PathLike`, *optional*):
-            A path to a directory for storing logs of locally-compatible loggers.
+            A path to a directory for storing logs of locally-compatible loggers. If not passed will save in
+            `project_dir` by default.
+        project_dir (`str`, `os.PathLike`, *optional*):
+            A path to a directory for storing data such as logs of locally-compatible loggers and potentially saved
+            checkpoints.
+        automatic_checkpoint_naming (`bool`, *optional*, defaults to `False`):
+            Whether saved states should be stored in `project_location` and be automatically iteratively named.
+        save_total_limit (`int`, *optional*):
+            The maximum number of checkpoints to keep if performing `automatic_checkpoint_naming`, will default to all
+            of them.
Fixed with a `SaveConfiguration` class
@@ -158,7 +159,16 @@ class Accelerator:
             If `"all"` is selected, will pick up all available trackers in the environment and initialize them. Can
             also accept implementations of `GeneralTracker` for custom trackers, and can be combined with `"all"`.
         logging_dir (`str`, `os.PathLike`, *optional*):
-            A path to a directory for storing logs of locally-compatible loggers.
+            A path to a directory for storing logs of locally-compatible loggers. If not passed will save in
+            `project_dir` by default.
+        project_dir (`str`, `os.PathLike`, *optional*):
+            A path to a directory for storing data such as logs of locally-compatible loggers and potentially saved
+            checkpoints.
+        automatic_checkpoint_naming (`bool`, *optional*, defaults to `False`):
+            Whether saved states should be stored in `project_location` and be automatically iteratively named.
+        save_total_limit (`int`, *optional*):
+            The maximum number of checkpoints to keep if performing `automatic_checkpoint_naming`, will default to all
+            of them.
Let's also deprecate logging_dir to include it in the SaveConfig? Wdyt?
@@ -205,13 +213,23 @@ def __init__(
         megatron_lm_plugin: MegatronLMPlugin = None,
         rng_types: Optional[List[Union[str, RNGType]]] = None,
         log_with: Optional[List[Union[str, LoggerType, GeneralTracker]]] = None,
+        project_dir: Optional[Union[str, os.PathLike]] = None,
+        save_config: Optional[SaveConfiguration] = None,
         logging_dir: Optional[Union[str, os.PathLike]] = None,
         dispatch_batches: Optional[bool] = None,
         even_batches: bool = True,
         step_scheduler_with_optimizer: bool = True,
         kwargs_handlers: Optional[List[KwargsHandler]] = None,
         dynamo_backend: Union[DynamoBackend, str] = None,
     ):
+        if save_config is not None:
+            self.save_configuration = save_config
+        else:
+            self.save_configuration = SaveConfiguration()
+        if project_dir is not None and self.project_dir is None:
+            self.save_configuration.project_dir = project_dir
+        if self.project_dir is not None and logging_dir is None:
I was debating on that, if we did so I think it'd be better to name it `ProjectConfiguration` instead since it's all relative to the total "project" at that point. But can make adjustments for it. The only "con" I see is given the conversation above for the user the code would change, but the end-result would not so I think it'd be fine
@@ -205,13 +213,23 @@ def __init__(
         megatron_lm_plugin: MegatronLMPlugin = None,
         rng_types: Optional[List[Union[str, RNGType]]] = None,
         log_with: Optional[List[Union[str, LoggerType, GeneralTracker]]] = None,
+        project_dir: Optional[Union[str, os.PathLike]] = None,
+        save_config: Optional[SaveConfiguration] = None,
         logging_dir: Optional[Union[str, os.PathLike]] = None,
         dispatch_batches: Optional[bool] = None,
         even_batches: bool = True,
         step_scheduler_with_optimizer: bool = True,
         kwargs_handlers: Optional[List[KwargsHandler]] = None,
         dynamo_backend: Union[DynamoBackend, str] = None,
     ):
+        if save_config is not None:
+            self.save_configuration = save_config
+        else:
+            self.save_configuration = SaveConfiguration()
+        if project_dir is not None and self.project_dir is None:
+            self.save_configuration.project_dir = project_dir
+        if self.project_dir is not None and logging_dir is None:
```suggestion
if there are other further configurations such as those which can be combined with the [`~utils.ProjectConfiguration`] dataclass.
```
@@ -83,6 +83,12 @@ for iteration in config["num_iterations"]:
 accelerator.end_training()
 ```
 
+If a tracker requires a directory to save data to such as `TensorBoard` then a `logging_dir` or `project_dir` can be passed in. `project_dir` is useful 
+for if there are other further configurations such as those which can be combined with the [`~utils.ProjectConfiguration`] dataclass.
_The documentation is not available anymore as the PR was closed or merged._
In my setup I have a VM for training and a LXC that hosts TensorBoard linked by a NFS share. I set `logging_dir` to the NFS share so the VM pushes its logs directly to TB as they are generated. Since this PR centralizes logging and output I would no longer be able to dump my logs directly to the TB server and instead have to do some complicated rsync watchdog service.

Can we keep `logging_dir` as an option? Just set it to `project_dir` by default, no need to depreciate it.
@Cyberes there is a deprecation cycle for `logging_dir` of two full versions and when it comes to trackers `project_dir` acts the exact same as the `logging_dir` parameter, so for your particular setup if I understand correctly would just need to change `logging_dir` to `project_dir` and nothing else.  

Is this reasonable enough for you? Or is there something more specific/drastic I've misunderstood.
Just to make sure I'm understanding the new functionality, the checkpoints are saved to `project_dir/checkpoints`  and the logs are saved to `project_dir/logs/[name]`?
Checkpoints are saved to `project_dir/checkpoints`, logs are saved to `project_dir/[name]`
Yeah, I'd like to have the option to put them in completely separate directories. For example, save the checkpoint states to `project_dir/checkpoints` and TB logs to `/mnt/tensorboard`.

If you're against the seperate directories thing I could always override `TensorBoardTracker`, I guess.
That use case makes sense, will do as you suggested and override only if not passed. Thanks @Cyberes :) 
Sweeeeeeeeeet, thanks @muellerzr.
https://api.github.com/repos/huggingface/accelerate/pulls/916
diff --git a/docs/source/usage_guides/checkpoint.mdx b/docs/source/usage_guides/checkpoint.mdx
index 7d6bbbf99..f2684971b 100644
--- a/docs/source/usage_guides/checkpoint.mdx
+++ b/docs/source/usage_guides/checkpoint.mdx
@@ -31,13 +31,13 @@ import torch
 accelerator = Accelerator()
 
 my_scheduler = torch.optim.lr_scheduler.StepLR(my_optimizer, step_size=1, gamma=0.99)
-my_model, my_optimizer, my_training_dataloader = accelerate.prepare(my_model, my_optimizer, my_training_dataloader)
+my_model, my_optimizer, my_training_dataloader = accelerator.prepare(my_model, my_optimizer, my_training_dataloader)
 
 # Register the LR scheduler
-accelerate.register_for_checkpointing(my_scheduler)
+accelerator.register_for_checkpointing(my_scheduler)
 
 # Save the starting state
-accelerate.save_state("my/save/path")
+accelerator.save_state("my/save/path")
 
 device = accelerator.device
 my_model.to(device)
@@ -56,5 +56,5 @@ for epoch in range(num_epochs):
     my_scheduler.step()
 
 # Restore previous state
-accelerate.load_state("my/save/path")
+accelerator.load_state("my/save/path")
 ```

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_915). All of your documentation changes will be reflected on that endpoint.
Thanks for the fixes! 🤗
https://api.github.com/repos/huggingface/accelerate/pulls/915
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 85fc1b4d2..1c3782067 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -939,9 +939,9 @@ def prepare_model(self, model: torch.nn.Module, device_placement=None):
                     backward_prefetch=fsdp_plugin.backward_prefetch,
                     mixed_precision=fsdp_plugin.mixed_precision_policy,
                     ignored_modules=fsdp_plugin.ignored_modules,
+                    device_id=self.device,
+                    limit_all_gathers=fsdp_plugin.limit_all_gathers,
                 )
-                if not fsdp_plugin.cpu_offload.offload_params:
-                    model.to(self.device)
             self._models[-1] = model
         elif self.distributed_type == DistributedType.MULTI_CPU:
             kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {}
diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py
index 2e286b657..eafad8a09 100644
--- a/src/accelerate/utils/dataclasses.py
+++ b/src/accelerate/utils/dataclasses.py
@@ -554,6 +554,16 @@ class FullyShardedDataParallelPlugin:
         },
     )
 
+    limit_all_gathers: bool = field(
+        default=False,
+        metadata={
+            "help": "If False, then FSDP allows the CPU thread to schedule all-gathers "
+            "without any extra synchronization. If True, then FSDP explicitly synchronizes the CPU thread to prevent "
+            "too many in-flight all-gathers. This bool only affects the sharded strategies that schedule all-gathers. "
+            "Enabling this can help lower the number of CUDA malloc retries."
+        },
+    )
+
     def __post_init__(self):
         from torch.distributed.fsdp.fully_sharded_data_parallel import (
             BackwardPrefetch,

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/911
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 85fc1b4d2..0808ca462 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -349,7 +349,11 @@ def __init__(
         self.scaler = None
         self.native_amp = False
         err = "{mode} mixed precision requires {requirement}"
-        if self.state.mixed_precision == "fp16" and self.distributed_type != DistributedType.MEGATRON_LM:
+        if (
+            self.state.mixed_precision == "fp16"
+            and self.device.type != "cpu"
+            and self.distributed_type != DistributedType.MEGATRON_LM
+        ):
             self.native_amp = True
             if not torch.cuda.is_available() and not parse_flag_from_env("ACCELERATE_USE_MPS_DEVICE"):
                 raise ValueError(err.format(mode="fp16", requirement="a GPU"))
@@ -1847,7 +1851,7 @@ def autocast(self):
         if self.native_amp:
             if self.mixed_precision == "fp16" and is_torch_version(">=", "1.10"):
                 autocast_context = torch.cuda.amp.autocast(dtype=torch.float16)
-            elif self.mixed_precision == "bf16" and is_bf16_available():
+            elif self.mixed_precision == "bf16":
                 if self.distributed_type in [DistributedType.NO, DistributedType.MULTI_CPU, DistributedType.MULTI_GPU]:
                     autocast_context = torch.autocast(dtype=torch.bfloat16, device_type=self.device.type)
             else:

@muellerzr @sgugger @yao-matrix @jiqing-feng.
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/909
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 12f67088a..97bd75833 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -1006,7 +1006,7 @@ def launch_command(args):
         raise ValueError(
             "You can only use one of `--cpu`, `--multi_gpu`, `--mps`, `--tpu`, `--use_deepspeed`, `--use_fsdp` at a time."
         )
-    if args.multi_gpu and args.num_processes < 2:
+    if args.multi_gpu and (args.num_processes is not None) and (args.num_processes < 2):
         raise ValueError("You need to use at least 2 processes to use `--multi_gpu`.")
 
     defaults = None

Can all be grouped in the same if ;-)
@@ -1006,8 +1006,9 @@ def launch_command(args):
         raise ValueError(
             "You can only use one of `--cpu`, `--multi_gpu`, `--mps`, `--tpu`, `--use_deepspeed`, `--use_fsdp` at a time."
         )
-    if args.multi_gpu and args.num_processes < 2:
-        raise ValueError("You need to use at least 2 processes to use `--multi_gpu`.")
+    if args.num_processes is not None:
+        if args.multi_gpu and args.num_processes < 2:
+            raise ValueError("You need to use at least 2 processes to use `--multi_gpu`.")
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/907
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 1d6f7ee10..85fc1b4d2 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -365,7 +365,10 @@ def __init__(
             and self.distributed_type != DistributedType.FSDP
             and self.distributed_type != DistributedType.MEGATRON_LM
         ):
-            self.native_amp = is_bf16_available(True)
+            if self.device.type == "cpu":
+                self.native_amp = is_torch_version(">=", "1.10")
+            else:
+                self.native_amp = is_bf16_available(True)
             if mixed_precision == "bf16" and not self.native_amp and not is_tpu_available():
                 raise ValueError(err.format(mode="bf16", requirement="PyTorch >= 1.10 and a supported device."))
 
@@ -948,8 +951,7 @@ def prepare_model(self, model: torch.nn.Module, device_placement=None):
             if self.mixed_precision == "fp16" and is_torch_version(">=", "1.10"):
                 model.forward = torch.cuda.amp.autocast(dtype=torch.float16)(model.forward)
             elif self.mixed_precision == "bf16" and self.distributed_type != DistributedType.TPU:
-                device_type = "cuda" if torch.cuda.is_available() else "cpu"
-                model.forward = torch.autocast(device_type=device_type, dtype=torch.bfloat16)(model.forward)
+                model.forward = torch.autocast(device_type=self.device.type, dtype=torch.bfloat16)(model.forward)
             else:
                 model.forward = torch.cuda.amp.autocast()(model.forward)
             model.forward = convert_outputs_to_fp32(model.forward)
@@ -1847,8 +1849,7 @@ def autocast(self):
                 autocast_context = torch.cuda.amp.autocast(dtype=torch.float16)
             elif self.mixed_precision == "bf16" and is_bf16_available():
                 if self.distributed_type in [DistributedType.NO, DistributedType.MULTI_CPU, DistributedType.MULTI_GPU]:
-                    device_type = "cpu" if not torch.cuda.is_available() else "cuda"
-                    autocast_context = torch.autocast(dtype=torch.bfloat16, device_type=device_type)
+                    autocast_context = torch.autocast(dtype=torch.bfloat16, device_type=self.device.type)
             else:
                 autocast_context = torch.cuda.amp.autocast()
 
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 12f67088a..af3ba0691 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -1089,7 +1089,7 @@ def launch_command(args):
                 ["MPI_LOCALNRANKS", "OMPI_COMM_WORLD_LOCAL_SIZE", "MV2_COMM_WORLD_LOCAL_SIZE"], 1
             )
             threads_per_process = int(psutil.cpu_count(logical=False) / local_size)
-            if args.num_cpu_threads_per_process > 1:
+            if threads_per_process > 1:
                 args.num_cpu_threads_per_process = threads_per_process
                 warned.append(
                     f"\t`--num_cpu_threads_per_process` was set to `{args.num_cpu_threads_per_process}` to improve out-of-box performance when training on CPUs"

This line (and the next occurrence) can be removed entirely and just pass in `self.device.type` to `torch.autocast`
@@ -948,7 +951,7 @@ def prepare_model(self, model: torch.nn.Module, device_placement=None):
             if self.mixed_precision == "fp16" and is_torch_version(">=", "1.10"):
                 model.forward = torch.cuda.amp.autocast(dtype=torch.float16)(model.forward)
             elif self.mixed_precision == "bf16" and self.distributed_type != DistributedType.TPU:
-                device_type = "cuda" if torch.cuda.is_available() else "cpu"
+                device_type = "cuda" if self.device.type != "cpu" and torch.cuda.is_available() else "cpu"
good point. done
@@ -948,7 +951,7 @@ def prepare_model(self, model: torch.nn.Module, device_placement=None):
             if self.mixed_precision == "fp16" and is_torch_version(">=", "1.10"):
                 model.forward = torch.cuda.amp.autocast(dtype=torch.float16)(model.forward)
             elif self.mixed_precision == "bf16" and self.distributed_type != DistributedType.TPU:
-                device_type = "cuda" if torch.cuda.is_available() else "cpu"
+                device_type = "cuda" if self.device.type != "cpu" and torch.cuda.is_available() else "cpu"
done
@@ -948,7 +951,7 @@ def prepare_model(self, model: torch.nn.Module, device_placement=None):
             if self.mixed_precision == "fp16" and is_torch_version(">=", "1.10"):
                 model.forward = torch.cuda.amp.autocast(dtype=torch.float16)(model.forward)
             elif self.mixed_precision == "bf16" and self.distributed_type != DistributedType.TPU:
-                device_type = "cuda" if torch.cuda.is_available() else "cpu"
+                device_type = "cuda" if self.device.type != "cpu" and torch.cuda.is_available() else "cpu"
@sgugger @yao-matrix  @jiqing-feng  please notice the issue
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/906
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 081380331..12f67088a 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -1002,8 +1002,12 @@ def sagemaker_launcher(sagemaker_config: SageMakerConfig, args):
 
 def launch_command(args):
     # Sanity checks
-    if sum([args.multi_gpu, args.tpu, args.use_deepspeed, args.use_fsdp]) > 1:
-        raise ValueError("You can only pick one between `--multi_gpu`, `--use_deepspeed`, `--tpu`, `--use_fsdp`.")
+    if sum([args.multi_gpu, args.cpu, args.tpu, args.mps, args.use_deepspeed, args.use_fsdp]) > 1:
+        raise ValueError(
+            "You can only use one of `--cpu`, `--multi_gpu`, `--mps`, `--tpu`, `--use_deepspeed`, `--use_fsdp` at a time."
+        )
+    if args.multi_gpu and args.num_processes < 2:
+        raise ValueError("You need to use at least 2 processes to use `--multi_gpu`.")
 
     defaults = None
     warned = []

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/905
diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py
index 490ea3f4f..2e286b657 100644
--- a/src/accelerate/utils/dataclasses.py
+++ b/src/accelerate/utils/dataclasses.py
@@ -65,7 +65,18 @@ class DistributedDataParallelKwargs(KwargsHandler):
 
     `static_graph` is only available in PyTorch 1.11.0 and later versions.
 
-    </Tip>"""
+    </Tip>
+
+    Example:
+
+    ```python
+    from accelerate import Accelerator
+    from accelerate.utils import DistributedDataParallelKwargs
+
+    kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
+    accelerator = Accelerator(kwargs_handlers=[kwargs])
+    ```
+    """
 
     dim: int = 0
     broadcast_buffers: bool = True
@@ -87,7 +98,18 @@ class GradScalerKwargs(KwargsHandler):
 
     `GradScaler` is only available in PyTorch 1.5.0 and later versions.
 
-    </Tip>"""
+    </Tip>
+
+    Example:
+
+    ```python
+    from accelerate import Accelerator
+    from accelerate.utils import GradScalerKwargs
+
+    kwargs = GradScalerKwargs(backoff_filter=0.25)
+    accelerator = Accelerator(kwargs_handlers=[kwargs])
+    ```
+    """
 
     init_scale: float = 65536.0
     growth_factor: float = 2.0
@@ -103,6 +125,15 @@ class InitProcessGroupKwargs(KwargsHandler):
     to the documentation of this
     [method](https://pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group) for more
     information on each argument.
+
+    ```python
+    from datetime import timedelta
+    from accelerate import Accelerator
+    from accelerate.utils import InitProcessGroupKwargs
+
+    kwargs = InitProcessGroupKwargs(timeout=timedelta(seconds=800))
+    accelerator = Accelerator(kwargs_handlers=[kwargs])
+    ```
     """
 
     init_method: Optional[str] = None

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/904
diff --git a/docs/source/package_reference/cli.mdx b/docs/source/package_reference/cli.mdx
index 3d72aada6..402d57781 100644
--- a/docs/source/package_reference/cli.mdx
+++ b/docs/source/package_reference/cli.mdx
@@ -125,6 +125,8 @@ accelerate launch [arguments] {training_script} --{training_script-argument-1} -
 * `-m`, `--module` (`bool`) -- Change each process to interpret the launch script as a Python module, executing with the same behavior as 'python -m'.
 * `--no_python` (`bool`) -- Skip prepending the training script with 'python' - just execute it directly. Useful when the script is not a Python script.
 * `--debug` (`bool`) -- Whether to print out the torch.distributed stack trace when something fails.
+* `-q`, `--quiet` (`bool`) -- Silence subprocess errors from the launch stack trace to only show the relevant tracebacks. (Only applicable to DeepSpeed and single-process configurations).
+
 
 The rest of these arguments are configured through `accelerate config` and are read in from the specified `--config_file` (or default configuration) for their 
 values. They can also be passed in manually.
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index d778177c6..081380331 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -137,6 +137,12 @@ def launch_command_parser(subparsers=None):
     parser.add_argument(
         "--config_file", default=None, help="The config file to use for the default values in the launching script."
     )
+    parser.add_argument(
+        "--quiet",
+        "-q",
+        action="store_true",
+        help="Silence subprocess errors from the launch stack trace and only show the relevant tracebacks. (Only applicable to DeepSpeed and single-process configurations)",
+    )
     # Hardware selection arguments
     hardware_args = parser.add_argument_group(
         "Hardware Selection Arguments", "Arguments for selecting the hardware to be used."
@@ -564,7 +570,10 @@ def simple_launcher(args):
     process = subprocess.Popen(cmd, env=current_env)
     process.wait()
     if process.returncode != 0:
-        raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)
+        if not args.quiet:
+            raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)
+        else:
+            sys.exit(1)
 
 
 def multi_gpu_launcher(args):
@@ -770,7 +779,10 @@ def deepspeed_launcher(args):
         process = subprocess.Popen(cmd, env=current_env)
         process.wait()
         if process.returncode != 0:
-            raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)
+            if not args.quiet:
+                raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)
+            else:
+                sys.exit(1)
     else:
         if is_torch_version("<", "1.9.0"):
             raise NotImplementedError("Multi-node training requires pytorch>=1.9.0")

```suggestion
        help="Silence subprocess errors from the launch stack trace and only show the relevant tracebacks. (Only applicable to DeepSpeed and single-process configurations)",
```
@@ -137,6 +137,12 @@ def launch_command_parser(subparsers=None):
     parser.add_argument(
         "--config_file", default=None, help="The config file to use for the default values in the launching script."
     )
+    parser.add_argument(
+        "--quiet",
+        "-q",
+        action="store_true",
+        help="Don't print an error message if an error return code is returned from launch.",
_The documentation is not available anymore as the PR was closed or merged._
I followed the instructions in https://github.com/huggingface/accelerate/pull/895#issuecomment-1333213726
```
black tests src examples benchmarks
reformatted src/accelerate/commands/launch.py

All done! ✨ 🍰 ✨
1 file reformatted, 104 files left unchanged.
isort tests src examples benchmarks
python utils/style_doc.py src/accelerate docs/source --max_len 119
black --check tests src examples benchmarks
All done! ✨ 🍰 ✨
105 files would be left unchanged.
isort --check-only tests src examples benchmarks
flake8 tests src examples benchmarks
python utils/style_doc.py src/accelerate docs/source --max_len 119 --check_only

```
I'll make those doc changes. Let me know if I should change it to `--silent`/`-s`.
TODO: update docs if it's changed to `--silent`
No strong opinion on the name on my side.
We'll keep it as `-q` then, thanks for your contribution @Cyberes! :) 
https://api.github.com/repos/huggingface/accelerate/pulls/902
diff --git a/docs/source/package_reference/utilities.mdx b/docs/source/package_reference/utilities.mdx
index 0c64953fc..3c1a1065e 100644
--- a/docs/source/package_reference/utilities.mdx
+++ b/docs/source/package_reference/utilities.mdx
@@ -93,3 +93,10 @@ These utilities relate to setting and synchronizing of all the random states.
 [[autodoc]] utils.synchronize_rng_state
 
 [[autodoc]] utils.synchronize_rng_states
+
+
+## PyTorch XLA
+
+These include utilities that are useful while using PyTorch with XLA.
+
+[[autodoc]] utils.install_xla
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
index f46bce267..be78d0089 100644
--- a/src/accelerate/utils/__init__.py
+++ b/src/accelerate/utils/__init__.py
@@ -128,4 +128,5 @@
     write_basic_config,
 )
 from .random import set_seed, synchronize_rng_state, synchronize_rng_states
+from .torch_xla import install_xla
 from .tqdm import tqdm
diff --git a/src/accelerate/utils/torch_xla.py b/src/accelerate/utils/torch_xla.py
new file mode 100644
index 000000000..59fe7970c
--- /dev/null
+++ b/src/accelerate/utils/torch_xla.py
@@ -0,0 +1,44 @@
+# 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 subprocess
+import sys
+
+import pkg_resources
+
+
+def install_xla(upgrade: bool = False):
+    """
+    Helper function to install appropriate xla wheels based on the `torch` version.
+
+    Args:
+        upgrade (`bool`, *optional*, defaults to `False`):
+            Whether to upgrade `torch` and install the latest `torch_xla` wheels.
+    """
+    in_colab = False
+    if "IPython" in sys.modules:
+        in_colab = "google.colab" in str(sys.modules["IPython"].get_ipython())
+
+    if in_colab:
+        if upgrade:
+            torch_install_cmd = ["pip", "install", "-U", "torch"]
+            subprocess.run(torch_install_cmd, check=True)
+        # get the current version of torch
+        torch_version = pkg_resources.get_distribution("torch").version
+        torch_version_trunc = torch_version[: torch_version.rindex(".")]
+        xla_wheel = f"https://storage.googleapis.com/tpu-pytorch/wheels/colab/torch_xla-{torch_version_trunc}-cp37-cp37m-linux_x86_64.whl"
+        xla_install_cmd = ["pip", "install", xla_wheel]
+        subprocess.run(xla_install_cmd, check=True)
+    else:
+        raise RuntimeError("`install_xla` utility works only on google colab.")

Is there a reason to limit this? Everything should work fine in any other env no? Or is it that the wheels can only work on Colab?
@@ -0,0 +1,43 @@
+# 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 sys
+import subprocess
+import pkg_resources
+
+
+def install_xla(upgrade: bool = False):
+    """
+    Helper function to install appropriate xla wheels based on the `torch` version.
+
+    Args:
+        upgrade (`bool`, *optional*, defaults to `False`):
+            Whether to upgrade `torch` and install the latest `torch_xla` wheels.
+    """
+    in_colab = False
+    if "IPython" in sys.modules:
+        in_colab = "google.colab" in str(sys.modules["IPython"].get_ipython())
+
+    if in_colab:
The FR was for colab specifically because they have their own wheels for it https://github.com/huggingface/accelerate/issues/586 (otherwise things get more complicated such as trying to get the right python version, etc. Things that can be added later but for rn just an easy way in colab to skip a line :) )
@@ -0,0 +1,43 @@
+# 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 sys
+import subprocess
+import pkg_resources
+
+
+def install_xla(upgrade: bool = False):
+    """
+    Helper function to install appropriate xla wheels based on the `torch` version.
+
+    Args:
+        upgrade (`bool`, *optional*, defaults to `False`):
+            Whether to upgrade `torch` and install the latest `torch_xla` wheels.
+    """
+    in_colab = False
+    if "IPython" in sys.modules:
+        in_colab = "google.colab" in str(sys.modules["IPython"].get_ipython())
+
+    if in_colab:
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/896
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 4e23c18af..1d6f7ee10 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -237,7 +237,7 @@ def __init__(
         else:
             assert isinstance(
                 deepspeed_plugin, DeepSpeedPlugin
-            ), "`deepspeed_plugin` must be a DeepSpeedPlugin object."
+            ), "`deepspeed_plugin` must be an `accelerate.utils.DeepSpeedPlugin` object."
             os.environ["ACCELERATE_USE_DEEPSPEED"] = "true"  # use DeepSpeed if plugin is provided
         if deepspeed_plugin:
             if not is_deepspeed_available():
@@ -285,7 +285,9 @@ def __init__(
         self.init_handler = None
         if kwargs_handlers is not None:
             for handler in kwargs_handlers:
-                assert isinstance(handler, KwargsHandler), f"Unsupported kwargs handler passed: {handler}."
+                assert isinstance(
+                    handler, KwargsHandler
+                ), f"Unsupported kwargs handler passed: {handler}, must be one that inherits `accelerate.utils.KwargsHandler`."
                 if isinstance(handler, DistributedDataParallelKwargs):
                     if self.ddp_handler is not None:
                         raise ValueError("You can only pass one `DistributedDataParallelKwargs` in `kwargs_handler`.")

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/894
diff --git a/src/accelerate/commands/menu/keymap.py b/src/accelerate/commands/menu/keymap.py
index 7ce6c0637..3f08236b3 100644
--- a/src/accelerate/commands/menu/keymap.py
+++ b/src/accelerate/commands/menu/keymap.py
@@ -17,10 +17,9 @@
 """
 
 
+import os
 import string
 import sys
-import termios
-import tty
 
 
 ARROW_KEY_FLAG = 1 << 8
@@ -36,24 +35,77 @@
     "mod_int": 91,
     "undefined": sys.maxsize,
     "interrupt": 3,
+    "insert": 50,
+    "delete": 51,
+    "pg_up": 53,
+    "pg_down": 54,
 }
 
 KEYMAP["arrow_begin"] = KEYMAP["up"]
 KEYMAP["arrow_end"] = KEYMAP["left"]
 
+if sys.platform == "win32":
+    WIN_CH_BUFFER = []
+    WIN_KEYMAP = {
+        b"\xe0H": KEYMAP["up"] - ARROW_KEY_FLAG,
+        b"\x00H": KEYMAP["up"] - ARROW_KEY_FLAG,
+        b"\xe0P": KEYMAP["down"] - ARROW_KEY_FLAG,
+        b"\x00P": KEYMAP["down"] - ARROW_KEY_FLAG,
+        b"\xe0M": KEYMAP["right"] - ARROW_KEY_FLAG,
+        b"\x00M": KEYMAP["right"] - ARROW_KEY_FLAG,
+        b"\xe0K": KEYMAP["left"] - ARROW_KEY_FLAG,
+        b"\x00K": KEYMAP["left"] - ARROW_KEY_FLAG,
+    }
+
 for i in range(10):
     KEYMAP[str(i)] = ord(str(i))
 
 
 def get_raw_chars():
     "Gets raw characters from inputs"
-    fd = sys.stdin.fileno()
-    old_settings = termios.tcgetattr(fd)
-    try:
-        tty.setraw(fd)
-        ch = sys.stdin.read(1)
-    finally:
-        termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
+    if os.name == "nt":
+        import msvcrt
+
+        encoding = "mbcs"
+        # Flush the keyboard buffer
+        while msvcrt.kbhit():
+            msvcrt.getwch()
+        if len(WIN_CH_BUFFER) == 0:
+            # Read the keystroke
+            ch = msvcrt.getwch()
+            # If it is a prefix char, get second part
+            if ch.encode(encoding) in (b"\x00", b"\xe0"):
+                ch2 = ch + msvcrt.getwch()
+                # Translate actual Win chars to bullet char types
+                try:
+                    chx = chr(WIN_KEYMAP[ch2.encode(encoding)])
+                    WIN_CH_BUFFER.append(chr(KEYMAP["mod_int"]))
+                    WIN_CH_BUFFER.append(chx)
+                    if ord(chx) in (
+                        KEYMAP["insert"] - 1 << 9,
+                        KEYMAP["delete"] - 1 << 9,
+                        KEYMAP["pg_up"] - 1 << 9,
+                        KEYMAP["pg_down"] - 1 << 9,
+                    ):
+                        WIN_CH_BUFFER.append(chr(126))
+                    ch = chr(KEYMAP["esc"])
+                except KeyError:
+                    ch = ch2[1]
+            else:
+                pass
+        else:
+            ch = WIN_CH_BUFFER.pop(0)
+    elif os.name == "posix":
+        import termios
+        import tty
+
+        fd = sys.stdin.fileno()
+        old_settings = termios.tcgetattr(fd)
+        try:
+            tty.setraw(fd)
+            ch = sys.stdin.read(1)
+        finally:
+            termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
     return ch
 
 
diff --git a/src/accelerate/commands/menu/selection_menu.py b/src/accelerate/commands/menu/selection_menu.py
index 751f6818a..d5a1c8e35 100644
--- a/src/accelerate/commands/menu/selection_menu.py
+++ b/src/accelerate/commands/menu/selection_menu.py
@@ -15,6 +15,8 @@
 """
 Main driver for the selection menu, based on https://github.com/bchao1/bullet
 """
+import sys
+
 from . import cursor, input
 from .helpers import Direction, clear_line, forceWrite, linebreak, move_cursor, reset_cursor, writeColor
 from .keymap import KEYMAP
@@ -30,12 +32,22 @@ def __init__(self, prompt: str = None, choices: list = []):
         self.position = 0
         self.choices = choices
         self.prompt = prompt
+        if sys.platform == "win32":
+            self.arrow_char = "*"
+        else:
+            self.arrow_char = "➔ "
+
+    def write_choice(self, index, end: str = ""):
+        if sys.platform != "win32":
+            writeColor(self.choices[index], 32, end)
+        else:
+            forceWrite(self.choices[index], end)
 
     def print_choice(self, index: int):
         "Prints the choice at the given index"
         if index == self.position:
-            forceWrite(" ➔  ")
-            writeColor(self.choices[index], 32)
+            forceWrite(f" {self.arrow_char} ")
+            self.write_choice(index)
         else:
             forceWrite(f"    {self.choices[index]}")
         reset_cursor()
@@ -109,6 +121,5 @@ def run(self, default_choice: int = 0):
                     for _ in range(len(self.choices) + 1):
                         move_cursor(1, "UP")
                         clear_line()
-                    forceWrite(" ➔  ")
-                    writeColor(self.choices[choice], 32, "\n")
+                    self.write_choice(choice, "\n")
                     return choice

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/893
diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
index e471731a4..ff3e95ead 100644
--- a/src/accelerate/utils/modeling.py
+++ b/src/accelerate/utils/modeling.py
@@ -21,11 +21,10 @@
 from collections import defaultdict
 from typing import Dict, List, Optional, Tuple, Union
 
-import numpy as np
 import torch
 import torch.nn as nn
 
-from .offload import offload_weight, save_offload_index
+from .offload import load_offloaded_weight, offload_weight, save_offload_index
 
 
 WEIGHTS_INDEX_NAME = "pytorch_model.bin.index.json"
@@ -324,9 +323,8 @@ def load_offloaded_weights(model, index, offload_folder):
 
     for param_name, metadata in index.items():
         tensor_file = os.path.join(offload_folder, f"{param_name}.dat")
-        shape = tuple(metadata["shape"])
-        weight = np.memmap(tensor_file, dtype=metadata["dtype"], mode="r", shape=shape)
-        set_module_tensor_to_device(model, param_name, "cpu", value=torch.tensor(weight))
+        weight = load_offloaded_weight(tensor_file, metadata)
+        set_module_tensor_to_device(model, param_name, "cpu", value=weight)
 
 
 def get_balanced_memory(

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/892
diff --git a/tests/deepspeed/test_deepspeed.py b/tests/deepspeed/test_deepspeed.py
index 86c9766a1..eca75c1ed 100644
--- a/tests/deepspeed/test_deepspeed.py
+++ b/tests/deepspeed/test_deepspeed.py
@@ -119,7 +119,7 @@ def setUp(self):
         self.ds_config_dict = dict(zero2=config_zero2, zero3=config_zero3)
 
         self.dist_env = dict(
-            USE_DEEPSPEED="true",
+            ACCELERATE_USE_DEEPSPEED="true",
             MASTER_ADDR="localhost",
             MASTER_PORT="10999",
             RANK="0",
diff --git a/tests/fsdp/test_fsdp.py b/tests/fsdp/test_fsdp.py
index e90106c17..19917c671 100644
--- a/tests/fsdp/test_fsdp.py
+++ b/tests/fsdp/test_fsdp.py
@@ -58,7 +58,7 @@ def setUp(self):
         super().setUp()
 
         self.dist_env = dict(
-            USE_FSDP="true",
+            ACCELERATE_USE_FSDP="true",
             MASTER_ADDR="localhost",
             MASTER_PORT="10999",
             RANK="0",
@@ -147,7 +147,7 @@ def test_mixed_precision(self):
 
         for mp_dtype in dtypes:
             env = self.dist_env.copy()
-            env["MIXED_PRECISION"] = mp_dtype
+            env["ACCELERATE_MIXED_PRECISION"] = mp_dtype
             with mockenv_context(**env):
                 accelerator = Accelerator()
                 if mp_dtype == "fp16":

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/891
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 0d0fe0aa6..fabf4991e 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -119,9 +119,9 @@ class Accelerator:
             in your script multiplied by the number of processes.
         mixed_precision (`str`, *optional*):
             Whether or not to use mixed precision training (fp16 or bfloat16). Choose from 'no','fp16','bf16'. Will
-            default to the value in the environment variable `MIXED_PRECISION`, which will use the default value in the
-            accelerate config of the current system or the flag passed with the `accelerate.launch` command. 'fp16'
-            requires pytorch 1.6 or higher. 'bf16' requires pytorch 1.10 or higher.
+            default to the value in the environment variable `ACCELERATE_MIXED_PRECISION`, which will use the default
+            value in the accelerate config of the current system or the flag passed with the `accelerate.launch`
+            command. 'fp16' requires pytorch 1.6 or higher. 'bf16' requires pytorch 1.10 or higher.
         gradient_accumulation_steps (`int`, *optional*, default to 1):
             The number of steps that should pass before gradients are accumulated. A number > 1 should be combined with
             `Accelerator.accumulate`.
@@ -231,39 +231,49 @@ def __init__(
             dynamo_backend = DynamoBackend(dynamo_backend.upper())
 
         if deepspeed_plugin is None:  # init from env variables
-            deepspeed_plugin = DeepSpeedPlugin() if os.environ.get("USE_DEEPSPEED", "false") == "true" else None
+            deepspeed_plugin = (
+                DeepSpeedPlugin() if os.environ.get("ACCELERATE_USE_DEEPSPEED", "false") == "true" else None
+            )
         else:
             assert isinstance(
                 deepspeed_plugin, DeepSpeedPlugin
             ), "`deepspeed_plugin` must be a DeepSpeedPlugin object."
-            os.environ["USE_DEEPSPEED"] = "true"  # use DeepSpeed if plugin is provided
+            os.environ["ACCELERATE_USE_DEEPSPEED"] = "true"  # use DeepSpeed if plugin is provided
         if deepspeed_plugin:
             if not is_deepspeed_available():
                 raise ImportError("DeepSpeed is not installed => run `pip install deepspeed` or build it from source.")
             if compare_versions("deepspeed", "<", "0.6.5"):
                 raise ImportError("DeepSpeed version must be >= 0.6.5. Please update DeepSpeed.")
 
-            mixed_precision = os.environ.get("MIXED_PRECISION", "no") if mixed_precision is None else mixed_precision
+            mixed_precision = (
+                os.environ.get("ACCELERATE_MIXED_PRECISION", "no") if mixed_precision is None else mixed_precision
+            )
             deepspeed_plugin.set_mixed_precision(mixed_precision)
             deepspeed_plugin.set_deepspeed_weakref()
 
-        if os.environ.get("USE_FSDP", "false") == "true" or isinstance(fsdp_plugin, FullyShardedDataParallelPlugin):
+        if os.environ.get("ACCELERATE_USE_FSDP", "false") == "true" or isinstance(
+            fsdp_plugin, FullyShardedDataParallelPlugin
+        ):
             if is_torch_version("<", "1.12.0"):
                 raise ValueError("FSDP requires PyTorch >= 1.12.0")
 
         if fsdp_plugin is None:  # init from env variables
-            fsdp_plugin = FullyShardedDataParallelPlugin() if os.environ.get("USE_FSDP", "false") == "true" else None
+            fsdp_plugin = (
+                FullyShardedDataParallelPlugin() if os.environ.get("ACCELERATE_USE_FSDP", "false") == "true" else None
+            )
         else:
             if not isinstance(fsdp_plugin, FullyShardedDataParallelPlugin):
                 raise TypeError("`fsdp_plugin` must be a FullyShardedDataParallelPlugin object.")
-            os.environ["USE_FSDP"] = "true"  # use FSDP if plugin is provided
+            os.environ["ACCELERATE_USE_FSDP"] = "true"  # use FSDP if plugin is provided
 
         if megatron_lm_plugin is None:  # init from env variables
-            megatron_lm_plugin = MegatronLMPlugin() if os.environ.get("USE_MEGATRON_LM", "false") == "true" else None
+            megatron_lm_plugin = (
+                MegatronLMPlugin() if os.environ.get("ACCELERATE_USE_MEGATRON_LM", "false") == "true" else None
+            )
         else:
             if not isinstance(megatron_lm_plugin, MegatronLMPlugin):
                 raise TypeError("`megatron_lm_plugin` must be a MegatronLMPlugin object.")
-            os.environ["USE_MEGATRON_LM"] = "true"  # use MegatronLM if plugin is provided
+            os.environ["ACCELERATE_USE_MEGATRON_LM"] = "true"  # use MegatronLM if plugin is provided
 
         if megatron_lm_plugin:
             if not is_megatron_lm_available():
@@ -339,7 +349,7 @@ def __init__(
         err = "{mode} mixed precision requires {requirement}"
         if self.state.mixed_precision == "fp16" and self.distributed_type != DistributedType.MEGATRON_LM:
             self.native_amp = True
-            if not torch.cuda.is_available() and not parse_flag_from_env("USE_MPS_DEVICE"):
+            if not torch.cuda.is_available() and not parse_flag_from_env("ACCELERATE_USE_MPS_DEVICE"):
                 raise ValueError(err.format(mode="fp16", requirement="a GPU"))
             kwargs = self.scaler_handler.to_kwargs() if self.scaler_handler is not None else {}
             if self.distributed_type == DistributedType.FSDP:
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 9e37bf0db..d778177c6 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -518,14 +518,14 @@ def simple_launcher(args):
     cmd.extend(args.training_script_args)
 
     current_env = os.environ.copy()
-    current_env["USE_CPU"] = str(args.cpu or args.use_cpu)
+    current_env["ACCELERATE_USE_CPU"] = str(args.cpu or args.use_cpu)
     if args.use_mps_device:
         warnings.warn(
             '`use_mps_device` flag is deprecated and will be removed in version 0.15.0 of 🤗 Accelerate. Use "--mps" instead.',
             FutureWarning,
         )
         args.mps = True
-    current_env["USE_MPS_DEVICE"] = str(args.mps)
+    current_env["ACCELERATE_USE_MPS_DEVICE"] = str(args.mps)
     if args.mps:
         current_env["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
     elif args.gpu_ids != "all" and args.gpu_ids is not None:
@@ -551,13 +551,13 @@ def simple_launcher(args):
         )
         mixed_precision = "fp16"
 
-    current_env["MIXED_PRECISION"] = str(mixed_precision)
+    current_env["ACCELERATE_MIXED_PRECISION"] = str(mixed_precision)
 
     try:
         dynamo_backend = DynamoBackend(args.dynamo_backend.upper())
     except ValueError:
         raise ValueError(f"Unknown dynamo backend: {args.dynamo_backend.upper()}. Choose between {DYNAMO_BACKENDS}.")
-    current_env["DYNAMO_BACKEND"] = dynamo_backend.value
+    current_env["ACCELERATE_DYNAMO_BACKEND"] = dynamo_backend.value
 
     current_env["OMP_NUM_THREADS"] = str(args.num_cpu_threads_per_process)
 
@@ -612,16 +612,16 @@ def multi_gpu_launcher(args):
         )
         mixed_precision = "fp16"
 
-    current_env["MIXED_PRECISION"] = str(mixed_precision)
+    current_env["ACCELERATE_MIXED_PRECISION"] = str(mixed_precision)
 
     try:
         dynamo_backend = DynamoBackend(args.dynamo_backend.upper())
     except ValueError:
         raise ValueError(f"Unknown dynamo backend: {args.dynamo_backend.upper()}. Choose between {DYNAMO_BACKENDS}.")
-    current_env["DYNAMO_BACKEND"] = dynamo_backend.value
+    current_env["ACCELERATE_DYNAMO_BACKEND"] = dynamo_backend.value
 
     if args.use_fsdp:
-        current_env["USE_FSDP"] = "true"
+        current_env["ACCELERATE_USE_FSDP"] = "true"
         current_env["FSDP_SHARDING_STRATEGY"] = str(args.fsdp_sharding_strategy)
         current_env["FSDP_OFFLOAD_PARAMS"] = str(args.fsdp_offload_params).lower()
         current_env["FSDP_MIN_NUM_PARAMS"] = str(args.fsdp_min_num_params)
@@ -636,7 +636,7 @@ def multi_gpu_launcher(args):
 
     if args.use_megatron_lm:
         prefix = "MEGATRON_LM_"
-        current_env["USE_MEGATRON_LM"] = "true"
+        current_env["ACCELERATE_USE_MEGATRON_LM"] = "true"
         current_env[prefix + "TP_DEGREE"] = str(args.megatron_lm_tp_degree)
         current_env[prefix + "PP_DEGREE"] = str(args.megatron_lm_pp_degree)
         current_env[prefix + "GRADIENT_CLIPPING"] = str(args.megatron_lm_gradient_clipping)
@@ -748,8 +748,8 @@ def deepspeed_launcher(args):
         mixed_precision = "fp16"
 
     current_env["PYTHONPATH"] = env_var_path_add("PYTHONPATH", os.path.abspath("."))
-    current_env["MIXED_PRECISION"] = str(mixed_precision)
-    current_env["USE_DEEPSPEED"] = "true"
+    current_env["ACCELERATE_MIXED_PRECISION"] = str(mixed_precision)
+    current_env["ACCELERATE_USE_DEEPSPEED"] = "true"
     current_env["DEEPSPEED_ZERO_STAGE"] = str(args.zero_stage)
     current_env["GRADIENT_ACCUMULATION_STEPS"] = str(args.gradient_accumulation_steps)
     current_env["GRADIENT_CLIPPING"] = str(args.gradient_clipping).lower()
@@ -924,10 +924,10 @@ def sagemaker_launcher(sagemaker_config: SageMakerConfig, args):
 
     # Environment variables to be set for use during training job
     environment = {
-        "USE_SAGEMAKER": "true",
-        "MIXED_PRECISION": str(mixed_precision),
-        "DYNAMO_BACKEND": dynamo_backend.value,
-        "SAGEMAKER_DISTRIBUTED_TYPE": sagemaker_config.distributed_type.value,
+        "ACCELERATE_USE_SAGEMAKER": "true",
+        "ACCELERATE_MIXED_PRECISION": str(mixed_precision),
+        "ACCELERATE_DYNAMO_BACKEND": dynamo_backend.value,
+        "ACCELERATE_SAGEMAKER_DISTRIBUTED_TYPE": sagemaker_config.distributed_type.value,
     }
     # configure distribution set up
     distribution = None
diff --git a/src/accelerate/launchers.py b/src/accelerate/launchers.py
index 4de8b5858..5768ec06b 100644
--- a/src/accelerate/launchers.py
+++ b/src/accelerate/launchers.py
@@ -163,9 +163,9 @@ def debug_launcher(function, args=(), num_processes=2):
             world_size=num_processes,
             master_addr="127.0.01",
             master_port="29500",
-            mixed_precision="no",
+            accelerate_mixed_precision="no",
             accelerate_debug_rdv_file=tmp_file.name,
-            use_cpu="yes",
+            accelerate_use_cpu="yes",
         ):
             launcher = PrepareForLaunch(function, debug=True)
             start_processes(launcher, args=args, nprocs=num_processes, start_method="fork")
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index 18df61bbc..a8a5e7960 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -66,7 +66,7 @@ def __init__(
         **kwargs,
     ):
         self.__dict__ = self._shared_state
-        if parse_flag_from_env("USE_CPU"):
+        if parse_flag_from_env("ACCELERATE_USE_CPU"):
             cpu = True
         self._check_initialized(mixed_precision, cpu)
         self.fork_launched = parse_flag_from_env("FORK_LAUNCHED", 0)
@@ -74,10 +74,12 @@ def __init__(
             self.backend = None
             self.deepspeed_plugin = None
             mixed_precision = (
-                parse_choice_from_env("MIXED_PRECISION", "no") if mixed_precision is None else mixed_precision.lower()
+                parse_choice_from_env("ACCELERATE_MIXED_PRECISION", "no")
+                if mixed_precision is None
+                else mixed_precision.lower()
             )
             dynamo_backend = (
-                parse_choice_from_env("DYNAMO_BACKEND", "no") if dynamo_backend is None else dynamo_backend
+                parse_choice_from_env("ACCELERATE_DYNAMO_BACKEND", "no") if dynamo_backend is None else dynamo_backend
             )
             self.dynamo_backend = DynamoBackend(dynamo_backend.upper())
             if not _from_accelerator:
@@ -86,11 +88,11 @@ def __init__(
                     "before using any functionality from the `accelerate` library."
                 )
             if (
-                os.environ.get("USE_SAGEMAKER", "false") == "true"
-                and os.environ.get("SAGEMAKER_DISTRIBUTED_TYPE") != SageMakerDistributedType.NO
+                os.environ.get("ACCELERATE_USE_SAGEMAKER", "false") == "true"
+                and os.environ.get("ACCELERATE_SAGEMAKER_DISTRIBUTED_TYPE") != SageMakerDistributedType.NO
                 and not cpu
             ):
-                if os.environ.get("SAGEMAKER_DISTRIBUTED_TYPE") == SageMakerDistributedType.DATA_PARALLEL:
+                if os.environ.get("ACCELERATE_SAGEMAKER_DISTRIBUTED_TYPE") == SageMakerDistributedType.DATA_PARALLEL:
                     self.distributed_type = DistributedType.MULTI_GPU
                     import smdistributed.dataparallel.torch.torch_smddp  # noqa
 
@@ -110,7 +112,7 @@ def __init__(
                 self.local_process_index = xm.get_local_ordinal()
                 self.device = xm.xla_device()
                 if mixed_precision == "bf16":
-                    if os.environ.get("DOWNCAST_BF16"):
+                    if os.environ.get("ACCELERATE_DOWNCAST_BF16"):
                         os.environ["XLA_USE_BF16"] = str(0)
                         os.environ["XLA_DOWNCAST_BF16"] = str(1)
                         self.downcast_bfloat = True
@@ -119,7 +121,7 @@ def __init__(
                         os.environ["XLA_DOWNCAST_BF16"] = str(0)
                         self.downcast_bfloat = False
                 self.mixed_precision = mixed_precision
-            elif os.environ.get("USE_DEEPSPEED", "false") == "true" and not cpu:
+            elif os.environ.get("ACCELERATE_USE_DEEPSPEED", "false") == "true" and not cpu:
                 assert (
                     is_deepspeed_available()
                 ), "DeepSpeed is not available => install it using `pip3 install deepspeed` or build it from source"
@@ -153,12 +155,12 @@ def __init__(
                 self.device = torch.device("cuda", self.local_process_index)
                 torch.cuda.set_device(self.device)
                 self.mixed_precision = mixed_precision
-                if os.environ.get("USE_FSDP", "false") == "true":
+                if os.environ.get("ACCELERATE_USE_FSDP", "false") == "true":
                     self.distributed_type = DistributedType.FSDP
                     if self.mixed_precision != "no":
                         fsdp_plugin.set_mixed_precision(self.mixed_precision)
                     self.fsdp_plugin = fsdp_plugin
-                if os.environ.get("USE_MEGATRON_LM", "false") == "true":
+                if os.environ.get("ACCELERATE_USE_MEGATRON_LM", "false") == "true":
                     self.distributed_type = DistributedType.MEGATRON_LM
                     megatron_lm_plugin.set_mixed_precision(self.mixed_precision)
                     self.megatron_lm_plugin = megatron_lm_plugin
@@ -207,7 +209,7 @@ def __init__(
                 self.distributed_type = DistributedType.NO
                 self.num_processes = 1
                 self.process_index = self.local_process_index = 0
-                if parse_flag_from_env("USE_MPS_DEVICE") and not cpu:
+                if parse_flag_from_env("ACCELERATE_USE_MPS_DEVICE") and not cpu:
                     if not torch.backends.mps.is_available():
                         if not torch.backends.mps.is_built():
                             raise AssertionError(
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index bfedee162..4ac250e6e 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -15,6 +15,7 @@
 import importlib
 import os
 import sys
+import warnings
 from distutils.util import strtobool
 from functools import lru_cache
 
@@ -93,7 +94,7 @@ def is_bf16_available(ignore_tpu=False):
 
 
 def is_megatron_lm_available():
-    if strtobool(os.environ.get("USE_MEGATRON_LM", "False")) == 1:
+    if strtobool(os.environ.get("ACCELERATE_USE_MEGATRON_LM", "False")) == 1:
         package_exists = importlib.util.find_spec("megatron") is not None
         if package_exists:
             megatron_version = parse(importlib_metadata.version("megatron-lm"))
@@ -130,7 +131,15 @@ def is_boto3_available():
 
 
 def is_rich_available():
-    return (importlib.util.find_spec("rich") is not None) and (not parse_flag_from_env("DISABLE_RICH"))
+    if importlib.util.find_spec("rich") is not None:
+        if parse_flag_from_env("DISABLE_RICH"):
+            warnings.warn(
+                "The `DISABLE_RICH` flag is deprecated and will be removed in version 0.17.0 of 🤗 Accelerate. Use `ACCELERATE_DISABLE_RICH` instead.",
+                FutureWarning,
+            )
+            return not parse_flag_from_env("DISABLE_RICH")
+        return not parse_flag_from_env("ACCELERATE_DISABLE_RICH")
+    return False
 
 
 def is_sagemaker_available():

This feels over-engineered, and outside the scope of this PR.
@@ -0,0 +1,68 @@
+# 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 warnings
+
+
+def deprecate_value(old=None, new=None, version=None, warn=True, deprecate_str=None):
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/890
diff --git a/examples/by_feature/deepspeed_with_config_support.py b/examples/by_feature/deepspeed_with_config_support.py
index 0fdff58ae..6cfee9dda 100755
--- a/examples/by_feature/deepspeed_with_config_support.py
+++ b/examples/by_feature/deepspeed_with_config_support.py
@@ -642,7 +642,7 @@ def group_texts(examples):
                 total_loss += loss.detach().float()
             loss = loss / args.gradient_accumulation_steps
             accelerator.backward(loss)
-            if step % args.gradient_accumulation_steps == 0 or step == len(train_dataloader) - 1:
+            if (step + 1) % args.gradient_accumulation_steps == 0 or step == len(train_dataloader) - 1:
                 optimizer.step()
                 lr_scheduler.step()
                 optimizer.zero_grad()

_The documentation is not available anymore as the PR was closed or merged._
Thanks! 🤗
https://api.github.com/repos/huggingface/accelerate/pulls/887
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 0d0fe0aa6..0541ce506 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -795,7 +795,7 @@ def prepare(self, *args, device_placement=None):
                 - `torch.utils.data.DataLoader`: PyTorch Dataloader
                 - `torch.nn.Module`: PyTorch Module
                 - `torch.optim.Optimizer`: PyTorch Optimizer
-                - `torch.optim.lr_scheduler._LRScheduler`: PyTorch LR Scheduler
+                - `torch.optim.lr_scheduler.LRScheduler`: PyTorch LR Scheduler
 
             device_placement (`List[bool]`, *optional*):
                 Used to customize whether automatic device placement should be performed for each object passed. Needs
@@ -998,7 +998,7 @@ def _prepare_deepspeed(self, *args):
                 model = obj
             elif isinstance(obj, (torch.optim.Optimizer, DummyOptim)):
                 optimizer = obj
-            elif (isinstance(obj, (torch.optim.lr_scheduler._LRScheduler, DummyScheduler))) or (
+            elif (isinstance(obj, (LRScheduler, DummyScheduler))) or (
                 type(obj).__name__ in deepspeed.runtime.lr_schedules.VALID_LR_SCHEDULES
             ):
                 scheduler = obj
@@ -1097,7 +1097,7 @@ def _prepare_deepspeed(self, *args):
                     result[i] = engine
                 elif isinstance(result[i], (torch.optim.Optimizer, DummyOptim)):
                     result[i] = optimizer
-                elif (isinstance(result[i], (torch.optim.lr_scheduler._LRScheduler, DummyScheduler))) or (
+                elif (isinstance(result[i], (LRScheduler, DummyScheduler))) or (
                     type(result[i]).__name__ in deepspeed.runtime.lr_schedules.VALID_LR_SCHEDULES
                 ):
                     result[i] = scheduler
@@ -1150,7 +1150,7 @@ def _prepare_megatron_lm(self, *args):
                 model = obj
             elif isinstance(obj, (torch.optim.Optimizer)):
                 optimizer = obj
-            elif isinstance(obj, (torch.optim.lr_scheduler._LRScheduler, MegatronLMDummyScheduler)):
+            elif isinstance(obj, (LRScheduler, MegatronLMDummyScheduler)):
                 scheduler = obj
 
         if model is not None:
@@ -1260,13 +1260,13 @@ def prepare_optimizer(self, optimizer: torch.optim.Optimizer, device_placement=N
         self._optimizers.append(optimizer)
         return optimizer
 
-    def prepare_scheduler(self, scheduler: torch.optim.lr_scheduler._LRScheduler):
+    def prepare_scheduler(self, scheduler: LRScheduler):
         """
         Prepares a PyTorch Scheduler for training in any distributed setup. It is recommended to use
         [`Accelerator.prepare`] instead.
 
         Args:
-            scheduler (`torch.optim.lr_scheduler._LRScheduler`):
+            scheduler (`torch.optim.lr_scheduler.LRScheduler`):
                 A vanilla PyTorch scheduler to prepare
         """
         # We try to find the optimizer associated with `scheduler`, the default is the full list.

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/884
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index a29794fc4..0d0fe0aa6 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -751,7 +751,7 @@ def _prepare_fsdp(self, *args):
         for obj in args:
             if isinstance(obj, torch.optim.Optimizer):
                 if len(obj.param_groups) > 1:
-                    logger.warn(
+                    logger.warning(
                         "FSDP Warning: When using FSDP, several parameter groups will be conflated into "
                         "a single one due to nested module wrapping and parameter flattening."
                     )
@@ -831,7 +831,7 @@ def prepare(self, *args, device_placement=None):
                     "Then pass the optimizers to the prepare call in the same order as corresponding models."
                 )
             elif model_count == 1 and optimizer_present:
-                logger.warn(
+                logger.warning(
                     "FSDP Warning: When using FSDP, "
                     "it is efficient and recommended to call prepare for the model before creating the optimizer"
                 )
@@ -1713,7 +1713,7 @@ def load_state(self, input_dir: str):
             err = "Warning! Number of found checkpoints does not match the number of registered objects:"
             err += f"\n\tFound checkpoints: {len(custom_checkpoints)}"
             err += f"\n\tRegistered objects: {len(self._custom_objects)}\nSkipping."
-            logger.warn(err)
+            logger.warning(err)
         else:
             logger.info(f"Loading in {len(custom_checkpoints)} custom states")
             for index, obj in enumerate(self._custom_objects):
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index b2d95bad8..9e37bf0db 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -1085,7 +1085,7 @@ def launch_command(args):
         message += (
             "\nTo avoid this warning pass in values for each of the problematic parameters or run `accelerate config`."
         )
-        logger.warn(message)
+        logger.warning(message)
 
     # Use the proper launcher
     if args.use_deepspeed and not args.cpu:

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/881
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index d6a323559..be366c6e3 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -95,6 +95,12 @@
 if is_tpu_available(check_device=False):
     import torch_xla.distributed.xla_multiprocessing as xmp
 
+
+if is_torch_version("<=", "1.13.5"):
+    from torch.optim.lr_scheduler import _LRScheduler as LRScheduler
+else:
+    from torch.optim.lr_scheduler import LRScheduler as LRScheduler
+
 logger = get_logger(__name__)
 
 
@@ -725,7 +731,7 @@ def _prepare_one(self, obj, first_pass=False, device_placement=None):
                 optimizer = self.prepare_optimizer(obj, device_placement=device_placement)
                 return optimizer
         # Second pass of preparation: LR scheduler (which need the full list of optimizers)
-        elif isinstance(obj, torch.optim.lr_scheduler._LRScheduler):
+        elif isinstance(obj, LRScheduler):
             scheduler = self.prepare_scheduler(obj)
             return scheduler
         # Return the unprocessed object if previous criteria was not met

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/878
diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py
index bb14078d4..442b63cf5 100644
--- a/src/accelerate/utils/dataclasses.py
+++ b/src/accelerate/utils/dataclasses.py
@@ -31,6 +31,7 @@
 import torch
 
 from .constants import FSDP_AUTO_WRAP_POLICY, FSDP_BACKWARD_PREFETCH, FSDP_STATE_DICT_TYPE, MODEL_NAME, OPTIMIZER_NAME
+from .versions import is_torch_version
 
 
 class KwargsHandler:
@@ -526,9 +527,9 @@ def __post_init__(self):
         from torch.distributed.fsdp.fully_sharded_data_parallel import (
             BackwardPrefetch,
             CPUOffload,
+            FullStateDictConfig,
             ShardingStrategy,
             StateDictType,
-            _state_dict_type_to_config,
         )
 
         if self.sharding_strategy is None:
@@ -549,12 +550,8 @@ def __post_init__(self):
             state_dict_type_policy = os.environ.get("FSDP_STATE_DICT_TYPE", "FULL_STATE_DICT")
             self.state_dict_type = StateDictType(FSDP_STATE_DICT_TYPE.index(state_dict_type_policy) + 1)
 
-            if self.state_dict_type == StateDictType.FULL_STATE_DICT:
-                self.state_dict_config = _state_dict_type_to_config[self.state_dict_type](
-                    offload_to_cpu=True, rank0_only=True
-                )
-            else:
-                self.state_dict_config = _state_dict_type_to_config[self.state_dict_type]()
+            if self.state_dict_type == StateDictType.FULL_STATE_DICT and self.state_dict_config is None:
+                self.state_dict_config = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
 
     @staticmethod
     def get_module_class_from_name(module, name):
@@ -616,9 +613,14 @@ def save_model(self, accelerator, model, output_dir, model_index=0):
         from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP
         from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType
 
-        if self.state_dict_type == StateDictType.FULL_STATE_DICT:
+        if is_torch_version("<=", "1.13.5"):
             with FSDP.state_dict_type(model, self.state_dict_type, self.state_dict_config):
                 state_dict = model.state_dict()
+        else:
+            FSDP.set_state_dict_type(model, self.state_dict_type, self.state_dict_config)
+            state_dict = model.state_dict()
+
+        if self.state_dict_type == StateDictType.FULL_STATE_DICT:
             weights_name = f"{MODEL_NAME}.bin" if model_index == 0 else f"{MODEL_NAME}_{model_index}.bin"
             output_model_file = os.path.join(output_dir, weights_name)
             if accelerator.process_index == 0:
@@ -626,8 +628,6 @@ def save_model(self, accelerator, model, output_dir, model_index=0):
                 torch.save(state_dict, output_model_file)
                 print(f"Model saved to {output_model_file}")
         else:
-            with FSDP.state_dict_type(model, self.state_dict_type, self.state_dict_config):
-                state_dict = model.state_dict()
             weights_name = (
                 f"{MODEL_NAME}_rank{accelerator.process_index}.bin"
                 if model_index == 0
@@ -660,7 +660,12 @@ def load_model(self, accelerator, model, input_dir, model_index=0):
             print(f"Loading model from {input_model_file}")
             state_dict = torch.load(input_model_file)
             print(f"Model loaded from {input_model_file}")
-        with FSDP.state_dict_type(model, self.state_dict_type, self.state_dict_config):
+
+        if is_torch_version("<=", "1.13.5"):
+            with FSDP.state_dict_type(model, self.state_dict_type, self.state_dict_config):
+                model.load_state_dict(state_dict)
+        else:
+            FSDP.set_state_dict_type(model, self.state_dict_type, self.state_dict_config)
             model.load_state_dict(state_dict)
 
     def save_optimizer(self, accelerator, optimizer, model, output_dir, optimizer_index=0, optim_input=None):
diff --git a/tests/fsdp/test_fsdp.py b/tests/fsdp/test_fsdp.py
index 249d2b692..e90106c17 100644
--- a/tests/fsdp/test_fsdp.py
+++ b/tests/fsdp/test_fsdp.py
@@ -95,7 +95,7 @@ def test_backward_prefetch(self):
                     self.assertEqual(fsdp_plugin.backward_prefetch, BackwardPrefetch(i + 1))
 
     def test_state_dict_type(self):
-        from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType, _state_dict_type_to_config
+        from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType
 
         for i, state_dict_type in enumerate(FSDP_STATE_DICT_TYPE):
             env = self.dist_env.copy()
@@ -103,9 +103,6 @@ def test_state_dict_type(self):
             with mockenv_context(**env):
                 fsdp_plugin = FullyShardedDataParallelPlugin()
                 self.assertEqual(fsdp_plugin.state_dict_type, StateDictType(i + 1))
-                self.assertEqual(
-                    type(fsdp_plugin.state_dict_config), type(_state_dict_type_to_config[StateDictType(i + 1)]())
-                )
                 if state_dict_type == "FULL_STATE_DICT":
                     self.assertTrue(fsdp_plugin.state_dict_config.offload_to_cpu)
                     self.assertTrue(fsdp_plugin.state_dict_config.rank0_only)

```suggestion
        if is_torch_version("<=", "1.13.5"):
```
Just in case they make a patch release before 1.14 is out.
@@ -616,18 +613,21 @@ def save_model(self, accelerator, model, output_dir, model_index=0):
         from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP
         from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType
 
-        if self.state_dict_type == StateDictType.FULL_STATE_DICT:
+        if is_torch_version("<=", "1.13.0"):
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/877
diff --git a/src/accelerate/big_modeling.py b/src/accelerate/big_modeling.py
index e277e89a2..24857ad25 100644
--- a/src/accelerate/big_modeling.py
+++ b/src/accelerate/big_modeling.py
@@ -210,7 +210,8 @@ def dispatch_model(
     device_map: Dict[str, Union[str, int, torch.device]],
     main_device: Optional[torch.device] = None,
     state_dict: Optional[Dict[str, torch.Tensor]] = None,
-    offload_dir: Union[str, os.PathLike] = None,
+    offload_dir: Optional[Union[str, os.PathLike]] = None,
+    offload_index: Optional[Dict[str, str]] = None,
     offload_buffers: bool = False,
     preload_module_classes: Optional[List[str]] = None,
 ):
@@ -231,6 +232,9 @@ def dispatch_model(
             The state dict of the part of the model that will be kept on CPU.
         offload_dir (`str` or `os.PathLike`):
             The folder in which to offload the model weights (or where the model weights are already offloaded).
+        offload_index (`Dict`, *optional*):
+            A dictionary from weight name to their information (`dtype`/ `shape` or safetensors filename). Will default
+            to the index saved in `save_folder`.
         offload_buffers (`bool`, *optional*, defaults to `False`):
             Whether or not to offload the buffers with the model parameters.
         preload_module_classes (`List[str]`, *optional*):
@@ -256,13 +260,15 @@ def dispatch_model(
             state_dict = extract_submodules_state_dict(model.state_dict(), cpu_modules)
 
     disk_modules = [name for name, device in device_map.items() if device == "disk"]
-    if offload_dir is None and len(disk_modules) > 0:
+    if offload_dir is None and offload_index is None and len(disk_modules) > 0:
         raise ValueError(
             "We need an `offload_dir` to dispatch this model according to this `device_map`, the following submodules "
             f"need to be offloaded: {', '.join(disk_modules)}."
         )
-    if len(disk_modules) > 0 and (
-        not os.path.isdir(offload_dir) or not os.path.isfile(os.path.join(offload_dir, "index.json"))
+    if (
+        len(disk_modules) > 0
+        and offload_index is None
+        and (not os.path.isdir(offload_dir) or not os.path.isfile(os.path.join(offload_dir, "index.json")))
     ):
         disk_state_dict = extract_submodules_state_dict(model.state_dict(), disk_modules)
         offload_state_dict(offload_dir, disk_state_dict)
@@ -273,8 +279,11 @@ def dispatch_model(
     offloaded_devices = ["disk"] if main_device == "cpu" else ["cpu", "disk"]
     offload = {name: device in offloaded_devices for name, device in device_map.items()}
     save_folder = offload_dir if len(disk_modules) > 0 else None
-    if state_dict is not None or save_folder is not None:
-        weights_map = OffloadedWeightsLoader(state_dict=state_dict, save_folder=save_folder)
+    if state_dict is not None or save_folder is not None or offload_index is not None:
+        device = main_device if offload_index is not None else None
+        weights_map = OffloadedWeightsLoader(
+            state_dict=state_dict, save_folder=save_folder, index=offload_index, device=device
+        )
     else:
         weights_map = None
 
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
index 87f891021..f46bce267 100644
--- a/src/accelerate/utils/__init__.py
+++ b/src/accelerate/utils/__init__.py
@@ -34,6 +34,7 @@
     is_megatron_lm_available,
     is_mlflow_available,
     is_rich_available,
+    is_safetensors_available,
     is_sagemaker_available,
     is_tensorboard_available,
     is_tpu_available,
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index bfedee162..72d6202cd 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -101,6 +101,10 @@ def is_megatron_lm_available():
     return False
 
 
+def is_safetensors_available():
+    return importlib.util.find_spec("safetensors") is not None
+
+
 def is_transformers_available():
     return importlib.util.find_spec("transformers") is not None
 
diff --git a/src/accelerate/utils/offload.py b/src/accelerate/utils/offload.py
index bc0fcdef1..22f51cce1 100644
--- a/src/accelerate/utils/offload.py
+++ b/src/accelerate/utils/offload.py
@@ -20,6 +20,12 @@
 import numpy as np
 import torch
 
+from ..logging import get_logger
+from .imports import is_safetensors_available
+
+
+logger = get_logger(__name__)
+
 
 def offload_weight(weight, weight_name, offload_folder, index=None):
     dtype = None
@@ -131,8 +137,8 @@ class OffloadedWeightsLoader(Mapping):
         save_folder (`str` or `os.PathLike`, *optional*):
             The directory in which the weights are stored (by `offload_state_dict` for instance).
         index (`Dict`, *optional*):
-            A dictionary from weight name to their information (`dtype` and `shape`). Will default to the index saved
-            in `save_folder`.
+            A dictionary from weight name to their information (`dtype`/ `shape` or safetensors filename). Will default
+            to the index saved in `save_folder`.
     """
 
     def __init__(
@@ -140,6 +146,7 @@ def __init__(
         state_dict: Dict[str, torch.Tensor] = None,
         save_folder: Optional[Union[str, os.PathLike]] = None,
         index: Mapping = None,
+        device=None,
     ):
         if state_dict is None and save_folder is None:
             raise ValueError("Need either a `state_dict` or a `save_folder` containing offloaded weights.")
@@ -152,12 +159,32 @@ def __init__(
         self.index = {} if index is None else index
         self.all_keys = list(self.state_dict.keys())
         self.all_keys.extend([key for key in self.index if key not in self.all_keys])
+        self.device = device
 
     def __getitem__(self, key: str):
         # State dict gets priority
         if key in self.state_dict:
             return self.state_dict[key]
         weight_info = self.index[key]
+        if weight_info.get("safetensors_file") is not None:
+            if not is_safetensors_available():
+                raise ImportError("These offloaded weights require the use of safetensors: `pip install safetensors`.")
+
+            if "SAFETENSORS_FAST_GPU" not in os.environ:
+                logger.info("Enabling fast loading with safetensors by setting `SAFETENSORS_FAST_GPU` to 1.")
+                os.environ["SAFETENSORS_FAST_GPU"] = "1"
+
+            from safetensors import safe_open
+
+            device = "cpu" if self.device is None else self.device
+            with safe_open(weight_info["safetensors_file"], framework="pt", device=device) as f:
+                tensor = f.get_tensor(weight_info.get("weight_name", key))
+
+            if "dtype" in weight_info:
+                return tensor.to(getattr(torch, weight_info["dtype"]))
+            else:
+                return tensor
+
         weight_file = os.path.join(self.save_folder, f"{key}.dat")
         return load_offloaded_weight(weight_file, weight_info)
 

We need a guard here for `is_safetensors_available` and raise an error if not :) (aka make one :) )
@@ -152,12 +153,25 @@ def __init__(
         self.index = {} if index is None else index
         self.all_keys = list(self.state_dict.keys())
         self.all_keys.extend([key for key in self.index if key not in self.all_keys])
+        self.device = device
 
     def __getitem__(self, key: str):
         # State dict gets priority
         if key in self.state_dict:
             return self.state_dict[key]
         weight_info = self.index[key]
+        if weight_info.get("safetensors_file") is not None:
+            from safetensors import safe_open
Indeed, will clean that up on Monday!
@@ -152,12 +153,25 @@ def __init__(
         self.index = {} if index is None else index
         self.all_keys = list(self.state_dict.keys())
         self.all_keys.extend([key for key in self.index if key not in self.all_keys])
+        self.device = device
 
     def __getitem__(self, key: str):
         # State dict gets priority
         if key in self.state_dict:
             return self.state_dict[key]
         weight_info = self.index[key]
+        if weight_info.get("safetensors_file") is not None:
+            from safetensors import safe_open
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/873
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index d6a323559..d614b47e2 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -926,6 +926,7 @@ def prepare_model(self, model: torch.nn.Module, device_placement=None):
             kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {}
             model = torch.nn.parallel.DistributedDataParallel(model, **kwargs)
         if self.native_amp:
+            model._original_forward = model.forward
             if self.mixed_precision == "fp16" and is_torch_version(">=", "1.10"):
                 model.forward = torch.cuda.amp.autocast(dtype=torch.float16)(model.forward)
             elif self.mixed_precision == "bf16" and self.distributed_type != DistributedType.TPU:
diff --git a/src/accelerate/utils/operations.py b/src/accelerate/utils/operations.py
index 7ba4482cc..3ebea2eef 100644
--- a/src/accelerate/utils/operations.py
+++ b/src/accelerate/utils/operations.py
@@ -17,6 +17,7 @@
 """
 
 
+import pickle
 from functools import update_wrapper
 from typing import Any, Mapping
 
@@ -473,8 +474,6 @@ class ConvertOutputsToFp32:
     Decorator to apply to a function outputing tensors (like a model forward pass) that ensures the outputs in FP16
     precision will be convert back to FP32.
 
-    Use a class instead of a decorator because otherwise, the prepared model can no longer be pickled (issue #273).
-
     Args:
         model_forward (`Callable`):
             The function which outputs we want to treat.
@@ -490,6 +489,11 @@ def __init__(self, model_forward):
     def __call__(self, *args, **kwargs):
         return convert_to_fp32(self.model_forward(*args, **kwargs))
 
+    def __getstate__(self):
+        raise pickle.PicklingError(
+            "Cannot pickle a prepared model with automatic mixed precision, please unwrap the model with `Accelerator.unwrap_model(model)` before pickling it."
+        )
+
 
 convert_outputs_to_fp32 = ConvertOutputsToFp32
 
diff --git a/src/accelerate/utils/other.py b/src/accelerate/utils/other.py
index a4a7b6312..1eb7cca88 100644
--- a/src/accelerate/utils/other.py
+++ b/src/accelerate/utils/other.py
@@ -21,7 +21,6 @@
 from ..state import AcceleratorState
 from .dataclasses import DistributedType
 from .imports import is_deepspeed_available, is_tpu_available
-from .operations import ConvertOutputsToFp32
 
 
 if is_deepspeed_available():
@@ -53,8 +52,13 @@ def extract_model_from_parallel(model, keep_fp32_wrapper: bool = False):
 
     if not keep_fp32_wrapper:
         forward = getattr(model, "forward")
-        if isinstance(forward, ConvertOutputsToFp32):
-            setattr(model, "forward", forward.model_forward)
+        original_forward = model.__dict__.pop("_original_forward", None)
+        if original_forward is not None:
+            while hasattr(forward, "__wrapped__"):
+                forward = forward.__wrapped__
+                if forward == original_forward:
+                    break
+            model.forward = forward
     return model
 
 
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 1e9d18c17..7c7629a1a 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -19,8 +19,15 @@
 
 import torch
 
+from accelerate.test_utils.testing import require_cuda
 from accelerate.test_utils.training import RegressionModel
-from accelerate.utils import convert_outputs_to_fp32, find_device, patch_environment, send_to_device
+from accelerate.utils import (
+    convert_outputs_to_fp32,
+    extract_model_from_parallel,
+    find_device,
+    patch_environment,
+    send_to_device,
+)
 
 
 ExampleNamedTuple = namedtuple("ExampleNamedTuple", "a b c")
@@ -74,9 +81,20 @@ def test_patch_environment(self):
         self.assertNotIn("AA", os.environ)
         self.assertNotIn("BB", os.environ)
 
-    def test_convert_to_32_lets_model_pickle(self):
+    def test_can_undo_convert_outputs(self):
         model = RegressionModel()
+        model._original_forward = model.forward
         model.forward = convert_outputs_to_fp32(model.forward)
+        model = extract_model_from_parallel(model)
+        _ = pickle.dumps(model)
+
+    @require_cuda
+    def test_can_undo_fp16_conversion(self):
+        model = RegressionModel()
+        model._original_forward = model.forward
+        model.forward = torch.cuda.amp.autocast(dtype=torch.float16)(model.forward)
+        model.forward = convert_outputs_to_fp32(model.forward)
+        model = extract_model_from_parallel(model)
         _ = pickle.dumps(model)
 
     def test_find_device(self):

@sgugger I know how you feel about `_` but I felt this was a decent time for one because the user really *shouldn't* need to know this exists at all or try to accidentally call it in any way
@@ -926,6 +926,7 @@ def prepare_model(self, model: torch.nn.Module, device_placement=None):
             kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {}
             model = torch.nn.parallel.DistributedDataParallel(model, **kwargs)
         if self.native_amp:
+            model._original_forward = model.forward
Agree in this case :-)
@@ -926,6 +926,7 @@ def prepare_model(self, model: torch.nn.Module, device_placement=None):
             kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {}
             model = torch.nn.parallel.DistributedDataParallel(model, **kwargs)
         if self.native_amp:
+            model._original_forward = model.forward
Are we sure this does not make a regression with this issue?
@@ -468,13 +468,11 @@ def _is_fp16_bf16_tensor(tensor):
     return recursively_apply(_convert_to_fp32, tensor, test_type=_is_fp16_bf16_tensor)
 
 
-class ConvertOutputsToFp32:
+def convert_outputs_to_fp32(model_forward):
     """
     Decorator to apply to a function outputing tensors (like a model forward pass) that ensures the outputs in FP16
     precision will be convert back to FP32.
 
-    Use a class instead of a decorator because otherwise, the prepared model can no longer be pickled (issue #273).
-
The issue will always exist because it can never be pickled if we've wrapped it in a mixed precision decorator. So it will always have to go through `unwrap_model`. And this only ever comes up if we do mixed precision. See the nb shown in the bug report here: https://colab.research.google.com/drive/11fvrk1Jslw2VIRTF6h0pdGgJx5AkAIMv?usp=sharing

Where specifically our explicit wrapping around `torch.cuda.amp` causes issues, and technically was the original issue too.

Or, phrasing it another way, the original issue doesn't come from Accelerate actually it comes from torch and using the autocast decorator. To keep things consistent this just makes them all use decorators so that we can quickly just unwrap all of them. 
@@ -468,13 +468,11 @@ def _is_fp16_bf16_tensor(tensor):
     return recursively_apply(_convert_to_fp32, tensor, test_type=_is_fp16_bf16_tensor)
 
 
-class ConvertOutputsToFp32:
+def convert_outputs_to_fp32(model_forward):
     """
     Decorator to apply to a function outputing tensors (like a model forward pass) that ensures the outputs in FP16
     precision will be convert back to FP32.
 
-    Use a class instead of a decorator because otherwise, the prepared model can no longer be pickled (issue #273).
-
I can try to see if it's possible to inject a warning before trying to pickle/save it and tell the user they need to unwrap the model first. 
@@ -468,13 +468,11 @@ def _is_fp16_bf16_tensor(tensor):
     return recursively_apply(_convert_to_fp32, tensor, test_type=_is_fp16_bf16_tensor)
 
 
-class ConvertOutputsToFp32:
+def convert_outputs_to_fp32(model_forward):
     """
     Decorator to apply to a function outputing tensors (like a model forward pass) that ensures the outputs in FP16
     precision will be convert back to FP32.
 
-    Use a class instead of a decorator because otherwise, the prepared model can no longer be pickled (issue #273).
-
Let's put your comment on the PR here as well, to explain what this code is doing 🙏 
@@ -53,8 +52,13 @@ def extract_model_from_parallel(model, keep_fp32_wrapper: bool = False):
 
     if not keep_fp32_wrapper:
         forward = getattr(model, "forward")
-        if isinstance(forward, ConvertOutputsToFp32):
-            setattr(model, "forward", forward.model_forward)
+        original_forward = model.__dict__.pop("_original_forward", None)
+        if original_forward is not None:
+            while hasattr(forward, "__wrapped__"):
+                forward = forward.__wrapped__
+                if forward == original_forward:
+                    break
A custom `__getstate__` func let's us raise an error when the user tries to prepare a model that had automatic mixed precision performed, ensuring with a clear error what the proper behavior needs to be, and the user doesn't have to know that this is due to PyTorch *and* accelerate and their behaviors
@@ -490,6 +488,11 @@ def __init__(self, model_forward):
     def __call__(self, *args, **kwargs):
         return convert_to_fp32(self.model_forward(*args, **kwargs))
 
+    def __getstate__(self):
+        raise pickle.PicklingError(
+            "Cannot pickle a prepared model with automatic mixed precision, please unwrap the model with `Accelerator.unwrap_model(model)` before pickling it."
+        )
+
When a function is wrapped using functools (such that we have `torch.cuda.amp.autocast(dtype=torch.float16)(model.forward)` the wrapped function (the original `model.forward`) exists in the `__wrapped__` attribute. When using `functools.update_wrapper` (such as in `ConvertOutputsToFp32`) this will also in turn add a new `__wrapped__` function, allowing us to follow the chain of `__wrapped__` until we return to the original model forward function, which was stored away earlier.
@@ -53,8 +52,13 @@ def extract_model_from_parallel(model, keep_fp32_wrapper: bool = False):
 
     if not keep_fp32_wrapper:
         forward = getattr(model, "forward")
-        if isinstance(forward, ConvertOutputsToFp32):
-            setattr(model, "forward", forward.model_forward)
+        original_forward = model.__dict__.pop("_original_forward", None)
+        if original_forward is not None:
+            while hasattr(forward, "__wrapped__"):
+                forward = forward.__wrapped__
+                if forward == original_forward:
+                    break
_The documentation is not available anymore as the PR was closed or merged._
Some notes on the final version (which will also be trickled into the code of the PR):

- I kept this as a class so that we can implement a custom `__getstate__` which will return an explicit pickling error if someone tries to pickle a wrapped model with automatic mixed precision on it
- When a function is wrapped using functools (such that we have `torch.cuda.amp.autocast(dtype=torch.float16)(model.forward)` the wrapped function (the original `model.forward`) exists in the `__wrapped__` attribute. When using `functools.update_wrapper` (such as in `ConvertOutputsToFp32`) this will also in turn add a new `__wrapped__` function, allowing us to follow the chain of `__wrapped__` until we return to the original model forward function, which was stored away earlier. 
https://api.github.com/repos/huggingface/accelerate/pulls/872
diff --git a/src/accelerate/logging.py b/src/accelerate/logging.py
index 2eceb223e..679f2da5c 100644
--- a/src/accelerate/logging.py
+++ b/src/accelerate/logging.py
@@ -78,7 +78,8 @@ def get_logger(name: str, log_level: str = None):
     ```
     """
     if log_level is None:
-        log_level = os.environ.get("ACCELERATE_LOG_LEVEL", "WARNING")
+        log_level = os.environ.get("ACCELERATE_LOG_LEVEL", None)
     logger = logging.getLogger(name)
-    logging.basicConfig(level=log_level.upper())
+    if log_level is not None:
+        logger.setLevel(log_level.upper())
     return MultiProcessAdapter(logger, {})

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/871
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 6ac8d8d3a..d6a323559 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -953,7 +953,9 @@ def _prepare_deepspeed(self, *args):
                 batch_sizes = [batch_size // self.num_processes for batch_size in batch_sizes]
             if len(batch_sizes) == 0:
                 raise ValueError(
-                    "You must specify a training or evaluation dataloader in `accelerate.prepare()` when using DeepSpeed."
+                    "When using DeepSpeed `accelerate.prepare()` requires you to pass at least one of training or evaluation dataloaders "
+                    "or alternatively set an integer value in `train_micro_batch_size_per_gpu` in the deepspeed config file"
+                    "or assign integer value to `AcceleratorState().deepspeed_plugin.deepspeed_config['train_micro_batch_size_per_gpu']`."
                 )
 
             batch_size_per_device = min(batch_sizes) if deepspeed_plugin.is_train_batch_min else max(batch_sizes)
diff --git a/tests/deepspeed/test_deepspeed.py b/tests/deepspeed/test_deepspeed.py
index 1be794421..86c9766a1 100644
--- a/tests/deepspeed/test_deepspeed.py
+++ b/tests/deepspeed/test_deepspeed.py
@@ -350,7 +350,9 @@ def test_prepare_deepspeed(self, optim_type, scheduler_type):
                 with self.assertRaises(ValueError) as cm:
                     model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler)
                 self.assertTrue(
-                    "You must specify a training or evaluation dataloader in `accelerate.prepare()` when using DeepSpeed."
+                    "When using DeepSpeed `accelerate.prepare()` requires you to pass at least one of training or evaluation dataloaders "
+                    "or alternatively set an integer value in `train_micro_batch_size_per_gpu` in the deepspeed config file"
+                    "or assign integer value to `AcceleratorState().deepspeed_plugin.deepspeed_config['train_micro_batch_size_per_gpu']`."
                     in str(cm.exception)
                 )
 

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/868
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index a653b9d6a..a3dcd2dcb 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -63,21 +63,21 @@ def get_cluster_input():
     if distributed_type in [DistributedType.MULTI_GPU, DistributedType.MULTI_CPU]:
         num_machines = _ask_field(
             "How many different machines will you use (use more than 1 for multi-node training)? [1]: ",
-            lambda x: int(x),
+            int,
             default=1,
         )
         if num_machines > 1:
             machine_rank = _ask_options(
                 "What is the rank of this machine?",
                 list(range(num_machines)),
-                lambda x: int(x),
+                int,
             )
             main_process_ip = _ask_field(
                 "What is the IP address of the machine that will host the main process? ",
             )
             main_process_port = _ask_field(
                 "What is the port you will use to communicate with the main process? ",
-                lambda x: int(x),
+                int,
             )
             same_network = _ask_field(
                 "Are all the machines on the same local network? Answer `no` if nodes are on the cloud and/or on different network hosts [YES/no]: ",
@@ -153,14 +153,14 @@ def get_cluster_input():
             if use_deepspeed_config:
                 deepspeed_config["deepspeed_config_file"] = _ask_field(
                     "Please enter the path to the json DeepSpeed config file: ",
-                    lambda x: str(x),
+                    str,
                     default="none",
                 )
             else:
                 deepspeed_config["zero_stage"] = _ask_options(
                     "What should be your DeepSpeed's ZeRO optimization stage?",
                     [0, 1, 2, 3],
-                    lambda x: int(x),
+                    int,
                     default=2,
                 )
 
@@ -174,7 +174,7 @@ def get_cluster_input():
                     )
                 deepspeed_config["gradient_accumulation_steps"] = _ask_field(
                     "How many gradient accumulation steps you're passing in your script? [1]: ",
-                    lambda x: int(x),
+                    int,
                     default=1,
                 )
                 use_gradient_clipping = _ask_field(
@@ -186,7 +186,7 @@ def get_cluster_input():
                 if use_gradient_clipping:
                     deepspeed_config["gradient_clipping"] = _ask_field(
                         "What is the gradient clipping value? [1.0]: ",
-                        lambda x: float(x),
+                        float,
                         default=1.0,
                     )
                 if deepspeed_config["zero_stage"] == 3:
@@ -224,7 +224,7 @@ def get_cluster_input():
                         "for more information please refer official [documentation]"
                         "(https://www.deepspeed.ai/getting-started/#resource-configuration-multi-node). "
                         "Please specify the location of hostfile: ",
-                        lambda x: str(x),
+                        str,
                     )
 
                     is_exclusion_filter = _ask_field(
@@ -236,7 +236,7 @@ def get_cluster_input():
                     if is_exclusion_filter:
                         deepspeed_config["deepspeed_exclusion_filter"] = _ask_field(
                             "DeepSpeed exclusion filter string: ",
-                            lambda x: str(x),
+                            str,
                         )
 
                     is_inclusion_filter = _ask_field(
@@ -248,7 +248,7 @@ def get_cluster_input():
                     if is_inclusion_filter:
                         deepspeed_config["deepspeed_inclusion_filter"] = _ask_field(
                             "DeepSpeed inclusion filter string: ",
-                            lambda x: str(x),
+                            str,
                         )
 
     fsdp_config = {}
@@ -284,12 +284,12 @@ def get_cluster_input():
             if fsdp_config["fsdp_auto_wrap_policy"] == FSDP_AUTO_WRAP_POLICY[0]:
                 fsdp_config["fsdp_transformer_layer_cls_to_wrap"] = _ask_field(
                     "What is the transformer layer class name (case-sensitive) to wrap ,e.g, `BertLayer`, `GPTJBlock`, `T5Block` ...? : ",
-                    lambda x: str(x),
+                    str,
                 )
             elif fsdp_config["fsdp_auto_wrap_policy"] == FSDP_AUTO_WRAP_POLICY[1]:
                 fsdp_config["fsdp_min_num_params"] = _ask_field(
                     "What should be your FSDP's minimum number of parameters for Default Auto Wrapping Policy? [1e8]: ",
-                    lambda x: int(x),
+                    int,
                     default=1e8,
                 )
             fsdp_backward_prefetch_query = "What should be your FSDP's backward prefetch policy?"
@@ -319,7 +319,7 @@ def get_cluster_input():
             prefix = "megatron_lm_"
             megatron_lm_config[prefix + "tp_degree"] = _ask_field(
                 "What is the Tensor Parallelism degree/size? [1]:",
-                lambda x: int(x),
+                int,
                 default=1,
                 error_message="Please enter an integer.",
             )
@@ -333,14 +333,14 @@ def get_cluster_input():
 
             megatron_lm_config[prefix + "pp_degree"] = _ask_field(
                 "What is the Pipeline Parallelism degree/size? [1]:",
-                lambda x: int(x),
+                int,
                 default=1,
                 error_message="Please enter an integer.",
             )
             if megatron_lm_config[prefix + "pp_degree"] > 1:
                 megatron_lm_config[prefix + "num_micro_batches"] = _ask_field(
                     "What is the number of micro-batches? [1]:",
-                    lambda x: int(x),
+                    int,
                     default=1,
                     error_message="Please enter an integer.",
                 )
@@ -362,7 +362,7 @@ def get_cluster_input():
 
             megatron_lm_config[prefix + "gradient_clipping"] = _ask_field(
                 "What is the gradient clipping value based on global L2 Norm (0 to disable)? [1.0]: ",
-                lambda x: float(x),
+                float,
                 default=1.0,
             )
 
@@ -438,14 +438,14 @@ def get_cluster_input():
             machine_type += "(s)"
         num_processes = _ask_field(
             f"How many {machine_type} should be used for distributed training? [1]:",
-            lambda x: int(x),
+            int,
             default=1,
             error_message="Please enter an integer.",
         )
     elif distributed_type in [DistributedType.FSDP, DistributedType.DEEPSPEED, DistributedType.MEGATRON_LM]:
         num_processes = _ask_field(
             "How many GPU(s) should be used for distributed training? [1]:",
-            lambda x: int(x),
+            int,
             default=1,
             error_message="Please enter an integer.",
         )
diff --git a/src/accelerate/commands/config/sagemaker.py b/src/accelerate/commands/config/sagemaker.py
index 22264b41f..af4195f29 100644
--- a/src/accelerate/commands/config/sagemaker.py
+++ b/src/accelerate/commands/config/sagemaker.py
@@ -97,7 +97,7 @@ def get_sagemaker_input():
     credentials_configuration = _ask_options(
         "How do you want to authorize?",
         ["AWS Profile", "Credentials (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY) "],
-        lambda x: int(x),
+        int,
     )
     aws_profile = None
     if credentials_configuration == 0:
@@ -120,7 +120,7 @@ def get_sagemaker_input():
     role_management = _ask_options(
         "Do you already have an IAM Role for executing Amazon SageMaker Training Jobs?",
         ["Provide IAM Role name", "Create new IAM role using credentials"],
-        lambda x: int(x),
+        int,
     )
     if role_management == 0:
         iam_role_name = _ask_field("Enter your IAM role name: ")
@@ -206,13 +206,10 @@ def get_sagemaker_input():
         ec2_instance_type = _ask_field(ec2_instance_query, lambda x: str(x).lower(), default="ml.p3.2xlarge")
 
     num_machines = 1
-    if (
-        distributed_type == SageMakerDistributedType.DATA_PARALLEL
-        or distributed_type == SageMakerDistributedType.MODEL_PARALLEL
-    ):
+    if distributed_type in (SageMakerDistributedType.DATA_PARALLEL, SageMakerDistributedType.MODEL_PARALLEL):
         num_machines = _ask_field(
             "How many machines do you want use? [1]: ",
-            lambda x: int(x),
+            int,
             default=1,
         )
 
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 8eefb9d7f..b2d95bad8 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -863,7 +863,7 @@ def _infer_type(s):
                 parser.add_argument(argument, action=action)
 
     return {
-        key: (literal_eval(value) if value == "True" or value == "False" else value)
+        key: (literal_eval(value) if value in ("True", "False") else value)
         for key, value in parser.parse_args(nargs).__dict__.items()
     }
 
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index 4f3930e1c..8a45856cd 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import math
+from contextlib import suppress
 from typing import List, Optional, Union
 
 import torch
@@ -364,12 +365,10 @@ def __iter__(self):
         if self.rng_types is not None:
             synchronize_rng_states(self.rng_types, self.synchronized_generator)
         self.gradient_state._set_end_of_dataloader(False)
-        try:
+        # We can safely pass because the default is -1
+        with suppress(Exception):
             length = getattr(self.dataset, "total_dataset_length", len(self.dataset))
             self.gradient_state._set_remainder(length % self.total_batch_size)
-        except Exception:
-            # We can safely pass because the default is -1
-            pass
         dataloader_iter = super().__iter__()
         # We iterate one batch ahead to check when we are at the end
         try:
@@ -448,12 +447,10 @@ def __init__(self, dataset, split_batches: bool = False, _drop_last: bool = Fals
         self.gradient_state = GradientState()
         self.state = AcceleratorState()
         self._drop_last = _drop_last
-        try:
+        # We can safely pass because the default is -1
+        with suppress(Exception):
             length = getattr(self.dataset, "total_dataset_length", len(self.dataset))
             self.gradient_state._set_remainder(length % self.total_batch_size)
-        except Exception:
-            # We can safely pass because the default is -1
-            pass
 
     def _fetch_batches(self, iterator):
         batches, batch = None, None
diff --git a/src/accelerate/tracking.py b/src/accelerate/tracking.py
index 689ee5033..75694ba72 100644
--- a/src/accelerate/tracking.py
+++ b/src/accelerate/tracking.py
@@ -593,7 +593,7 @@ def filter_trackers(
                             if getattr(tracker_init, "requires_logging_directory"):
                                 if logging_dir is None:
                                     raise ValueError(
-                                        f"Logging with `{str(log_type)}` requires a `logging_dir` to be passed in."
+                                        f"Logging with `{log_type}` requires a `logging_dir` to be passed in."
                                     )
                             loggers.append(log_type)
                         else:
diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py
index bb14078d4..50013c4d5 100644
--- a/src/accelerate/utils/dataclasses.py
+++ b/src/accelerate/utils/dataclasses.py
@@ -228,7 +228,7 @@ def __str__(self):
     @classmethod
     def list(cls):
         "Method to list all the possible items in `cls`"
-        return list(map(lambda item: str(item), cls))
+        return list(map(str, cls))
 
 
 class LoggerType(BaseEnum):
diff --git a/src/accelerate/utils/launch.py b/src/accelerate/utils/launch.py
index c65df08dc..1a4dcb4eb 100644
--- a/src/accelerate/utils/launch.py
+++ b/src/accelerate/utils/launch.py
@@ -88,7 +88,7 @@ def __call__(self, index, *args):
                 store=torch.distributed.FileStore(rdv_file, world_size),
                 world_size=world_size,
             )
-        elif self.distributed_type == DistributedType.MULTI_GPU or self.distributed_type == DistributedType.MULTI_CPU:
+        elif self.distributed_type in (DistributedType.MULTI_GPU, DistributedType.MULTI_CPU):
             # Prepare the environment for torch.distributed
             os.environ["LOCAL_RANK"] = str(index)
             os.environ["RANK"] = str(index)
diff --git a/src/accelerate/utils/megatron_lm.py b/src/accelerate/utils/megatron_lm.py
index afee2fcca..ab3bd8c37 100644
--- a/src/accelerate/utils/megatron_lm.py
+++ b/src/accelerate/utils/megatron_lm.py
@@ -132,7 +132,7 @@ def prepare_model(accelerator):
         custom_model_provider_func = accelerator.state.megatron_lm_plugin.custom_model_provider_function
         model = accelerator.state.megatron_lm_plugin.custom_prepare_model_function(custom_model_provider_func)
     else:
-        if args.model_type_name == "bert" or args.model_type_name == "gpt":
+        if args.model_type_name in ("bert", "gpt"):
             model_type = ModelType.encoder_or_decoder
         elif args.model_type_name == "t5":
             model_type = ModelType.encoder_and_decoder
@@ -566,7 +566,7 @@ def loss_func_finetune(labels, logits):
                 #  We are doing regression
                 loss_fct = MSELoss()
                 loss = loss_fct(logits.view(-1), labels.view(-1))
-            elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+            elif self.num_labels > 1 and (labels.dtype in (torch.long, torch.int)):
                 loss_fct = CrossEntropyLoss()
                 loss = loss_fct(logits.view(-1, num_labels), labels.view(-1))
             else:

Disagree with this change here. It's more readable to explicitly have the `.keys()`.
@@ -98,7 +98,7 @@ def __call__(self, parser, namespace, values, option_string=None):
             "optional arguments",
         ]
         if len(args) > 1:
-            used_platforms = [arg for arg in args if arg in options_to_group.keys()]
+            used_platforms = [arg for arg in args if arg in options_to_group]
Not convinced this is more readable.
@@ -245,7 +245,7 @@ def dispatch_model(
     check_device_map(model, device_map)
 
     if main_device is None:
-        if set(device_map.values()) == {"cpu"} or set(device_map.values()) == {"cpu", "disk"}:
+        if set(device_map.values()) in ({"cpu"}, {"cpu", "disk"}):
While this is technically equivalent, I prefer the explicit `else`.
@@ -77,8 +77,7 @@ def handle_input(cls):
         if handler:
             cls.current_selection = char
             return handler(cls)
-        else:
-            return None
+        return None
I prefer the explicit `else`.
@@ -164,9 +165,8 @@ def __len__(self):
         elif self.even_batches:
             # When we even batches we always get +1
             return length + 1
-        else:
-            # Otherwise it depends on the process index.
-            return length + 1 if self.process_index < len(self.batch_sampler) % self.num_processes else length
+        # Otherwise it depends on the process index.
+        return length + 1 if self.process_index < len(self.batch_sampler) % self.num_processes else length
Same here.
@@ -402,8 +400,7 @@ def total_batch_size(self):
     def total_dataset_length(self):
         if hasattr("total_length", self.dataset):
             return self.dataset.total_length
-        else:
-            return len(self.dataset)
+        return len(self.dataset)
Same here.
@@ -548,8 +543,7 @@ def __len__(self):
             return whole_length
         elif self._drop_last:
             return whole_length // self.state.num_processes
-        else:
-            return math.ceil(whole_length / self.state.num_processes)
+        return math.ceil(whole_length / self.state.num_processes)
While this is equivalent, `del values[names]` is explicit about what it's doing. `values.pop(name)` is not, so this should really be reverted.
@@ -500,7 +500,7 @@ def store_init_configuration(self, values: dict):
                     f'Trainer is attempting to log a value of "{value}" for key "{name}" as a parameter. MLflow\'s'
                     f" log_param() only accepts values no longer than {mlflow.utils.validation.MAX_PARAM_VAL_LENGTH} characters so we dropped this attribute."
                 )
-                del values[name]
+                values.pop(name)
This is not equivalent, so this change should be reverted (f-strings use the repr, `str` does not).
@@ -593,7 +593,7 @@ def filter_trackers(
                             if getattr(tracker_init, "requires_logging_directory"):
                                 if logging_dir is None:
                                     raise ValueError(
-                                        f"Logging with `{str(log_type)}` requires a `logging_dir` to be passed in."
+                                        f"Logging with `{log_type}` requires a `logging_dir` to be passed in."
Leaving the else and return is clearer.
@@ -538,8 +538,7 @@ def get_batch_transformer(data_iterator):
 
         if megatron_dataset_flag:
             return get_batch_megatron
-        else:
-            return get_batch_transformer
+        return get_batch_transformer
Same here.
@@ -577,8 +576,7 @@ def loss_func_finetune(labels, logits):
 
         if pretraining_flag:
             return loss_func_pretrain
-        else:
-            return loss_func_finetune
+        return loss_func_finetune
Same here.
@@ -666,8 +664,7 @@ def get_batch_transformer(data_iterator):
 
         if megatron_dataset_flag:
             return get_batch_megatron
-        else:
-            return get_batch_transformer
+        return get_batch_transformer
Same here.
@@ -803,8 +800,7 @@ def get_batch_transformer(data_iterator):
 
         if megatron_dataset_flag:
             return get_batch_megatron
-        else:
-            return get_batch_transformer
+        return get_batch_transformer
Same here.
@@ -1118,8 +1114,7 @@ def eval_step(self, **batch_data):
                 else:
                     loss_reduced[key] = torch.concat(losses_reduced_for_key)
             return loss_reduced
-        else:
-            return {}
+        return {}
As above, this should be reverted.
@@ -304,7 +304,7 @@ def clean_device_map(device_map: Dict[str, Union[int, str, torch.device]], modul
     values = [v for k, v in device_map.items() if k.startswith(prefix)]
     if len(set(values)) == 1 and len(values) > 1:
         for k in [k for k in device_map if k.startswith(prefix)]:
-            del device_map[k]
+            device_map.pop(k)
Let's leave the else return.
@@ -225,8 +225,7 @@ def gather(tensor):
         return _gpu_gather(tensor)
     elif AcceleratorState().distributed_type == DistributedType.MULTI_CPU:
         return _cpu_gather(tensor)
-    else:
-        return tensor
+    return tensor
Same here.
@@ -297,8 +296,7 @@ def broadcast(tensor, from_process: int = 0):
         return _gpu_broadcast(tensor, src=from_process)
     elif AcceleratorState().distributed_type == DistributedType.MULTI_CPU:
         return _gpu_broadcast(tensor, src=from_process)
-    else:
-        return tensor
+    return tensor
Same here.
@@ -439,8 +437,7 @@ def _reduce_across_processes(tensor, reduction="mean"):
         else:
             if reduction == "sum":
                 return cloned_tensor.sum()
-            else:
-                return cloned_tensor.mean()
+            return cloned_tensor.mean()
I'd leave the `.keys()` as it's more readable this way.
@@ -40,7 +40,7 @@ def compare_versions(library_or_version: Union[str, Version], operation: str, re
         requirement_version (`str`):
             The version to compare the library version against
     """
-    if operation not in STR_OPERATION_TO_FUNC.keys():
+    if operation not in STR_OPERATION_TO_FUNC:
For the record I quickly tested it and it seems that it actually uses `__str__`, not `__rep__`. 

```py
class A:
    def __str__(self):
        return "Use str"

    def __repr__(self):
        return "Use repr"

a = A()
print(a)
print("%s" % a)
print(f"{a}")
print(f"{str(a)}")
print("{}".format(a))
```

Output (both in python 3.7 and 3.10)
```
Use str
Use str
Use str
Use str
Use str
```
@@ -593,7 +593,7 @@ def filter_trackers(
                             if getattr(tracker_init, "requires_logging_directory"):
                                 if logging_dir is None:
                                     raise ValueError(
-                                        f"Logging with `{str(log_type)}` requires a `logging_dir` to be passed in."
+                                        f"Logging with `{log_type}` requires a `logging_dir` to be passed in."
Oh! My bad then.
@@ -593,7 +593,7 @@ def filter_trackers(
                             if getattr(tracker_init, "requires_logging_directory"):
                                 if logging_dir is None:
                                     raise ValueError(
-                                        f"Logging with `{str(log_type)}` requires a `logging_dir` to be passed in."
+                                        f"Logging with `{log_type}` requires a `logging_dir` to be passed in."
It is actually a TIL for me, didn't know what to expect :)
@@ -593,7 +593,7 @@ def filter_trackers(
                             if getattr(tracker_init, "requires_logging_directory"):
                                 if logging_dir is None:
                                     raise ValueError(
-                                        f"Logging with `{str(log_type)}` requires a `logging_dir` to be passed in."
+                                        f"Logging with `{log_type}` requires a `logging_dir` to be passed in."
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/865
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 6ac8d8d3a..d6a323559 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -953,7 +953,9 @@ def _prepare_deepspeed(self, *args):
                 batch_sizes = [batch_size // self.num_processes for batch_size in batch_sizes]
             if len(batch_sizes) == 0:
                 raise ValueError(
-                    "You must specify a training or evaluation dataloader in `accelerate.prepare()` when using DeepSpeed."
+                    "When using DeepSpeed `accelerate.prepare()` requires you to pass at least one of training or evaluation dataloaders "
+                    "or alternatively set an integer value in `train_micro_batch_size_per_gpu` in the deepspeed config file"
+                    "or assign integer value to `AcceleratorState().deepspeed_plugin.deepspeed_config['train_micro_batch_size_per_gpu']`."
                 )
 
             batch_size_per_device = min(batch_sizes) if deepspeed_plugin.is_train_batch_min else max(batch_sizes)

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/861
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 7a9ac92f4..6ac8d8d3a 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -1464,7 +1464,7 @@ def pad_across_processes(self, tensor, dim=0, pad_index=0, pad_first=False):
         """
         return pad_across_processes(tensor, dim=dim, pad_index=pad_index, pad_first=pad_first)
 
-    def unwrap_model(self, model):
+    def unwrap_model(self, model, keep_fp32_wrapper: bool = False):
         """
         Unwraps the `model` from the additional layer possible added by [`~Accelerator.prepare`]. Useful before saving
         the model.
@@ -1472,8 +1472,10 @@ def unwrap_model(self, model):
         Args:
             model (`torch.nn.Module`):
                 The model to unwrap.
+            keep_fp32_wrapper (`bool`, *optional*, defaults to `False`):
+                Whether to not remove the mixed precision hook if it was added.
         """
-        return extract_model_from_parallel(model)
+        return extract_model_from_parallel(model, keep_fp32_wrapper)
 
     def wait_for_everyone(self):
         """
@@ -1760,7 +1762,7 @@ def get_state_dict(self, model, unwrap=True):
         Args:
             model (`torch.nn.Module`):
                 A PyTorch model sent through [`Accelerator.prepare`]
-            unwrap (`bool`, *optional*, defaults to True):
+            unwrap (`bool`, *optional*, defaults to `True`):
                 Whether to return the original underlying state_dict of `model` or to return the wrapped state_dict
         """
         is_zero_3 = False
diff --git a/src/accelerate/utils/other.py b/src/accelerate/utils/other.py
index a9b6f9d4b..a4a7b6312 100644
--- a/src/accelerate/utils/other.py
+++ b/src/accelerate/utils/other.py
@@ -21,6 +21,7 @@
 from ..state import AcceleratorState
 from .dataclasses import DistributedType
 from .imports import is_deepspeed_available, is_tpu_available
+from .operations import ConvertOutputsToFp32
 
 
 if is_deepspeed_available():
@@ -30,12 +31,15 @@
     import torch_xla.core.xla_model as xm
 
 
-def extract_model_from_parallel(model):
+def extract_model_from_parallel(model, keep_fp32_wrapper: bool = False):
     """
     Extract a model from its distributed containers.
 
     Args:
-        model (`torch.nn.Module`): The model to extract.
+        model (`torch.nn.Module`):
+            The model to extract.
+        keep_fp32_wrapper (`bool`, *optional*):
+            Whether to remove mixed precision hooks from the model.
 
     Returns:
         `torch.nn.Module`: The extracted model.
@@ -46,6 +50,11 @@ def extract_model_from_parallel(model):
 
     while isinstance(model, options):
         model = model.module
+
+    if not keep_fp32_wrapper:
+        forward = getattr(model, "forward")
+        if isinstance(forward, ConvertOutputsToFp32):
+            setattr(model, "forward", forward.model_forward)
     return model
 
 

```suggestion
        the model.
```
Add this sentence when documenting the argument below instead.
@@ -1464,16 +1464,16 @@ def pad_across_processes(self, tensor, dim=0, pad_index=0, pad_first=False):
         """
         return pad_across_processes(tensor, dim=dim, pad_index=pad_index, pad_first=pad_first)
 
-    def unwrap_model(self, model):
+    def unwrap_model(self, model, keep_fp32_wrapper: bool = False):
         """
         Unwraps the `model` from the additional layer possible added by [`~Accelerator.prepare`]. Useful before saving
-        the model.
+        the model. Will also remove the mixed precision hook if it was added by default.
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_860). All of your documentation changes will be reflected on that endpoint.
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_860). All of your documentation changes will be reflected on that endpoint.
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_860). All of your documentation changes will be reflected on that endpoint.
https://api.github.com/repos/huggingface/accelerate/pulls/860
diff --git a/src/accelerate/logging.py b/src/accelerate/logging.py
index 46118f7c6..2eceb223e 100644
--- a/src/accelerate/logging.py
+++ b/src/accelerate/logging.py
@@ -78,7 +78,7 @@ def get_logger(name: str, log_level: str = None):
     ```
     """
     if log_level is None:
-        log_level = os.environ.get("ACCELERATE_LOG_LEVEL", "INFO")
+        log_level = os.environ.get("ACCELERATE_LOG_LEVEL", "WARNING")
     logger = logging.getLogger(name)
     logging.basicConfig(level=log_level.upper())
     return MultiProcessAdapter(logger, {})

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_859). All of your documentation changes will be reflected on that endpoint.
https://api.github.com/repos/huggingface/accelerate/pulls/859
diff --git a/.github/workflows/build_pr_documentation.yml b/.github/workflows/build_pr_documentation.yml
index 12c600cb7..dc56751c6 100644
--- a/.github/workflows/build_pr_documentation.yml
+++ b/.github/workflows/build_pr_documentation.yml
@@ -9,11 +9,8 @@ concurrency:
 
 jobs:
   build:
-    uses: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml@use_hf_hub
+    uses: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml@main
     with:
       commit_sha: ${{ github.event.pull_request.head.sha }}
       pr_number: ${{ github.event.number }}
       package: accelerate
-    secrets:
-      token: ${{ secrets.HF_DOC_PUSH }}
-      comment_bot_token: ${{ secrets.HUGGINGFACE_PUSH }}
diff --git a/.github/workflows/delete_doc_comment.yml b/.github/workflows/delete_doc_comment.yml
index 973c2a8b3..da61d21df 100644
--- a/.github/workflows/delete_doc_comment.yml
+++ b/.github/workflows/delete_doc_comment.yml
@@ -7,10 +7,7 @@ on:
 
 jobs:
   delete:
-    uses: huggingface/doc-builder/.github/workflows/delete_doc_comment.yml@use_hf_hub
+    uses: huggingface/doc-builder/.github/workflows/delete_doc_comment.yml@main
     with:
       pr_number: ${{ github.event.number }}
       package: accelerate
-    secrets:
-      token: ${{ secrets.HF_DOC_PUSH }}
-      comment_bot_token: ${{ secrets.HUGGINGFACE_PUSH }}

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_857). All of your documentation changes will be reflected on that endpoint.
https://api.github.com/repos/huggingface/accelerate/pulls/857
diff --git a/docs/source/package_reference/cli.mdx b/docs/source/package_reference/cli.mdx
index eafec816e..3d72aada6 100644
--- a/docs/source/package_reference/cli.mdx
+++ b/docs/source/package_reference/cli.mdx
@@ -46,7 +46,7 @@ Create a default config file for Accelerate with only a few flags set.
 **Usage**: 
 
 ```bash
-accelerate default-config [arguments]
+accelerate config default [arguments]
 ```
 
 **Optional Arguments**:
@@ -57,6 +57,28 @@ accelerate default-config [arguments]
 * `-h`, `--help` (`bool`) -- Show a help message and exit
 * `--mixed_precision {no,fp16,bf16}` (`str`) -- Whether or not to use mixed precision training. Choose between FP16 and BF16 (bfloat16) training. BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.
 
+## accelerate config update
+
+**Command**:
+
+`accelerate config update` or `accelerate-config update`
+
+Update an existing config file with the latest defaults while maintaining the old configuration.
+
+**Usage**: 
+
+```bash
+accelerate config update [arguments]
+```
+
+**Optional Arguments**:
+* `--config_file CONFIG_FILE` (`str`) -- The path to the config file to update. Will default to a file named default_config.yaml in the cache location, which is the content
+                        of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have such an environment variable, your cache directory
+                        (`~/.cache` or the content of `XDG_CACHE_HOME`) suffixed with `huggingface`.
+
+* `-h`, `--help` (`bool`) -- Show a help message and exit
+
+
 ## accelerate env
 
 **Command**:
diff --git a/src/accelerate/commands/config/__init__.py b/src/accelerate/commands/config/__init__.py
index ba9d6cf30..9b1545982 100644
--- a/src/accelerate/commands/config/__init__.py
+++ b/src/accelerate/commands/config/__init__.py
@@ -16,27 +16,22 @@
 
 import argparse
 
-from .config import config_command, config_command_parser
+from .config import config_command_parser
 from .config_args import default_config_file, load_config_from_file  # noqa: F401
-from .default import default_command_parser, default_config_command
-
-
-def filter_command_args(args: dict, args_prefix: str):
-    "Filters args while only keeping ones that are prefixed with `{args_prefix}.`"
-    new_args = argparse.Namespace()
-    for key, value in vars(args).items():
-        if key.startswith(args_prefix):
-            setattr(new_args, key.replace(f"{args_prefix}.", ""), value)
-    return new_args
+from .default import default_command_parser
+from .update import update_command_parser
 
 
 def get_config_parser(subparsers=None):
     parent_parser = argparse.ArgumentParser(add_help=False)
     # The main config parser
     config_parser = config_command_parser(subparsers)
+    # The subparser to add commands to
+    subcommands = config_parser.add_subparsers(title="subcommands", dest="subcommand")
 
     # Then add other parsers with the parent parser
-    default_parser = default_command_parser(config_parser, parents=[parent_parser])  # noqa: F841
+    default_command_parser(subcommands, parents=[parent_parser])
+    update_command_parser(subcommands, parents=[parent_parser])
 
     return config_parser
 
@@ -44,12 +39,13 @@ def get_config_parser(subparsers=None):
 def main():
     config_parser = get_config_parser()
     args = config_parser.parse_args()
-    if not args.default:
-        args = filter_command_args(args, "config_args")
-        config_command(args)
-    elif args.default:
-        args = filter_command_args(args, "default_args")
-        default_config_command(args)
+
+    if not hasattr(args, "func"):
+        config_parser.print_help()
+        exit(1)
+
+    # Run
+    args.func(args)
 
 
 if __name__ == "__main__":
diff --git a/src/accelerate/commands/config/config.py b/src/accelerate/commands/config/config.py
index b504f07ad..72414f2ab 100644
--- a/src/accelerate/commands/config/config.py
+++ b/src/accelerate/commands/config/config.py
@@ -21,13 +21,7 @@
 
 from .cluster import get_cluster_input
 from .config_args import cache_dir, default_config_file, default_yaml_config_file, load_config_from_file  # noqa: F401
-from .config_utils import (  # noqa: F401
-    GroupedAction,
-    SubcommandHelpFormatter,
-    _ask_field,
-    _ask_options,
-    _convert_compute_environment,
-)
+from .config_utils import _ask_field, _ask_options, _convert_compute_environment  # noqa: F401
 from .sagemaker import get_sagemaker_input
 
 
@@ -49,18 +43,13 @@ def get_user_input():
 
 def config_command_parser(subparsers=None):
     if subparsers is not None:
-        parser = subparsers.add_parser("config", description=description, formatter_class=SubcommandHelpFormatter)
+        parser = subparsers.add_parser("config", description=description)
     else:
-        parser = argparse.ArgumentParser(
-            "Accelerate config command", description=description, formatter_class=SubcommandHelpFormatter
-        )
+        parser = argparse.ArgumentParser("Accelerate config command", description=description)
 
     parser.add_argument(
         "--config_file",
         default=None,
-        dest="config_args.config_file",
-        metavar="CONFIG_FILE",
-        action=GroupedAction,
         help=(
             "The path to use to store the config file. Will default to a file named default_config.yaml in the cache "
             "location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have "
@@ -87,6 +76,7 @@ def config_command(args):
         config.to_json_file(config_file)
     else:
         config.to_yaml_file(config_file)
+    print(f"accelerate configuration saved at {config_file}")
 
 
 def main():
diff --git a/src/accelerate/commands/config/config_utils.py b/src/accelerate/commands/config/config_utils.py
index 1deee9f32..736c3a3d5 100644
--- a/src/accelerate/commands/config/config_utils.py
+++ b/src/accelerate/commands/config/config_utils.py
@@ -88,29 +88,11 @@ def _convert_yes_no_to_bool(value):
     return {"yes": True, "no": False}[value.lower()]
 
 
-class GroupedAction(argparse.Action):
-    """
-    Filters arguments into seperate namespace groups based on the first part of the argument name.
-    """
-
-    def __call__(self, parser, namespace, values, option_string=None):
-        group, dest = self.dest.split(".", 2)
-        groupspace = getattr(namespace, group, argparse.Namespace())
-        setattr(groupspace, dest, values)
-        setattr(namespace, group, groupspace)
-
-
 class SubcommandHelpFormatter(argparse.RawDescriptionHelpFormatter):
     """
     A custom formatter that will remove the usage line from the help message for subcommands.
     """
 
-    def _format_action(self, action):
-        parts = super()._format_action(action)
-        if action.nargs == argparse.PARSER:
-            parts = "\n".join(parts.split("\n")[1:])
-        return parts
-
     def _format_usage(self, usage, actions, groups, prefix):
         usage = super()._format_usage(usage, actions, groups, prefix)
         usage = usage.replace("<command> [<args>] ", "")
diff --git a/src/accelerate/commands/config/default.py b/src/accelerate/commands/config/default.py
index b87a1d45b..f7ceb84ce 100644
--- a/src/accelerate/commands/config/default.py
+++ b/src/accelerate/commands/config/default.py
@@ -14,13 +14,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import argparse
 from pathlib import Path
 
 import torch
 
 from .config_args import ClusterConfig, default_json_config_file
-from .config_utils import GroupedAction
+from .config_utils import SubcommandHelpFormatter
+
+
+description = "Create a default config file for Accelerate with only a few flags set."
 
 
 def write_basic_config(mixed_precision="no", save_location: str = default_json_config_file, dynamo_backend="no"):
@@ -42,7 +44,7 @@ def write_basic_config(mixed_precision="no", save_location: str = default_json_c
         print(
             f"Configuration already exists at {save_location}, will not override. Run `accelerate config` manually or pass a different `save_location`."
         )
-        return
+        return False
     mixed_precision = mixed_precision.lower()
     if mixed_precision not in ["no", "fp16", "bf16"]:
         raise ValueError(f"`mixed_precision` should be one of 'no', 'fp16', or 'bf16'. Received {mixed_precision}")
@@ -64,20 +66,13 @@ def write_basic_config(mixed_precision="no", save_location: str = default_json_c
         config["use_cpu"] = True
         config["num_processes"] = 1
         config["distributed_type"] = "NO"
-    if not path.exists():
-        config = ClusterConfig(**config)
-        config.to_json_file(path)
+    config = ClusterConfig(**config)
+    config.to_json_file(path)
+    return path
 
 
-description = "Create a default config file for Accelerate with only a few flags set."
-
-
-def default_command_parser(parser=None, parents=None):
-    if parser is None and parents is None:
-        parser = argparse.ArgumentParser(description=description)
-    else:
-        default_parser = parser.add_subparsers(title="subcommand {default}", dest="default", description=description)
-        parser = default_parser.add_parser("default", parents=parents)
+def default_command_parser(parser, parents):
+    parser = parser.add_parser("default", parents=parents, help=description, formatter_class=SubcommandHelpFormatter)
     parser.add_argument(
         "--config_file",
         default=default_json_config_file,
@@ -87,9 +82,7 @@ def default_command_parser(parser=None, parents=None):
             "such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed "
             "with 'huggingface'."
         ),
-        dest="default_args.save_location",
-        metavar="CONFIG_FILE",
-        action=GroupedAction,
+        dest="save_location",
     )
 
     parser.add_argument(
@@ -100,14 +93,12 @@ def default_command_parser(parser=None, parents=None):
         "Choose between FP16 and BF16 (bfloat16) training. "
         "BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.",
         default="no",
-        dest="default_args.mixed_precision",
-        action=GroupedAction,
     )
     parser.set_defaults(func=default_config_command)
     return parser
 
 
 def default_config_command(args):
-    args = vars(args)
-    args.pop("func", None)
-    write_basic_config(**args)
+    config_file = write_basic_config(args.mixed_precision, args.save_location)
+    if config_file:
+        print(f"accelerate configuration saved at {config_file}")
diff --git a/src/accelerate/commands/config/update.py b/src/accelerate/commands/config/update.py
new file mode 100644
index 000000000..5f025594b
--- /dev/null
+++ b/src/accelerate/commands/config/update.py
@@ -0,0 +1,63 @@
+#!/usr/bin/env python
+
+# 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.
+
+from pathlib import Path
+
+from .config_args import default_config_file, load_config_from_file
+from .config_utils import SubcommandHelpFormatter
+
+
+description = "Update an existing config file with the latest defaults while maintaining the old configuration."
+
+
+def update_config(args):
+    """
+    Update an existing config file with the latest defaults while maintaining the old configuration.
+    """
+    config_file = args.config_file
+    if config_file is None and Path(default_config_file).exists():
+        config_file = default_config_file
+    elif not Path(config_file).exists():
+        raise ValueError(f"The passed config file located at {config_file} doesn't exist.")
+    config = load_config_from_file(config_file)
+
+    if config_file.endswith(".json"):
+        config.to_json_file(config_file)
+    else:
+        config.to_yaml_file(config_file)
+    return config_file
+
+
+def update_command_parser(parser, parents):
+    parser = parser.add_parser("update", parents=parents, help=description, formatter_class=SubcommandHelpFormatter)
+    parser.add_argument(
+        "--config_file",
+        default=None,
+        help=(
+            "The path to the config file to update. Will default to a file named default_config.yaml in the cache "
+            "location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have "
+            "such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed "
+            "with 'huggingface'."
+        ),
+    )
+
+    parser.set_defaults(func=update_config_command)
+    return parser
+
+
+def update_config_command(args):
+    config_file = update_config(args)
+    print(f"Sucessfully updated the configuration file at {config_file}.")

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_855). All of your documentation changes will be reflected on that endpoint.
https://api.github.com/repos/huggingface/accelerate/pulls/855
diff --git a/docs/source/package_reference/cli.mdx b/docs/source/package_reference/cli.mdx
index eafec816e..3d72aada6 100644
--- a/docs/source/package_reference/cli.mdx
+++ b/docs/source/package_reference/cli.mdx
@@ -46,7 +46,7 @@ Create a default config file for Accelerate with only a few flags set.
 **Usage**: 
 
 ```bash
-accelerate default-config [arguments]
+accelerate config default [arguments]
 ```
 
 **Optional Arguments**:
@@ -57,6 +57,28 @@ accelerate default-config [arguments]
 * `-h`, `--help` (`bool`) -- Show a help message and exit
 * `--mixed_precision {no,fp16,bf16}` (`str`) -- Whether or not to use mixed precision training. Choose between FP16 and BF16 (bfloat16) training. BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.
 
+## accelerate config update
+
+**Command**:
+
+`accelerate config update` or `accelerate-config update`
+
+Update an existing config file with the latest defaults while maintaining the old configuration.
+
+**Usage**: 
+
+```bash
+accelerate config update [arguments]
+```
+
+**Optional Arguments**:
+* `--config_file CONFIG_FILE` (`str`) -- The path to the config file to update. Will default to a file named default_config.yaml in the cache location, which is the content
+                        of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have such an environment variable, your cache directory
+                        (`~/.cache` or the content of `XDG_CACHE_HOME`) suffixed with `huggingface`.
+
+* `-h`, `--help` (`bool`) -- Show a help message and exit
+
+
 ## accelerate env
 
 **Command**:
diff --git a/src/accelerate/commands/config/__init__.py b/src/accelerate/commands/config/__init__.py
index ba9d6cf30..9b1545982 100644
--- a/src/accelerate/commands/config/__init__.py
+++ b/src/accelerate/commands/config/__init__.py
@@ -16,27 +16,22 @@
 
 import argparse
 
-from .config import config_command, config_command_parser
+from .config import config_command_parser
 from .config_args import default_config_file, load_config_from_file  # noqa: F401
-from .default import default_command_parser, default_config_command
-
-
-def filter_command_args(args: dict, args_prefix: str):
-    "Filters args while only keeping ones that are prefixed with `{args_prefix}.`"
-    new_args = argparse.Namespace()
-    for key, value in vars(args).items():
-        if key.startswith(args_prefix):
-            setattr(new_args, key.replace(f"{args_prefix}.", ""), value)
-    return new_args
+from .default import default_command_parser
+from .update import update_command_parser
 
 
 def get_config_parser(subparsers=None):
     parent_parser = argparse.ArgumentParser(add_help=False)
     # The main config parser
     config_parser = config_command_parser(subparsers)
+    # The subparser to add commands to
+    subcommands = config_parser.add_subparsers(title="subcommands", dest="subcommand")
 
     # Then add other parsers with the parent parser
-    default_parser = default_command_parser(config_parser, parents=[parent_parser])  # noqa: F841
+    default_command_parser(subcommands, parents=[parent_parser])
+    update_command_parser(subcommands, parents=[parent_parser])
 
     return config_parser
 
@@ -44,12 +39,13 @@ def get_config_parser(subparsers=None):
 def main():
     config_parser = get_config_parser()
     args = config_parser.parse_args()
-    if not args.default:
-        args = filter_command_args(args, "config_args")
-        config_command(args)
-    elif args.default:
-        args = filter_command_args(args, "default_args")
-        default_config_command(args)
+
+    if not hasattr(args, "func"):
+        config_parser.print_help()
+        exit(1)
+
+    # Run
+    args.func(args)
 
 
 if __name__ == "__main__":
diff --git a/src/accelerate/commands/config/config.py b/src/accelerate/commands/config/config.py
index b504f07ad..72414f2ab 100644
--- a/src/accelerate/commands/config/config.py
+++ b/src/accelerate/commands/config/config.py
@@ -21,13 +21,7 @@
 
 from .cluster import get_cluster_input
 from .config_args import cache_dir, default_config_file, default_yaml_config_file, load_config_from_file  # noqa: F401
-from .config_utils import (  # noqa: F401
-    GroupedAction,
-    SubcommandHelpFormatter,
-    _ask_field,
-    _ask_options,
-    _convert_compute_environment,
-)
+from .config_utils import _ask_field, _ask_options, _convert_compute_environment  # noqa: F401
 from .sagemaker import get_sagemaker_input
 
 
@@ -49,18 +43,13 @@ def get_user_input():
 
 def config_command_parser(subparsers=None):
     if subparsers is not None:
-        parser = subparsers.add_parser("config", description=description, formatter_class=SubcommandHelpFormatter)
+        parser = subparsers.add_parser("config", description=description)
     else:
-        parser = argparse.ArgumentParser(
-            "Accelerate config command", description=description, formatter_class=SubcommandHelpFormatter
-        )
+        parser = argparse.ArgumentParser("Accelerate config command", description=description)
 
     parser.add_argument(
         "--config_file",
         default=None,
-        dest="config_args.config_file",
-        metavar="CONFIG_FILE",
-        action=GroupedAction,
         help=(
             "The path to use to store the config file. Will default to a file named default_config.yaml in the cache "
             "location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have "
@@ -87,6 +76,7 @@ def config_command(args):
         config.to_json_file(config_file)
     else:
         config.to_yaml_file(config_file)
+    print(f"accelerate configuration saved at {config_file}")
 
 
 def main():
diff --git a/src/accelerate/commands/config/config_utils.py b/src/accelerate/commands/config/config_utils.py
index 1deee9f32..736c3a3d5 100644
--- a/src/accelerate/commands/config/config_utils.py
+++ b/src/accelerate/commands/config/config_utils.py
@@ -88,29 +88,11 @@ def _convert_yes_no_to_bool(value):
     return {"yes": True, "no": False}[value.lower()]
 
 
-class GroupedAction(argparse.Action):
-    """
-    Filters arguments into seperate namespace groups based on the first part of the argument name.
-    """
-
-    def __call__(self, parser, namespace, values, option_string=None):
-        group, dest = self.dest.split(".", 2)
-        groupspace = getattr(namespace, group, argparse.Namespace())
-        setattr(groupspace, dest, values)
-        setattr(namespace, group, groupspace)
-
-
 class SubcommandHelpFormatter(argparse.RawDescriptionHelpFormatter):
     """
     A custom formatter that will remove the usage line from the help message for subcommands.
     """
 
-    def _format_action(self, action):
-        parts = super()._format_action(action)
-        if action.nargs == argparse.PARSER:
-            parts = "\n".join(parts.split("\n")[1:])
-        return parts
-
     def _format_usage(self, usage, actions, groups, prefix):
         usage = super()._format_usage(usage, actions, groups, prefix)
         usage = usage.replace("<command> [<args>] ", "")
diff --git a/src/accelerate/commands/config/default.py b/src/accelerate/commands/config/default.py
index b87a1d45b..f7ceb84ce 100644
--- a/src/accelerate/commands/config/default.py
+++ b/src/accelerate/commands/config/default.py
@@ -14,13 +14,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import argparse
 from pathlib import Path
 
 import torch
 
 from .config_args import ClusterConfig, default_json_config_file
-from .config_utils import GroupedAction
+from .config_utils import SubcommandHelpFormatter
+
+
+description = "Create a default config file for Accelerate with only a few flags set."
 
 
 def write_basic_config(mixed_precision="no", save_location: str = default_json_config_file, dynamo_backend="no"):
@@ -42,7 +44,7 @@ def write_basic_config(mixed_precision="no", save_location: str = default_json_c
         print(
             f"Configuration already exists at {save_location}, will not override. Run `accelerate config` manually or pass a different `save_location`."
         )
-        return
+        return False
     mixed_precision = mixed_precision.lower()
     if mixed_precision not in ["no", "fp16", "bf16"]:
         raise ValueError(f"`mixed_precision` should be one of 'no', 'fp16', or 'bf16'. Received {mixed_precision}")
@@ -64,20 +66,13 @@ def write_basic_config(mixed_precision="no", save_location: str = default_json_c
         config["use_cpu"] = True
         config["num_processes"] = 1
         config["distributed_type"] = "NO"
-    if not path.exists():
-        config = ClusterConfig(**config)
-        config.to_json_file(path)
+    config = ClusterConfig(**config)
+    config.to_json_file(path)
+    return path
 
 
-description = "Create a default config file for Accelerate with only a few flags set."
-
-
-def default_command_parser(parser=None, parents=None):
-    if parser is None and parents is None:
-        parser = argparse.ArgumentParser(description=description)
-    else:
-        default_parser = parser.add_subparsers(title="subcommand {default}", dest="default", description=description)
-        parser = default_parser.add_parser("default", parents=parents)
+def default_command_parser(parser, parents):
+    parser = parser.add_parser("default", parents=parents, help=description, formatter_class=SubcommandHelpFormatter)
     parser.add_argument(
         "--config_file",
         default=default_json_config_file,
@@ -87,9 +82,7 @@ def default_command_parser(parser=None, parents=None):
             "such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed "
             "with 'huggingface'."
         ),
-        dest="default_args.save_location",
-        metavar="CONFIG_FILE",
-        action=GroupedAction,
+        dest="save_location",
     )
 
     parser.add_argument(
@@ -100,14 +93,12 @@ def default_command_parser(parser=None, parents=None):
         "Choose between FP16 and BF16 (bfloat16) training. "
         "BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.",
         default="no",
-        dest="default_args.mixed_precision",
-        action=GroupedAction,
     )
     parser.set_defaults(func=default_config_command)
     return parser
 
 
 def default_config_command(args):
-    args = vars(args)
-    args.pop("func", None)
-    write_basic_config(**args)
+    config_file = write_basic_config(args.mixed_precision, args.save_location)
+    if config_file:
+        print(f"accelerate configuration saved at {config_file}")
diff --git a/src/accelerate/commands/config/update.py b/src/accelerate/commands/config/update.py
new file mode 100644
index 000000000..5f025594b
--- /dev/null
+++ b/src/accelerate/commands/config/update.py
@@ -0,0 +1,63 @@
+#!/usr/bin/env python
+
+# 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.
+
+from pathlib import Path
+
+from .config_args import default_config_file, load_config_from_file
+from .config_utils import SubcommandHelpFormatter
+
+
+description = "Update an existing config file with the latest defaults while maintaining the old configuration."
+
+
+def update_config(args):
+    """
+    Update an existing config file with the latest defaults while maintaining the old configuration.
+    """
+    config_file = args.config_file
+    if config_file is None and Path(default_config_file).exists():
+        config_file = default_config_file
+    elif not Path(config_file).exists():
+        raise ValueError(f"The passed config file located at {config_file} doesn't exist.")
+    config = load_config_from_file(config_file)
+
+    if config_file.endswith(".json"):
+        config.to_json_file(config_file)
+    else:
+        config.to_yaml_file(config_file)
+    return config_file
+
+
+def update_command_parser(parser, parents):
+    parser = parser.add_parser("update", parents=parents, help=description, formatter_class=SubcommandHelpFormatter)
+    parser.add_argument(
+        "--config_file",
+        default=None,
+        help=(
+            "The path to the config file to update. Will default to a file named default_config.yaml in the cache "
+            "location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have "
+            "such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed "
+            "with 'huggingface'."
+        ),
+    )
+
+    parser.set_defaults(func=update_config_command)
+    return parser
+
+
+def update_config_command(args):
+    config_file = update_config(args)
+    print(f"Sucessfully updated the configuration file at {config_file}.")

_The documentation is not available anymore as the PR was closed or merged._
The DeepSpeed CI failures are unrelated and look like a startup issue
https://api.github.com/repos/huggingface/accelerate/pulls/853
diff --git a/docs/source/package_reference/cli.mdx b/docs/source/package_reference/cli.mdx
index 562e1f4a4..eafec816e 100644
--- a/docs/source/package_reference/cli.mdx
+++ b/docs/source/package_reference/cli.mdx
@@ -35,11 +35,11 @@ accelerate config [arguments]
                         (`~/.cache` or the content of `XDG_CACHE_HOME`) suffixed with `huggingface`.
 * `-h`, `--help` (`bool`) -- Show a help message and exit
 
-## accelerate default-config 
+## accelerate config default
 
 **Command**:
 
-`accelerate default-config` or `accelerate-default-config`
+`accelerate config default` or `accelerate-config default`
 
 Create a default config file for Accelerate with only a few flags set.
 
diff --git a/setup.py b/setup.py
index 622d6de69..f96ebc440 100644
--- a/setup.py
+++ b/setup.py
@@ -48,7 +48,6 @@
             "accelerate=accelerate.commands.accelerate_cli:main",
             "accelerate-config=accelerate.commands.config:main",
             "accelerate-launch=accelerate.commands.launch:main",
-            "accelerate-default-config=accelerate.commands.config.default:main",
         ]
     },
     python_requires=">=3.7.0",
diff --git a/src/accelerate/commands/accelerate_cli.py b/src/accelerate/commands/accelerate_cli.py
index 8ffda3572..7716526c5 100644
--- a/src/accelerate/commands/accelerate_cli.py
+++ b/src/accelerate/commands/accelerate_cli.py
@@ -16,8 +16,7 @@
 
 from argparse import ArgumentParser
 
-from accelerate.commands.config import config_command_parser
-from accelerate.commands.config.default import default_command_parser
+from accelerate.commands.config import get_config_parser
 from accelerate.commands.env import env_command_parser
 from accelerate.commands.launch import launch_command_parser
 from accelerate.commands.test import test_command_parser
@@ -29,8 +28,7 @@ def main():
     subparsers = parser.add_subparsers(help="accelerate command helpers")
 
     # Register commands
-    config_command_parser(subparsers=subparsers)
-    default_command_parser(subparsers=subparsers)
+    get_config_parser(subparsers=subparsers)
     env_command_parser(subparsers=subparsers)
     launch_command_parser(subparsers=subparsers)
     tpu_command_parser(subparsers=subparsers)
diff --git a/src/accelerate/commands/config/__init__.py b/src/accelerate/commands/config/__init__.py
index 2400555c4..ba9d6cf30 100644
--- a/src/accelerate/commands/config/__init__.py
+++ b/src/accelerate/commands/config/__init__.py
@@ -15,73 +15,41 @@
 # limitations under the License.
 
 import argparse
-import os
 
-from accelerate.utils import ComputeEnvironment
+from .config import config_command, config_command_parser
+from .config_args import default_config_file, load_config_from_file  # noqa: F401
+from .default import default_command_parser, default_config_command
 
-from .cluster import get_cluster_input
-from .config_args import cache_dir, default_config_file, default_yaml_config_file, load_config_from_file  # noqa: F401
-from .config_utils import _ask_field, _ask_options, _convert_compute_environment  # noqa: F401
-from .sagemaker import get_sagemaker_input
 
+def filter_command_args(args: dict, args_prefix: str):
+    "Filters args while only keeping ones that are prefixed with `{args_prefix}.`"
+    new_args = argparse.Namespace()
+    for key, value in vars(args).items():
+        if key.startswith(args_prefix):
+            setattr(new_args, key.replace(f"{args_prefix}.", ""), value)
+    return new_args
 
-description = "Launches a series of prompts to create and save a `default_config.yaml` configuration file for your training system. Should always be ran first on your machine"
 
+def get_config_parser(subparsers=None):
+    parent_parser = argparse.ArgumentParser(add_help=False)
+    # The main config parser
+    config_parser = config_command_parser(subparsers)
 
-def get_user_input():
-    compute_environment = _ask_options(
-        "In which compute environment are you running?",
-        ["This machine", "AWS (Amazon SageMaker)"],
-        _convert_compute_environment,
-    )
-    if compute_environment == ComputeEnvironment.AMAZON_SAGEMAKER:
-        config = get_sagemaker_input()
-    else:
-        config = get_cluster_input()
-    return config
+    # Then add other parsers with the parent parser
+    default_parser = default_command_parser(config_parser, parents=[parent_parser])  # noqa: F841
 
-
-def config_command_parser(subparsers=None):
-    if subparsers is not None:
-        parser = subparsers.add_parser("config", description=description)
-    else:
-        parser = argparse.ArgumentParser("Accelerate config command", description=description)
-
-    parser.add_argument(
-        "--config_file",
-        default=None,
-        help=(
-            "The path to use to store the config file. Will default to a file named default_config.yaml in the cache "
-            "location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have "
-            "such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed "
-            "with 'huggingface'."
-        ),
-    )
-
-    if subparsers is not None:
-        parser.set_defaults(func=config_command)
-    return parser
-
-
-def config_command(args):
-    config = get_user_input()
-    if args.config_file is not None:
-        config_file = args.config_file
-    else:
-        if not os.path.isdir(cache_dir):
-            os.makedirs(cache_dir)
-        config_file = default_yaml_config_file
-
-    if config_file.endswith(".json"):
-        config.to_json_file(config_file)
-    else:
-        config.to_yaml_file(config_file)
+    return config_parser
 
 
 def main():
-    parser = config_command_parser()
-    args = parser.parse_args()
-    config_command(args)
+    config_parser = get_config_parser()
+    args = config_parser.parse_args()
+    if not args.default:
+        args = filter_command_args(args, "config_args")
+        config_command(args)
+    elif args.default:
+        args = filter_command_args(args, "default_args")
+        default_config_command(args)
 
 
 if __name__ == "__main__":
diff --git a/src/accelerate/commands/config/config.py b/src/accelerate/commands/config/config.py
new file mode 100644
index 000000000..b504f07ad
--- /dev/null
+++ b/src/accelerate/commands/config/config.py
@@ -0,0 +1,99 @@
+#!/usr/bin/env python
+
+# 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 argparse
+import os
+
+from accelerate.utils import ComputeEnvironment
+
+from .cluster import get_cluster_input
+from .config_args import cache_dir, default_config_file, default_yaml_config_file, load_config_from_file  # noqa: F401
+from .config_utils import (  # noqa: F401
+    GroupedAction,
+    SubcommandHelpFormatter,
+    _ask_field,
+    _ask_options,
+    _convert_compute_environment,
+)
+from .sagemaker import get_sagemaker_input
+
+
+description = "Launches a series of prompts to create and save a `default_config.yaml` configuration file for your training system. Should always be ran first on your machine"
+
+
+def get_user_input():
+    compute_environment = _ask_options(
+        "In which compute environment are you running?",
+        ["This machine", "AWS (Amazon SageMaker)"],
+        _convert_compute_environment,
+    )
+    if compute_environment == ComputeEnvironment.AMAZON_SAGEMAKER:
+        config = get_sagemaker_input()
+    else:
+        config = get_cluster_input()
+    return config
+
+
+def config_command_parser(subparsers=None):
+    if subparsers is not None:
+        parser = subparsers.add_parser("config", description=description, formatter_class=SubcommandHelpFormatter)
+    else:
+        parser = argparse.ArgumentParser(
+            "Accelerate config command", description=description, formatter_class=SubcommandHelpFormatter
+        )
+
+    parser.add_argument(
+        "--config_file",
+        default=None,
+        dest="config_args.config_file",
+        metavar="CONFIG_FILE",
+        action=GroupedAction,
+        help=(
+            "The path to use to store the config file. Will default to a file named default_config.yaml in the cache "
+            "location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have "
+            "such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed "
+            "with 'huggingface'."
+        ),
+    )
+
+    if subparsers is not None:
+        parser.set_defaults(func=config_command)
+    return parser
+
+
+def config_command(args):
+    config = get_user_input()
+    if args.config_file is not None:
+        config_file = args.config_file
+    else:
+        if not os.path.isdir(cache_dir):
+            os.makedirs(cache_dir)
+        config_file = default_yaml_config_file
+
+    if config_file.endswith(".json"):
+        config.to_json_file(config_file)
+    else:
+        config.to_yaml_file(config_file)
+
+
+def main():
+    parser = config_command_parser()
+    args = parser.parse_args()
+    config_command(args)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/accelerate/commands/config/config_utils.py b/src/accelerate/commands/config/config_utils.py
index def29fe06..1deee9f32 100644
--- a/src/accelerate/commands/config/config_utils.py
+++ b/src/accelerate/commands/config/config_utils.py
@@ -14,6 +14,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import argparse
+
 from ...utils.dataclasses import (
     ComputeEnvironment,
     DistributedType,
@@ -84,3 +86,32 @@ def _convert_sagemaker_distributed_mode(value):
 
 def _convert_yes_no_to_bool(value):
     return {"yes": True, "no": False}[value.lower()]
+
+
+class GroupedAction(argparse.Action):
+    """
+    Filters arguments into seperate namespace groups based on the first part of the argument name.
+    """
+
+    def __call__(self, parser, namespace, values, option_string=None):
+        group, dest = self.dest.split(".", 2)
+        groupspace = getattr(namespace, group, argparse.Namespace())
+        setattr(groupspace, dest, values)
+        setattr(namespace, group, groupspace)
+
+
+class SubcommandHelpFormatter(argparse.RawDescriptionHelpFormatter):
+    """
+    A custom formatter that will remove the usage line from the help message for subcommands.
+    """
+
+    def _format_action(self, action):
+        parts = super()._format_action(action)
+        if action.nargs == argparse.PARSER:
+            parts = "\n".join(parts.split("\n")[1:])
+        return parts
+
+    def _format_usage(self, usage, actions, groups, prefix):
+        usage = super()._format_usage(usage, actions, groups, prefix)
+        usage = usage.replace("<command> [<args>] ", "")
+        return usage
diff --git a/src/accelerate/commands/config/default.py b/src/accelerate/commands/config/default.py
index 956c3d719..b87a1d45b 100644
--- a/src/accelerate/commands/config/default.py
+++ b/src/accelerate/commands/config/default.py
@@ -15,21 +15,69 @@
 # limitations under the License.
 
 import argparse
-
-from accelerate.utils import write_basic_config
-
-from .config_args import default_json_config_file
+from pathlib import Path
+
+import torch
+
+from .config_args import ClusterConfig, default_json_config_file
+from .config_utils import GroupedAction
+
+
+def write_basic_config(mixed_precision="no", save_location: str = default_json_config_file, dynamo_backend="no"):
+    """
+    Creates and saves a basic cluster config to be used on a local machine with potentially multiple GPUs. Will also
+    set CPU if it is a CPU-only machine.
+
+    Args:
+        mixed_precision (`str`, *optional*, defaults to "no"):
+            Mixed Precision to use. Should be one of "no", "fp16", or "bf16"
+        save_location (`str`, *optional*, defaults to `default_json_config_file`):
+            Optional custom save location. Should be passed to `--config_file` when using `accelerate launch`. Default
+            location is inside the huggingface cache folder (`~/.cache/huggingface`) but can be overriden by setting
+            the `HF_HOME` environmental variable, followed by `accelerate/default_config.yaml`.
+    """
+    path = Path(save_location)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    if path.exists():
+        print(
+            f"Configuration already exists at {save_location}, will not override. Run `accelerate config` manually or pass a different `save_location`."
+        )
+        return
+    mixed_precision = mixed_precision.lower()
+    if mixed_precision not in ["no", "fp16", "bf16"]:
+        raise ValueError(f"`mixed_precision` should be one of 'no', 'fp16', or 'bf16'. Received {mixed_precision}")
+    config = {
+        "compute_environment": "LOCAL_MACHINE",
+        "mixed_precision": mixed_precision,
+        "dynamo_backend": dynamo_backend,
+    }
+    if torch.cuda.is_available():
+        num_gpus = torch.cuda.device_count()
+        config["num_processes"] = num_gpus
+        config["use_cpu"] = False
+        if num_gpus > 1:
+            config["distributed_type"] = "MULTI_GPU"
+        else:
+            config["distributed_type"] = "NO"
+    else:
+        num_gpus = 0
+        config["use_cpu"] = True
+        config["num_processes"] = 1
+        config["distributed_type"] = "NO"
+    if not path.exists():
+        config = ClusterConfig(**config)
+        config.to_json_file(path)
 
 
 description = "Create a default config file for Accelerate with only a few flags set."
 
 
-def default_command_parser(subparsers=None):
-    if subparsers is not None:
-        parser = subparsers.add_parser("default-config", description=description)
+def default_command_parser(parser=None, parents=None):
+    if parser is None and parents is None:
+        parser = argparse.ArgumentParser(description=description)
     else:
-        parser = argparse.ArgumentParser("Accelerate default-config command", description=description)
-
+        default_parser = parser.add_subparsers(title="subcommand {default}", dest="default", description=description)
+        parser = default_parser.add_parser("default", parents=parents)
     parser.add_argument(
         "--config_file",
         default=default_json_config_file,
@@ -39,7 +87,9 @@ def default_command_parser(subparsers=None):
             "such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed "
             "with 'huggingface'."
         ),
-        dest="save_location",
+        dest="default_args.save_location",
+        metavar="CONFIG_FILE",
+        action=GroupedAction,
     )
 
     parser.add_argument(
@@ -50,24 +100,14 @@ def default_command_parser(subparsers=None):
         "Choose between FP16 and BF16 (bfloat16) training. "
         "BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.",
         default="no",
+        dest="default_args.mixed_precision",
+        action=GroupedAction,
     )
-
-    if subparsers is not None:
-        parser.set_defaults(func=config_command)
+    parser.set_defaults(func=default_config_command)
     return parser
 
 
-def config_command(args):
+def default_config_command(args):
     args = vars(args)
     args.pop("func", None)
     write_basic_config(**args)
-
-
-def main():
-    parser = default_command_parser()
-    args = parser.parse_args()
-    config_command(args)
-
-
-if __name__ == "__main__":
-    main()
diff --git a/src/accelerate/utils/other.py b/src/accelerate/utils/other.py
index f1af373b6..a9b6f9d4b 100644
--- a/src/accelerate/utils/other.py
+++ b/src/accelerate/utils/other.py
@@ -14,12 +14,10 @@
 
 import os
 from contextlib import contextmanager
-from pathlib import Path
 
 import torch
 
-from ..commands.config.cluster import ClusterConfig
-from ..commands.config.config_args import default_json_config_file
+from ..commands.config.default import write_basic_config  # noqa: F401
 from ..state import AcceleratorState
 from .dataclasses import DistributedType
 from .imports import is_deepspeed_available, is_tpu_available
@@ -113,49 +111,3 @@ def get_pretty_name(obj):
     if hasattr(obj, "__name__"):
         return obj.__name__
     return str(obj)
-
-
-def write_basic_config(mixed_precision="no", save_location: str = default_json_config_file, dynamo_backend="no"):
-    """
-    Creates and saves a basic cluster config to be used on a local machine with potentially multiple GPUs. Will also
-    set CPU if it is a CPU-only machine.
-
-    Args:
-        mixed_precision (`str`, *optional*, defaults to "no"):
-            Mixed Precision to use. Should be one of "no", "fp16", or "bf16"
-        save_location (`str`, *optional*, defaults to `default_json_config_file`):
-            Optional custom save location. Should be passed to `--config_file` when using `accelerate launch`. Default
-            location is inside the huggingface cache folder (`~/.cache/huggingface`) but can be overriden by setting
-            the `HF_HOME` environmental variable, followed by `accelerate/default_config.yaml`.
-    """
-    path = Path(save_location)
-    path.parent.mkdir(parents=True, exist_ok=True)
-    if path.exists():
-        print(
-            f"Configuration already exists at {save_location}, will not override. Run `accelerate config` manually or pass a different `save_location`."
-        )
-        return
-    mixed_precision = mixed_precision.lower()
-    if mixed_precision not in ["no", "fp16", "bf16"]:
-        raise ValueError(f"`mixed_precision` should be one of 'no', 'fp16', or 'bf16'. Received {mixed_precision}")
-    config = {
-        "compute_environment": "LOCAL_MACHINE",
-        "mixed_precision": mixed_precision,
-        "dynamo_backend": dynamo_backend,
-    }
-    if torch.cuda.is_available():
-        num_gpus = torch.cuda.device_count()
-        config["num_processes"] = num_gpus
-        config["use_cpu"] = False
-        if num_gpus > 1:
-            config["distributed_type"] = "MULTI_GPU"
-        else:
-            config["distributed_type"] = "NO"
-    else:
-        num_gpus = 0
-        config["use_cpu"] = True
-        config["num_processes"] = 1
-        config["distributed_type"] = "NO"
-    if not path.exists():
-        config = ClusterConfig(**config)
-        config.to_json_file(path)

```suggestion
```
@@ -15,73 +15,42 @@
 # limitations under the License.
 
 import argparse
-import os
 
-from accelerate.utils import ComputeEnvironment
+from .config import config_command, config_command_parser
+from .config_args import default_config_file, load_config_from_file  # noqa: F401
+from .default import default_command_parser, default_config_command
 
-from .cluster import get_cluster_input
-from .config_args import cache_dir, default_config_file, default_yaml_config_file, load_config_from_file  # noqa: F401
-from .config_utils import _ask_field, _ask_options, _convert_compute_environment  # noqa: F401
-from .sagemaker import get_sagemaker_input
 
+def filter_command_args(args: dict, args_prefix: str):
+    "Filters args while only keeping ones that are prefixed with `{args_prefix}.`"
+    new_args = argparse.Namespace()
+    for key, value in vars(args).items():
+        print(key, value)
is this required here anymore?
@@ -14,12 +14,10 @@
 
 import os
 from contextlib import contextmanager
-from pathlib import Path
 
 import torch
 
-from ..commands.config.cluster import ClusterConfig
-from ..commands.config.config_args import default_json_config_file
+from ..commands.config.default import write_basic_config  # noqa: F401
Yep, this is the nit about making sure we can keep `from accelerate.utils import write_basic_config` 
@@ -14,12 +14,10 @@
 
 import os
 from contextlib import contextmanager
-from pathlib import Path
 
 import torch
 
-from ..commands.config.cluster import ClusterConfig
-from ..commands.config.config_args import default_json_config_file
+from ..commands.config.default import write_basic_config  # noqa: F401
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/851
diff --git a/src/accelerate/scheduler.py b/src/accelerate/scheduler.py
index 835d4e0d9..2652a2105 100644
--- a/src/accelerate/scheduler.py
+++ b/src/accelerate/scheduler.py
@@ -69,8 +69,9 @@ def step(self, *args, **kwargs):
             num_processes = AcceleratorState().num_processes
             for _ in range(num_processes):
                 # Special case when using OneCycle and `drop_last` was not used
-                if hasattr(self.scheduler, "total_steps") and self.scheduler._step_count <= self.scheduler.total_steps:
-                    self.scheduler.step(*args, **kwargs)
+                if hasattr(self.scheduler, "total_steps"):
+                    if self.scheduler._step_count <= self.scheduler.total_steps:
+                        self.scheduler.step(*args, **kwargs)
                 else:
                     self.scheduler.step(*args, **kwargs)
 

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/849
diff --git a/examples/complete_cv_example.py b/examples/complete_cv_example.py
index 880981594..bca55aa87 100644
--- a/examples/complete_cv_example.py
+++ b/examples/complete_cv_example.py
@@ -230,6 +230,7 @@ def training_function(config, args):
                     accelerator.save_state(output_dir)
         model.eval()
         accurate = 0
+        num_elems = 0
         for step, batch in enumerate(eval_dataloader):
             # We could avoid this line since we set the accelerator with `device_placement=True`.
             batch = {k: v.to(accelerator.device) for k, v in batch.items()}
@@ -239,9 +240,10 @@ def training_function(config, args):
             predictions = outputs.argmax(dim=-1)
             predictions, references = accelerator.gather_for_metrics((predictions, batch["label"]))
             accurate_preds = predictions == references
+            num_elems += accurate_preds.shape[0]
             accurate += accurate_preds.long().sum()
 
-        eval_metric = accurate.item() / accelerator.gradient_state.samples_seen
+        eval_metric = accurate.item() / num_elems
         # Use accelerator.print to print only on the main process.
         accelerator.print(f"epoch {epoch}: {100 * eval_metric:.2f}")
         if args.with_tracking:

https://api.github.com/repos/huggingface/accelerate/pulls/848
diff --git a/docs/source/package_reference/logging.mdx b/docs/source/package_reference/logging.mdx
index 675af41ee..85e844690 100644
--- a/docs/source/package_reference/logging.mdx
+++ b/docs/source/package_reference/logging.mdx
@@ -21,4 +21,14 @@ To utilize this replace cases of `logging` with `accelerate.logging`:
 + logger = get_logger(__name__)
 ```
 
+## Setting the log level
+
+The log level can be set with the `ACCELERATE_LOG_LEVEL` environment variable or by passing 
+`log_level` to `get_logger`:
+```python
+from accelerate.logging import get_logger
+
+logger = get_logger(__name__, log_level="INFO")
+```
+
 [[autodoc]] logging.get_logger
\ No newline at end of file
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index f595c08cc..7a9ac92f4 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -207,11 +207,6 @@ def __init__(
         dynamo_backend: Union[DynamoBackend, str] = None,
     ):
         self.logging_dir = logging_dir
-        trackers = filter_trackers(log_with, self.logging_dir)
-        if len(trackers) < 1 and log_with is not None:
-            warnings.warn(f"`log_with={log_with}` was passed but no supported trackers are currently installed.")
-        self.log_with = trackers
-
         if mixed_precision is not None:
             mixed_precision = str(mixed_precision)
             if mixed_precision not in PrecisionType:
@@ -303,6 +298,11 @@ def __init__(
             **kwargs,
         )
 
+        trackers = filter_trackers(log_with, self.logging_dir)
+        if len(trackers) < 1 and log_with is not None:
+            warnings.warn(f"`log_with={log_with}` was passed but no supported trackers are currently installed.")
+        self.log_with = trackers
+
         if (
             (mixed_precision != "bf16")
             and getattr(self.state, "downcast_bfloat", False)
diff --git a/src/accelerate/logging.py b/src/accelerate/logging.py
index dca093215..46118f7c6 100644
--- a/src/accelerate/logging.py
+++ b/src/accelerate/logging.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import logging
+import os
 
 from .state import AcceleratorState
 from .utils import DistributedType
@@ -49,7 +50,7 @@ def log(self, level, msg, *args, **kwargs):
             self.logger.log(level, msg, *args, **kwargs)
 
 
-def get_logger(name: str):
+def get_logger(name: str, log_level: str = None):
     """
     Returns a `logging.Logger` for `name` that can handle multiprocessing.
 
@@ -58,6 +59,8 @@ def get_logger(name: str):
     Args:
         name (`str`):
             The name for the logger, such as `__file__`
+        log_level (`str`, *optional*):
+            The log level to use. If not passed, will default to the `LOG_LEVEL` environment variable, or `INFO` if not
 
     Example:
 
@@ -68,7 +71,14 @@ def get_logger(name: str):
 
     >>> logger.info("My log", main_process_only=False)
     >>> logger.debug("My log", main_process_only=True)
+
+    >>> logger = get_logger(__name__, accelerate_log_level="DEBUG")
+    >>> logger.info("My log")
+    >>> logger.debug("My second log")
     ```
     """
+    if log_level is None:
+        log_level = os.environ.get("ACCELERATE_LOG_LEVEL", "INFO")
     logger = logging.getLogger(name)
+    logging.basicConfig(level=log_level.upper())
     return MultiProcessAdapter(logger, {})
diff --git a/src/accelerate/tracking.py b/src/accelerate/tracking.py
index 3a31fed93..689ee5033 100644
--- a/src/accelerate/tracking.py
+++ b/src/accelerate/tracking.py
@@ -577,7 +577,6 @@ def filter_trackers(
     if log_with is not None:
         if not isinstance(log_with, (list, tuple)):
             log_with = [log_with]
-            logger.debug(f"{log_with}")
         if "all" in log_with or LoggerType.ALL in log_with:
             loggers = [o for o in log_with if issubclass(type(o), GeneralTracker)] + get_available_trackers()
         else:

Is this a standard name that comes from `logging`? Should we also look for an `ACCELERATE_LOG_LEVEL` if users want to use a different one from the standard one? (So priority would be to look for the accelerate log level, then the general one).
@@ -21,4 +21,14 @@ To utilize this replace cases of `logging` with `accelerate.logging`:
 + logger = get_logger(__name__)
 ```
 
+## Setting the log level
+
+The log level can be set with the `LOG_LEVEL` environment variable or by passing 
```suggestion
        log_level (`str`, *optional*):
```
Make sure to re-read the doc-writing guide ;-)
@@ -58,6 +59,8 @@ def get_logger(name: str):
     Args:
         name (`str`):
             The name for the logger, such as `__file__`
+        log_level (`str`, `optional`, defaults to `None`):
There is not, generally I've seen it as `LOGLEVEL` for the recommended usage but then it's advised to pull it in from `os.environ`. However you're right having something Accelerate specific would be much better so I will opt for `ACCELERATE_LOG_LEVEL` instead
@@ -21,4 +21,14 @@ To utilize this replace cases of `logging` with `accelerate.logging`:
 + logger = get_logger(__name__)
 ```
 
+## Setting the log level
+
+The log level can be set with the `LOG_LEVEL` environment variable or by passing 
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/842
diff --git a/docs/source/package_reference/cli.mdx b/docs/source/package_reference/cli.mdx
index 9643d3b4f..562e1f4a4 100644
--- a/docs/source/package_reference/cli.mdx
+++ b/docs/source/package_reference/cli.mdx
@@ -35,6 +35,28 @@ accelerate config [arguments]
                         (`~/.cache` or the content of `XDG_CACHE_HOME`) suffixed with `huggingface`.
 * `-h`, `--help` (`bool`) -- Show a help message and exit
 
+## accelerate default-config 
+
+**Command**:
+
+`accelerate default-config` or `accelerate-default-config`
+
+Create a default config file for Accelerate with only a few flags set.
+
+**Usage**: 
+
+```bash
+accelerate default-config [arguments]
+```
+
+**Optional Arguments**:
+* `--config_file CONFIG_FILE` (`str`) -- The path to use to store the config file. Will default to a file named default_config.yaml in the cache location, which is the content
+                        of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have such an environment variable, your cache directory
+                        (`~/.cache` or the content of `XDG_CACHE_HOME`) suffixed with `huggingface`.
+
+* `-h`, `--help` (`bool`) -- Show a help message and exit
+* `--mixed_precision {no,fp16,bf16}` (`str`) -- Whether or not to use mixed precision training. Choose between FP16 and BF16 (bfloat16) training. BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.
+
 ## accelerate env
 
 **Command**:
diff --git a/setup.py b/setup.py
index f96ebc440..622d6de69 100644
--- a/setup.py
+++ b/setup.py
@@ -48,6 +48,7 @@
             "accelerate=accelerate.commands.accelerate_cli:main",
             "accelerate-config=accelerate.commands.config:main",
             "accelerate-launch=accelerate.commands.launch:main",
+            "accelerate-default-config=accelerate.commands.config.default:main",
         ]
     },
     python_requires=">=3.7.0",
diff --git a/src/accelerate/commands/accelerate_cli.py b/src/accelerate/commands/accelerate_cli.py
index f0e76fd2c..8ffda3572 100644
--- a/src/accelerate/commands/accelerate_cli.py
+++ b/src/accelerate/commands/accelerate_cli.py
@@ -17,6 +17,7 @@
 from argparse import ArgumentParser
 
 from accelerate.commands.config import config_command_parser
+from accelerate.commands.config.default import default_command_parser
 from accelerate.commands.env import env_command_parser
 from accelerate.commands.launch import launch_command_parser
 from accelerate.commands.test import test_command_parser
@@ -29,6 +30,7 @@ def main():
 
     # Register commands
     config_command_parser(subparsers=subparsers)
+    default_command_parser(subparsers=subparsers)
     env_command_parser(subparsers=subparsers)
     launch_command_parser(subparsers=subparsers)
     tpu_command_parser(subparsers=subparsers)
diff --git a/src/accelerate/commands/config/__init__.py b/src/accelerate/commands/config/__init__.py
index 4b02ac40c..2400555c4 100644
--- a/src/accelerate/commands/config/__init__.py
+++ b/src/accelerate/commands/config/__init__.py
@@ -25,6 +25,9 @@
 from .sagemaker import get_sagemaker_input
 
 
+description = "Launches a series of prompts to create and save a `default_config.yaml` configuration file for your training system. Should always be ran first on your machine"
+
+
 def get_user_input():
     compute_environment = _ask_options(
         "In which compute environment are you running?",
@@ -40,9 +43,9 @@ def get_user_input():
 
 def config_command_parser(subparsers=None):
     if subparsers is not None:
-        parser = subparsers.add_parser("config")
+        parser = subparsers.add_parser("config", description=description)
     else:
-        parser = argparse.ArgumentParser("Accelerate config command")
+        parser = argparse.ArgumentParser("Accelerate config command", description=description)
 
     parser.add_argument(
         "--config_file",
diff --git a/src/accelerate/commands/config/default.py b/src/accelerate/commands/config/default.py
new file mode 100644
index 000000000..956c3d719
--- /dev/null
+++ b/src/accelerate/commands/config/default.py
@@ -0,0 +1,73 @@
+#!/usr/bin/env python
+
+# 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 argparse
+
+from accelerate.utils import write_basic_config
+
+from .config_args import default_json_config_file
+
+
+description = "Create a default config file for Accelerate with only a few flags set."
+
+
+def default_command_parser(subparsers=None):
+    if subparsers is not None:
+        parser = subparsers.add_parser("default-config", description=description)
+    else:
+        parser = argparse.ArgumentParser("Accelerate default-config command", description=description)
+
+    parser.add_argument(
+        "--config_file",
+        default=default_json_config_file,
+        help=(
+            "The path to use to store the config file. Will default to a file named default_config.yaml in the cache "
+            "location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have "
+            "such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed "
+            "with 'huggingface'."
+        ),
+        dest="save_location",
+    )
+
+    parser.add_argument(
+        "--mixed_precision",
+        choices=["no", "fp16", "bf16"],
+        type=str,
+        help="Whether or not to use mixed precision training. "
+        "Choose between FP16 and BF16 (bfloat16) training. "
+        "BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.",
+        default="no",
+    )
+
+    if subparsers is not None:
+        parser.set_defaults(func=config_command)
+    return parser
+
+
+def config_command(args):
+    args = vars(args)
+    args.pop("func", None)
+    write_basic_config(**args)
+
+
+def main():
+    parser = default_command_parser()
+    args = parser.parse_args()
+    config_command(args)
+
+
+if __name__ == "__main__":
+    main()

I wouldn't include this in the default.
@@ -0,0 +1,82 @@
+#!/usr/bin/env python
+
+# 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 argparse
+
+from accelerate.utils import write_basic_config
+
+from .config_utils import DYNAMO_BACKENDS
+
+
+description = "Create a default config file for Accelerate with only a few flags set."
+
+
+def default_command_parser(subparsers=None):
+    if subparsers is not None:
+        parser = subparsers.add_parser("default-config", description=description)
+    else:
+        parser = argparse.ArgumentParser("Accelerate default-config command", description=description)
+
+    parser.add_argument(
+        "--config_file",
+        default=None,
+        help=(
+            "The path to use to store the config file. Will default to a file named default_config.yaml in the cache "
+            "location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have "
+            "such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed "
+            "with 'huggingface'."
+        ),
+        dest="save_location",
+    )
+
+    parser.add_argument(
+        "--mixed_precision",
+        choices=["no", "fp16", "bf16"],
+        type=str,
+        help="Whether or not to use mixed precision training. "
+        "Choose between FP16 and BF16 (bfloat16) training. "
+        "BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.",
+        default="no",
+    )
+
+    parser.add_argument(
+        "--dynamo_backend",
+        type=str,
+        choices=["no"] + [b.lower() for b in DYNAMO_BACKENDS],
+        help="Choose a backend to optimize your training with dynamo, see more at "
+        "https://github.com/pytorch/torchdynamo.",
+        default="no",
+    )
+
+    if subparsers is not None:
+        parser.set_defaults(func=config_command)
+    return parser
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/840
diff --git a/src/accelerate/commands/menu/helpers.py b/src/accelerate/commands/menu/helpers.py
index 2cc2ece3f..687f05c50 100644
--- a/src/accelerate/commands/menu/helpers.py
+++ b/src/accelerate/commands/menu/helpers.py
@@ -37,6 +37,10 @@ def forceWrite(content, end=""):
     sys.stdout.flush()
 
 
+def writeColor(content, color, end=""):
+    forceWrite(f"\u001b[{color}m{content}\u001b[0m", end)
+
+
 def reset_cursor():
     forceWrite("\r")
 
diff --git a/src/accelerate/commands/menu/selection_menu.py b/src/accelerate/commands/menu/selection_menu.py
index d4af34c8f..751f6818a 100644
--- a/src/accelerate/commands/menu/selection_menu.py
+++ b/src/accelerate/commands/menu/selection_menu.py
@@ -16,7 +16,7 @@
 Main driver for the selection menu, based on https://github.com/bchao1/bullet
 """
 from . import cursor, input
-from .helpers import Direction, clear_line, forceWrite, linebreak, move_cursor, reset_cursor
+from .helpers import Direction, clear_line, forceWrite, linebreak, move_cursor, reset_cursor, writeColor
 from .keymap import KEYMAP
 
 
@@ -34,7 +34,8 @@ def __init__(self, prompt: str = None, choices: list = []):
     def print_choice(self, index: int):
         "Prints the choice at the given index"
         if index == self.position:
-            forceWrite(f" ➔  {self.choices[index]}")
+            forceWrite(" ➔  ")
+            writeColor(self.choices[index], 32)
         else:
             forceWrite(f"    {self.choices[index]}")
         reset_cursor()
@@ -108,5 +109,6 @@ def run(self, default_choice: int = 0):
                     for _ in range(len(self.choices) + 1):
                         move_cursor(1, "UP")
                         clear_line()
-                    forceWrite(f" * {self.choices[choice]}", "\n")
+                    forceWrite(" ➔  ")
+                    writeColor(self.choices[choice], 32, "\n")
                     return choice

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/839
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index 15996bfd0..a653b9d6a 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -164,14 +164,13 @@ def get_cluster_input():
                     default=2,
                 )
 
+                deepspeed_devices = ["none", "cpu", "nvme"]
                 if deepspeed_config["zero_stage"] >= 2:
                     deepspeed_config["offload_optimizer_device"] = _ask_options(
-                        "Where to offload optimizer states?",
-                        ["none", "cpu", "nvme"],
+                        "Where to offload optimizer states?", deepspeed_devices, lambda x: deepspeed_devices[int(x)]
                     )
                     deepspeed_config["offload_param_device"] = _ask_options(
-                        "Where to offload parameters?",
-                        ["none", "cpu", "nvme"],
+                        "Where to offload parameters?", deepspeed_devices, lambda x: deepspeed_devices[int(x)]
                     )
                 deepspeed_config["gradient_accumulation_steps"] = _ask_field(
                     "How many gradient accumulation steps you're passing in your script? [1]: ",
@@ -294,7 +293,6 @@ def get_cluster_input():
                     default=1e8,
                 )
             fsdp_backward_prefetch_query = "What should be your FSDP's backward prefetch policy?"
-            fsdp_backward_prefetch_query = fsdp_backward_prefetch_query[:-2] + ")? [0]: "
             fsdp_config["fsdp_backward_prefetch_policy"] = _ask_options(
                 fsdp_backward_prefetch_query,
                 FSDP_BACKWARD_PREFETCH,

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_836). All of your documentation changes will be reflected on that endpoint.
https://api.github.com/repos/huggingface/accelerate/pulls/836
diff --git a/src/accelerate/commands/config/__init__.py b/src/accelerate/commands/config/__init__.py
index 11719835a..4b02ac40c 100644
--- a/src/accelerate/commands/config/__init__.py
+++ b/src/accelerate/commands/config/__init__.py
@@ -21,15 +21,15 @@
 
 from .cluster import get_cluster_input
 from .config_args import cache_dir, default_config_file, default_yaml_config_file, load_config_from_file  # noqa: F401
-from .config_utils import _ask_field, _convert_compute_environment
+from .config_utils import _ask_field, _ask_options, _convert_compute_environment  # noqa: F401
 from .sagemaker import get_sagemaker_input
 
 
 def get_user_input():
-    compute_environment = _ask_field(
-        "In which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)): ",
+    compute_environment = _ask_options(
+        "In which compute environment are you running?",
+        ["This machine", "AWS (Amazon SageMaker)"],
         _convert_compute_environment,
-        error_message="Please enter 0 or 1",
     )
     if compute_environment == ComputeEnvironment.AMAZON_SAGEMAKER:
         config = get_sagemaker_input()
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index 5f032debc..15996bfd0 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -31,14 +31,21 @@
     FSDP_STATE_DICT_TYPE,
 )
 from .config_args import ClusterConfig
-from .config_utils import _ask_field, _convert_distributed_mode, _convert_dynamo_backend, _convert_yes_no_to_bool
+from .config_utils import (
+    _ask_field,
+    _ask_options,
+    _convert_distributed_mode,
+    _convert_dynamo_backend,
+    _convert_mixed_precision,
+    _convert_yes_no_to_bool,
+)
 
 
 def get_cluster_input():
-    distributed_type = _ask_field(
-        "Which type of machine are you using? ([0] No distributed training, [1] multi-CPU, [2] multi-GPU, [3] TPU [4] MPS): ",
+    distributed_type = _ask_options(
+        "Which type of machine are you using?",
+        ["No distributed training", "multi-CPU", "multi-GPU", "TPU", "MPS"],
         _convert_distributed_mode,
-        error_message="Please enter 0, 1, 2, 3 or 4.",
     )
 
     machine_rank = 0
@@ -60,10 +67,10 @@ def get_cluster_input():
             default=1,
         )
         if num_machines > 1:
-            machine_rank = _ask_field(
-                "What is the rank of this machine (from 0 to the number of machines - 1 )? [0]: ",
+            machine_rank = _ask_options(
+                "What is the rank of this machine?",
+                list(range(num_machines)),
                 lambda x: int(x),
-                default=0,
             )
             main_process_ip = _ask_field(
                 "What is the IP address of the machine that will host the main process? ",
@@ -102,11 +109,22 @@ def get_cluster_input():
         error_message="Please enter yes or no.",
     )
     if use_dynamo:
-        dynamo_backend = _ask_field(
-            "Which dynamo backend would you like to use? ([0] eager, [1] aot_eager, [2] inductor, [3] nvfuser, [5] aot_nvfuser, [6] aot_cudagraphs, [7] ofi, [8] onnxrt, [9] ipex) [2]: ",
+        dynamo_backend = _ask_options(
+            "Which dynamo backend would you like to use?",
+            [
+                "eager",
+                "aot_eager",
+                "inductor",
+                "nvfuser",
+                "aot_nvfuser",
+                "aot_cudagraphs",
+                "ofi",
+                "fx2trt",
+                "onnxrt",
+                "ipex",
+            ],
             _convert_dynamo_backend,
-            default=DynamoBackend.INDUCTOR,
-            error_message="Please enter 0, 1, 2, 3, 4, 5, 6, 7, 8 or 9.",
+            default=2,
         )
     else:
         dynamo_backend = DynamoBackend.NO
@@ -139,22 +157,21 @@ def get_cluster_input():
                     default="none",
                 )
             else:
-                deepspeed_config["zero_stage"] = _ask_field(
-                    "What should be your DeepSpeed's ZeRO optimization stage (0, 1, 2, 3)? [2]: ",
+                deepspeed_config["zero_stage"] = _ask_options(
+                    "What should be your DeepSpeed's ZeRO optimization stage?",
+                    [0, 1, 2, 3],
                     lambda x: int(x),
                     default=2,
                 )
 
                 if deepspeed_config["zero_stage"] >= 2:
-                    deepspeed_config["offload_optimizer_device"] = _ask_field(
-                        "Where to offload optimizer states? [none/cpu/nvme]: ",
-                        lambda x: str(x),
-                        default="none",
+                    deepspeed_config["offload_optimizer_device"] = _ask_options(
+                        "Where to offload optimizer states?",
+                        ["none", "cpu", "nvme"],
                     )
-                    deepspeed_config["offload_param_device"] = _ask_field(
-                        "Where to offload parameters? [none/cpu/nvme]: ",
-                        lambda x: str(x),
-                        default="none",
+                    deepspeed_config["offload_param_device"] = _ask_options(
+                        "Where to offload parameters?",
+                        ["none", "cpu", "nvme"],
                     )
                 deepspeed_config["gradient_accumulation_steps"] = _ask_field(
                     "How many gradient accumulation steps you're passing in your script? [1]: ",
@@ -194,14 +211,11 @@ def get_cluster_input():
                     )
 
             if num_machines > 1:
-                launcher_query = "Which Type of launcher do you want to use "
-                for i, launcher in enumerate(DEEPSPEED_MULTINODE_LAUNCHERS):
-                    launcher_query += f"[{i}] {launcher}, "
-                launcher_query = launcher_query[:-2] + ")? [0]: "
-                deepspeed_config["deepspeed_multinode_launcher"] = _ask_field(
+                launcher_query = "Which Type of launcher do you want to use?"
+                deepspeed_config["deepspeed_multinode_launcher"] = _ask_options(
                     launcher_query,
+                    DEEPSPEED_MULTINODE_LAUNCHERS,
                     lambda x: DEEPSPEED_MULTINODE_LAUNCHERS[int(x)],
-                    default=DEEPSPEED_MULTINODE_LAUNCHERS[0],
                 )
 
                 if deepspeed_config["deepspeed_multinode_launcher"] != DEEPSPEED_MULTINODE_LAUNCHERS[1]:
@@ -249,13 +263,11 @@ def get_cluster_input():
         if use_fsdp:
             distributed_type = DistributedType.FSDP
         if distributed_type == DistributedType.FSDP:
-            sharding_strategy_query = "What should be your sharding strategy ("
-            for i, strategy in enumerate(FSDP_SHARDING_STRATEGY):
-                sharding_strategy_query += f"[{i+1}] {strategy}, "
-            sharding_strategy_query = sharding_strategy_query[:-2] + ")? [1]: "
-            fsdp_config["fsdp_sharding_strategy"] = _ask_field(
+            sharding_strategy_query = "What should be your sharding strategy?"
+            fsdp_config["fsdp_sharding_strategy"] = _ask_options(
                 sharding_strategy_query,
-                lambda x: int(x),
+                FSDP_SHARDING_STRATEGY,
+                lambda x: int(x) + 1,
                 default=1,
             )
             fsdp_config["fsdp_offload_params"] = _ask_field(
@@ -264,14 +276,11 @@ def get_cluster_input():
                 default=False,
                 error_message="Please enter yes or no.",
             )
-            fsdp_wrap_query = "What should be your auto wrap policy ("
-            for i, wrap_policy in enumerate(FSDP_AUTO_WRAP_POLICY):
-                fsdp_wrap_query += f"[{i}] {wrap_policy}, "
-            fsdp_wrap_query = fsdp_wrap_query[:-2] + ")? [0]: "
-            fsdp_config["fsdp_auto_wrap_policy"] = _ask_field(
+            fsdp_wrap_query = "What should be your auto wrap policy?"
+            fsdp_config["fsdp_auto_wrap_policy"] = _ask_options(
                 fsdp_wrap_query,
+                FSDP_AUTO_WRAP_POLICY,
                 lambda x: FSDP_AUTO_WRAP_POLICY[int(x)],
-                default="TRANSFORMER_BASED_WRAP",
             )
             if fsdp_config["fsdp_auto_wrap_policy"] == FSDP_AUTO_WRAP_POLICY[0]:
                 fsdp_config["fsdp_transformer_layer_cls_to_wrap"] = _ask_field(
@@ -284,23 +293,18 @@ def get_cluster_input():
                     lambda x: int(x),
                     default=1e8,
                 )
-            fsdp_backward_prefetch_query = "What should be your FSDP's backward prefetch policy ("
-            for i, backward_prefetch_policy in enumerate(FSDP_BACKWARD_PREFETCH):
-                fsdp_backward_prefetch_query += f"[{i}] {backward_prefetch_policy}, "
+            fsdp_backward_prefetch_query = "What should be your FSDP's backward prefetch policy?"
             fsdp_backward_prefetch_query = fsdp_backward_prefetch_query[:-2] + ")? [0]: "
-            fsdp_config["fsdp_backward_prefetch_policy"] = _ask_field(
+            fsdp_config["fsdp_backward_prefetch_policy"] = _ask_options(
                 fsdp_backward_prefetch_query,
+                FSDP_BACKWARD_PREFETCH,
                 lambda x: FSDP_BACKWARD_PREFETCH[int(x)],
-                default="BACKWARD_PRE",
             )
-            fsdp_state_dict_type_query = "What should be your FSDP's state dict type ("
-            for i, state_dict_type in enumerate(FSDP_STATE_DICT_TYPE):
-                fsdp_state_dict_type_query += f"[{i}] {state_dict_type}, "
-            fsdp_state_dict_type_query = fsdp_state_dict_type_query[:-2] + ")? [0]: "
-            fsdp_config["fsdp_state_dict_type"] = _ask_field(
+            fsdp_state_dict_type_query = "What should be your FSDP's state dict type?"
+            fsdp_config["fsdp_state_dict_type"] = _ask_options(
                 fsdp_state_dict_type_query,
+                FSDP_STATE_DICT_TYPE,
                 lambda x: FSDP_STATE_DICT_TYPE[int(x)],
-                default="FULL_STATE_DICT",
             )
 
     megatron_lm_config = {}
@@ -460,10 +464,10 @@ def get_cluster_input():
         if distributed_type == DistributedType.DEEPSPEED and use_deepspeed_config:
             mixed_precision = "no"
         else:
-            mixed_precision = _ask_field(
-                "Do you wish to use FP16 or BF16 (mixed precision)? [NO/fp16/bf16]: ",
-                lambda x: str(x).lower(),
-                default="no",
+            mixed_precision = _ask_options(
+                "Do you wish to use FP16 or BF16 (mixed precision)?",
+                ["no", "fp16", "bf16"],
+                _convert_mixed_precision,
             )
     else:
         mixed_precision = "no"
diff --git a/src/accelerate/commands/config/config_utils.py b/src/accelerate/commands/config/config_utils.py
index 8974cddc9..def29fe06 100644
--- a/src/accelerate/commands/config/config_utils.py
+++ b/src/accelerate/commands/config/config_utils.py
@@ -14,7 +14,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ...utils.dataclasses import ComputeEnvironment, DistributedType, DynamoBackend, SageMakerDistributedType
+from ...utils.dataclasses import (
+    ComputeEnvironment,
+    DistributedType,
+    DynamoBackend,
+    PrecisionType,
+    SageMakerDistributedType,
+)
+from ..menu import BulletMenu
 
 
 DYNAMO_BACKENDS = [
@@ -44,6 +51,12 @@ def _ask_field(input_text, convert_value=None, default=None, error_message=None)
                 print(error_message)
 
 
+def _ask_options(input_text, options=[], convert_value=None, default=0):
+    menu = BulletMenu(input_text, options)
+    result = menu.run(default_choice=default)
+    return convert_value(result) if convert_value is not None else result
+
+
 def _convert_compute_environment(value):
     value = int(value)
     return ComputeEnvironment(["LOCAL_MACHINE", "AMAZON_SAGEMAKER"][value])
@@ -56,7 +69,12 @@ def _convert_distributed_mode(value):
 
 def _convert_dynamo_backend(value):
     value = int(value)
-    return DynamoBackend(DYNAMO_BACKENDS[value + 1])
+    return DynamoBackend(DYNAMO_BACKENDS[value])
+
+
+def _convert_mixed_precision(value):
+    value = int(value)
+    return PrecisionType(["no", "fp16", "bf16"][value])
 
 
 def _convert_sagemaker_distributed_mode(value):
diff --git a/src/accelerate/commands/config/sagemaker.py b/src/accelerate/commands/config/sagemaker.py
index 71de4fea9..22264b41f 100644
--- a/src/accelerate/commands/config/sagemaker.py
+++ b/src/accelerate/commands/config/sagemaker.py
@@ -22,7 +22,9 @@
 from .config_args import SageMakerConfig
 from .config_utils import (
     _ask_field,
+    _ask_options,
     _convert_dynamo_backend,
+    _convert_mixed_precision,
     _convert_sagemaker_distributed_mode,
     _convert_yes_no_to_bool,
 )
@@ -92,8 +94,9 @@ def _get_iam_role_arn(role_name):
 
 
 def get_sagemaker_input():
-    credentials_configuration = _ask_field(
-        "How do you want to authorize? ([0] AWS Profile, [1] Credentials (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY)): ",
+    credentials_configuration = _ask_options(
+        "How do you want to authorize?",
+        ["AWS Profile", "Credentials (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY) "],
         lambda x: int(x),
     )
     aws_profile = None
@@ -114,8 +117,9 @@ def get_sagemaker_input():
     aws_region = _ask_field("Enter your AWS Region: [us-east-1]", default="us-east-1")
     os.environ["AWS_DEFAULT_REGION"] = aws_region
 
-    role_management = _ask_field(
-        "Do you already have an IAM Role for executing Amazon SageMaker Training Jobs? ([0] provide IAM Role name, [1] create new IAM role using credentials: ",
+    role_management = _ask_options(
+        "Do you already have an IAM Role for executing Amazon SageMaker Training Jobs?",
+        ["Provide IAM Role name", "Create new IAM role using credentials"],
         lambda x: int(x),
     )
     if role_management == 0:
@@ -161,12 +165,11 @@ def get_sagemaker_input():
             lambda x: str(x).lower(),
         )
 
-    distributed_type = _ask_field(
-        "What is the distributed mode? ([0] No distributed training, [1] data parallelism): ",
+    distributed_type = _ask_options(
+        "What is the distributed mode?",
+        ["No distributed training", "Data parallelism"],
         _convert_sagemaker_distributed_mode,
-        error_message="Please enter 0 or 1",
     )
-
     use_dynamo = _ask_field(
         "Do you wish to optimize your script with torch dynamo?[yes/NO]:",
         _convert_yes_no_to_bool,
@@ -174,22 +177,30 @@ def get_sagemaker_input():
         error_message="Please enter yes or no.",
     )
     if use_dynamo:
-        dynamo_backend = _ask_field(
-            "Which dynamo backend would you like to use? ([0] eager, [1] aot_eager, [2] inductor, [3] nvfuser, [5] aot_nvfuser, [6] aot_cudagraphs, [7] ofi, [8] onnxrt, [9] ipex) [2]: ",
+        dynamo_backend = _ask_options(
+            "Which dynamo backend would you like to use?",
+            [
+                "eager",
+                "aot_eager",
+                "inductor",
+                "nvfuser",
+                "aot_nvfuser",
+                "aot_cudagraphs",
+                "ofi",
+                "fx2trt",
+                "onnxrt",
+                "ipex",
+            ],
             _convert_dynamo_backend,
-            default=DynamoBackend.INDUCTOR,
-            error_message="Please enter 0, 1, 2, 3, 4, 5, 6, 7, 8 or 9.",
+            default=2,
         )
     else:
         dynamo_backend = DynamoBackend.NO
-
-    ec2_instance_query = "Which EC2 instance type you want to use for your training "
+    ec2_instance_query = "Which EC2 instance type you want to use for your training?"
     if distributed_type != SageMakerDistributedType.NO:
-        ec2_instance_query += "("
-        for i, instance_type in enumerate(SAGEMAKER_PARALLEL_EC2_INSTANCES):
-            ec2_instance_query += f"[{i}] {instance_type}, "
-        ec2_instance_query = ec2_instance_query[:-2] + ")? [0]: "
-        ec2_instance_type = _ask_field(ec2_instance_query, lambda x: SAGEMAKER_PARALLEL_EC2_INSTANCES[int(x)])
+        ec2_instance_type = _ask_options(
+            ec2_instance_query, SAGEMAKER_PARALLEL_EC2_INSTANCES, lambda x: SAGEMAKER_PARALLEL_EC2_INSTANCES[int(x)]
+        )
     else:
         ec2_instance_query += "? [ml.p3.2xlarge]:"
         ec2_instance_type = _ask_field(ec2_instance_query, lambda x: str(x).lower(), default="ml.p3.2xlarge")
@@ -205,10 +216,10 @@ def get_sagemaker_input():
             default=1,
         )
 
-    mixed_precision = _ask_field(
-        "Do you wish to use FP16 or BF16 (mixed precision)? [No/FP16/BF16]: ",
-        lambda x: str(x).lower(),
-        default="no",
+    mixed_precision = _ask_options(
+        "Do you wish to use FP16 or BF16 (mixed precision)?",
+        ["no", "fp16", "bf16"],
+        _convert_mixed_precision,
     )
 
     if use_dynamo and mixed_precision == "no":
diff --git a/src/accelerate/commands/menu/__init__.py b/src/accelerate/commands/menu/__init__.py
new file mode 100644
index 000000000..ec17fba58
--- /dev/null
+++ b/src/accelerate/commands/menu/__init__.py
@@ -0,0 +1,5 @@
+# flake8: noqa
+# There's no way to ignore "F401 '...' imported but unused" warnings in this
+# module, but to preserve other warnings. So, don't check this module at all
+
+from .selection_menu import BulletMenu
diff --git a/src/accelerate/commands/menu/cursor.py b/src/accelerate/commands/menu/cursor.py
new file mode 100644
index 000000000..c1f0bb7b6
--- /dev/null
+++ b/src/accelerate/commands/menu/cursor.py
@@ -0,0 +1,65 @@
+# Copyright 2022 The HuggingFace Team and Brian Chao. 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.
+
+"""
+A utility for showing and hiding the terminal cursor on Windows and Linux, based on https://github.com/bchao1/bullet
+"""
+
+import os
+import sys
+from contextlib import contextmanager
+
+
+# Windows only
+if os.name == "nt":
+    import ctypes
+    import msvcrt  # noqa
+
+    class CursorInfo(ctypes.Structure):
+        # _fields is a specific attr expected by ctypes
+        _fields_ = [("size", ctypes.c_int), ("visible", ctypes.c_byte)]
+
+
+def hide_cursor():
+    if os.name == "nt":
+        ci = CursorInfo()
+        handle = ctypes.windll.kernel32.GetStdHandle(-11)
+        ctypes.windll.kernel32.GetConsoleCursorInfo(handle, ctypes.byref(ci))
+        ci.visible = False
+        ctypes.windll.kernel32.SetConsoleCursorInfo(handle, ctypes.byref(ci))
+    elif os.name == "posix":
+        sys.stdout.write("\033[?25l")
+        sys.stdout.flush()
+
+
+def show_cursor():
+    if os.name == "nt":
+        ci = CursorInfo()
+        handle = ctypes.windll.kernel32.GetStdHandle(-11)
+        ctypes.windll.kernel32.GetConsoleCursorInfo(handle, ctypes.byref(ci))
+        ci.visible = True
+        ctypes.windll.kernel32.SetConsoleCursorInfo(handle, ctypes.byref(ci))
+    elif os.name == "posix":
+        sys.stdout.write("\033[?25h")
+        sys.stdout.flush()
+
+
+@contextmanager
+def hide():
+    "Context manager to hide the terminal cursor"
+    try:
+        hide_cursor()
+        yield
+    finally:
+        show_cursor()
diff --git a/src/accelerate/commands/menu/helpers.py b/src/accelerate/commands/menu/helpers.py
new file mode 100644
index 000000000..2cc2ece3f
--- /dev/null
+++ b/src/accelerate/commands/menu/helpers.py
@@ -0,0 +1,55 @@
+# Copyright 2022 The HuggingFace Team and Brian Chao. 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.
+
+"""
+A variety of helper functions and constants when dealing with terminal menu choices, based on
+https://github.com/bchao1/bullet
+"""
+
+import enum
+import shutil
+import sys
+
+
+TERMINAL_WIDTH, _ = shutil.get_terminal_size()
+
+CURSOR_TO_CHAR = {"UP": "A", "DOWN": "B", "RIGHT": "C", "LEFT": "D"}
+
+
+class Direction(enum.Enum):
+    UP = 0
+    DOWN = 1
+
+
+def forceWrite(content, end=""):
+    sys.stdout.write(content + end)
+    sys.stdout.flush()
+
+
+def reset_cursor():
+    forceWrite("\r")
+
+
+def move_cursor(num_lines: int, direction: str):
+    forceWrite(f"\033[{num_lines}{CURSOR_TO_CHAR[direction.upper()]}")
+
+
+def clear_line():
+    forceWrite(" " * TERMINAL_WIDTH)
+    reset_cursor()
+
+
+def linebreak():
+    reset_cursor()
+    forceWrite("-" * TERMINAL_WIDTH)
diff --git a/src/accelerate/commands/menu/input.py b/src/accelerate/commands/menu/input.py
new file mode 100644
index 000000000..266f7e7db
--- /dev/null
+++ b/src/accelerate/commands/menu/input.py
@@ -0,0 +1,86 @@
+# Copyright 2022 The HuggingFace Team and Brian Chao. 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.
+
+"""
+This file contains utilities for handling input from the user and registering specific keys to specific functions,
+based on https://github.com/bchao1/bullet
+"""
+
+from typing import List
+
+from .keymap import KEYMAP, get_character
+
+
+def mark(key: str):
+    """
+    Mark the function with the key code so it can be handled in the register
+    """
+
+    def decorator(func):
+        handle = getattr(func, "handle_key", [])
+        handle += [key]
+        setattr(func, "handle_key", handle)
+        return func
+
+    return decorator
+
+
+def mark_multiple(*keys: List[str]):
+    """
+    Mark the function with the key codes so it can be handled in the register
+    """
+
+    def decorator(func):
+        handle = getattr(func, "handle_key", [])
+        handle += keys
+        setattr(func, "handle_key", handle)
+        return func
+
+    return decorator
+
+
+class KeyHandler(type):
+    """
+    Metaclass that adds the key handlers to the class
+    """
+
+    def __new__(cls, name, bases, attrs):
+        new_cls = super().__new__(cls, name, bases, attrs)
+        if not hasattr(new_cls, "key_handler"):
+            setattr(new_cls, "key_handler", {})
+        setattr(new_cls, "handle_input", KeyHandler.handle_input)
+
+        for value in attrs.values():
+            handled_keys = getattr(value, "handle_key", [])
+            for key in handled_keys:
+                new_cls.key_handler[key] = value
+        return new_cls
+
+    @staticmethod
+    def handle_input(cls):
+        "Finds and returns the selected character if it exists in the handler"
+        char = get_character()
+        if char != KEYMAP["undefined"]:
+            char = ord(char)
+        handler = cls.key_handler.get(char)
+        if handler:
+            cls.current_selection = char
+            return handler(cls)
+        else:
+            return None
+
+
+def register(cls):
+    """Adds KeyHandler metaclass to the class"""
+    return KeyHandler(cls.__name__, cls.__bases__, cls.__dict__.copy())
diff --git a/src/accelerate/commands/menu/keymap.py b/src/accelerate/commands/menu/keymap.py
new file mode 100644
index 000000000..7ce6c0637
--- /dev/null
+++ b/src/accelerate/commands/menu/keymap.py
@@ -0,0 +1,81 @@
+# Copyright 2022 The HuggingFace Team and Brian Chao. 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.
+
+"""
+Utilities relating to parsing raw characters from the keyboard, based on https://github.com/bchao1/bullet
+"""
+
+
+import string
+import sys
+import termios
+import tty
+
+
+ARROW_KEY_FLAG = 1 << 8
+
+KEYMAP = {
+    "tab": ord("\t"),
+    "newline": ord("\r"),
+    "esc": 27,
+    "up": 65 + ARROW_KEY_FLAG,
+    "down": 66 + ARROW_KEY_FLAG,
+    "right": 67 + ARROW_KEY_FLAG,
+    "left": 68 + ARROW_KEY_FLAG,
+    "mod_int": 91,
+    "undefined": sys.maxsize,
+    "interrupt": 3,
+}
+
+KEYMAP["arrow_begin"] = KEYMAP["up"]
+KEYMAP["arrow_end"] = KEYMAP["left"]
+
+for i in range(10):
+    KEYMAP[str(i)] = ord(str(i))
+
+
+def get_raw_chars():
+    "Gets raw characters from inputs"
+    fd = sys.stdin.fileno()
+    old_settings = termios.tcgetattr(fd)
+    try:
+        tty.setraw(fd)
+        ch = sys.stdin.read(1)
+    finally:
+        termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
+    return ch
+
+
+def get_character():
+    "Gets a character from the keyboard and returns the key code"
+    char = get_raw_chars()
+    if ord(char) in [KEYMAP["interrupt"], KEYMAP["newline"]]:
+        return char
+
+    elif ord(char) == KEYMAP["esc"]:
+        combo = get_raw_chars()
+        if ord(combo) == KEYMAP["mod_int"]:
+            key = get_raw_chars()
+            if ord(key) >= KEYMAP["arrow_begin"] - ARROW_KEY_FLAG and ord(key) <= KEYMAP["arrow_end"] - ARROW_KEY_FLAG:
+                return chr(ord(key) + ARROW_KEY_FLAG)
+            else:
+                return KEYMAP["undefined"]
+        else:
+            return get_raw_chars()
+
+    else:
+        if char in string.printable:
+            return char
+        else:
+            return KEYMAP["undefined"]
diff --git a/src/accelerate/commands/menu/selection_menu.py b/src/accelerate/commands/menu/selection_menu.py
new file mode 100644
index 000000000..d4af34c8f
--- /dev/null
+++ b/src/accelerate/commands/menu/selection_menu.py
@@ -0,0 +1,112 @@
+# Copyright 2022 The HuggingFace Team and Brian Chao. 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.
+
+"""
+Main driver for the selection menu, based on https://github.com/bchao1/bullet
+"""
+from . import cursor, input
+from .helpers import Direction, clear_line, forceWrite, linebreak, move_cursor, reset_cursor
+from .keymap import KEYMAP
+
+
+@input.register
+class BulletMenu:
+    """
+    A CLI menu to select a choice from a list of choices using the keyboard.
+    """
+
+    def __init__(self, prompt: str = None, choices: list = []):
+        self.position = 0
+        self.choices = choices
+        self.prompt = prompt
+
+    def print_choice(self, index: int):
+        "Prints the choice at the given index"
+        if index == self.position:
+            forceWrite(f" ➔  {self.choices[index]}")
+        else:
+            forceWrite(f"    {self.choices[index]}")
+        reset_cursor()
+
+    def move_direction(self, direction: Direction, num_spaces: int = 1):
+        "Should not be directly called, used to move a direction of either up or down"
+        old_position = self.position
+        if direction == Direction.DOWN:
+            if self.position + 1 >= len(self.choices):
+                return
+            self.position += num_spaces
+        else:
+            if self.position - 1 < 0:
+                return
+            self.position -= num_spaces
+        clear_line()
+        self.print_choice(old_position)
+        move_cursor(num_spaces, direction.name)
+        self.print_choice(self.position)
+
+    @input.mark(KEYMAP["up"])
+    def move_up(self):
+        self.move_direction(Direction.UP)
+
+    @input.mark(KEYMAP["down"])
+    def move_down(self):
+        self.move_direction(Direction.DOWN)
+
+    @input.mark(KEYMAP["newline"])
+    def select(self):
+        move_cursor(len(self.choices) - self.position, "DOWN")
+        return self.position
+
+    @input.mark(KEYMAP["interrupt"])
+    def interrupt(self):
+        move_cursor(len(self.choices) - self.position, "DOWN")
+        raise KeyboardInterrupt
+
+    @input.mark_multiple(*[KEYMAP[str(number)] for number in range(10)])
+    def select_row(self):
+        index = int(chr(self.current_selection))
+        movement = index - self.position
+        if index == self.position:
+            return
+        if index < len(self.choices):
+            if self.position > index:
+                self.move_direction(Direction.UP, -movement)
+            elif self.position < index:
+                self.move_direction(Direction.DOWN, movement)
+            else:
+                return
+        else:
+            return
+
+    def run(self, default_choice: int = 0):
+        "Start the menu and return the selected choice"
+        if self.prompt:
+            linebreak()
+            forceWrite(self.prompt, "\n")
+            forceWrite("Please select a choice using the arrow or number keys, and selecting with enter", "\n")
+        self.position = default_choice
+        for i in range(len(self.choices)):
+            self.print_choice(i)
+            forceWrite("\n")
+        move_cursor(len(self.choices) - self.position, "UP")
+        with cursor.hide():
+            while True:
+                choice = self.handle_input()
+                if choice is not None:
+                    reset_cursor()
+                    for _ in range(len(self.choices) + 1):
+                        move_cursor(1, "UP")
+                        clear_line()
+                    forceWrite(f" * {self.choices[choice]}", "\n")
+                    return choice

Let's name this submodule `menu` since it's about that?
@@ -0,0 +1,5 @@
+# flake8: noqa
You know how I feel about _ prefixing names.
@@ -0,0 +1,64 @@
+# 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.
+
+"""
+A utility for showing and hiding the terminal cursor on Windows and Linux, based on https://github.com/bchao1/bullet
+"""
+
+import os
+import sys
+from contextlib import contextmanager
+
+
+# Windows only
+if os.name == "nt":
+    import ctypes
+    import msvcrt  # noqa
+
+    class _CursorInfo(ctypes.Structure):
+        _fields_ = [("size", ctypes.c_int), ("visible", ctypes.c_byte)]
If this is heavily inspired by another lib, it needs to be reflected in the copyright probably.
@@ -0,0 +1,64 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
`_fields_` is very specific to `ctypes.Structure`, but will rename `CursorInfo`
@@ -0,0 +1,64 @@
+# 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.
+
+"""
+A utility for showing and hiding the terminal cursor on Windows and Linux, based on https://github.com/bchao1/bullet
+"""
+
+import os
+import sys
+from contextlib import contextmanager
+
+
+# Windows only
+if os.name == "nt":
+    import ctypes
+    import msvcrt  # noqa
+
+    class _CursorInfo(ctypes.Structure):
+        _fields_ = [("size", ctypes.c_int), ("visible", ctypes.c_byte)]
This and every function in this module pretty please!
@@ -0,0 +1,64 @@
+# 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.
+
+"""
+A utility for showing and hiding the terminal cursor on Windows and Linux, based on https://github.com/bchao1/bullet
+"""
+
+import os
+import sys
+from contextlib import contextmanager
+
+
+# Windows only
+if os.name == "nt":
+    import ctypes
+    import msvcrt  # noqa
+
+    class _CursorInfo(ctypes.Structure):
+        _fields_ = [("size", ctypes.c_int), ("visible", ctypes.c_byte)]
```suggestion
                lambda x: int(x)+1,
```
@@ -227,12 +223,10 @@ def get_cluster_input():
         if use_fsdp:
             distributed_type = DistributedType.FSDP
         if distributed_type == DistributedType.FSDP:
-            sharding_strategy_query = "What should be your sharding strategy ("
-            for i, strategy in enumerate(FSDP_SHARDING_STRATEGY):
-                sharding_strategy_query += f"[{i+1}] {strategy}, "
-            sharding_strategy_query = sharding_strategy_query[:-2] + ")? [1]: "
-            fsdp_config["fsdp_sharding_strategy"] = _ask_field(
+            sharding_strategy_query = "What should be your sharding strategy?"
+            fsdp_config["fsdp_sharding_strategy"] = _ask_options(
                 sharding_strategy_query,
+                FSDP_SHARDING_STRATEGY,
                 lambda x: int(x),
Because it's 1 based numbering
@@ -227,12 +223,10 @@ def get_cluster_input():
         if use_fsdp:
             distributed_type = DistributedType.FSDP
         if distributed_type == DistributedType.FSDP:
-            sharding_strategy_query = "What should be your sharding strategy ("
-            for i, strategy in enumerate(FSDP_SHARDING_STRATEGY):
-                sharding_strategy_query += f"[{i+1}] {strategy}, "
-            sharding_strategy_query = sharding_strategy_query[:-2] + ")? [1]: "
-            fsdp_config["fsdp_sharding_strategy"] = _ask_field(
+            sharding_strategy_query = "What should be your sharding strategy?"
+            fsdp_config["fsdp_sharding_strategy"] = _ask_options(
                 sharding_strategy_query,
+                FSDP_SHARDING_STRATEGY,
                 lambda x: int(x),
Missing the conversion function here.
@@ -460,10 +463,9 @@ def get_cluster_input():
         if distributed_type == DistributedType.DEEPSPEED and use_deepspeed_config:
             mixed_precision = "no"
         else:
-            mixed_precision = _ask_field(
-                "Do you wish to use FP16 or BF16 (mixed precision)? [NO/fp16/bf16]: ",
-                lambda x: str(x).lower(),
-                default="no",
+            mixed_precision = _ask_options(
+                "Do you wish to use FP16 or BF16 (mixed precision)?",
+                ["no", "fp16", "bf16"],
Thanks! Fixed here and for sagemaker
@@ -460,10 +463,9 @@ def get_cluster_input():
         if distributed_type == DistributedType.DEEPSPEED and use_deepspeed_config:
             mixed_precision = "no"
         else:
-            mixed_precision = _ask_field(
-                "Do you wish to use FP16 or BF16 (mixed precision)? [NO/fp16/bf16]: ",
-                lambda x: str(x).lower(),
-                default="no",
+            mixed_precision = _ask_options(
+                "Do you wish to use FP16 or BF16 (mixed precision)?",
+                ["no", "fp16", "bf16"],
This line (297) should have been removed
@@ -284,23 +293,18 @@ def get_cluster_input():
                     lambda x: int(x),
                     default=1e8,
                 )
-            fsdp_backward_prefetch_query = "What should be your FSDP's backward prefetch policy ("
-            for i, backward_prefetch_policy in enumerate(FSDP_BACKWARD_PREFETCH):
-                fsdp_backward_prefetch_query += f"[{i}] {backward_prefetch_policy}, "
+            fsdp_backward_prefetch_query = "What should be your FSDP's backward prefetch policy?"
             fsdp_backward_prefetch_query = fsdp_backward_prefetch_query[:-2] + ")? [0]: "
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_830). All of your documentation changes will be reflected on that endpoint.
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_830). All of your documentation changes will be reflected on that endpoint.
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_830). All of your documentation changes will be reflected on that endpoint.
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_830). All of your documentation changes will be reflected on that endpoint.
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_830). All of your documentation changes will be reflected on that endpoint.
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_830). All of your documentation changes will be reflected on that endpoint.
https://api.github.com/repos/huggingface/accelerate/pulls/830
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 63ea4a71b..ad22c8ee4 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -36,6 +36,7 @@
     DeepSpeedPlugin,
     DistributedDataParallelKwargs,
     DistributedType,
+    DynamoBackend,
     FullyShardedDataParallelPlugin,
     GradScalerKwargs,
     InitProcessGroupKwargs,
@@ -163,6 +164,8 @@ class Accelerator:
         kwargs_handlers (`List[KwargHandler]`, *optional*)
             A list of `KwargHandler` to customize how the objects related to distributed training or mixed precision
             are created. See [kwargs](kwargs) for more information.
+        dynamo_backend (`str` or `DynamoBackend`, *optional*, defaults to `"no"`):
+            Set to one of the possible dynamo backends to optimize your training with torch dynamo.
 
     **Available attributes:**
 
@@ -198,6 +201,7 @@ def __init__(
         even_batches: bool = True,
         step_scheduler_with_optimizer: bool = True,
         kwargs_handlers: Optional[List[KwargsHandler]] = None,
+        dynamo_backend: Union[DynamoBackend, str] = None,
     ):
         self.logging_dir = logging_dir
         trackers = filter_trackers(log_with, self.logging_dir)
@@ -219,6 +223,9 @@ def __init__(
             )
             mixed_precision = "fp16"
 
+        if dynamo_backend is not None:
+            dynamo_backend = DynamoBackend(dynamo_backend.upper())
+
         if deepspeed_plugin is None:  # init from env variables
             deepspeed_plugin = DeepSpeedPlugin() if os.environ.get("USE_DEEPSPEED", "false") == "true" else None
         else:
@@ -285,6 +292,7 @@ def __init__(
         self.state = AcceleratorState(
             mixed_precision=mixed_precision,
             cpu=cpu,
+            dynamo_backend=dynamo_backend,
             deepspeed_plugin=deepspeed_plugin,
             fsdp_plugin=fsdp_plugin,
             megatron_lm_plugin=megatron_lm_plugin,
@@ -793,6 +801,10 @@ def prepare_model(self, model: torch.nn.Module, device_placement=None):
         self._models.append(model)
         if device_placement:
             model = model.to(self.device)
+        if self.state.dynamo_backend != DynamoBackend.NO:
+            import torch._dynamo as dynamo
+
+            model = dynamo.optimize(self.state.dynamo_backend.value.lower())(model)
         if self.distributed_type == DistributedType.MULTI_GPU:
             if any(p.requires_grad for p in model.parameters()):
                 kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {}
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index aa61a6857..5f032debc 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -16,7 +16,13 @@
 
 import os
 
-from ...utils import ComputeEnvironment, DistributedType, is_deepspeed_available, is_transformers_available
+from ...utils import (
+    ComputeEnvironment,
+    DistributedType,
+    DynamoBackend,
+    is_deepspeed_available,
+    is_transformers_available,
+)
 from ...utils.constants import (
     DEEPSPEED_MULTINODE_LAUNCHERS,
     FSDP_AUTO_WRAP_POLICY,
@@ -25,7 +31,7 @@
     FSDP_STATE_DICT_TYPE,
 )
 from .config_args import ClusterConfig
-from .config_utils import _ask_field, _convert_distributed_mode, _convert_yes_no_to_bool
+from .config_utils import _ask_field, _convert_distributed_mode, _convert_dynamo_backend, _convert_yes_no_to_bool
 
 
 def get_cluster_input():
@@ -89,6 +95,22 @@ def get_cluster_input():
     else:
         use_cpu = False
 
+    use_dynamo = _ask_field(
+        "Do you wish to optimize your script with torch dynamo?[yes/NO]:",
+        _convert_yes_no_to_bool,
+        default=False,
+        error_message="Please enter yes or no.",
+    )
+    if use_dynamo:
+        dynamo_backend = _ask_field(
+            "Which dynamo backend would you like to use? ([0] eager, [1] aot_eager, [2] inductor, [3] nvfuser, [5] aot_nvfuser, [6] aot_cudagraphs, [7] ofi, [8] onnxrt, [9] ipex) [2]: ",
+            _convert_dynamo_backend,
+            default=DynamoBackend.INDUCTOR,
+            error_message="Please enter 0, 1, 2, 3, 4, 5, 6, 7, 8 or 9.",
+        )
+    else:
+        dynamo_backend = DynamoBackend.NO
+
     deepspeed_config = {}
     if distributed_type in [DistributedType.MULTI_GPU, DistributedType.NO]:
         use_deepspeed = _ask_field(
@@ -446,6 +468,11 @@ def get_cluster_input():
     else:
         mixed_precision = "no"
 
+    if use_dynamo and mixed_precision == "no" and not use_cpu:
+        print(
+            "Torch dynamo used without mixed precision requires TF32 to be efficient. Accelerate will enable it by default when launching your scripts."
+        )
+
     downcast_bf16 = "no"
     if distributed_type == DistributedType.TPU and mixed_precision == "bf16":
         downcast_bf16 = _ask_field(
@@ -474,4 +501,5 @@ def get_cluster_input():
         tpu_zone=tpu_zone,
         commands=commands,
         command_file=command_file,
+        dynamo_backend=dynamo_backend,
     )
diff --git a/src/accelerate/commands/config/config_args.py b/src/accelerate/commands/config/config_args.py
index a5a841900..ba492802e 100644
--- a/src/accelerate/commands/config/config_args.py
+++ b/src/accelerate/commands/config/config_args.py
@@ -22,7 +22,7 @@
 
 import yaml
 
-from ...utils import ComputeEnvironment, DistributedType, SageMakerDistributedType
+from ...utils import ComputeEnvironment, DistributedType, DynamoBackend, SageMakerDistributedType
 from ...utils.constants import SAGEMAKER_PYTHON_VERSION, SAGEMAKER_PYTORCH_VERSION, SAGEMAKER_TRANSFORMERS_VERSION
 
 
@@ -70,6 +70,7 @@ class BaseConfig:
     distributed_type: Union[DistributedType, SageMakerDistributedType]
     mixed_precision: str
     use_cpu: bool
+    dynamo_backend: DynamoBackend
 
     def to_dict(self):
         result = self.__dict__
@@ -92,6 +93,8 @@ def from_json_file(cls, json_file=None):
             del config_dict["fp16"]
         if "use_cpu" not in config_dict:
             config_dict["use_cpu"] = False
+        if "dynamo_backend" not in config_dict:
+            config_dict["dynamo_backend"] = DynamoBackend.NO
         return cls(**config_dict)
 
     def to_json_file(self, json_file):
@@ -113,6 +116,8 @@ def from_yaml_file(cls, yaml_file=None):
             del config_dict["fp16"]
         if "use_cpu" not in config_dict:
             config_dict["use_cpu"] = False
+        if "dynamo_backend" not in config_dict:
+            config_dict["dynamo_backend"] = DynamoBackend.NO
 
         return cls(**config_dict)
 
@@ -128,6 +133,8 @@ def __post_init__(self):
                 self.distributed_type = SageMakerDistributedType(self.distributed_type)
             else:
                 self.distributed_type = DistributedType(self.distributed_type)
+        if isinstance(self.dynamo_backend, str):
+            self.dynamo_backend = DynamoBackend(self.dynamo_backend.upper())
 
 
 @dataclass
diff --git a/src/accelerate/commands/config/config_utils.py b/src/accelerate/commands/config/config_utils.py
index 9dd1f4c99..8974cddc9 100644
--- a/src/accelerate/commands/config/config_utils.py
+++ b/src/accelerate/commands/config/config_utils.py
@@ -14,7 +14,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ...utils.dataclasses import ComputeEnvironment, DistributedType, SageMakerDistributedType
+from ...utils.dataclasses import ComputeEnvironment, DistributedType, DynamoBackend, SageMakerDistributedType
+
+
+DYNAMO_BACKENDS = [
+    "EAGER",
+    "AOT_EAGER",
+    "INDUCTOR",
+    "NVFUSER",
+    "AOT_NVFUSER",
+    "AOT_CUDAGRAPHS",
+    "OFI",
+    "FX2TRT",
+    "ONNXRT",
+    "IPEX",
+]
 
 
 def _ask_field(input_text, convert_value=None, default=None, error_message=None):
@@ -40,6 +54,11 @@ def _convert_distributed_mode(value):
     return DistributedType(["NO", "MULTI_CPU", "MULTI_GPU", "TPU", "MPS"][value])
 
 
+def _convert_dynamo_backend(value):
+    value = int(value)
+    return DynamoBackend(DYNAMO_BACKENDS[value + 1])
+
+
 def _convert_sagemaker_distributed_mode(value):
     value = int(value)
     return SageMakerDistributedType(["NO", "DATA_PARALLEL", "MODEL_PARALLEL"][value])
diff --git a/src/accelerate/commands/config/sagemaker.py b/src/accelerate/commands/config/sagemaker.py
index b3a45c9e4..71de4fea9 100644
--- a/src/accelerate/commands/config/sagemaker.py
+++ b/src/accelerate/commands/config/sagemaker.py
@@ -17,10 +17,15 @@
 import os
 
 from ...utils.constants import SAGEMAKER_PARALLEL_EC2_INSTANCES
-from ...utils.dataclasses import ComputeEnvironment, SageMakerDistributedType
+from ...utils.dataclasses import ComputeEnvironment, DynamoBackend, SageMakerDistributedType
 from ...utils.imports import is_boto3_available
 from .config_args import SageMakerConfig
-from .config_utils import _ask_field, _convert_sagemaker_distributed_mode, _convert_yes_no_to_bool
+from .config_utils import (
+    _ask_field,
+    _convert_dynamo_backend,
+    _convert_sagemaker_distributed_mode,
+    _convert_yes_no_to_bool,
+)
 
 
 if is_boto3_available():
@@ -162,6 +167,22 @@ def get_sagemaker_input():
         error_message="Please enter 0 or 1",
     )
 
+    use_dynamo = _ask_field(
+        "Do you wish to optimize your script with torch dynamo?[yes/NO]:",
+        _convert_yes_no_to_bool,
+        default=False,
+        error_message="Please enter yes or no.",
+    )
+    if use_dynamo:
+        dynamo_backend = _ask_field(
+            "Which dynamo backend would you like to use? ([0] eager, [1] aot_eager, [2] inductor, [3] nvfuser, [5] aot_nvfuser, [6] aot_cudagraphs, [7] ofi, [8] onnxrt, [9] ipex) [2]: ",
+            _convert_dynamo_backend,
+            default=DynamoBackend.INDUCTOR,
+            error_message="Please enter 0, 1, 2, 3, 4, 5, 6, 7, 8 or 9.",
+        )
+    else:
+        dynamo_backend = DynamoBackend.NO
+
     ec2_instance_query = "Which EC2 instance type you want to use for your training "
     if distributed_type != SageMakerDistributedType.NO:
         ec2_instance_query += "("
@@ -186,15 +207,21 @@ def get_sagemaker_input():
 
     mixed_precision = _ask_field(
         "Do you wish to use FP16 or BF16 (mixed precision)? [No/FP16/BF16]: ",
-        lambda x: str(x),
-        default="No",
+        lambda x: str(x).lower(),
+        default="no",
     )
 
+    if use_dynamo and mixed_precision == "no":
+        print(
+            "Torch dynamo used without mixed precision requires TF32 to be efficient. Accelerate will enable it by default when launching your scripts."
+        )
+
     return SageMakerConfig(
         image_uri=docker_image,
         compute_environment=ComputeEnvironment.AMAZON_SAGEMAKER,
         distributed_type=distributed_type,
         use_cpu=False,
+        dynamo_backend=dynamo_backend,
         ec2_instance_type=ec2_instance_type,
         profile=aws_profile,
         region=aws_region,
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 46d95c713..d96ab2f45 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -30,10 +30,12 @@
 import psutil
 from accelerate.commands.config import default_config_file, load_config_from_file
 from accelerate.commands.config.config_args import SageMakerConfig
+from accelerate.commands.config.config_utils import DYNAMO_BACKENDS
 from accelerate.state import get_int_from_env
 from accelerate.utils import (
     ComputeEnvironment,
     DistributedType,
+    DynamoBackend,
     PrecisionType,
     PrepareForLaunch,
     _filter_args,
@@ -171,6 +173,13 @@ def launch_command_parser(subparsers=None):
     resource_args = parser.add_argument_group(
         "Resource Selection Arguments", "Arguments for fine-tuning how available hardware should be used."
     )
+    resource_args.add_argument(
+        "--dynamo_backend",
+        type=str,
+        choices=["no"] + [b.lower() for b in DYNAMO_BACKENDS],
+        help="Choose a backend to optimize your training with dynamo, see more at "
+        "https://github.com/pytorch/torchdynamo.",
+    )
     resource_args.add_argument(
         "--mixed_precision",
         type=str,
@@ -546,6 +555,13 @@ def simple_launcher(args):
         mixed_precision = "fp16"
 
     current_env["MIXED_PRECISION"] = str(mixed_precision)
+
+    try:
+        dynamo_backend = DynamoBackend(args.dynamo_backend.upper())
+    except ValueError:
+        raise ValueError(f"Unknown dynamo backend: {args.dynamo_backend.upper()}. Choose between {DYNAMO_BACKENDS}.")
+    current_env["DYNAMO_BACKEND"] = dynamo_backend.value
+
     current_env["OMP_NUM_THREADS"] = str(args.num_cpu_threads_per_process)
 
     process = subprocess.Popen(cmd, env=current_env)
@@ -598,6 +614,13 @@ def multi_gpu_launcher(args):
         mixed_precision = "fp16"
 
     current_env["MIXED_PRECISION"] = str(mixed_precision)
+
+    try:
+        dynamo_backend = DynamoBackend(args.dynamo_backend.upper())
+    except ValueError:
+        raise ValueError(f"Unknown dynamo backend: {args.dynamo_backend.upper()}. Choose between {DYNAMO_BACKENDS}.")
+    current_env["DYNAMO_BACKEND"] = dynamo_backend.value
+
     if args.use_fsdp:
         current_env["USE_FSDP"] = "true"
         current_env["FSDP_SHARDING_STRATEGY"] = str(args.fsdp_sharding_strategy)
@@ -893,10 +916,16 @@ def sagemaker_launcher(sagemaker_config: SageMakerConfig, args):
         warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', FutureWarning)
         mixed_precision = "fp16"
 
+    try:
+        dynamo_backend = DynamoBackend(args.dynamo_backend.upper())
+    except ValueError:
+        raise ValueError(f"Unknown dynamo backend: {args.dynamo_backend.upper()}. Choose between {DYNAMO_BACKENDS}.")
+
     # Environment variables to be set for use during training job
     environment = {
         "USE_SAGEMAKER": "true",
         "MIXED_PRECISION": str(mixed_precision),
+        "DYNAMO_BACKEND": dynamo_backend.value,
         "SAGEMAKER_DISTRIBUTED_TYPE": sagemaker_config.distributed_type.value,
     }
     # configure distribution set up
@@ -1017,6 +1046,9 @@ def launch_command(args):
                 args.mixed_precision = "fp16"
             else:
                 args.mixed_precision = defaults.mixed_precision
+        if args.dynamo_backend is None:
+            warned.append("\t`--dynamo_backend` was set to a value of `'no'`")
+            args.dynamo_backend = "no"
     else:
         if args.num_processes is None:
             args.num_processes = torch.cuda.device_count() if args.multi_gpu else 1
@@ -1029,6 +1061,9 @@ def launch_command(args):
             args.mixed_precision = "no"
         if not hasattr(args, "use_cpu"):
             args.use_cpu = args.cpu
+        if args.dynamo_backend is None:
+            warned.append("\t`--dynamo_backend` was set to a value of `'no'`")
+            args.dynamo_backend = "no"
 
     if args.num_cpu_threads_per_process is None:
         args.num_cpu_threads_per_process = 1
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index 00ba3ab67..18df61bbc 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -19,6 +19,7 @@
 
 from .utils import (
     DistributedType,
+    DynamoBackend,
     get_ccl_version,
     get_int_from_env,
     is_ccl_available,
@@ -57,6 +58,7 @@ def __init__(
         self,
         mixed_precision: str = None,
         cpu: bool = False,
+        dynamo_backend=None,
         deepspeed_plugin=None,
         fsdp_plugin=None,
         megatron_lm_plugin=None,
@@ -74,6 +76,10 @@ def __init__(
             mixed_precision = (
                 parse_choice_from_env("MIXED_PRECISION", "no") if mixed_precision is None else mixed_precision.lower()
             )
+            dynamo_backend = (
+                parse_choice_from_env("DYNAMO_BACKEND", "no") if dynamo_backend is None else dynamo_backend
+            )
+            self.dynamo_backend = DynamoBackend(dynamo_backend.upper())
             if not _from_accelerator:
                 raise ValueError(
                     "Please make sure to properly initialize your accelerator via `accelerator = Accelerator()` "
@@ -230,6 +236,9 @@ def __init__(
                 else:
                     self.device = torch.device("cuda")
                 self.mixed_precision = mixed_precision
+
+            if self.dynamo_backend != DynamoBackend.NO and self.mixed_precision == "no" and self.device.type == "cuda":
+                torch.backends.cuda.matmul.allow_tf32 = True
             self.initialized = True
 
     def __repr__(self):
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
index a017a2ad5..3ad09ecab 100644
--- a/src/accelerate/utils/__init__.py
+++ b/src/accelerate/utils/__init__.py
@@ -8,6 +8,7 @@
     DeepSpeedPlugin,
     DistributedDataParallelKwargs,
     DistributedType,
+    DynamoBackend,
     FullyShardedDataParallelPlugin,
     GradScalerKwargs,
     InitProcessGroupKwargs,
diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py
index 7dd8798e9..e51576865 100644
--- a/src/accelerate/utils/dataclasses.py
+++ b/src/accelerate/utils/dataclasses.py
@@ -164,6 +164,50 @@ class ComputeEnvironment(str, enum.Enum):
     AMAZON_SAGEMAKER = "AMAZON_SAGEMAKER"
 
 
+class DynamoBackend(str, enum.Enum):
+    """
+    Represents a dynamo backend (see https://github.com/pytorch/torchdynamo).
+
+    Values:
+
+        - **NO** -- Do not use torch dynamo.
+        - **EAGER** -- Uses PyTorch to run the extracted GraphModule. This is quite useful in debugging TorchDynamo
+          issues.
+        - **AOT_EAGER** -- Uses AotAutograd with no compiler, i.e, just using PyTorch eager for the AotAutograd's
+          extracted forward and backward graphs. This is useful for debugging, and unlikely to give speedups.
+        - **INDUCTOR** -- Uses TorchInductor backend with AotAutograd and cudagraphs by leveraging codegened Triton
+          kernels. [Read
+          more](https://dev-discuss.pytorch.org/t/torchinductor-a-pytorch-native-compiler-with-define-by-run-ir-and-symbolic-shapes/747)
+        - **NVFUSER** -- nvFuser with TorchScript. [Read
+          more](https://dev-discuss.pytorch.org/t/tracing-with-primitives-update-1-nvfuser-and-its-primitives/593)
+        - **AOT_NVFUSER** -- nvFuser with AotAutograd. [Read
+          more](https://dev-discuss.pytorch.org/t/tracing-with-primitives-update-1-nvfuser-and-its-primitives/593)
+        - **AOT_CUDAGRAPHS** -- cudagraphs with AotAutograd. [Read
+          more](https://github.com/pytorch/torchdynamo/pull/757)
+        - **OFI** -- Uses Torchscript optimize_for_inference. Inference only. [Read
+          more](https://pytorch.org/docs/stable/generated/torch.jit.optimize_for_inference.html)
+        - **FX2TRT** -- Uses Nvidia TensorRT for inference optimizations. Inference only. [Read
+          more](https://github.com/pytorch/TensorRT/blob/master/docsrc/tutorials/getting_started_with_fx_path.rst)
+        - **ONNXRT** -- Uses ONNXRT for inference on CPU/GPU. Inference only. [Read more](https://onnxruntime.ai/)
+        - **IPEX** -- Uses IPEX for inference on CPU. Inference only. [Read
+          more](https://github.com/intel/intel-extension-for-pytorch).
+
+    """
+
+    # Subclassing str as well as Enum allows the `SageMakerDistributedType` to be JSON-serializable out of the box.
+    NO = "NO"
+    EAGER = "EAGER"
+    AOT_EAGER = "AOT_EAGER"
+    INDUCTOR = "INDUCTOR"
+    NVFUSER = "NVFUSER"
+    AOT_NVFUSER = "AOT_NVFUSER"
+    AOT_CUDAGRAPHS = "AOT_CUDAGRAPHS"
+    OFI = "OFI"
+    FX2TRT = "FX2TRT"
+    ONNXRT = "ONNXRT"
+    IPEX = "IPEX"
+
+
 class EnumWithContains(enum.EnumMeta):
     "A metaclass that adds the ability to check if `self` contains an item with the `in` operator"
 
diff --git a/src/accelerate/utils/other.py b/src/accelerate/utils/other.py
index ff360038d..f1af373b6 100644
--- a/src/accelerate/utils/other.py
+++ b/src/accelerate/utils/other.py
@@ -115,7 +115,7 @@ def get_pretty_name(obj):
     return str(obj)
 
 
-def write_basic_config(mixed_precision="no", save_location: str = default_json_config_file):
+def write_basic_config(mixed_precision="no", save_location: str = default_json_config_file, dynamo_backend="no"):
     """
     Creates and saves a basic cluster config to be used on a local machine with potentially multiple GPUs. Will also
     set CPU if it is a CPU-only machine.
@@ -138,7 +138,11 @@ def write_basic_config(mixed_precision="no", save_location: str = default_json_c
     mixed_precision = mixed_precision.lower()
     if mixed_precision not in ["no", "fp16", "bf16"]:
         raise ValueError(f"`mixed_precision` should be one of 'no', 'fp16', or 'bf16'. Received {mixed_precision}")
-    config = {"compute_environment": "LOCAL_MACHINE", "mixed_precision": mixed_precision}
+    config = {
+        "compute_environment": "LOCAL_MACHINE",
+        "mixed_precision": mixed_precision,
+        "dynamo_backend": dynamo_backend,
+    }
     if torch.cuda.is_available():
         num_gpus = torch.cuda.device_count()
         config["num_processes"] = num_gpus

```suggestion
            Set to one of the possible dynamo backends to optimize your training with torch dynamo.
```

@@ -163,6 +164,8 @@ class Accelerator:
         kwargs_handlers (`List[KwargHandler]`, *optional*)
             A list of `KwargHandler` to customize how the objects related to distributed training or mixed precision
             are created. See [kwargs](kwargs) for more information.
+        dynamo_backend (`str` or `DynamoBackend`, *optional*, defaults to `"no"`):
+            Set to one of the possible dynamo backends to optimizer your training with torch dynamo.
```suggestion
            if self.dynamo_backend != DynamoBackend.NO and self.mixed_precision == "no" and self.device.type == "cuda":
```

@@ -230,6 +236,9 @@ def __init__(
                 else:
                     self.device = torch.device("cuda")
                 self.mixed_precision = mixed_precision
+
+            if self.dynamo_backend != DynamoBackend.NO and self.mixed_precision != "no" and self.device.type == "cuda":
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_829). All of your documentation changes will be reflected on that endpoint.
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_829). All of your documentation changes will be reflected on that endpoint.
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_829). All of your documentation changes will be reflected on that endpoint.
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_829). All of your documentation changes will be reflected on that endpoint.
https://api.github.com/repos/huggingface/accelerate/pulls/829
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 46d95c713..9238fbca9 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -56,9 +56,6 @@
     logging.basicConfig(format=FORMAT, datefmt="[%X]", handlers=[RichHandler()])
 
 
-if is_torch_version(">=", "1.9.0"):
-    import torch.distributed.run as distrib_run
-
 logger = logging.getLogger(__name__)
 
 options_to_group = {
@@ -555,6 +552,8 @@ def simple_launcher(args):
 
 
 def multi_gpu_launcher(args):
+    if is_torch_version(">=", "1.9.0"):
+        import torch.distributed.run as distrib_run
     num_processes = getattr(args, "num_processes")
     num_machines = getattr(args, "num_machines")
     main_process_ip = getattr(args, "main_process_ip")
@@ -644,6 +643,8 @@ def multi_gpu_launcher(args):
 
 
 def deepspeed_launcher(args):
+    if is_torch_version(">=", "1.9.0"):
+        import torch.distributed.run as distrib_run
     if not is_deepspeed_available():
         raise ImportError("DeepSpeed is not installed => run `pip3 install deepspeed` or build it from source.")
     num_processes = getattr(args, "num_processes")
diff --git a/src/accelerate/utils/launch.py b/src/accelerate/utils/launch.py
index 8642a441a..c65df08dc 100644
--- a/src/accelerate/utils/launch.py
+++ b/src/accelerate/utils/launch.py
@@ -21,10 +21,6 @@
 from .dataclasses import DistributedType
 
 
-if is_torch_version(">=", "1.9.0"):
-    import torch.distributed.run as distrib_run
-
-
 def get_launch_prefix():
     """
     Grabs the correct launcher for starting a distributed command, such as either `torchrun`, `python -m
@@ -43,6 +39,8 @@ def _filter_args(args):
     """
     Filters out all `accelerate` specific args
     """
+    if is_torch_version(">=", "1.9.0"):
+        import torch.distributed.run as distrib_run
     distrib_args = distrib_run.get_args_parser()
     new_args, _ = distrib_args.parse_known_args()
 

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_828). All of your documentation changes will be reflected on that endpoint.
https://api.github.com/repos/huggingface/accelerate/pulls/828
diff --git a/.github/workflows/build_pr_documentation.yml b/.github/workflows/build_pr_documentation.yml
index dc56751c6..12c600cb7 100644
--- a/.github/workflows/build_pr_documentation.yml
+++ b/.github/workflows/build_pr_documentation.yml
@@ -9,8 +9,11 @@ concurrency:
 
 jobs:
   build:
-    uses: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml@main
+    uses: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml@use_hf_hub
     with:
       commit_sha: ${{ github.event.pull_request.head.sha }}
       pr_number: ${{ github.event.number }}
       package: accelerate
+    secrets:
+      token: ${{ secrets.HF_DOC_PUSH }}
+      comment_bot_token: ${{ secrets.HUGGINGFACE_PUSH }}
diff --git a/.github/workflows/delete_doc_comment.yml b/.github/workflows/delete_doc_comment.yml
index da61d21df..973c2a8b3 100644
--- a/.github/workflows/delete_doc_comment.yml
+++ b/.github/workflows/delete_doc_comment.yml
@@ -7,7 +7,10 @@ on:
 
 jobs:
   delete:
-    uses: huggingface/doc-builder/.github/workflows/delete_doc_comment.yml@main
+    uses: huggingface/doc-builder/.github/workflows/delete_doc_comment.yml@use_hf_hub
     with:
       pr_number: ${{ github.event.number }}
       package: accelerate
+    secrets:
+      token: ${{ secrets.HF_DOC_PUSH }}
+      comment_bot_token: ${{ secrets.HUGGINGFACE_PUSH }}

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_827). All of your documentation changes will be reflected on that endpoint.
https://api.github.com/repos/huggingface/accelerate/pulls/827
diff --git a/setup.py b/setup.py
index 22d7cfb8d..b4708353c 100644
--- a/setup.py
+++ b/setup.py
@@ -19,7 +19,7 @@
 extras["quality"] = ["black ~= 22.0", "isort >= 5.5.4", "flake8 >= 3.8.3", "hf-doc-builder >= 0.3.0"]
 extras["docs"] = []
 extras["test_prod"] = ["pytest", "pytest-xdist", "pytest-subtests", "parameterized"]
-extras["test_dev"] = ["datasets", "evaluate", "transformers", "scipy", "sklearn", "deepspeed<0.7.0", "tqdm"]
+extras["test_dev"] = ["datasets", "evaluate", "transformers", "scipy", "scikit-learn", "deepspeed<0.7.0", "tqdm"]
 extras["testing"] = extras["test_prod"] + extras["test_dev"]
 extras["rich"] = ["rich"]
 

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/825
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index acc4a1828..aa61a6857 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -385,11 +385,23 @@ def get_cluster_input():
                     )
                     command_file = os.path.abspath(command_file)
                 else:
-                    commands = _ask_field(
-                        "What commands do you wish to run on startup in each pod? ",
-                        default=None,
-                        error_message="Please enter the commands you wish to run on startup in each pod as a single string.",
-                    )
+                    print("Please enter each command seperately you wish to run on startup in each pod.")
+                    commands = []
+                    another_command = True
+                    while another_command:
+                        commands.append(
+                            _ask_field(
+                                "Please enter a single command to be ran ",
+                                default=None,
+                                error_message="Please enter the commands you wish to run on startup in each pod as a single string.",
+                            )
+                        )
+                        another_command = _ask_field(
+                            "Do you wish to add another command? [yes/NO]: ",
+                            _convert_yes_no_to_bool,
+                            default=False,
+                            error_message="Please enter yes or no.",
+                        )
 
     else:
         main_training_function = "main"
diff --git a/src/accelerate/commands/config/config_args.py b/src/accelerate/commands/config/config_args.py
index 9a1247c55..a5a841900 100644
--- a/src/accelerate/commands/config/config_args.py
+++ b/src/accelerate/commands/config/config_args.py
@@ -155,7 +155,7 @@ class ClusterConfig(BaseConfig):
     tpu_name: str = None
     tpu_zone: str = None
     command_file: str = None
-    command: List[str] = None
+    commands: List[str] = None
 
     def __post_init__(self):
         if self.deepspeed_config is None:
diff --git a/src/accelerate/commands/tpu.py b/src/accelerate/commands/tpu.py
index 59bbb08e9..6b90770c7 100644
--- a/src/accelerate/commands/tpu.py
+++ b/src/accelerate/commands/tpu.py
@@ -89,8 +89,8 @@ def tpu_command_launcher(args):
         defaults = load_config_from_file(args.config_file)
         if not args.command_file and defaults.command_file is not None and not args.command:
             args.command_file = defaults.command_file
-        if not args.command and defaults.command is not None:
-            args.command = defaults.command
+        if not args.command and defaults.commands is not None:
+            args.command = defaults.commands
         if not args.tpu_name:
             args.tpu_name = defaults.tpu_name
         if not args.tpu_zone:
@@ -110,7 +110,8 @@ def tpu_command_launcher(args):
             args.command = [f.read().splitlines()]
 
     # To turn list of lists into list of strings
-    args.command = [line for cmd in args.command for line in cmd]
+    if isinstance(args.command[0], list):
+        args.command = [line for cmd in args.command for line in cmd]
     # Default to the shared folder and install accelerate
     new_cmd = ["cd /usr/share"]
     if args.install_accelerate:
diff --git a/tests/test_configs/latest.yaml b/tests/test_configs/latest.yaml
index de6be03a4..456348ef9 100644
--- a/tests/test_configs/latest.yaml
+++ b/tests/test_configs/latest.yaml
@@ -17,5 +17,5 @@ same_network: true
 use_cpu: false
 tpu_name: 'test-tpu'
 tpu_zone: 'us-central1-a'
-command: null
+commands: null
 command_file: tests/test_samples/test_command_file.sh
\ No newline at end of file

_The documentation is not available anymore as the PR was closed or merged._
@sgugger fails are unrelated, see https://github.com/huggingface/accelerate/pull/825
https://api.github.com/repos/huggingface/accelerate/pulls/824
diff --git a/docs/source/usage_guides/tracking.mdx b/docs/source/usage_guides/tracking.mdx
index cc5c17418..b22a4dd14 100644
--- a/docs/source/usage_guides/tracking.mdx
+++ b/docs/source/usage_guides/tracking.mdx
@@ -17,11 +17,12 @@ There are a large number of experiment tracking API's available, however getting
 
 ## Integrated Trackers
 
-Currently `Accelerate` supports three trackers out-of-the-box:
+Currently `Accelerate` supports four trackers out-of-the-box:
 
 - TensorBoard
 - WandB
 - CometML
+- MLFlow
 
 To use any of them, pass in the selected type(s) to the `log_with` parameter in [`Accelerate`]:
 ```python
diff --git a/docs/source/usage_guides/training_zoo.mdx b/docs/source/usage_guides/training_zoo.mdx
index 444784b49..176d97d2b 100644
--- a/docs/source/usage_guides/training_zoo.mdx
+++ b/docs/source/usage_guides/training_zoo.mdx
@@ -21,7 +21,10 @@ Below contains a non-exhuastive list of tutorials and scripts showcasing Acceler
 These examples showcase the base features of Accelerate and are a great starting point
 
 - [Barebones NLP example](https://github.com/huggingface/accelerate/blob/main/examples/nlp_example.py)
+- [Barebones distributed NLP example in a Jupyter Notebook](https://github.com/huggingface/notebooks/blob/main/examples/accelerate_examples/simple_nlp_example.ipynb)
 - [Barebones computer vision example](https://github.com/huggingface/accelerate/blob/main/examples/cv_example.py)
+- [Barebones distributed computer vision example in a Jupyter Notebook](https://github.com/huggingface/notebooks/blob/main/examples/accelerate_examples/simple_cv_example.ipynb)
+- [Using Accelerate in Kaggle](https://www.kaggle.com/code/muellerzr/multi-gpu-and-accelerate)
 
 ### Feature Specific Examples
 
@@ -101,6 +104,10 @@ These are tutorials from libraries that integrate with 🤗 Accelerate:
 
 - [Perform Deep Learning with 3D data](https://pytorch3d.org/tutorials/)
 
+### Stable-Dreamfusion
+
+- [Training with Stable-Dreamfusion to convert text to a 3D model](https://colab.research.google.com/drive/1MXT3yfOFvO0ooKEfiUUvTKwUkrrlCHpF?usp=sharing)
+
 ### Tez 
 
 - [Leaf disease detection with Tez and Accelerate](https://www.kaggle.com/code/abhishek/tez-faster-and-easier-training-for-leaf-detection/notebook)

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/823
diff --git a/examples/by_feature/deepspeed_with_config_support.py b/examples/by_feature/deepspeed_with_config_support.py
index 36ace84b9..0fdff58ae 100755
--- a/examples/by_feature/deepspeed_with_config_support.py
+++ b/examples/by_feature/deepspeed_with_config_support.py
@@ -285,10 +285,9 @@ def evaluate(args, model, eval_dataloader, accelerator, eval_dataset):
             outputs = model(**batch)
 
         loss = outputs.loss
-        losses.append(accelerator.gather(loss.repeat(args.per_device_eval_batch_size)))
+        losses.append(accelerator.gather_for_metrics(loss.repeat(args.per_device_eval_batch_size)))
 
     losses = torch.cat(losses)
-    losses = losses[: len(eval_dataset)]
     try:
         eval_loss = torch.mean(losses)
         perplexity = math.exp(eval_loss)

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

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

One thing that I was thinking about whilst working on this though is that, although uneven batches is likely only going to be needed during evaluation, the approach that we have taken with setting a variable at the accelerator level is that it will apply to both training and validation. So an example would look something like:

```
accelerator = Accelerator(even_batches=True)

model, optimizer, train_dataloader, val_dataloader = create_components()

ddp_model, optimizer, train_dataloader = accelerator.prepare(model, optimizer, train_dataloader)

# training loop as normal
...


accelerator.even_batches = False

val_dataloader = accelerator.prepare(val_dataloader)

# run eval loop
with accelerator.join_uneven_inputs([ddp_model]):
    ...  
```
It's a bit clunky, but as it is an advanced use case, it may be acceptable.

I'm interested to hear any thoughts that you have!
@Chris-hughes10 @sgugger what do you think about something like this. Here's an MVP example showing the basic behavior:

```python
from contextlib import contextmanager

class A:
    "A mock version of what we would call the dataloader"
    def __init__(self, a:int):
        self.a = a

class B:
    "A mock version of what we would call the Accelerator"
    @contextmanager
    def do_the_thing(self, a):
        old_a = a.a
        a.a = 2
        try:
            yield
        finally:
            a.a = old_a

a = A(1)
b = B()

with b.do_the_thing(a):
    print(a.a)
print(a.a)
```
You'd find that `a.a` prints 2 in the first and 1 in the second. We can apply that similarly here since we expect if the user is calling this that they will be training/evaluating/doing things on uneven inputs, leading to an API such as:

```python
with accelerator.join_uneven_inputs([model], dataloader=eval_dataloader):
   ...
```
Inside this ctx manager `a.a` would be `eval_dataloader.even_batches`, and if it wasn't disabled originally we can print for the user that using this context manager will temporarily change this value or something. This let's us get rid of the explicit of setting in the accelerator and having to prepare later, so the only code change is the context manager. 

Lmk what you both think

Otherwise I think you've done a great job with it! Super small API design nit of my own 😄 

(specifically this would do `eval_dataloader.even_batches=False` during the `a.a` bits in the context manager example)

It could probably go as far as accept a tuple of models and dataloaders and Accelerate could figure out which one is which and apply them rather than having to pass in a list 🤔 Decisions decisions :) 

If wanted I could do some of this in a follow up PR, as this PR is great already
Agreed, maybe remove some of the magic with an `override` arg as well (that defaults to `False`) and a very explicit docstring as middle ground. (Also would love others opinions 😄 )
Hi both, I have given this some thought and have a couple of ideas. 

My first reaction was to agree with @sgugger, changing the `even_batches` value in the dataloader seems a little too magical, but it definitely makes the particular use case smoother. I think the main irk that I have with this is that it is changing a value on an unrelated object; if it was changing the value on the accelerator itself, I think it would be quite clean. Additionally, it seems a bit like a side effect, as it isn't really part of the same operation as the join.

One thought that I had, is that we could expose a method on the dataloader to make it easy to toggle this behaviour. Perhaps something like:
```
accelerator = Accelerator()

model, optimizer, train_dataloader, val_dataloader = create_components()

ddp_model, optimizer, train_dataloader, val_dataloader = accelerator.prepare(model, optimizer, train_dataloader, val_dataloader)

# training loop as normal
...


val_dataloader.ensure_even_batches(enabled=False)

# run eval loop
with accelerator.join_uneven_inputs([ddp_model]):
    ...  
```
There is still an extra step, but then it is very explicit what is going on. Conversely, this does mean that the flag in the accelerator is a little redundant. Also, it probably means that, in most cases, setting this flag and the context manager will have to be used together.

As an alternative, if we would like to keep everything accelerator-centric, we could potentially have the accelerator maintain a reference to any prepared dataloaders. This way, if we change the setting on the accelerator, we could change the value on all dataloaders. This could look something like:

```
accelerator = Accelerator(even_batches=True)

model, optimizer, train_dataloader, val_dataloader = create_components()

ddp_model, optimizer, train_dataloader, val_dataloader = accelerator.prepare(model, optimizer, train_dataloader, val_dataloader)

# training loop as normal
...

# run eval loop
with accelerator.join_uneven_inputs([ddp_model], even_batches=False):
    ...  
```
Personally, I think that this is quite clean, but it does mean adding additional logic to the accelerator, and perhaps the actual difference between this solution and the one proposed by @muellerzr is a bit academic.

EDIT: cc @pacman100 , I think we commented essentially at the same time. What do you think about the second approach proposed here?
Yes, I like the second approach proposed above ✨
@Chris-hughes10 I too like that solution, as it keeps in line with how we design bits like this such as the gradient accumulation wrapper 😄 It also keeps the dataloader preparation the same as optimizer and model in the sense of pointers. This just means that there would also need to be a change to `accelerator.free_memory()` to also free the dataloader references. 

I also think it should just be an `Accelerator` only attribute, there wouldn't be a need to put this in `AcceleratorState` since we deem that too magical already 😃 
Works for me too!
It looks like there are failing tests that are unrelated to this PR. Perhaps https://github.com/huggingface/accelerate/pull/825 will solve the problem.
@Chris-hughes10 correct, you can ignore the "checkpoint" and "rest" failures here. 
I have made some more updates, so I think it should be feature complete now.

Unfortunately, I am getting a couple of failing tests on the CI, which look like they occur in the bf16 section of the test script, comparing training on CPU and GPU. I am unable to reproduce locally as I don't have a device with bf16 available, and the rest of the tests in that script seem to pass. 

It is quite strange though, because looking at what it is failing, it doesn't really interact with anything in this PR other than the accelerator maintaining references to the dataloader; which doesn't seem to be a problem for the other training configurations in the script. After noticing that the accelerator state is cleared multiple times in the test - which does not include the dataloaders, as discussed above - I also tried adding resetting the accelerator manually before the bf16 section, using `accelerator.free_memory`, but this did not make any difference. 

Honesty, I am stumped at this point. Do you guys have any ideas on what could be the issue?
@Chris-hughes10 I'll be able to dig into this further tomorrow and see if I can't root out what's going on, so don't stress on it today 😃 
Hello @Chris-hughes10 , I went over the PR and tests. I found the cause of the tests failing. In the below screenshot please find the changes required in `test_accelerator.py` to resolve failing tests.

<img width="764" alt="Screenshot 2022-11-10 at 7 28 32 PM" src="https://user-images.githubusercontent.com/13534540/201111269-367a3afa-5957-4b3e-8288-bbdd11107393.png">

> Hello @Chris-hughes10 , I went over the PR and tests. I found the cause of the tests failing. In the below screenshot please find the changes required in `test_accelerator.py` to resolve failing tests.
> 


Awesome, thanks for this! I didn't anticipate that the state would carry over outside of each test, but I suppose it makes sense for a singleton. I will make the change.
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_820). All of your documentation changes will be reflected on that endpoint.
There seems to be a problem with building the PR documentation, is this something that I can fix?
@Chris-hughes10 nope don't worry about that, it's a problem on our side :)
Thanks a bunch @Chris-hughes10, excellent work! I'll be keeping an eye on our CI that runs multi-gpu tests post this merge to see if anything is up and I'll let you know if there is and I can't seem to figure out what may be failing quickly :) 
Thanks for helping to get this merged @muellerzr @sgugger @pacman100! 
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_820). All of your documentation changes will be reflected on that endpoint.
https://api.github.com/repos/huggingface/accelerate/pulls/820
diff --git a/docs/source/usage_guides/megatron_lm.mdx b/docs/source/usage_guides/megatron_lm.mdx
index d26ecd4b5..23b024f75 100644
--- a/docs/source/usage_guides/megatron_lm.mdx
+++ b/docs/source/usage_guides/megatron_lm.mdx
@@ -435,6 +435,103 @@ python checkpoint_utils/megatgron_gpt2/checkpoint_reshaping_and_interoperability
 --print-checkpoint-structure
 ```
 
+## Megatron-LM GPT models support returning logits and `megatron_generate` function for text generation
+
+1. Returning logits require setting `require_logits=True` in MegatronLMPlugin as shown below. 
+These would be available on the in the last stage of pipeline.
+```python
+megatron_lm_plugin = MegatronLMPlugin(return_logits=True)
+```
+
+2. `megatron_generate` method for Megatron-LM GPT model: This will use Tensor and Pipeline Parallelism to complete 
+generations for a batch of inputs when using greedy with/without top_k/top_p sampling and for individual prompt inputs when using beam search decoding. 
+Only a subset of features of transformers generate is supported. This will help in using large models via tensor and pipeline parallelism 
+for generation (already does key-value caching and uses fused kernels by default).
+This requires data parallel size to be 1, sequence parallelism and activation checkpointing to be disabled.
+It also requires specifying path to tokenizer's vocab file and merges file. 
+Below example shows how to configure and use `megatron_generate` method for Megatron-LM GPT model.
+```python
+# specifying tokenizer's vocab and merges file
+vocab_file = os.path.join(args.resume_from_checkpoint, "vocab.json")
+merge_file = os.path.join(args.resume_from_checkpoint, "merges.txt")
+other_megatron_args = {"vocab_file": vocab_file, "merge_file": merge_file}
+megatron_lm_plugin = MegatronLMPlugin(other_megatron_args=other_megatron_args)
+
+# inference using `megatron_generate` functionality
+tokenizer.pad_token = tokenizer.eos_token
+max_new_tokens = 64
+batch_texts = [
+    "Are you human?",
+    "The purpose of life is",
+    "The arsenal was constructed at the request of",
+    "How are you doing these days?",
+]
+batch_encodings = tokenizer(batch_texts, return_tensors="pt", padding=True)
+
+# top-p sampling
+generated_tokens = model.megatron_generate(
+    batch_encodings["input_ids"],
+    batch_encodings["attention_mask"],
+    max_new_tokens=max_new_tokens,
+    top_p=0.8,
+    top_p_decay=0.5,
+    temperature=0.9,
+)
+decoded_preds = tokenizer.batch_decode(generated_tokens.cpu().numpy())
+accelerator.print(decoded_preds)
+
+# top-k sampling
+generated_tokens = model.megatron_generate(
+    batch_encodings["input_ids"],
+    batch_encodings["attention_mask"],
+    max_new_tokens=max_new_tokens,
+    top_k=50,
+    temperature=0.9,
+)
+decoded_preds = tokenizer.batch_decode(generated_tokens.cpu().numpy())
+accelerator.print(decoded_preds)
+
+# adding `bos` token at the start
+generated_tokens = model.megatron_generate(
+    batch_encodings["input_ids"], batch_encodings["attention_mask"], max_new_tokens=max_new_tokens, add_BOS=True
+)
+decoded_preds = tokenizer.batch_decode(generated_tokens.cpu().numpy())
+accelerator.print(decoded_preds)
+
+# beam search => only takes single prompt
+batch_texts = ["The purpose of life is"]
+batch_encodings = tokenizer(batch_texts, return_tensors="pt", padding=True)
+generated_tokens = model.megatron_generate(
+    batch_encodings["input_ids"],
+    batch_encodings["attention_mask"],
+    max_new_tokens=max_new_tokens,
+    num_beams=20,
+    length_penalty=1.5,
+)
+decoded_preds = tokenizer.batch_decode(generated_tokens.cpu().numpy())
+accelerator.print(decoded_preds)
+```
+
+3. An end-to-end example of using `megatron_generate` method for Megatron-LM GPT model is available at
+[megatron_gpt2_generation.py](https://github.com/pacman100/accelerate-megatron-test/blob/main/src/inference/megatron_gpt2_generation.py) with 
+config file [megatron_lm_gpt_generate_config.yaml](https://github.com/pacman100/accelerate-megatron-test/blob/main/src/Configs/megatron_lm_gpt_generate_config.yaml).
+The bash script with accelerate launch command is available at [megatron_lm_gpt_generate.sh](https://github.com/pacman100/accelerate-megatron-test/blob/main/megatron_lm_gpt_generate.sh).
+The output logs of the script are available at [megatron_lm_gpt_generate.log](https://github.com/pacman100/accelerate-megatron-test/blob/main/output_logs/megatron_lm_gpt_generate.log).
+
+## Support for ROPE and ALiBi Positional embeddings and Multi-Query Attention
+
+1. For ROPE/ALiBi attention, pass `position_embedding_type` with `("absolute" | "rotary" | "alibi")` to `MegatronLMPlugin` as shown below.
+```python
+other_megatron_args = {"position_embedding_type": "alibi"}
+megatron_lm_plugin = MegatronLMPlugin(other_megatron_args=other_megatron_args)
+```
+
+2. For Multi-Query Attention, pass `attention_head_type` with `("multihead" | "multiquery")` to `MegatronLMPlugin` as shown below.
+```python
+other_megatron_args = {"attention_head_type": "multiquery"}
+megatron_lm_plugin = MegatronLMPlugin(other_megatron_args=other_megatron_args)
+```
+
 ## Caveats
 
 1. Supports Transformers GPT2, Megatron-BERT and T5 models.
@@ -445,8 +542,12 @@ there is quite complex interplay of pipeline, tensor and data parallelsim behind
 The `model(**batch_data)` call return loss(es) averaged across the data parallel ranks.
 This is fine for most cases wherein pre-training jobs are run using Megatron-LM features and
 you can easily compute the `perplexity` using the loss. 
+For GPT model, returning logits in addition to loss(es) is supported. 
+These logits aren't gathered across data prallel ranks. Use `accelerator.utils.gather_across_data_parallel_groups`
+to gather logits across data parallel ranks. These logits along with labels can be used for computing various 
+performance metrics. 
 
-3. The main process is the last rank as the losses are available in the last stage of pipeline.
+3. The main process is the last rank as the losses/logits are available in the last stage of pipeline.
 `accelerator.is_main_process` and `accelerator.is_local_main_process` return `True` for last rank when using 
 Megatron-LM integration.
 
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
index b472ec556..0e67daf4b 100644
--- a/src/accelerate/utils/__init__.py
+++ b/src/accelerate/utils/__init__.py
@@ -107,6 +107,7 @@
     MegatronLMSchedulerWrapper,
     T5TrainStep,
     avg_losses_across_data_parallel_group,
+    gather_across_data_parallel_groups,
 )
 from .megatron_lm import initialize as megatron_lm_initialize
 from .megatron_lm import prepare_data_loader as megatron_lm_prepare_data_loader
diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py
index 8227fa5bc..5fcb5e68f 100644
--- a/src/accelerate/utils/dataclasses.py
+++ b/src/accelerate/utils/dataclasses.py
@@ -771,6 +771,18 @@ class MegatronLMPlugin:
         default=False,
         metadata={"help": "Whether to set all logging options."},
     )
+    eval_iters: int = field(
+        default=100, metadata={"help": "Number of iterations to run for evaluation validation/test for."}
+    )
+    eval_interval: int = field(
+        default=1000, metadata={"help": "Interval between running evaluation on validation set."}
+    )
+    return_logits: bool = field(
+        default=False,
+        metadata={"help": "Whether to return logits from the model."},
+    )
+
+    # custom train step args
     custom_train_step_class: Optional[Any] = field(
         default=None,
         metadata={"help": "Custom train step class."},
@@ -779,11 +791,22 @@ class MegatronLMPlugin:
         default=None,
         metadata={"help": "Custom train step kwargs."},
     )
-    eval_iters: int = field(
-        default=100, metadata={"help": "Number of iterations to run for evaluation validation/test for."}
+
+    # custom model args
+    custom_model_provider_function: Optional[Callable] = field(
+        default=None,
+        metadata={"help": "Custom model provider function."},
     )
-    eval_interval: int = field(
-        default=1000, metadata={"help": "Interval between running evaluation on validation set."}
+    custom_prepare_model_function: Optional[Callable] = field(
+        default=None,
+        metadata={"help": "Custom prepare model function."},
+    )
+
+    # remaining args such as enabling Alibi/ROPE positional embeddings,
+    # wandb logging, Multi-Query Attention, etc.
+    other_megatron_args: Optional[Dict[str, Any]] = field(
+        default=None,
+        metadata={"help": "Other Megatron-LM arguments. Please refer Megatron-LM"},
     )
 
     def __post_init__(self):
@@ -840,6 +863,8 @@ def __post_init__(self):
             self.megatron_lm_default_args["tensorboard_dir"] = self.tensorboard_dir
             if self.set_all_logging_options:
                 self.set_tensorboard_logging_options()
+        if self.other_megatron_args is not None:
+            self.megatron_lm_default_args.update(self.other_megatron_args)
 
     def set_network_size_args(self, model, batch_data=None):
         # Check if the model is either BERT, GPT or T5 else raise error
@@ -884,6 +909,8 @@ def set_network_size_args(self, model, batch_data=None):
             else:
                 self.seq_length = max_position_embeddings
             self.megatron_lm_default_args["seq_length"] = self.seq_length
+            self.megatron_lm_default_args["return_logits"] = self.return_logits
+            self.megatron_lm_default_args["tokenizer_type"] = "GPT2BPETokenizer"
         elif "t5" in model.config.model_type.lower():
             model_type_name = "t5"
             num_layers = model.config.num_layers
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index bdb53988a..85e20a746 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -13,7 +13,9 @@
 # limitations under the License.
 
 import importlib
+import os
 import sys
+from distutils.util import strtobool
 from functools import lru_cache
 
 import torch
@@ -90,10 +92,11 @@ def is_bf16_available(ignore_tpu=False):
 
 
 def is_megatron_lm_available():
-    package_exists = importlib.util.find_spec("megatron") is not None
-    if package_exists:
-        megatron_version = parse(importlib_metadata.version("megatron-lm"))
-        return compare_versions(megatron_version, ">=", "2.2.0")
+    if strtobool(os.environ.get("USE_MEGATRON_LM", "False")) == 1:
+        package_exists = importlib.util.find_spec("megatron") is not None
+        if package_exists:
+            megatron_version = parse(importlib_metadata.version("megatron-lm"))
+            return compare_versions(megatron_version, ">=", "2.2.0")
     return False
 
 
diff --git a/src/accelerate/utils/megatron_lm.py b/src/accelerate/utils/megatron_lm.py
index cb988e360..afee2fcca 100644
--- a/src/accelerate/utils/megatron_lm.py
+++ b/src/accelerate/utils/megatron_lm.py
@@ -20,11 +20,12 @@
 import torch
 import torch.nn.functional as F
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+from torch.nn.parallel.distributed import DistributedDataParallel as torchDDP
 
 from ..optimizer import AcceleratedOptimizer
 from ..scheduler import AcceleratedScheduler
 from .imports import is_megatron_lm_available, is_transformers_available
-from .operations import send_to_device
+from .operations import recursively_apply, send_to_device
 
 
 if is_transformers_available():
@@ -41,6 +42,7 @@
         get_num_microbatches,
         get_tensorboard_writer,
         get_timers,
+        get_tokenizer,
         mpu,
         print_rank_0,
         print_rank_last,
@@ -56,16 +58,24 @@
         set_jit_fusion_options,
         write_args_to_tensorboard,
     )
-    from megatron.model import BertModel, GPTModel, ModelType, T5Model
+    from megatron.model import BertModel
+    from megatron.model import DistributedDataParallel as LocalDDP
+    from megatron.model import Float16Module, GPTModel, ModelType, T5Model
     from megatron.model.classification import Classification
     from megatron.optimizer import get_megatron_optimizer
     from megatron.schedules import get_forward_backward_func
+    from megatron.text_generation.communication import broadcast_int_list, broadcast_tensor
+    from megatron.text_generation.generation import (
+        beam_search_and_return_on_first_stage,
+        generate_tokens_probs_and_return_on_first_stage,
+    )
     from megatron.tokenizer.tokenizer import _vocab_size_with_padding
     from megatron.training import get_model, get_optimizer_param_scheduler, training_log
     from megatron.utils import (
         average_losses_across_data_parallel_group,
         calc_params_l2_norm,
         get_ltor_masks_and_position_ids,
+        unwrap_model,
     )
 
 
@@ -114,13 +124,21 @@ def model_provider_func(pre_process=True, post_process=True, add_encoder=True, a
 def prepare_model(accelerator):
     accelerator.print("Preparing model")
     args = get_args()
-    if args.model_type_name == "bert" or args.model_type_name == "gpt":
-        model_type = ModelType.encoder_or_decoder
-    elif args.model_type_name == "t5":
-        model_type = ModelType.encoder_and_decoder
-        if args.pipeline_model_parallel_split_rank is None and args.pipeline_model_parallel_size > 1:
-            args.pipeline_model_parallel_split_rank = args.pipeline_model_parallel_size // 2
-    model = get_model(model_provider_func, model_type)
+    if accelerator.state.megatron_lm_plugin.custom_prepare_model_function is not None:
+        if accelerator.state.megatron_lm_plugin.custom_model_provider_function is None:
+            raise ValueError(
+                "You must provide a `custom_model_provider_function` when using a `custom_prepare_model_function`."
+            )
+        custom_model_provider_func = accelerator.state.megatron_lm_plugin.custom_model_provider_function
+        model = accelerator.state.megatron_lm_plugin.custom_prepare_model_function(custom_model_provider_func)
+    else:
+        if args.model_type_name == "bert" or args.model_type_name == "gpt":
+            model_type = ModelType.encoder_or_decoder
+        elif args.model_type_name == "t5":
+            model_type = ModelType.encoder_and_decoder
+            if args.pipeline_model_parallel_split_rank is None and args.pipeline_model_parallel_size > 1:
+                args.pipeline_model_parallel_split_rank = args.pipeline_model_parallel_size // 2
+        model = get_model(model_provider_func, model_type)
     return model
 
 
@@ -593,6 +611,12 @@ def __init__(self, args):
         self.loss_func = self.get_loss_func()
         self.forward_step = self.get_forward_step_func()
         self.eod_token = args.padded_vocab_size - 1
+        if args.vocab_file is not None:
+            tokenizer = get_tokenizer()
+            self.eod_token = tokenizer.eod
+        self.reset_position_ids = args.reset_position_ids
+        self.reset_attention_mask = args.reset_attention_mask
+        self.eod_mask_loss = args.eod_mask_loss
         if not args.model_return_dict:
             self.model_output_class = None
         else:
@@ -601,8 +625,6 @@ def __init__(self, args):
     def get_batch_func(self, megatron_dataset_flag):
         def get_batch_megatron(data_iterator):
             """Generate a batch"""
-            args = get_args()
-
             # Items and their type.
             keys = ["text"]
             datatype = torch.int64
@@ -621,7 +643,7 @@ def get_batch_megatron(data_iterator):
 
             # Get the masks and postition ids.
             attention_mask, loss_mask, position_ids = get_ltor_masks_and_position_ids(
-                tokens, self.eod_token, args.reset_position_ids, args.reset_attention_mask, args.eod_mask_loss
+                tokens, self.eod_token, self.reset_position_ids, self.reset_attention_mask, self.eod_mask_loss
             )
 
             return tokens, labels, loss_mask, attention_mask, position_ids
@@ -638,7 +660,7 @@ def get_batch_transformer(data_iterator):
             tokens = tokens_[:, :-1].contiguous()
             # Get the masks and postition ids.
             attention_mask, loss_mask, position_ids = get_ltor_masks_and_position_ids(
-                tokens, self.eod_token, False, False, True
+                tokens, self.eod_token, self.reset_position_ids, self.reset_attention_mask, True
             )
             return tokens, labels, loss_mask, attention_mask, position_ids
 
@@ -648,15 +670,24 @@ def get_batch_transformer(data_iterator):
             return get_batch_transformer
 
     def get_loss_func(self):
+        args = get_args()
+
         def loss_func(loss_mask, output_tensor):
-            losses = output_tensor.float()
+            if args.return_logits:
+                losses, logits = output_tensor
+            else:
+                losses = output_tensor
+            losses = losses.float()
             loss_mask = loss_mask.view(-1).float()
             loss = torch.sum(losses.view(-1) * loss_mask) / loss_mask.sum()
 
             # Reduce loss for logging.
             averaged_loss = average_losses_across_data_parallel_group([loss])
 
-            return loss, {"lm loss": averaged_loss[0]}
+            output_dict = {"lm loss": averaged_loss[0]}
+            if args.return_logits:
+                output_dict.update({"logits": logits})
+            return loss, output_dict
 
         return loss_func
 
@@ -1031,7 +1062,10 @@ def train_step(self, **batch_data):
             loss_reduced = {}
             for key in losses_reduced[0]:
                 losses_reduced_for_key = [x[key] for x in losses_reduced]
-                loss_reduced[key] = sum(losses_reduced_for_key) / len(losses_reduced_for_key)
+                if len(losses_reduced_for_key[0].shape) == 0:
+                    loss_reduced[key] = sum(losses_reduced_for_key) / len(losses_reduced_for_key)
+                else:
+                    loss_reduced[key] = torch.concat(losses_reduced_for_key)
             return loss_reduced, skipped_iter, grad_norm, num_zeros_in_grad
         return {}, skipped_iter, grad_norm, num_zeros_in_grad
 
@@ -1079,7 +1113,10 @@ def eval_step(self, **batch_data):
             loss_reduced = {}
             for key in loss_dicts[0]:
                 losses_reduced_for_key = [x[key] for x in loss_dicts]
-                loss_reduced[key] = sum(losses_reduced_for_key) / len(losses_reduced_for_key)
+                if len(losses_reduced_for_key[0].shape) == 0:
+                    loss_reduced[key] = sum(losses_reduced_for_key) / len(losses_reduced_for_key)
+                else:
+                    loss_reduced[key] = torch.concat(losses_reduced_for_key)
             return loss_reduced
         else:
             return {}
@@ -1134,10 +1171,15 @@ def forward(self, **batch_data):
 
         loss = torch.tensor(0.0, device=args.local_rank)
         for key in loss_dict:
-            loss += loss_dict[key]
+            if len(loss_dict[key].shape) == 0:
+                loss += loss_dict[key]
+
+        logits = None
+        if "logits" in loss_dict:
+            logits = loss_dict["logits"]
         # loss = reduce(loss)
         if self.train_step_handler.model_output_class is not None:
-            return self.train_step_handler.model_output_class(loss=loss)
+            return self.train_step_handler.model_output_class(loss=loss, logits=logits)
         return loss
 
     def log_eval_results(self):
@@ -1186,6 +1228,184 @@ def load_checkpoint(self, input_dir):
         if args.fp16 and self.iteration == 0:
             self.optimizer.reload_model_params()
 
+    def megatron_generate(
+        self,
+        inputs,
+        attention_mask=None,
+        max_length=None,
+        max_new_tokens=None,
+        num_beams=None,
+        temperature=None,
+        top_k=None,
+        top_p=None,
+        length_penalty=None,
+        **kwargs,
+    ):
+        """
+        Generate method for GPT2 model. This method is used for inference. Supports both greedy and beam search along
+        with sampling. Refer the Megatron-LM repo for more details
+
+        Args:
+            inputs (torch.Tensor): input ids
+            attention_mask (torch.Tensor, optional): attention mask. Defaults to None.
+            max_length (int, optional): max length of the generated sequence. Defaults to None.
+            Either this or max_new_tokens should be provided.
+            max_new_tokens (int, optional): max number of tokens to be generated. Defaults to None.
+            Either this or max_length should be provided.
+            num_beams (int, optional): number of beams to use for beam search. Defaults to None.
+            temperature (float, optional): temperature for sampling. Defaults to 1.0.
+            top_k (int, optional): top k tokens to consider for sampling. Defaults to 0.0.
+            top_p (float, optional): tokens in top p probability are considered for sampling. Defaults to 0.0.
+            length_penalty (float, optional): length penalty for beam search. Defaults to None.
+            kwargs: additional key-value arguments
+        """
+
+        # checking if required arguments are passed
+        args = get_args()
+        if args.model_type_name != "gpt":
+            raise NotImplementedError("Generate method is not implemented for this model")
+
+        if args.data_parallel_size > 1:
+            raise ValueError("Generate method requires data parallelism to be 1")
+
+        if args.sequence_parallel:
+            raise ValueError("Generate method requires sequence parallelism to be False")
+
+        if args.recompute_granularity is not None:
+            raise ValueError("Checkpoint activations cannot be set for inference")
+
+        if args.vocab_file is None:
+            raise ValueError("Vocab file is required for inference")
+
+        # Prepare inputs
+        if max_length is None and max_new_tokens is None:
+            raise ValueError("`max_length` or `max_new_tokens` are required for inference")
+
+        if temperature is None:
+            temperature = 1.0
+        elif not (0.0 < temperature <= 100.0):
+            raise ValueError("temperature must be a positive number less than or equal to 100.0")
+
+        if top_k is None:
+            top_k = 0
+        elif not (0 <= top_k <= 1000):
+            raise ValueError("top_k must be a positive number less than or equal to 1000")
+
+        if top_p is None:
+            top_p = 0.0
+        elif top_p > 0.0 and top_k > 0.0:
+            raise ValueError("top_p and top_k sampling cannot be set together")
+        else:
+            if not (0.0 <= top_p <= 1.0):
+                raise ValueError("top_p must be less than or equal to 1.0")
+
+        top_p_decay = kwargs.get("top_p_decay", 0.0)
+        if not (0.0 <= top_p_decay <= 1.0):
+            raise ValueError("top_p_decay must be less than or equal to 1.0")
+
+        top_p_bound = kwargs.get("top_p_bound", 0.0)
+        if not (0.0 <= top_p_bound <= 1.0):
+            raise ValueError("top_p_bound must be less than or equal to 1.0")
+
+        add_BOS = kwargs.get("add_BOS", False)
+        if not (isinstance(add_BOS, bool)):
+            raise ValueError("add_BOS must be a boolean")
+
+        beam_width = num_beams
+        if beam_width is not None:
+            if not isinstance(beam_width, int):
+                raise ValueError("beam_width must be an integer")
+            if beam_width < 1:
+                raise ValueError("beam_width must be greater than 0")
+            if inputs.shape[0] > 1:
+                return "When doing beam_search, batch size must be 1"
+
+        tokenizer = get_tokenizer()
+
+        stop_token = kwargs.get("stop_token", tokenizer.eod)
+        if stop_token is not None:
+            if not isinstance(stop_token, int):
+                raise ValueError("stop_token must be an integer")
+
+        if length_penalty is None:
+            length_penalty = 1.0
+
+        sizes_list = None
+        prompts_tokens_tensor = None
+        prompts_length_tensor = None
+        if torch.distributed.get_rank() == 0:
+            # Get the prompts length.
+            if attention_mask is None:
+                prompts_length_tensor = torch.cuda.LongTensor([inputs.shape[1]] * inputs.shape[0])
+            else:
+                prompts_length_tensor = attention_mask.sum(axis=-1).cuda()
+
+            if max_new_tokens is None:
+                max_new_tokens = max_length - inputs.shape[1]
+            if max_new_tokens <= 0:
+                raise ValueError("max_new_tokens must be greater than 0")
+
+            if add_BOS:
+                max_length = max_new_tokens + inputs.shape[1] + 1
+                # making sure that `max_length` is a multiple of 4 to leverage fused kernels
+                max_length = 4 * math.ceil(max_length / 4)
+                max_new_tokens = max_length - (inputs.shape[1] + 1)
+                padding = torch.cuda.LongTensor([[tokenizer.eod] * max_new_tokens] * inputs.shape[0])
+                prompts_tokens_tensor = torch.concat(
+                    [torch.unsqueeze(padding[:, 0], axis=-1), inputs.cuda(), padding], axis=-1
+                )
+            else:
+                # making sure that `max_length` is a multiple of 4 to leverage fused kernels
+                max_length = max_new_tokens + inputs.shape[1]
+                max_length = 4 * math.ceil(max_length / 4)
+                max_new_tokens = max_length - inputs.shape[1]
+                padding = torch.cuda.LongTensor([[tokenizer.eod] * max_new_tokens] * inputs.shape[0])
+                prompts_tokens_tensor = torch.concat([inputs.cuda(), padding], axis=-1)
+
+            # We need the sizes of these tensors for the boradcast
+            sizes_list = [
+                prompts_tokens_tensor.size(0),  # Batch size
+                prompts_tokens_tensor.size(1),
+            ]  # Sequence lenght
+
+        # First, broadcast the sizes.
+        sizes_tensor = broadcast_int_list(2, int_list=sizes_list, rank=0)
+
+        # Now that we have the sizes, we can boradcast the tokens
+        # and length tensors.
+        sizes = sizes_tensor.tolist()
+        context_tokens_tensor = broadcast_tensor(sizes, torch.int64, tensor=prompts_tokens_tensor, rank=0)
+        context_length_tensor = broadcast_tensor(sizes[0], torch.int64, tensor=prompts_length_tensor, rank=0)
+
+        # Run the inference
+        random_seed = kwargs.get("random_seed", 0)
+        torch.random.manual_seed(random_seed)
+        unwrapped_model = unwrap_model(self.base_model, (torchDDP, LocalDDP, Float16Module))
+        if beam_width is not None:
+            tokens, _ = beam_search_and_return_on_first_stage(
+                unwrapped_model,
+                context_tokens_tensor,
+                context_length_tensor,
+                beam_width,
+                stop_token=stop_token,
+                num_return_gen=1,
+                length_penalty=length_penalty,
+            )
+        else:
+            tokens, _, _ = generate_tokens_probs_and_return_on_first_stage(
+                unwrapped_model,
+                context_tokens_tensor,
+                context_length_tensor,
+                return_output_log_probs=False,
+                top_k=top_k,
+                top_p=top_p,
+                top_p_decay=top_p_decay,
+                top_p_bound=top_p_bound,
+                temperature=temperature,
+                use_eod_token_for_early_termination=True,
+            )
+        return tokens
+
 
 # other utilities
 def avg_losses_across_data_parallel_group(losses):
@@ -1197,3 +1417,25 @@ def avg_losses_across_data_parallel_group(losses):
     """
 
     return average_losses_across_data_parallel_group(losses)
+
+
+def gather_across_data_parallel_groups(tensor):
+    """
+    Recursively gather tensor in a nested list/tuple/dictionary of tensors from data parallel ranks.
+
+    Args:
+        tensor (nested list/tuple/dictionary of `torch.Tensor`):
+            The data to gather across data parallel ranks.
+
+    """
+
+    def _gpu_gather_one(tensor):
+        if tensor.ndim == 0:
+            tensor = tensor.clone()[None]
+        output_tensors = [
+            tensor.clone() for _ in range(torch.distributed.get_world_size(group=mpu.get_data_parallel_group()))
+        ]
+        torch.distributed.all_gather(output_tensors, tensor, group=mpu.get_data_parallel_group())
+        return torch.cat(output_tensors, dim=0)
+
+    return recursively_apply(_gpu_gather_one, tensor, error_on_other_type=True)

Here you should lowercase, or use the str2bool conversion, just to make sure to catch variants.
@@ -90,10 +91,11 @@ def is_bf16_available(ignore_tpu=False):
 
 
 def is_megatron_lm_available():
-    package_exists = importlib.util.find_spec("megatron") is not None
-    if package_exists:
-        megatron_version = parse(importlib_metadata.version("megatron-lm"))
-        return compare_versions(megatron_version, ">=", "2.2.0")
+    if os.environ.get("USE_MEGATRON_LM", "false") == "true":
Fits on one line.
@@ -1134,10 +1171,18 @@ def forward(self, **batch_data):
 
         loss = torch.tensor(0.0, device=args.local_rank)
         for key in loss_dict:
-            loss += loss_dict[key]
+            if len(loss_dict[key].shape) == 0:
+                loss += loss_dict[key]
+
+        logits = None
+        if "logits" in loss_dict:
+            logits = loss_dict["logits"]
         # loss = reduce(loss)
         if self.train_step_handler.model_output_class is not None:
-            return self.train_step_handler.model_output_class(loss=loss)
+            return self.train_step_handler.model_output_class(
+                loss=loss,
+                logits=logits,
+            )
Very reluctant to add this here with the same name as in Transformers. This method does not have the same signature/behavior so users will be very confused to get something different. 
@@ -1186,6 +1231,184 @@ def load_checkpoint(self, input_dir):
         if args.fp16 and self.iteration == 0:
             self.optimizer.reload_model_params()
 
+    def generate(
Done.
@@ -90,10 +91,11 @@ def is_bf16_available(ignore_tpu=False):
 
 
 def is_megatron_lm_available():
-    package_exists = importlib.util.find_spec("megatron") is not None
-    if package_exists:
-        megatron_version = parse(importlib_metadata.version("megatron-lm"))
-        return compare_versions(megatron_version, ">=", "2.2.0")
+    if os.environ.get("USE_MEGATRON_LM", "false") == "true":
Done.
@@ -1134,10 +1171,18 @@ def forward(self, **batch_data):
 
         loss = torch.tensor(0.0, device=args.local_rank)
         for key in loss_dict:
-            loss += loss_dict[key]
+            if len(loss_dict[key].shape) == 0:
+                loss += loss_dict[key]
+
+        logits = None
+        if "logits" in loss_dict:
+            logits = loss_dict["logits"]
         # loss = reduce(loss)
         if self.train_step_handler.model_output_class is not None:
-            return self.train_step_handler.model_output_class(loss=loss)
+            return self.train_step_handler.model_output_class(
+                loss=loss,
+                logits=logits,
+            )
renamed to `megatron_generate`
@@ -1186,6 +1231,184 @@ def load_checkpoint(self, input_dir):
         if args.fp16 and self.iteration == 0:
             self.optimizer.reload_model_params()
 
+    def generate(
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/819
diff --git a/docs/source/package_reference/cli.mdx b/docs/source/package_reference/cli.mdx
index f7424d656..9643d3b4f 100644
--- a/docs/source/package_reference/cli.mdx
+++ b/docs/source/package_reference/cli.mdx
@@ -177,6 +177,37 @@ The following arguments are only useful when training in SageMaker
 * `--aws_access_key_id AWS_ACCESS_KEY_ID` (`str`) -- The AWS_ACCESS_KEY_ID used to launch the Amazon SageMaker training job
 * `--aws_secret_access_key AWS_SECRET_ACCESS_KEY` (`str`) -- The AWS_SECRET_ACCESS_KEY used to launch the Amazon SageMaker training job
 
+## accelerate tpu-config
+
+`accelerate tpu-config`
+
+**Usage**:
+
+```bash
+accelerate tpu-config [arguments]
+```
+
+**Optional Arguments**:
+* `-h`, `--help` (`bool`) -- Show a help message and exit
+
+**Config Arguments**:
+
+Arguments that can be configured through `accelerate config`.
+
+* `--config_file` (`str`) -- Path to the config file to use for accelerate.
+* `--tpu_name` (`str`) -- The name of the TPU to use. If not specified, will use the TPU specified in the config file.
+* `--tpu_zone` (`str`) -- The zone of the TPU to use. If not specified, will use the zone specified in the config file.
+
+**TPU Arguments**:
+
+Arguments for options ran inside the TPU.
+
+* `--command_file` (`str`) -- The path to the file containing the commands to run on the pod on startup.
+* `--command` (`str`) -- A command to run on the pod. Can be passed multiple times.
+* `--install_accelerate` (`bool`) -- Whether to install accelerate on the pod. Defaults to False.
+* `--accelerate_version` (`str`) -- The version of accelerate to install on the pod. If not specified, will use the latest pypi version. Specify 'dev' to install from GitHub.
+* `--debug` (`bool`) -- If set, will print the command that would be run instead of running it.
+
 ## accelerate test
 
 `accelerate test` or `accelerate-test`
diff --git a/src/accelerate/commands/accelerate_cli.py b/src/accelerate/commands/accelerate_cli.py
index 515a66d1a..f0e76fd2c 100644
--- a/src/accelerate/commands/accelerate_cli.py
+++ b/src/accelerate/commands/accelerate_cli.py
@@ -19,8 +19,8 @@
 from accelerate.commands.config import config_command_parser
 from accelerate.commands.env import env_command_parser
 from accelerate.commands.launch import launch_command_parser
-from accelerate.commands.pod import pod_command_parser
 from accelerate.commands.test import test_command_parser
+from accelerate.commands.tpu import tpu_command_parser
 
 
 def main():
@@ -31,7 +31,7 @@ def main():
     config_command_parser(subparsers=subparsers)
     env_command_parser(subparsers=subparsers)
     launch_command_parser(subparsers=subparsers)
-    pod_command_parser(subparsers=subparsers)
+    tpu_command_parser(subparsers=subparsers)
     test_command_parser(subparsers=subparsers)
 
     # Let's go
diff --git a/src/accelerate/commands/pod.py b/src/accelerate/commands/tpu.py
similarity index 81%
rename from src/accelerate/commands/pod.py
rename to src/accelerate/commands/tpu.py
index 87cd4f904..59bbb08e9 100644
--- a/src/accelerate/commands/pod.py
+++ b/src/accelerate/commands/tpu.py
@@ -22,70 +22,66 @@
 from packaging.version import Version, parse
 
 
-_description = "Run commands across a pod of TPU VMs for initial setup before running `accelerate launch`. Will also install Accelerate on the pod."
+_description = "Run commands across TPU VMs for initial setup before running `accelerate launch`."
 
 
-def pod_command_parser(subparsers=None):
+def tpu_command_parser(subparsers=None):
     if subparsers is not None:
-        parser = subparsers.add_parser("pod-config", description=_description)
+        parser = subparsers.add_parser("tpu-config", description=_description)
     else:
-        parser = argparse.ArgumentParser("Accelerate pod-config command", description=_description)
-
-    parser.add_argument(
+        parser = argparse.ArgumentParser("Accelerate tpu-config command", description=_description)
+    # Core arguments
+    config_args = parser.add_argument_group(
+        "Config Arguments", "Arguments that can be configured through `accelerate config`."
+    )
+    config_args.add_argument(
         "--config_file",
         type=str,
         default=None,
         help="Path to the config file to use for accelerate.",
     )
-
-    parser.add_argument(
-        "--pod_config_file",
-        type=str,
+    config_args.add_argument(
+        "--tpu_name",
         default=None,
-        help="Path to the config file to use for the pod.",
+        help="The name of the TPU to use. If not specified, will use the TPU specified in the config file.",
     )
-
-    parser.add_argument(
+    config_args.add_argument(
+        "--tpu_zone",
+        default=None,
+        help="The zone of the TPU to use. If not specified, will use the zone specified in the config file.",
+    )
+    pod_args = parser.add_argument_group("TPU Arguments", "Arguments for options ran inside the TPU.")
+    pod_args.add_argument(
         "--command_file",
         default=None,
         help="The path to the file containing the commands to run on the pod on startup.",
     )
-    parser.add_argument(
+    pod_args.add_argument(
         "--command",
         action="append",
         nargs="+",
-        help="A command to run on the pod. If not specified, will use the command specified in the command file.",
-    )
-    parser.add_argument(
-        "--tpu_name",
-        default=None,
-        help="The name of the TPU to use. If not specified, will use the TPU specified in the config file.",
-    )
-    parser.add_argument(
-        "--tpu_zone",
-        default=None,
-        help="The zone of the TPU to use. If not specified, will use the zone specified in the config file.",
+        help="A command to run on the pod. Can be passed multiple times.",
     )
-    parser.add_argument(
+    pod_args.add_argument(
         "--install_accelerate",
         action="store_true",
         help="Whether to install accelerate on the pod. Defaults to False.",
     )
-    parser.add_argument(
+    pod_args.add_argument(
         "--accelerate_version",
         default="latest",
         help="The version of accelerate to install on the pod. If not specified, will use the latest pypi version. Specify 'dev' to install from GitHub.",
     )
-    parser.add_argument(
+    pod_args.add_argument(
         "--debug", action="store_true", help="If set, will print the command that would be run instead of running it."
     )
 
     if subparsers is not None:
-        parser.set_defaults(func=pod_launcher)
+        parser.set_defaults(func=tpu_command_launcher)
     return parser
 
 
-def pod_launcher(args):
+def tpu_command_launcher(args):
     defaults = None
 
     # Get the default from the config file if it exists.
@@ -146,7 +142,7 @@ def pod_launcher(args):
 
 
 def main():
-    parser = pod_command_parser()
+    parser = tpu_command_parser()
     args = parser.parse_args()
 
-    pod_launcher(args)
+    tpu_command_launcher(args)
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 2479f736d..50ed87d9f 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -66,15 +66,15 @@ def test_config_compatibility(self):
                 )
 
 
-class PodConfigTester(unittest.TestCase):
+class TpuConfigTester(unittest.TestCase):
     """
-    Test case for verifying the `accelerate pod-config` CLI passes the right `gcloud` command.
+    Test case for verifying the `accelerate tpu-config` CLI passes the right `gcloud` command.
     """
 
     tpu_name = "test-tpu"
     tpu_zone = "us-central1-a"
     command = "ls"
-    cmd = ["accelerate", "pod-config"]
+    cmd = ["accelerate", "tpu-config"]
     base_output = "cd /usr/share"
     command_file = "tests/test_samples/test_command_file.sh"
     gcloud = "Running gcloud compute tpus tpu-vm ssh"

This arg was never used in the end, but not removed. 
@@ -22,70 +22,66 @@
 from packaging.version import Version, parse
 
 
-_description = "Run commands across a pod of TPU VMs for initial setup before running `accelerate launch`. Will also install Accelerate on the pod."
+_description = "Run commands across TPU VMs for initial setup before running `accelerate launch`."
 
 
-def pod_command_parser(subparsers=None):
+def tpu_command_parser(subparsers=None):
     if subparsers is not None:
-        parser = subparsers.add_parser("pod-config", description=_description)
+        parser = subparsers.add_parser("tpu-config", description=_description)
     else:
-        parser = argparse.ArgumentParser("Accelerate pod-config command", description=_description)
-
-    parser.add_argument(
+        parser = argparse.ArgumentParser("Accelerate tpu-config command", description=_description)
+    # Core arguments
+    config_args = parser.add_argument_group(
+        "Config Arguments", "Arguments that can be configured through `accelerate config`."
+    )
+    config_args.add_argument(
         "--config_file",
         type=str,
         default=None,
         help="Path to the config file to use for accelerate.",
     )
-
-    parser.add_argument(
-        "--pod_config_file",
-        type=str,
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/818
diff --git a/docs/source/package_reference/cli.mdx b/docs/source/package_reference/cli.mdx
index 52e8ecfb6..f7424d656 100644
--- a/docs/source/package_reference/cli.mdx
+++ b/docs/source/package_reference/cli.mdx
@@ -78,63 +78,102 @@ accelerate launch [arguments] {training_script} --{training_script-argument-1} -
 
 * `-h`, `--help` (`bool`) -- Show a help message and exit
 * `--config_file CONFIG_FILE` (`str`)-- The config file to use for the default values in the launching script.
-* `--cpu` (`bool`) -- Whether or not to force the training on the CPU.
-* `--mixed_precision {no,fp16,bf16}` (`str`) -- Whether or not to use mixed precision training. Choose between FP16 and BF16 (bfloat16) training. BF16 training is only supported on
-                    Nvidia Ampere GPUs and PyTorch 1.10 or later.
-* `--multi_gpu` (`bool`, defaults to `False`) -- Whether or not this should launch a distributed GPU training.
 * `-m`, `--module` (`bool`) -- Change each process to interpret the launch script as a Python module, executing with the same behavior as 'python -m'.
 * `--no_python` (`bool`) -- Skip prepending the training script with 'python' - just execute it directly. Useful when the script is not a Python script.
+* `--debug` (`bool`) -- Whether to print out the torch.distributed stack trace when something fails.
 
 The rest of these arguments are configured through `accelerate config` and are read in from the specified `--config_file` (or default configuration) for their 
 values. They can also be passed in manually.
 
-**Machine Configuration Arguments**:
+**Hardware Selection Arguments**:
 
-The following arguments are useful for customization of worker machines
-* `--machine_rank MACHINE_RANK` (`int`) -- The rank of the machine on which this script is launched.
-* `--num_machines NUM_MACHINES` (`int`) -- The total number of machines used in this training.
+* `--cpu` (`bool`) -- Whether or not to force the training on the CPU.
+* `--multi_gpu` (`bool`) -- Whether or not this should launch a distributed GPU training.
+* `--mps` (`bool`) -- Whether or not this should use MPS-enabled GPU device on MacOS machines.
+* `--tpu` (`bool`) -- Whether or not this should launch a TPU training.
+
+**Resource Selection Arguments**:
+
+The following arguments are useful for fine-tuning how available hardware should be used
+
+* `--mixed_precision {no,fp16,bf16}` (`str`) -- Whether or not to use mixed precision training. Choose between FP16 and BF16 (bfloat16) training. BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.
 * `--num_processes NUM_PROCESSES` (`int`) -- The total number of processes to be launched in parallel.
+* `--num_machines NUM_MACHINES` (`int`) -- The total number of machines used in this training.
+* `--num_cpu_threads_per_process NUM_CPU_THREADS_PER_PROCESS` (`int`) -- The number of CPU threads per process. Can be tuned for optimal performance.
+
+**Training Paradigm Arguments**:
+
+The following arguments are useful for selecting which training paradigm to use.
+
+* `--use_deepspeed` (`bool`) -- Whether or not to use DeepSpeed for training.
+* `--use_fsdp` (`bool`) -- Whether or not to use FullyShardedDataParallel for training.
+* `--use_megatron_lm` (`bool`) -- Whether or not to use Megatron-LM for training.
+
+**Distributed GPU Arguments**:
+
+The following arguments are only useful when `multi_gpu` is passed or multi-gpu training is configured through `accelerate config`: 
+
 * `--gpu_ids` (`str`) -- What GPUs (by id) should be used for training on this machine as a comma-seperated list
 * `--same_network` (`bool`) -- Whether all machines used for multinode training exist on the same local network.
+* `--machine_rank MACHINE_RANK` (`int`) -- The rank of the machine on which this script is launched.
 * `--main_process_ip MAIN_PROCESS_IP` (`str`) -- The IP address of the machine of rank 0.
 * `--main_process_port MAIN_PROCESS_PORT` (`int`) -- The port to use to communicate with the machine of rank 0.
 * `--rdzv_conf` (`str`) -- Additional rendezvous configuration (<key1>=<value1>,<key2>=<value2>,...).
-* `--num_cpu_threads_per_process NUM_CPU_THREADS_PER_PROCESS` (`int`) -- The number of CPU threads per process. Can be tuned for optimal performance.
 * `--max_restarts` (`int`) -- Maximum number of worker group restarts before failing.
 * `--monitor_interval` (`float`) -- Interval, in seconds, to monitor the state of workers.
 
+**TPU Arguments**:
+
+The following arguments are only useful when `tpu` is passed or TPU training is configured through `accelerate config`: 
+
+* `--main_training_function MAIN_TRAINING_FUNCTION` (`str`) -- The name of the main function to be executed in your script.
+* `--downcast_bf16` (`bool`) -- Whether when using bf16 precision on TPUs if both float and double tensors are cast to bfloat16 or if double tensors remain as float32.
+
 **DeepSpeed Arguments**:
 
-The following arguments are only useful when `use_deepspeed` is passed: 
-* `--use_deepspeed` (`bool`) -- Whether to use deepspeed.
-* `--deepspeed_config_file DEEPSPEED_CONFIG_FILE` (`str`) -- DeepSpeed config file.
-* `--zero_stage ZERO_STAGE` (`str`) -- DeepSpeed's ZeRO optimization stage
-* `--offload_optimizer_device OFFLOAD_OPTIMIZER_DEVICE` (`str`) -- Decides where (none|cpu|nvme) to offload optimizer states
-* `--offload_param_device OFFLOAD_PARAM_DEVICE` (`str`) -- Decides where (none|cpu|nvme) to offload parameters
-* `--gradient_accumulation_steps GRADIENT_ACCUMULATION_STEPS` (`int`) -- Number of gradient_accumulation_steps used in your training script
-* `--gradient_clipping GRADIENT_CLIPPING` (`float`) -- gradient clipping value used in your training script
-The following arguments are related to using ZeRO Stage-3
-* `--zero3_init_flag ZERO3_INIT_FLAG` (`bool`) -- Decides Whether (true|false) to enable `deepspeed.zero.Init` for constructing massive models
-* `--zero3_save_16bit_model ZERO3_SAVE_16BIT_MODEL` (`bool`) -- Decides Whether (true|false) to save 16-bit model weights when using ZeRO Stage-3
+The following arguments are only useful when `use_deepspeed` is passed or `deepspeed` is configured through `accelerate config`: 
+
+* `--deepspeed_config_file` (`str`) -- DeepSpeed config file.
+* `--zero_stage` (`int`) -- DeepSpeed's ZeRO optimization stage.
+* `--offload_optimizer_device` (`str`) -- Decides where (none|cpu|nvme) to offload optimizer states.
+* `--offload_param_device` (`str`) -- Decides where (none|cpu|nvme) to offload parameters.
+* `--gradient_accumulation_steps` (`int`) -- No of gradient_accumulation_steps used in your training script.
+* `--gradient_clipping` (`float`) -- Gradient clipping value used in your training script.
+* `--zero3_init_flag` (`str`) -- Decides Whether (true|false) to enable `deepspeed.zero.Init` for constructing massive models. Only applicable with DeepSpeed ZeRO Stage-3.
+* `--zero3_save_16bit_model` (`str`) -- Decides Whether (true|false) to save 16-bit model weights when using ZeRO Stage-3. Only applicable with DeepSpeed ZeRO Stage-3.
+* `--deepspeed_hostfile` (`str`) -- DeepSpeed hostfile for configuring multi-node compute resources.
+* `--deepspeed_exclusion_filter` (`str`) -- DeepSpeed exclusion filter string when using mutli-node setup.
+* `--deepspeed_inclusion_filter` (`str`) -- DeepSpeed inclusion filter string when using mutli-node setup.
+* `--deepspeed_multinode_launcher` (`str`) -- DeepSpeed multi-node launcher to use.
 
 **Fully Sharded Data Parallelism Arguments**:
 
-The following arguments are only useful when `use_fdsp` is passed:
-* `--use_fsdp` (`bool`) -- Whether to use fsdp.
-* `--offload_params OFFLOAD_PARAMS` (`bool`) -- Decides Whether (true|false) to offload parameters and gradients to CPU.
-* `--min_num_params MIN_NUM_PARAMS` (`int`) -- FSDP's minimum number of parameters for Default Auto Wrapping.
-* `--sharding_strategy SHARDING_STRATEGY` (`str`) -- FSDP's Sharding Strategy.
+The following arguments are only useful when `use_fdsp` is passed or Fully Sharded Data Parallelism is configured through `accelerate config`:
 
-**TPU Arguments**:
+* `--fsdp_offload_params` (`str`) -- Decides Whether (true|false) to offload parameters and gradients to CPU.
+* `--fsdp_min_num_params` (`int`) -- FSDP's minimum number of parameters for Default Auto Wrapping.
+* `--fsdp_sharding_strategy` (`int`) -- FSDP's Sharding Strategy.
+* `--fsdp_auto_wrap_policy` (`str`) -- FSDP's auto wrap policy.
+* `--fsdp_transformer_layer_cls_to_wrap` (`str`) -- Transformer layer class name (case-sensitive) to wrap, e.g, `BertLayer`, `GPTJBlock`, `T5Block` ...
+* `--fsdp_backward_prefetch_policy` (`str`) -- FSDP's backward prefetch policy.
+* `--fsdp_state_dict_type` (`str`) -- FSDP's state dict type.
 
-The following arguments are only useful when `tpu` is passed:
-* `--tpu` (`bool`) -- Whether or not this should launch a TPU training.
-* `--main_training_function MAIN_TRAINING_FUNCTION` (`str`) -- The name of the main function to be executed in your script.
-* `--downcast_bf16` (`bool`) -- Whether when using bf16 precision on TPUs if both float and double tensors are cast to bfloat16 or if double tensors remain as float32.
+**Megatron-LM Arguments**:
+
+The following arguments are only useful when `use_megatron_lm` is passed or Megatron-LM is configured through `accelerate config`:
+
+* `--megatron_lm_tp_degree` (``) -- Megatron-LM's Tensor Parallelism (TP) degree.
+* `--megatron_lm_pp_degree` (``) -- Megatron-LM's Pipeline Parallelism (PP) degree.
+* `--megatron_lm_num_micro_batches` (``) -- Megatron-LM's number of micro batches when PP degree > 1.
+* `--megatron_lm_sequence_parallelism` (``) -- Decides Whether (true|false) to enable Sequence Parallelism when TP degree > 1.
+* `--megatron_lm_recompute_activations` (``) -- Decides Whether (true|false) to enable Selective Activation Recomputation.
+* `--megatron_lm_use_distributed_optimizer` (``) -- Decides Whether (true|false) to use distributed optimizer which shards optimizer state and gradients across Data Pralellel (DP) ranks.
+* `--megatron_lm_gradient_clipping` (``) -- Megatron-LM's gradient clipping value based on global L2 Norm (0 to disable).
 
 **AWS SageMaker Arguments**:
 
 The following arguments are only useful when training in SageMaker
+
 * `--aws_access_key_id AWS_ACCESS_KEY_ID` (`str`) -- The AWS_ACCESS_KEY_ID used to launch the Amazon SageMaker training job
 * `--aws_secret_access_key AWS_SECRET_ACCESS_KEY` (`str`) -- The AWS_SECRET_ACCESS_KEY used to launch the Amazon SageMaker training job
 
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 415b06906..46d95c713 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -64,6 +64,7 @@
 options_to_group = {
     "--multi-gpu": "Distributed GPUs",
     "--tpu": "TPU",
+    "--mps": "MPS",
     "--use_mps_device": "MPS",
     "--use_deepspeed": "DeepSpeed Arguments",
     "--use_fsdp": "FSDP Arguments",
@@ -144,6 +145,12 @@ def launch_command_parser(subparsers=None):
     hardware_args.add_argument(
         "--cpu", default=False, action="store_true", help="Whether or not to force the training on the CPU."
     )
+    hardware_args.add_argument(
+        "--mps",
+        default=False,
+        action="store_true",
+        help="Whether or not this should use MPS-enabled GPU device on MacOS machines.",
+    )
     hardware_args.add_argument(
         "--multi_gpu",
         default=False,
@@ -157,7 +164,7 @@ def launch_command_parser(subparsers=None):
         "--use_mps_device",
         default=False,
         action="store_true",
-        help="Whether or not this should use MPS-enabled GPU device on MacOS machines.",
+        help="This argument is deprecated, use `--mps` instead.",
     )
 
     # Resource selection arguments
@@ -506,8 +513,14 @@ def simple_launcher(args):
 
     current_env = os.environ.copy()
     current_env["USE_CPU"] = str(args.cpu or args.use_cpu)
-    current_env["USE_MPS_DEVICE"] = str(args.use_mps_device)
     if args.use_mps_device:
+        warnings.warn(
+            '`use_mps_device` flag is deprecated and will be removed in version 0.15.0 of 🤗 Accelerate. Use "--mps" instead.',
+            FutureWarning,
+        )
+        args.mps = True
+    current_env["USE_MPS_DEVICE"] = str(args.mps)
+    if args.mps:
         current_env["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
     elif args.gpu_ids != "all" and args.gpu_ids is not None:
         current_env["CUDA_VISIBLE_DEVICES"] = args.gpu_ids
@@ -704,7 +717,10 @@ def deepspeed_launcher(args):
         )
 
     if args.fp16:
-        warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', FutureWarning)
+        warnings.warn(
+            '--fp16 flag is deprecated and will be removed in version 0.15.0 of 🤗 Accelerate. Use "--mixed_precision fp16" instead.',
+            FutureWarning,
+        )
         mixed_precision = "fp16"
 
     current_env["PYTHONPATH"] = env_var_path_add("PYTHONPATH", os.path.abspath("."))
@@ -955,18 +971,18 @@ def launch_command(args):
         if (
             not args.multi_gpu
             and not args.tpu
+            and not args.mps
             and not args.use_deepspeed
             and not args.use_fsdp
-            and not args.use_mps_device
             and not args.use_megatron_lm
         ):
             args.use_deepspeed = defaults.distributed_type == DistributedType.DEEPSPEED
             args.multi_gpu = defaults.distributed_type == DistributedType.MULTI_GPU
             args.tpu = defaults.distributed_type == DistributedType.TPU
             args.use_fsdp = defaults.distributed_type == DistributedType.FSDP
-            args.use_mps_device = defaults.distributed_type == DistributedType.MPS
+            args.mps = defaults.distributed_type == DistributedType.MPS
             args.use_megatron_lm = defaults.distributed_type == DistributedType.MEGATRON_LM
-        if not args.use_mps_device:
+        if not args.mps:
             if args.gpu_ids is None:
                 if defaults.gpu_ids is not None:
                     args.gpu_ids = defaults.gpu_ids

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/814
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 287e6899a..63ea4a71b 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -213,7 +213,10 @@ def __init__(
                 )
 
         if fp16:
-            warnings.warn('fp16=True is deprecated. Use mixed_precision="fp16" instead.', DeprecationWarning)
+            warnings.warn(
+                "`fp16=True` is deprecated and will be removed in version 0.15.0 of 🤗 Accelerate. Use `mixed_precision='fp16'` instead.",
+                FutureWarning,
+            )
             mixed_precision = "fp16"
 
         if deepspeed_plugin is None:  # init from env variables
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index df4f703c3..415b06906 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -173,7 +173,10 @@ def launch_command_parser(subparsers=None):
         "BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.",
     )
     resource_args.add_argument(
-        "--fp16", default=False, action="store_true", help="Whether or not to use mixed precision training."
+        "--fp16",
+        default=False,
+        action="store_true",
+        help="This argument is deprecated, use `--mixed_precision fp16` instead.",
     )
     resource_args.add_argument(
         "--num_processes", type=int, default=None, help="The total number of processes to be launched in parallel."
@@ -403,30 +406,6 @@ def launch_command_parser(subparsers=None):
         type=str,
         help="FSDP's state dict type. (useful only when `use_fsdp` flag is passed).",
     )
-    fsdp_args.add_argument(
-        "--offload_params",
-        default=None,
-        type=str,
-        help="This argument is deprecated. Use `fsdp_offload_params` instead.",
-    )
-    fsdp_args.add_argument(
-        "--min_num_params",
-        type=int,
-        default=None,
-        help="This argument is deprecated. Use `fsdp_min_num_params` instead.",
-    )
-    fsdp_args.add_argument(
-        "--sharding_strategy",
-        type=int,
-        default=None,
-        help="This argument is deprecated. Use `fsdp_sharding_strategy` instead.",
-    )
-    fsdp_args.add_argument(
-        "--transformer_layer_cls_to_wrap",
-        default=None,
-        type=str,
-        help="This argument is deprecated. Use `fsdp_transformer_layer_cls_to_wrap` instead.",
-    )
 
     # megatron_lm args
     megatron_lm_args = parser.add_argument_group("Megatron-LM Arguments", "Arguments related to Megatron-LM.")
@@ -547,7 +526,10 @@ def simple_launcher(args):
         )
 
     if args.fp16:
-        warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', DeprecationWarning)
+        warnings.warn(
+            "`fp16` is deprecated and will be removed in version 0.15.0 of 🤗 Accelerate. Use `mixed_precision fp16` instead.",
+            FutureWarning,
+        )
         mixed_precision = "fp16"
 
     current_env["MIXED_PRECISION"] = str(mixed_precision)
@@ -596,43 +578,14 @@ def multi_gpu_launcher(args):
         raise ValueError(f"Unknown mixed_precision mode: {mixed_precision}. Choose between {PrecisionType.list()}.")
 
     if args.fp16:
-        warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', DeprecationWarning)
+        warnings.warn(
+            "`fp16` is deprecated and will be removed in version 0.15.0 of 🤗 Accelerate. Use `mixed_precision fp16` instead.",
+            FutureWarning,
+        )
         mixed_precision = "fp16"
 
     current_env["MIXED_PRECISION"] = str(mixed_precision)
     if args.use_fsdp:
-        if args.sharding_strategy is not None:
-            warnings.warn(
-                "`sharding_strategy` is deprecated and will be removed in version 0.13.0 of 🤗 Accelerate. Use"
-                " `fsdp_sharding_strategy` instead",
-                FutureWarning,
-            )
-            args.fsdp_sharding_strategy = args.sharding_strategy
-
-        if args.offload_params is not None:
-            warnings.warn(
-                "`offload_params` is deprecated and will be removed in version 0.13.0 of 🤗 Accelerate. Use"
-                " `fsdp_offload_params` instead",
-                FutureWarning,
-            )
-            args.fsdp_offload_params = args.offload_params
-
-        if args.min_num_params is not None:
-            warnings.warn(
-                "`min_num_params` is deprecated and will be removed in version 0.13.0 of 🤗 Accelerate. Use"
-                " `fsdp_min_num_params` instead",
-                FutureWarning,
-            )
-            args.fsdp_min_num_params = args.min_num_params
-
-        if args.transformer_layer_cls_to_wrap is not None:
-            warnings.warn(
-                "`transformer_layer_cls_to_wrap` is deprecated and will be removed in version 0.13.0 of 🤗 Accelerate. Use"
-                " `fsdp_transformer_layer_cls_to_wrap` instead",
-                FutureWarning,
-            )
-            args.fsdp_transformer_layer_cls_to_wrap = args.transformer_layer_cls_to_wrap
-
         current_env["USE_FSDP"] = "true"
         current_env["FSDP_SHARDING_STRATEGY"] = str(args.fsdp_sharding_strategy)
         current_env["FSDP_OFFLOAD_PARAMS"] = str(args.fsdp_offload_params).lower()
@@ -751,7 +704,7 @@ def deepspeed_launcher(args):
         )
 
     if args.fp16:
-        warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', DeprecationWarning)
+        warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', FutureWarning)
         mixed_precision = "fp16"
 
     current_env["PYTHONPATH"] = env_var_path_add("PYTHONPATH", os.path.abspath("."))
@@ -921,7 +874,7 @@ def sagemaker_launcher(sagemaker_config: SageMakerConfig, args):
         )
 
     if args.fp16:
-        warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', DeprecationWarning)
+        warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', FutureWarning)
         mixed_precision = "fp16"
 
     # Environment variables to be set for use during training job
diff --git a/src/accelerate/launchers.py b/src/accelerate/launchers.py
index ea8fffa51..4de8b5858 100644
--- a/src/accelerate/launchers.py
+++ b/src/accelerate/launchers.py
@@ -105,7 +105,10 @@ def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, mix
                 )
 
             if use_fp16:
-                warnings.warn('use_fp16=True is deprecated. Use mixed_precision="fp16" instead.', DeprecationWarning)
+                warnings.warn(
+                    "`fp16=True` is deprecated and will be removed in version 0.15.0 of 🤗 Accelerate. Use `mixed_precision='fp16'` instead.",
+                    FutureWarning,
+                )
                 mixed_precision = "fp16"
 
             # torch.distributed will expect a few environment variable to be here. We set the ones common to each

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/813
diff --git a/src/accelerate/hooks.py b/src/accelerate/hooks.py
index 120bb7f42..ef429a369 100644
--- a/src/accelerate/hooks.py
+++ b/src/accelerate/hooks.py
@@ -160,17 +160,19 @@ def new_forward(*args, **kwargs):
     return module
 
 
-def remove_hook_from_module(module: nn.Module):
+def remove_hook_from_module(module: nn.Module, recurse=False):
     """
     Removes any hook attached to a module via `add_hook_to_module`.
 
     Args:
         module (`torch.nn.Module`): The module to attach a hook to.
+        recurse (`bool`, **optional**): Whether to remove the hooks recursively
 
     Returns:
         `torch.nn.Module`: The same module, with the hook detached (the module is modified in place, so the result can
         be discarded).
     """
+
     if hasattr(module, "_hf_hook"):
         module._hf_hook.detach_hook(module)
         delattr(module, "_hf_hook")
@@ -179,6 +181,10 @@ def remove_hook_from_module(module: nn.Module):
         module.forward = module._old_forward
         delattr(module, "_old_forward")
 
+    if recurse:
+        for child in module.children():
+            remove_hook_from_module(child, recurse)
+
     return module
 
 

Not sure the check is useful (should be >0 but the loop will just be empty)
```suggestion
    if recurse:
```
@@ -179,6 +181,10 @@ def remove_hook_from_module(module: nn.Module):
         module.forward = module._old_forward
         delattr(module, "_old_forward")
 
+    if recurse and len(list(module.children())) > 1:
```suggestion
        for child in module.children():
```
No need to ask for the names if you don't want them ;-)
@@ -179,6 +181,10 @@ def remove_hook_from_module(module: nn.Module):
         module.forward = module._old_forward
         delattr(module, "_old_forward")
 
+    if recurse and len(list(module.children())) > 1:
+        for _, child in module.named_children():
Nice catch! Yes indeed ;) 
@@ -179,6 +181,10 @@ def remove_hook_from_module(module: nn.Module):
         module.forward = module._old_forward
         delattr(module, "_old_forward")
 
+    if recurse and len(list(module.children())) > 1:
_The documentation is not available anymore as the PR was closed or merged._
Thanks for the suggestion! Can confirm the `BetterTransformers` tests pass on my side ;) 
https://api.github.com/repos/huggingface/accelerate/pulls/812
diff --git a/src/accelerate/optimizer.py b/src/accelerate/optimizer.py
index 4fad12c72..d5eeef99a 100644
--- a/src/accelerate/optimizer.py
+++ b/src/accelerate/optimizer.py
@@ -157,3 +157,9 @@ def is_overflow(self):
     def step_was_skipped(self):
         """Whether or not the optimizer step was skipped."""
         return self._is_overflow
+
+    def __getstate__(self):
+        return self.__dict__.copy()
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
new file mode 100644
index 000000000..15a095bf7
--- /dev/null
+++ b/tests/test_optimizer.py
@@ -0,0 +1,36 @@
+# 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 pickle
+import unittest
+
+import torch
+
+from accelerate import Accelerator
+from accelerate.state import AcceleratorState
+from accelerate.test_utils import require_cpu
+
+
+@require_cpu
+class OptimizerTester(unittest.TestCase):
+    def test_accelerated_optimizer_pickling(self):
+        model = torch.nn.Linear(10, 10)
+        optimizer = torch.optim.SGD(model.parameters(), 0.1)
+        accelerator = Accelerator()
+        optimizer = accelerator.prepare(optimizer)
+        try:
+            pickle.loads(pickle.dumps(optimizer))
+        except Exception as e:
+            self.fail(f"Accelerated optimizer pickling failed with {e}")
+        AcceleratorState._reset_state()

_The documentation is not available anymore as the PR was closed or merged._
Perfect, thanks!
https://api.github.com/repos/huggingface/accelerate/pulls/811
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 7c8534e93..df4f703c3 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -61,180 +61,367 @@
 
 logger = logging.getLogger(__name__)
 
+options_to_group = {
+    "--multi-gpu": "Distributed GPUs",
+    "--tpu": "TPU",
+    "--use_mps_device": "MPS",
+    "--use_deepspeed": "DeepSpeed Arguments",
+    "--use_fsdp": "FSDP Arguments",
+    "--use_megatron_lm": "Megatron-LM Arguments",
+}
+
+
+def clean_option(option):
+    "Finds all cases of - after the first two characters and changes them to _"
+    if option.startswith("--"):
+        return option[:3] + option[3:].replace("-", "_")
+
+
+class _CustomHelpAction(argparse._HelpAction):
+    """
+    This is a custom help action that will hide all arguments that are not used in the command line when the help is
+    called. This is useful for the case where the user is using a specific platform and only wants to see the arguments
+    for that platform.
+    """
+
+    def __call__(self, parser, namespace, values, option_string=None):
+        if "accelerate" in sys.argv[0] and "launch" in sys.argv[1:]:
+            args = sys.argv[2:]
+        else:
+            args = sys.argv[1:]
+        opts = parser._actions
+        titles = [
+            "Hardware Selection Arguments",
+            "Resource Selection Arguments",
+            "Training Paradigm Arguments",
+            "positional arguments",
+            "optional arguments",
+        ]
+        if len(args) > 1:
+            used_platforms = [arg for arg in args if arg in options_to_group.keys()]
+            args = list(map(clean_option, args))
+            used_titles = [options_to_group[o] for o in used_platforms]
+            for i, arg in enumerate(opts):
+                # If the argument's container is outside of the used titles, hide it
+                if arg.container.title not in titles + used_titles:
+                    setattr(opts[i], "help", argparse.SUPPRESS)
+                # If the argument is hardware selection, but not being passed, hide it
+                elif arg.container.title == "Hardware Selection Arguments":
+                    if set(arg.option_strings).isdisjoint(set(args)):
+                        setattr(opts[i], "help", argparse.SUPPRESS)
+                    else:
+                        setattr(opts[i], "help", arg.help + " (currently selected)")
+                # If the argument is a training paradigm, but not being passed, hide it
+                elif arg.container.title == "Training Paradigm Arguments":
+                    if set(arg.option_strings).isdisjoint(set(used_platforms)):
+                        setattr(opts[i], "help", argparse.SUPPRESS)
+                    else:
+                        setattr(opts[i], "help", arg.help + " (currently selected)")
+            for i, group in enumerate(list(parser._action_groups)):
+                # If all arguments in the group are hidden, hide the group
+                if all([arg.help == argparse.SUPPRESS for arg in group._group_actions]):
+                    parser._action_groups.remove(group)
+
+        super().__call__(parser, namespace, values, option_string)
+
 
 def launch_command_parser(subparsers=None):
     if subparsers is not None:
-        parser = subparsers.add_parser("launch")
+        parser = subparsers.add_parser("launch", add_help=False)
     else:
-        parser = argparse.ArgumentParser("Accelerate launch command")
+        parser = argparse.ArgumentParser("Accelerate launch command", add_help=False)
+
+    parser.register("action", "help", _CustomHelpAction)
+    parser.add_argument("-h", "--help", action="help", help="Show this help message and exit.")
 
     parser.add_argument(
         "--config_file", default=None, help="The config file to use for the default values in the launching script."
     )
-    parser.add_argument(
+    # Hardware selection arguments
+    hardware_args = parser.add_argument_group(
+        "Hardware Selection Arguments", "Arguments for selecting the hardware to be used."
+    )
+    hardware_args.add_argument(
+        "--cpu", default=False, action="store_true", help="Whether or not to force the training on the CPU."
+    )
+    hardware_args.add_argument(
         "--multi_gpu",
         default=False,
         action="store_true",
         help="Whether or not this should launch a distributed GPU training.",
     )
-    parser.add_argument(
+    hardware_args.add_argument(
+        "--tpu", default=False, action="store_true", help="Whether or not this should launch a TPU training."
+    )
+    hardware_args.add_argument(
         "--use_mps_device",
         default=False,
         action="store_true",
         help="Whether or not this should use MPS-enabled GPU device on MacOS machines.",
     )
 
-    # deepspeed args
-    parser.add_argument(
+    # Resource selection arguments
+    resource_args = parser.add_argument_group(
+        "Resource Selection Arguments", "Arguments for fine-tuning how available hardware should be used."
+    )
+    resource_args.add_argument(
+        "--mixed_precision",
+        type=str,
+        choices=["no", "fp16", "bf16"],
+        help="Whether or not to use mixed precision training. "
+        "Choose between FP16 and BF16 (bfloat16) training. "
+        "BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.",
+    )
+    resource_args.add_argument(
+        "--fp16", default=False, action="store_true", help="Whether or not to use mixed precision training."
+    )
+    resource_args.add_argument(
+        "--num_processes", type=int, default=None, help="The total number of processes to be launched in parallel."
+    )
+    resource_args.add_argument(
+        "--num_machines", type=int, default=None, help="The total number of machines used in this training."
+    )
+    resource_args.add_argument(
+        "--num_cpu_threads_per_process",
+        type=int,
+        default=None,
+        help="The number of CPU threads per process. Can be tuned for optimal performance.",
+    )
+
+    # Training Paradigm arguments
+    paradigm_args = parser.add_argument_group(
+        "Training Paradigm Arguments", "Arguments for selecting which training paradigm to be used."
+    )
+    paradigm_args.add_argument(
         "--use_deepspeed",
         default=False,
         action="store_true",
         help="Whether to use deepspeed.",
     )
+    paradigm_args.add_argument(
+        "--use_fsdp",
+        default=False,
+        action="store_true",
+        help="Whether to use fsdp.",
+    )
+    paradigm_args.add_argument(
+        "--use_megatron_lm",
+        default=False,
+        action="store_true",
+        help="Whether to use Megatron-LM.",
+    )
+
+    # distributed GPU training arguments
+    distributed_args = parser.add_argument_group("Distributed GPUs", "Arguments related to distributed GPU training.")
+    distributed_args.add_argument(
+        "--gpu_ids",
+        default=None,
+        help="What GPUs (by id) should be used for training on this machine as a comma-seperated list",
+    )
+    distributed_args.add_argument(
+        "--same_network",
+        default=False,
+        action="store_true",
+        help="Whether all machines used for multinode training exist on the same local network.",
+    )
+    distributed_args.add_argument(
+        "--machine_rank", type=int, default=None, help="The rank of the machine on which this script is launched."
+    )
+    distributed_args.add_argument(
+        "--main_process_ip", type=str, default=None, help="The IP address of the machine of rank 0."
+    )
+    distributed_args.add_argument(
+        "--main_process_port",
+        type=int,
+        default=None,
+        help="The port to use to communicate with the machine of rank 0.",
+    )
+    # Rendezvous related arguments
+    distributed_args.add_argument(
+        "--rdzv_conf",
+        type=str,
+        default="",
+        help="Additional rendezvous configuration (<key1>=<value1>,<key2>=<value2>,...).",
+    )
+    distributed_args.add_argument(
+        "--max_restarts",
+        type=int,
+        default=0,
+        help="Maximum number of worker group restarts before failing.",
+    )
+    distributed_args.add_argument(
+        "--monitor_interval",
+        type=float,
+        default=5,
+        help="Interval, in seconds, to monitor the state of workers.",
+    )
+    parser.add_argument(
+        "-m",
+        "--module",
+        action="store_true",
+        help="Change each process to interpret the launch script as a Python module, executing with the same behavior as 'python -m'.",
+    )
     parser.add_argument(
+        "--no_python",
+        action="store_true",
+        help="Skip prepending the training script with 'python' - just execute it directly. Useful when the script is not a Python script.",
+    )
+
+    # tpu arguments
+    tpu_args = parser.add_argument_group("TPU", "Arguments related to TPU.")
+    tpu_args.add_argument(
+        "--main_training_function",
+        type=str,
+        default=None,
+        help="The name of the main function to be executed in your script (only for TPU training).",
+    )
+    tpu_args.add_argument(
+        "--downcast_bf16",
+        action="store_true",
+        help="Whether when using bf16 precision on TPUs if both float and double tensors are cast to bfloat16 or if double tensors remain as float32.",
+    )
+
+    # DeepSpeed arguments
+    deepspeed_args = parser.add_argument_group("DeepSpeed Arguments", "Arguments related to DeepSpeed.")
+    deepspeed_args.add_argument(
         "--deepspeed_config_file",
         default=None,
         type=str,
         help="DeepSpeed config file.",
     )
-    parser.add_argument(
+    deepspeed_args.add_argument(
         "--zero_stage",
         default=None,
         type=int,
         help="DeepSpeed's ZeRO optimization stage (useful only when `use_deepspeed` flag is passed).",
     )
-    parser.add_argument(
+    deepspeed_args.add_argument(
         "--offload_optimizer_device",
         default=None,
         type=str,
         help="Decides where (none|cpu|nvme) to offload optimizer states (useful only when `use_deepspeed` flag is passed).",
     )
-    parser.add_argument(
+    deepspeed_args.add_argument(
         "--offload_param_device",
         default=None,
         type=str,
         help="Decides where (none|cpu|nvme) to offload parameters (useful only when `use_deepspeed` flag is passed).",
     )
-    parser.add_argument(
+    deepspeed_args.add_argument(
         "--gradient_accumulation_steps",
         default=None,
         type=int,
         help="No of gradient_accumulation_steps used in your training script (useful only when `use_deepspeed` flag is passed).",
     )
-    parser.add_argument(
+    deepspeed_args.add_argument(
         "--gradient_clipping",
         default=None,
         type=float,
         help="gradient clipping value used in your training script (useful only when `use_deepspeed` flag is passed).",
     )
-    parser.add_argument(
+    deepspeed_args.add_argument(
         "--zero3_init_flag",
         default=None,
         type=str,
         help="Decides Whether (true|false) to enable `deepspeed.zero.Init` for constructing massive models. "
         "Only applicable with DeepSpeed ZeRO Stage-3.",
     )
-    parser.add_argument(
+    deepspeed_args.add_argument(
         "--zero3_save_16bit_model",
         default=None,
         type=str,
         help="Decides Whether (true|false) to save 16-bit model weights when using ZeRO Stage-3. "
         "Only applicable with DeepSpeed ZeRO Stage-3.",
     )
-    parser.add_argument(
+    deepspeed_args.add_argument(
         "--deepspeed_hostfile",
         default=None,
         type=str,
         help="DeepSpeed hostfile for configuring multi-node compute resources.",
     )
-    parser.add_argument(
+    deepspeed_args.add_argument(
         "--deepspeed_exclusion_filter",
         default=None,
         type=str,
         help="DeepSpeed exclusion filter string when using mutli-node setup.",
     )
-    parser.add_argument(
+    deepspeed_args.add_argument(
         "--deepspeed_inclusion_filter",
         default=None,
         type=str,
         help="DeepSpeed inclusion filter string when using mutli-node setup.",
     )
-    parser.add_argument(
+    deepspeed_args.add_argument(
         "--deepspeed_multinode_launcher",
         default=None,
         type=str,
         help="DeepSpeed multi-node launcher to use.",
     )
 
-    # fsdp args
-    parser.add_argument(
-        "--use_fsdp",
-        default=False,
-        action="store_true",
-        help="Whether to use fsdp.",
-    )
-    parser.add_argument(
+    # fsdp arguments
+    fsdp_args = parser.add_argument_group("FSDP Arguments", "Arguments related to Fully Shared Data Parallelism.")
+    fsdp_args.add_argument(
         "--fsdp_offload_params",
         default="false",
         type=str,
         help="Decides Whether (true|false) to offload parameters and gradients to CPU. (useful only when `use_fsdp` flag is passed).",
     )
-    parser.add_argument(
+    fsdp_args.add_argument(
         "--fsdp_min_num_params",
         type=int,
         default=1e8,
         help="FSDP's minimum number of parameters for Default Auto Wrapping. (useful only when `use_fsdp` flag is passed).",
     )
-    parser.add_argument(
+    fsdp_args.add_argument(
         "--fsdp_sharding_strategy",
         type=int,
         default=1,
         help="FSDP's Sharding Strategy. (useful only when `use_fsdp` flag is passed).",
     )
-    parser.add_argument(
+    fsdp_args.add_argument(
         "--fsdp_auto_wrap_policy",
         type=str,
         default=None,
         help="FSDP's auto wrap policy. (useful only when `use_fsdp` flag is passed).",
     )
-    parser.add_argument(
+    fsdp_args.add_argument(
         "--fsdp_transformer_layer_cls_to_wrap",
         default=None,
         type=str,
         help="Transformer layer class name (case-sensitive) to wrap ,e.g, `BertLayer`, `GPTJBlock`, `T5Block` .... "
         "(useful only when `use_fsdp` flag is passed).",
     )
-    parser.add_argument(
+    fsdp_args.add_argument(
         "--fsdp_backward_prefetch_policy",
         default=None,
         type=str,
         help="FSDP's backward prefetch policy. (useful only when `use_fsdp` flag is passed).",
     )
-    parser.add_argument(
+    fsdp_args.add_argument(
         "--fsdp_state_dict_type",
         default=None,
         type=str,
         help="FSDP's state dict type. (useful only when `use_fsdp` flag is passed).",
     )
-    parser.add_argument(
+    fsdp_args.add_argument(
         "--offload_params",
         default=None,
         type=str,
         help="This argument is deprecated. Use `fsdp_offload_params` instead.",
     )
-    parser.add_argument(
+    fsdp_args.add_argument(
         "--min_num_params",
         type=int,
         default=None,
         help="This argument is deprecated. Use `fsdp_min_num_params` instead.",
     )
-    parser.add_argument(
+    fsdp_args.add_argument(
         "--sharding_strategy",
         type=int,
         default=None,
         help="This argument is deprecated. Use `fsdp_sharding_strategy` instead.",
     )
-    parser.add_argument(
+    fsdp_args.add_argument(
         "--transformer_layer_cls_to_wrap",
         default=None,
         type=str,
@@ -242,45 +429,40 @@ def launch_command_parser(subparsers=None):
     )
 
     # megatron_lm args
-    parser.add_argument(
-        "--use_megatron_lm",
-        default=False,
-        action="store_true",
-        help="Whether to use Megatron-LM.",
-    )
-    parser.add_argument(
+    megatron_lm_args = parser.add_argument_group("Megatron-LM Arguments", "Arguments related to Megatron-LM.")
+    megatron_lm_args.add_argument(
         "--megatron_lm_tp_degree",
         type=int,
         default=1,
         help="Megatron-LM's Tensor Parallelism (TP) degree. (useful only when `use_megatron_lm` flag is passed).",
     )
-    parser.add_argument(
+    megatron_lm_args.add_argument(
         "--megatron_lm_pp_degree",
         type=int,
         default=1,
         help="Megatron-LM's Pipeline Parallelism (PP) degree. (useful only when `use_megatron_lm` flag is passed).",
     )
-    parser.add_argument(
+    megatron_lm_args.add_argument(
         "--megatron_lm_num_micro_batches",
         type=int,
         default=None,
         help="Megatron-LM's number of micro batches when PP degree > 1. (useful only when `use_megatron_lm` flag is passed).",
     )
-    parser.add_argument(
+    megatron_lm_args.add_argument(
         "--megatron_lm_sequence_parallelism",
         default=None,
         type=str,
         help="Decides Whether (true|false) to enable Sequence Parallelism when TP degree > 1. "
         "(useful only when `use_megatron_lm` flag is passed).",
     )
-    parser.add_argument(
+    megatron_lm_args.add_argument(
         "--megatron_lm_recompute_activations",
         default=None,
         type=str,
         help="Decides Whether (true|false) to enable Selective Activation Recomputation. "
         "(useful only when `use_megatron_lm` flag is passed).",
     )
-    parser.add_argument(
+    megatron_lm_args.add_argument(
         "--megatron_lm_use_distributed_optimizer",
         default=None,
         type=str,
@@ -288,7 +470,7 @@ def launch_command_parser(subparsers=None):
         "which shards optimizer state and gradients across Data Pralellel (DP) ranks. "
         "(useful only when `use_megatron_lm` flag is passed).",
     )
-    parser.add_argument(
+    megatron_lm_args.add_argument(
         "--megatron_lm_gradient_clipping",
         default=1.0,
         type=float,
@@ -296,105 +478,15 @@ def launch_command_parser(subparsers=None):
         "(useful only when `use_megatron_lm` flag is passed).",
     )
 
-    parser.add_argument(
-        "--tpu", default=False, action="store_true", help="Whether or not this should launch a TPU training."
-    )
-    parser.add_argument(
-        "--mixed_precision",
-        type=str,
-        choices=["no", "fp16", "bf16"],
-        help="Whether or not to use mixed precision training. "
-        "Choose between FP16 and BF16 (bfloat16) training. "
-        "BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.",
-    )
-
-    parser.add_argument(
-        "--fp16", default=False, action="store_true", help="Whether or not to use mixed precision training."
-    )
-    parser.add_argument(
-        "--cpu", default=False, action="store_true", help="Whether or not to force the training on the CPU."
-    )
-    parser.add_argument(
-        "--num_processes", type=int, default=None, help="The total number of processes to be launched in parallel."
-    )
-    parser.add_argument(
-        "--num_machines", type=int, default=None, help="The total number of machines used in this training."
-    )
-    parser.add_argument(
-        "--gpu_ids",
-        default=None,
-        help="What GPUs (by id) should be used for training on this machine as a comma-seperated list",
-    )
-    parser.add_argument(
-        "--same_network",
-        default=False,
-        action="store_true",
-        help="Whether all machines used for multinode training exist on the same local network.",
-    )
-    parser.add_argument(
-        "--machine_rank", type=int, default=None, help="The rank of the machine on which this script is launched."
-    )
-    parser.add_argument("--main_process_ip", type=str, default=None, help="The IP address of the machine of rank 0.")
-    parser.add_argument(
-        "--main_process_port",
-        type=int,
-        default=None,
-        help="The port to use to communicate with the machine of rank 0.",
-    )
-    # Rendezvous related arguments
-    parser.add_argument(
-        "--rdzv_conf",
-        type=str,
-        default="",
-        help="Additional rendezvous configuration (<key1>=<value1>,<key2>=<value2>,...).",
-    )
-    parser.add_argument(
-        "--max_restarts",
-        type=int,
-        default=0,
-        help="Maximum number of worker group restarts before failing.",
-    )
-    parser.add_argument(
-        "--monitor_interval",
-        type=float,
-        default=5,
-        help="Interval, in seconds, to monitor the state of workers.",
-    )
-    parser.add_argument(
-        "--main_training_function",
-        type=str,
-        default=None,
-        help="The name of the main function to be executed in your script (only for TPU training).",
-    )
-    parser.add_argument(
-        "--downcast_bf16",
-        action="store_true",
-        help="Whether when using bf16 precision on TPUs if both float and double tensors are cast to bfloat16 or if double tensors remain as float32.",
-    )
-    parser.add_argument(
-        "-m",
-        "--module",
-        action="store_true",
-        help="Change each process to interpret the launch script as a Python module, executing with the same behavior as 'python -m'.",
-    )
-    parser.add_argument(
-        "--no_python",
-        action="store_true",
-        help="Skip prepending the training script with 'python' - just execute it directly. Useful when the script is not a Python script.",
-    )
-    parser.add_argument(
-        "--num_cpu_threads_per_process",
-        type=int,
-        default=None,
-        help="The number of CPU threads per process. Can be tuned for optimal performance.",
-    )
-    parser.add_argument(
+    # AWS arguments
+    aws_args = parser.add_argument_group("AWS Arguments", "Arguments related to AWS.")
+    aws_args.add_argument(
         "--aws_access_key_id",
         type=str,
         default=None,
         help="The AWS_ACCESS_KEY_ID used to launch the Amazon SageMaker training job",
     )
-    parser.add_argument(
+    aws_args.add_argument(
         "--aws_secret_access_key",
         type=str,
         default=None,

I hate this underscore convention in Python as it hurts readability and nothing is ever truly private. In Accelerate, like for Transformers, the convention is that things that are not in the main init are private. So no need to add all those _ :-)
@@ -61,340 +61,430 @@
 
 logger = logging.getLogger(__name__)
 
+_options_to_group = {
+    "--multi-gpu": "Distributed GPUs",
+    "--tpu": "TPU",
+    "--use_mps_device": "MPS",
+    "--use_deepspeed": "DeepSpeed",
+    "--use_fsdp": "FSDP",
+    "--use_megatron_lm": "Megatron-LM",
+}
+
+
+def _clean_option(option):
Made public! :) Will note down for next time
@@ -61,340 +61,430 @@
 
 logger = logging.getLogger(__name__)
 
+_options_to_group = {
+    "--multi-gpu": "Distributed GPUs",
+    "--tpu": "TPU",
+    "--use_mps_device": "MPS",
+    "--use_deepspeed": "DeepSpeed",
+    "--use_fsdp": "FSDP",
+    "--use_megatron_lm": "Megatron-LM",
+}
+
+
+def _clean_option(option):
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/810
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 81759962c..287e6899a 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -647,7 +647,15 @@ def _prepare_fsdp(self, *args):
                         "FSDP Warning: When using FSDP, several parameter groups will be conflated into "
                         "a single one due to nested module wrapping and parameter flattening."
                     )
-                optimizer = obj.optimizer.__class__(model.parameters(), **obj.optimizer.defaults)
+                try:
+                    optimizer = obj.optimizer.__class__(model.parameters(), **obj.optimizer.defaults)
+                except TypeError:
+                    if "differentiable" in obj.optimizer.defaults:
+                        # https://github.com/huggingface/accelerate/issues/801
+                        defaults = {k: v for k, v in obj.optimizer.defaults.items() if k != "differentiable"}
+                        optimizer = obj.optimizer.__class__(model.parameters(), **defaults)
+                    else:
+                        raise
                 obj = self.prepare_optimizer(optimizer)
                 optimizers.append(obj)
             elif isinstance(obj, torch.nn.Module):

This will delete inplace, we should probably create a new dict here
@@ -647,7 +647,15 @@ def _prepare_fsdp(self, *args):
                         "FSDP Warning: When using FSDP, several parameter groups will be conflated into "
                         "a single one due to nested module wrapping and parameter flattening."
                     )
-                optimizer = obj.optimizer.__class__(model.parameters(), **obj.optimizer.defaults)
+                try:
+                    optimizer = obj.optimizer.__class__(model.parameters(), **obj.optimizer.defaults)
+                except TypeError:
+                    if "differentiable" in obj.optimizer.defaults:
+                        # https://github.com/huggingface/accelerate/issues/801
+                        del obj.optimizer.defaults["differentiable"]
```suggestion
                        defaults = obj.optimizer.defaults.copy()
                        del defaults["differentiable"]
                        optimizer = obj.optimizer.__class__(model.parameters(), **defaults)
```
@@ -647,7 +647,15 @@ def _prepare_fsdp(self, *args):
                         "FSDP Warning: When using FSDP, several parameter groups will be conflated into "
                         "a single one due to nested module wrapping and parameter flattening."
                     )
-                optimizer = obj.optimizer.__class__(model.parameters(), **obj.optimizer.defaults)
+                try:
+                    optimizer = obj.optimizer.__class__(model.parameters(), **obj.optimizer.defaults)
+                except TypeError:
+                    if "differentiable" in obj.optimizer.defaults:
+                        # https://github.com/huggingface/accelerate/issues/801
+                        del obj.optimizer.defaults["differentiable"]
+                        optimizer = obj.optimizer.__class__(model.parameters(), **obj.optimizer.defaults)
Good point. Done. PTAL.
@@ -647,7 +647,15 @@ def _prepare_fsdp(self, *args):
                         "FSDP Warning: When using FSDP, several parameter groups will be conflated into "
                         "a single one due to nested module wrapping and parameter flattening."
                     )
-                optimizer = obj.optimizer.__class__(model.parameters(), **obj.optimizer.defaults)
+                try:
+                    optimizer = obj.optimizer.__class__(model.parameters(), **obj.optimizer.defaults)
+                except TypeError:
+                    if "differentiable" in obj.optimizer.defaults:
+                        # https://github.com/huggingface/accelerate/issues/801
+                        del obj.optimizer.defaults["differentiable"]
```suggestion
                        defaults = {k: v for k, v in obj.optimizer.defaults.items() if k != "differentiable"}
```
@@ -647,7 +647,16 @@ def _prepare_fsdp(self, *args):
                         "FSDP Warning: When using FSDP, several parameter groups will be conflated into "
                         "a single one due to nested module wrapping and parameter flattening."
                     )
-                optimizer = obj.optimizer.__class__(model.parameters(), **obj.optimizer.defaults)
+                try:
+                    optimizer = obj.optimizer.__class__(model.parameters(), **obj.optimizer.defaults)
+                except TypeError:
+                    if "differentiable" in obj.optimizer.defaults:
+                        # https://github.com/huggingface/accelerate/issues/801
+                        defaults = obj.optimizer.defaults.copy()
+                        del defaults["differentiable"]
closing #801
_The documentation is not available anymore as the PR was closed or merged._
Thank you for the quick fix! 🤗
https://api.github.com/repos/huggingface/accelerate/pulls/803
diff --git a/src/accelerate/commands/accelerate_cli.py b/src/accelerate/commands/accelerate_cli.py
index 6300c7c1d..515a66d1a 100644
--- a/src/accelerate/commands/accelerate_cli.py
+++ b/src/accelerate/commands/accelerate_cli.py
@@ -19,6 +19,7 @@
 from accelerate.commands.config import config_command_parser
 from accelerate.commands.env import env_command_parser
 from accelerate.commands.launch import launch_command_parser
+from accelerate.commands.pod import pod_command_parser
 from accelerate.commands.test import test_command_parser
 
 
@@ -28,9 +29,10 @@ def main():
 
     # Register commands
     config_command_parser(subparsers=subparsers)
+    env_command_parser(subparsers=subparsers)
     launch_command_parser(subparsers=subparsers)
+    pod_command_parser(subparsers=subparsers)
     test_command_parser(subparsers=subparsers)
-    env_command_parser(subparsers=subparsers)
 
     # Let's go
     args = parser.parse_args()
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index 0fa3ceab2..acc4a1828 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -14,6 +14,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+
 from ...utils import ComputeEnvironment, DistributedType, is_deepspeed_available, is_transformers_available
 from ...utils.constants import (
     DEEPSPEED_MULTINODE_LAUNCHERS,
@@ -41,6 +43,10 @@ def get_cluster_input():
     main_process_port = None
     rdzv_backend = "static"
     same_network = True
+    tpu_name = None
+    tpu_zone = None
+    commands = None
+    command_file = None
     if distributed_type in [DistributedType.MULTI_GPU, DistributedType.MULTI_CPU]:
         num_machines = _ask_field(
             "How many different machines will you use (use more than 1 for multi-node training)? [1]: ",
@@ -341,6 +347,50 @@ def get_cluster_input():
             "What is the name of the function in your script that should be launched in all parallel scripts? [main]: ",
             default="main",
         )
+        use_cluster = _ask_field(
+            "Are you using a TPU cluster? [yes/NO]: ",
+            _convert_yes_no_to_bool,
+            default=False,
+            error_message="Please enter yes or no.",
+        )
+        if use_cluster:
+            tpu_name = _ask_field(
+                "What is the name of your TPU cluster? ",
+                default=None,
+                error_message="Please enter the name of your TPU cluster.",
+            )
+            tpu_zone = _ask_field(
+                "What is the zone of your TPU cluster? ",
+                default=None,
+                error_message="Please enter the zone of your TPU cluster.",
+            )
+            run_commands = _ask_field(
+                "Do you have code you wish to run on startup in each pod? [yes/NO]: ",
+                _convert_yes_no_to_bool,
+                default=False,
+                error_message="Please enter yes or no.",
+            )
+            if run_commands:
+                use_command_file = _ask_field(
+                    "Is this code located in a bash script? [yes/NO]: ",
+                    _convert_yes_no_to_bool,
+                    default=False,
+                    error_message="Please enter yes or no.",
+                )
+                if use_command_file:
+                    command_file = _ask_field(
+                        "What is the path to your bash script? ",
+                        default=None,
+                        error_message="Please enter the path to your bash script.",
+                    )
+                    command_file = os.path.abspath(command_file)
+                else:
+                    commands = _ask_field(
+                        "What commands do you wish to run on startup in each pod? ",
+                        default=None,
+                        error_message="Please enter the commands you wish to run on startup in each pod as a single string.",
+                    )
+
     else:
         main_training_function = "main"
 
@@ -408,4 +458,8 @@ def get_cluster_input():
         use_cpu=use_cpu,
         rdzv_backend=rdzv_backend,
         same_network=same_network,
+        tpu_name=tpu_name,
+        tpu_zone=tpu_zone,
+        commands=commands,
+        command_file=command_file,
     )
diff --git a/src/accelerate/commands/config/config_args.py b/src/accelerate/commands/config/config_args.py
index 61f585858..9a1247c55 100644
--- a/src/accelerate/commands/config/config_args.py
+++ b/src/accelerate/commands/config/config_args.py
@@ -18,7 +18,7 @@
 import os
 from dataclasses import dataclass
 from enum import Enum
-from typing import Optional, Union
+from typing import List, Optional, Union
 
 import yaml
 
@@ -151,6 +151,12 @@ class ClusterConfig(BaseConfig):
     # args for TPU
     downcast_bf16: bool = False
 
+    # args for TPU pods
+    tpu_name: str = None
+    tpu_zone: str = None
+    command_file: str = None
+    command: List[str] = None
+
     def __post_init__(self):
         if self.deepspeed_config is None:
             self.deepspeed_config = {}
diff --git a/src/accelerate/commands/env.py b/src/accelerate/commands/env.py
index b66008e1b..a19c04d4a 100644
--- a/src/accelerate/commands/env.py
+++ b/src/accelerate/commands/env.py
@@ -1,3 +1,19 @@
+#!/usr/bin/env python
+
+# 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 platform
diff --git a/src/accelerate/commands/pod.py b/src/accelerate/commands/pod.py
new file mode 100644
index 000000000..87cd4f904
--- /dev/null
+++ b/src/accelerate/commands/pod.py
@@ -0,0 +1,152 @@
+#!/usr/bin/env python
+
+# 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 subprocess
+
+from accelerate.commands.config.config_args import default_config_file, load_config_from_file
+from packaging.version import Version, parse
+
+
+_description = "Run commands across a pod of TPU VMs for initial setup before running `accelerate launch`. Will also install Accelerate on the pod."
+
+
+def pod_command_parser(subparsers=None):
+    if subparsers is not None:
+        parser = subparsers.add_parser("pod-config", description=_description)
+    else:
+        parser = argparse.ArgumentParser("Accelerate pod-config command", description=_description)
+
+    parser.add_argument(
+        "--config_file",
+        type=str,
+        default=None,
+        help="Path to the config file to use for accelerate.",
+    )
+
+    parser.add_argument(
+        "--pod_config_file",
+        type=str,
+        default=None,
+        help="Path to the config file to use for the pod.",
+    )
+
+    parser.add_argument(
+        "--command_file",
+        default=None,
+        help="The path to the file containing the commands to run on the pod on startup.",
+    )
+    parser.add_argument(
+        "--command",
+        action="append",
+        nargs="+",
+        help="A command to run on the pod. If not specified, will use the command specified in the command file.",
+    )
+    parser.add_argument(
+        "--tpu_name",
+        default=None,
+        help="The name of the TPU to use. If not specified, will use the TPU specified in the config file.",
+    )
+    parser.add_argument(
+        "--tpu_zone",
+        default=None,
+        help="The zone of the TPU to use. If not specified, will use the zone specified in the config file.",
+    )
+    parser.add_argument(
+        "--install_accelerate",
+        action="store_true",
+        help="Whether to install accelerate on the pod. Defaults to False.",
+    )
+    parser.add_argument(
+        "--accelerate_version",
+        default="latest",
+        help="The version of accelerate to install on the pod. If not specified, will use the latest pypi version. Specify 'dev' to install from GitHub.",
+    )
+    parser.add_argument(
+        "--debug", action="store_true", help="If set, will print the command that would be run instead of running it."
+    )
+
+    if subparsers is not None:
+        parser.set_defaults(func=pod_launcher)
+    return parser
+
+
+def pod_launcher(args):
+    defaults = None
+
+    # Get the default from the config file if it exists.
+    if args.config_file is not None or os.path.isfile(default_config_file):
+        defaults = load_config_from_file(args.config_file)
+        if not args.command_file and defaults.command_file is not None and not args.command:
+            args.command_file = defaults.command_file
+        if not args.command and defaults.command is not None:
+            args.command = defaults.command
+        if not args.tpu_name:
+            args.tpu_name = defaults.tpu_name
+        if not args.tpu_zone:
+            args.tpu_zone = defaults.tpu_zone
+    if args.accelerate_version == "dev":
+        args.accelerate_version = "git+https://github.com/huggingface/accelerate.git"
+    elif args.accelerate_version == "latest":
+        args.accelerate_version = "accelerate -U"
+    elif isinstance(parse(args.accelerate_version), Version):
+        args.accelerate_version = f"accelerate=={args.accelerate_version}"
+
+    if not args.command_file and not args.command:
+        raise ValueError("You must specify either a command file or a command to run on the pod.")
+
+    if args.command_file:
+        with open(args.command_file, "r") as f:
+            args.command = [f.read().splitlines()]
+
+    # To turn list of lists into list of strings
+    args.command = [line for cmd in args.command for line in cmd]
+    # Default to the shared folder and install accelerate
+    new_cmd = ["cd /usr/share"]
+    if args.install_accelerate:
+        new_cmd += [f"pip install {args.accelerate_version}"]
+    new_cmd += args.command
+    args.command = "; ".join(new_cmd)
+
+    # Then send it to gcloud
+    # Eventually try to use google-api-core to do this instead of subprocess
+    cmd = [
+        "gcloud",
+        "compute",
+        "tpus",
+        "tpu-vm",
+        "ssh",
+        args.tpu_name,
+        "--zone",
+        args.tpu_zone,
+        "--command",
+        args.command,
+        "--worker",
+        "all",
+    ]
+    if args.debug:
+        print(f"Running {' '.join(cmd)}")
+        return
+    subprocess.run(cmd)
+    print("Successfully setup pod.")
+
+
+def main():
+    parser = pod_command_parser()
+    args = parser.parse_args()
+
+    pod_launcher(args)
diff --git a/tests/test_cli.py b/tests/test_cli.py
index ceed58290..2479f736d 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -21,6 +21,7 @@
 
 import accelerate
 from accelerate.test_utils import execute_subprocess_async
+from accelerate.test_utils.testing import run_command
 
 
 class AccelerateLauncherTester(unittest.TestCase):
@@ -63,3 +64,151 @@ def test_config_compatibility(self):
                 execute_subprocess_async(
                     self.base_cmd + ["--config_file", str(config), self.test_file_path], env=os.environ.copy()
                 )
+
+
+class PodConfigTester(unittest.TestCase):
+    """
+    Test case for verifying the `accelerate pod-config` CLI passes the right `gcloud` command.
+    """
+
+    tpu_name = "test-tpu"
+    tpu_zone = "us-central1-a"
+    command = "ls"
+    cmd = ["accelerate", "pod-config"]
+    base_output = "cd /usr/share"
+    command_file = "tests/test_samples/test_command_file.sh"
+    gcloud = "Running gcloud compute tpus tpu-vm ssh"
+
+    @staticmethod
+    def clean_output(output):
+        return "".join(output).rstrip()
+
+    def test_base(self):
+        output = run_command(
+            self.cmd
+            + ["--command", self.command, "--tpu_zone", self.tpu_zone, "--tpu_name", self.tpu_name, "--debug"],
+            return_stdout=True,
+        )
+        self.assertEqual(
+            self.clean_output(output),
+            f"{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls --worker all",
+        )
+
+    def test_base_backward_compatibility(self):
+        output = run_command(
+            self.cmd
+            + [
+                "--config_file",
+                "tests/test_configs/0_12_0.yaml",
+                "--command",
+                self.command,
+                "--tpu_zone",
+                self.tpu_zone,
+                "--tpu_name",
+                self.tpu_name,
+                "--debug",
+            ],
+            return_stdout=True,
+        )
+        self.assertEqual(
+            self.clean_output(output),
+            f"{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls --worker all",
+        )
+
+    def test_with_config_file(self):
+        output = run_command(
+            self.cmd + ["--config_file", "tests/test_configs/latest.yaml", "--debug"], return_stdout=True
+        )
+        self.assertEqual(
+            self.clean_output(output),
+            f'{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; echo "hello world"; echo "this is a second command" --worker all',
+        )
+
+    def test_with_config_file_and_command(self):
+        output = run_command(
+            self.cmd + ["--config_file", "tests/test_configs/latest.yaml", "--command", self.command, "--debug"],
+            return_stdout=True,
+        )
+        self.assertEqual(
+            self.clean_output(output),
+            f"{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls --worker all",
+        )
+
+    def test_with_config_file_and_multiple_command(self):
+        output = run_command(
+            self.cmd
+            + [
+                "--config_file",
+                "tests/test_configs/latest.yaml",
+                "--command",
+                self.command,
+                "--command",
+                'echo "Hello World"',
+                "--debug",
+            ],
+            return_stdout=True,
+        )
+        self.assertEqual(
+            self.clean_output(output),
+            f'{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls; echo "Hello World" --worker all',
+        )
+
+    def test_with_config_file_and_command_file(self):
+        output = run_command(
+            self.cmd
+            + ["--config_file", "tests/test_configs/latest.yaml", "--command_file", self.command_file, "--debug"],
+            return_stdout=True,
+        )
+        self.assertEqual(
+            self.clean_output(output),
+            f'{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; echo "hello world"; echo "this is a second command" --worker all',
+        )
+
+    def test_with_config_file_and_command_file_backward_compatibility(self):
+        output = run_command(
+            self.cmd
+            + [
+                "--config_file",
+                "tests/test_configs/0_12_0.yaml",
+                "--command_file",
+                self.command_file,
+                "--tpu_zone",
+                self.tpu_zone,
+                "--tpu_name",
+                self.tpu_name,
+                "--debug",
+            ],
+            return_stdout=True,
+        )
+        self.assertEqual(
+            self.clean_output(output),
+            f'{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; echo "hello world"; echo "this is a second command" --worker all',
+        )
+
+    def test_accelerate_install(self):
+        output = run_command(
+            self.cmd + ["--config_file", "tests/test_configs/latest.yaml", "--install_accelerate", "--debug"],
+            return_stdout=True,
+        )
+        self.assertEqual(
+            self.clean_output(output),
+            f'{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; pip install accelerate -U; echo "hello world"; echo "this is a second command" --worker all',
+        )
+
+    def test_accelerate_install_version(self):
+        output = run_command(
+            self.cmd
+            + [
+                "--config_file",
+                "tests/test_configs/latest.yaml",
+                "--install_accelerate",
+                "--accelerate_version",
+                "12.0.0",
+                "--debug",
+            ],
+            return_stdout=True,
+        )
+        self.assertEqual(
+            self.clean_output(output),
+            f'{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; pip install accelerate==12.0.0; echo "hello world"; echo "this is a second command" --worker all',
+        )
diff --git a/tests/test_configs/latest.yaml b/tests/test_configs/latest.yaml
index 87b294cf8..de6be03a4 100644
--- a/tests/test_configs/latest.yaml
+++ b/tests/test_configs/latest.yaml
@@ -15,3 +15,7 @@ num_processes: 1
 rdzv_backend: static
 same_network: true
 use_cpu: false
+tpu_name: 'test-tpu'
+tpu_zone: 'us-central1-a'
+command: null
+command_file: tests/test_samples/test_command_file.sh
\ No newline at end of file
diff --git a/tests/test_samples/test_command_file.sh b/tests/test_samples/test_command_file.sh
new file mode 100644
index 000000000..592a7d532
--- /dev/null
+++ b/tests/test_samples/test_command_file.sh
@@ -0,0 +1,2 @@
+echo "hello world"
+echo "this is a second command"
\ No newline at end of file

```suggestion
                "Do you have code you wish to run on startup in each pod? [yes/NO]: ",
```
@@ -341,6 +347,50 @@ def get_cluster_input():
             "What is the name of the function in your script that should be launched in all parallel scripts? [main]: ",
             default="main",
         )
+        use_cluster = _ask_field(
+            "Are you using a TPU cluster? [yes/NO]: ",
+            _convert_yes_no_to_bool,
+            default=False,
+            error_message="Please enter yes or no.",
+        )
+        if use_cluster:
+            tpu_name = _ask_field(
+                "What is the name of your TPU cluster? ",
+                default=None,
+                error_message="Please enter the name of your TPU cluster.",
+            )
+            tpu_zone = _ask_field(
+                "What is the zone of your TPU cluster? ",
+                default=None,
+                error_message="Please enter the zone of your TPU cluster.",
+            )
+            run_commands = _ask_field(
+                "Do you have code you wish to be ran on startup in each pod? [yes/NO]: ",
We can allow for mix and match by appending commands from the `args.command_file` to the `args.command`? Would it be better than raising error?
@@ -0,0 +1,152 @@
+#!/usr/bin/env python
+
+# 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 subprocess
+
+from accelerate.commands.config.config_args import default_config_file, load_config_from_file
+from packaging.version import Version, parse
+
+
+_description = "Run commands across a pod of TPU VMs for initial setup before running `accelerate launch`. Will also install Accelerate on the pod."
+
+
+def pod_command_parser(subparsers=None):
+    if subparsers is not None:
+        parser = subparsers.add_parser("pod-config", description=_description)
+    else:
+        parser = argparse.ArgumentParser("Accelerate pod-config command", description=_description)
+
+    parser.add_argument(
+        "--config_file",
+        type=str,
+        default=None,
+        help="Path to the config file to use for accelerate.",
+    )
+
+    parser.add_argument(
+        "--pod_config_file",
+        type=str,
+        default=None,
+        help="Path to the config file to use for the pod.",
+    )
+
+    parser.add_argument(
+        "--command_file",
+        default=None,
+        help="The path to the file containing the commands to run on the pod on startup.",
+    )
+    parser.add_argument(
+        "--command",
+        action="append",
+        nargs="+",
+        help="A command to run on the pod. If not specified, will use the command specified in the command file.",
+    )
+    parser.add_argument(
+        "--tpu_name",
+        default=None,
+        help="The name of the TPU to use. If not specified, will use the TPU specified in the config file.",
+    )
+    parser.add_argument(
+        "--tpu_zone",
+        default=None,
+        help="The zone of the TPU to use. If not specified, will use the zone specified in the config file.",
+    )
+    parser.add_argument(
+        "--install_accelerate",
+        action="store_true",
+        help="Whether to install accelerate on the pod. Defaults to False.",
+    )
+    parser.add_argument(
+        "--accelerate_version",
+        default="latest",
+        help="The version of accelerate to install on the pod. If not specified, will use the latest pypi version. Specify 'dev' to install from GitHub.",
+    )
+    parser.add_argument(
+        "--debug", action="store_true", help="If set, will print the command that would be run instead of running it."
+    )
+
+    if subparsers is not None:
+        parser.set_defaults(func=pod_launcher)
+    return parser
+
+
+def pod_launcher(args):
+    defaults = None
+
+    # Get the default from the config file if it exists.
+    if args.config_file is not None or os.path.isfile(default_config_file):
+        defaults = load_config_from_file(args.config_file)
+        if not args.command_file and defaults.command_file is not None and not args.command:
+            args.command_file = defaults.command_file
+        if not args.command and defaults.command is not None:
+            args.command = defaults.command
+        if not args.tpu_name:
+            args.tpu_name = defaults.tpu_name
+        if not args.tpu_zone:
+            args.tpu_zone = defaults.tpu_zone
+    if args.accelerate_version == "dev":
+        args.accelerate_version = "git+https://github.com/huggingface/accelerate.git"
+    elif args.accelerate_version == "latest":
+        args.accelerate_version = "accelerate -U"
+    elif isinstance(parse(args.accelerate_version), Version):
+        args.accelerate_version = f"accelerate=={args.accelerate_version}"
+
+    if not args.command_file and not args.command:
+        raise ValueError("You must specify either a command file or a command to run on the pod.")
```suggestion
        print(f"Running {' '.join(cmd)}")
```
@@ -0,0 +1,152 @@
+#!/usr/bin/env python
+
+# 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 subprocess
+
+from accelerate.commands.config.config_args import default_config_file, load_config_from_file
+from packaging.version import Version, parse
+
+
+_description = "Run commands across a pod of TPU VMs for initial setup before running `accelerate launch`. Will also install Accelerate on the pod."
+
+
+def pod_command_parser(subparsers=None):
+    if subparsers is not None:
+        parser = subparsers.add_parser("pod-config", description=_description)
+    else:
+        parser = argparse.ArgumentParser("Accelerate pod-config command", description=_description)
+
+    parser.add_argument(
+        "--config_file",
+        type=str,
+        default=None,
+        help="Path to the config file to use for accelerate.",
+    )
+
+    parser.add_argument(
+        "--pod_config_file",
+        type=str,
+        default=None,
+        help="Path to the config file to use for the pod.",
+    )
+
+    parser.add_argument(
+        "--command_file",
+        default=None,
+        help="The path to the file containing the commands to run on the pod on startup.",
+    )
+    parser.add_argument(
+        "--command",
+        action="append",
+        nargs="+",
+        help="A command to run on the pod. If not specified, will use the command specified in the command file.",
+    )
+    parser.add_argument(
+        "--tpu_name",
+        default=None,
+        help="The name of the TPU to use. If not specified, will use the TPU specified in the config file.",
+    )
+    parser.add_argument(
+        "--tpu_zone",
+        default=None,
+        help="The zone of the TPU to use. If not specified, will use the zone specified in the config file.",
+    )
+    parser.add_argument(
+        "--install_accelerate",
+        action="store_true",
+        help="Whether to install accelerate on the pod. Defaults to False.",
+    )
+    parser.add_argument(
+        "--accelerate_version",
+        default="latest",
+        help="The version of accelerate to install on the pod. If not specified, will use the latest pypi version. Specify 'dev' to install from GitHub.",
+    )
+    parser.add_argument(
+        "--debug", action="store_true", help="If set, will print the command that would be run instead of running it."
+    )
+
+    if subparsers is not None:
+        parser.set_defaults(func=pod_launcher)
+    return parser
+
+
+def pod_launcher(args):
+    defaults = None
+
+    # Get the default from the config file if it exists.
+    if args.config_file is not None or os.path.isfile(default_config_file):
+        defaults = load_config_from_file(args.config_file)
+        if not args.command_file and defaults.command_file is not None and not args.command:
+            args.command_file = defaults.command_file
+        if not args.command and defaults.command is not None:
+            args.command = defaults.command
+        if not args.tpu_name:
+            args.tpu_name = defaults.tpu_name
+        if not args.tpu_zone:
+            args.tpu_zone = defaults.tpu_zone
+    if args.accelerate_version == "dev":
+        args.accelerate_version = "git+https://github.com/huggingface/accelerate.git"
+    elif args.accelerate_version == "latest":
+        args.accelerate_version = "accelerate -U"
+    elif isinstance(parse(args.accelerate_version), Version):
+        args.accelerate_version = f"accelerate=={args.accelerate_version}"
+
+    if not args.command_file and not args.command:
+        raise ValueError("You must specify either a command file or a command to run on the pod.")
+
+    if args.command_file:
+        with open(args.command_file, "r") as f:
+            args.command = [f.read().splitlines()]
+
+    # To turn list of lists into list of strings
+    args.command = [line for cmd in args.command for line in cmd]
+    # Default to the shared folder and install accelerate
+    new_cmd = ["cd /usr/share"]
+    if args.install_accelerate:
+        new_cmd += [f"pip install {args.accelerate_version}"]
+    new_cmd += args.command
+    args.command = "; ".join(new_cmd)
+
+    # Then send it to gcloud
+    # Eventually try to use google-api-core to do this instead of subprocess
+    cmd = [
+        "gcloud",
+        "compute",
+        "tpus",
+        "tpu-vm",
+        "ssh",
+        args.tpu_name,
+        "--zone",
+        args.tpu_zone,
+        "--command",
+        args.command,
+        "--worker",
+        "all",
+    ]
+    if args.debug:
+        print(cmd)
That's a test class which is inherently private, thus we don't need to add all those `_` that hurt readability :-) 
@@ -63,3 +65,155 @@ def test_config_compatibility(self):
                 execute_subprocess_async(
                     self.base_cmd + ["--config_file", str(config), self.test_file_path], env=os.environ.copy()
                 )
+
+
+class PodConfigTester(unittest.TestCase):
+    """
+    Test case for verifying the `accelerate pod-config` CLI passes the right `gcloud` command.
+    """
+
+    _tpu_name = "test-tpu"
+    _tpu_zone = "us-central1-a"
+    _command = "ls"
+    cmd = ["accelerate", "pod-config"]
+    _base_output = "cd /usr/share"
+    _command_file = "tests/test_samples/test_command_file.sh"
With only allowing one vs the other it keeps the API simplistic, as otherwise we then have to worry about when in the order should the command come from in the CLI vs the bash script and that can be confusing to users. However in this particular case the check is just to ensure that you've passed *some* command to run in :) 
@@ -0,0 +1,152 @@
+#!/usr/bin/env python
+
+# 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 subprocess
+
+from accelerate.commands.config.config_args import default_config_file, load_config_from_file
+from packaging.version import Version, parse
+
+
+_description = "Run commands across a pod of TPU VMs for initial setup before running `accelerate launch`. Will also install Accelerate on the pod."
+
+
+def pod_command_parser(subparsers=None):
+    if subparsers is not None:
+        parser = subparsers.add_parser("pod-config", description=_description)
+    else:
+        parser = argparse.ArgumentParser("Accelerate pod-config command", description=_description)
+
+    parser.add_argument(
+        "--config_file",
+        type=str,
+        default=None,
+        help="Path to the config file to use for accelerate.",
+    )
+
+    parser.add_argument(
+        "--pod_config_file",
+        type=str,
+        default=None,
+        help="Path to the config file to use for the pod.",
+    )
+
+    parser.add_argument(
+        "--command_file",
+        default=None,
+        help="The path to the file containing the commands to run on the pod on startup.",
+    )
+    parser.add_argument(
+        "--command",
+        action="append",
+        nargs="+",
+        help="A command to run on the pod. If not specified, will use the command specified in the command file.",
+    )
+    parser.add_argument(
+        "--tpu_name",
+        default=None,
+        help="The name of the TPU to use. If not specified, will use the TPU specified in the config file.",
+    )
+    parser.add_argument(
+        "--tpu_zone",
+        default=None,
+        help="The zone of the TPU to use. If not specified, will use the zone specified in the config file.",
+    )
+    parser.add_argument(
+        "--install_accelerate",
+        action="store_true",
+        help="Whether to install accelerate on the pod. Defaults to False.",
+    )
+    parser.add_argument(
+        "--accelerate_version",
+        default="latest",
+        help="The version of accelerate to install on the pod. If not specified, will use the latest pypi version. Specify 'dev' to install from GitHub.",
+    )
+    parser.add_argument(
+        "--debug", action="store_true", help="If set, will print the command that would be run instead of running it."
+    )
+
+    if subparsers is not None:
+        parser.set_defaults(func=pod_launcher)
+    return parser
+
+
+def pod_launcher(args):
+    defaults = None
+
+    # Get the default from the config file if it exists.
+    if args.config_file is not None or os.path.isfile(default_config_file):
+        defaults = load_config_from_file(args.config_file)
+        if not args.command_file and defaults.command_file is not None and not args.command:
+            args.command_file = defaults.command_file
+        if not args.command and defaults.command is not None:
+            args.command = defaults.command
+        if not args.tpu_name:
+            args.tpu_name = defaults.tpu_name
+        if not args.tpu_zone:
+            args.tpu_zone = defaults.tpu_zone
+    if args.accelerate_version == "dev":
+        args.accelerate_version = "git+https://github.com/huggingface/accelerate.git"
+    elif args.accelerate_version == "latest":
+        args.accelerate_version = "accelerate -U"
+    elif isinstance(parse(args.accelerate_version), Version):
+        args.accelerate_version = f"accelerate=={args.accelerate_version}"
+
+    if not args.command_file and not args.command:
+        raise ValueError("You must specify either a command file or a command to run on the pod.")
We can keep this idea open though and if users want that we can enable it
@@ -0,0 +1,152 @@
+#!/usr/bin/env python
+
+# 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 subprocess
+
+from accelerate.commands.config.config_args import default_config_file, load_config_from_file
+from packaging.version import Version, parse
+
+
+_description = "Run commands across a pod of TPU VMs for initial setup before running `accelerate launch`. Will also install Accelerate on the pod."
+
+
+def pod_command_parser(subparsers=None):
+    if subparsers is not None:
+        parser = subparsers.add_parser("pod-config", description=_description)
+    else:
+        parser = argparse.ArgumentParser("Accelerate pod-config command", description=_description)
+
+    parser.add_argument(
+        "--config_file",
+        type=str,
+        default=None,
+        help="Path to the config file to use for accelerate.",
+    )
+
+    parser.add_argument(
+        "--pod_config_file",
+        type=str,
+        default=None,
+        help="Path to the config file to use for the pod.",
+    )
+
+    parser.add_argument(
+        "--command_file",
+        default=None,
+        help="The path to the file containing the commands to run on the pod on startup.",
+    )
+    parser.add_argument(
+        "--command",
+        action="append",
+        nargs="+",
+        help="A command to run on the pod. If not specified, will use the command specified in the command file.",
+    )
+    parser.add_argument(
+        "--tpu_name",
+        default=None,
+        help="The name of the TPU to use. If not specified, will use the TPU specified in the config file.",
+    )
+    parser.add_argument(
+        "--tpu_zone",
+        default=None,
+        help="The zone of the TPU to use. If not specified, will use the zone specified in the config file.",
+    )
+    parser.add_argument(
+        "--install_accelerate",
+        action="store_true",
+        help="Whether to install accelerate on the pod. Defaults to False.",
+    )
+    parser.add_argument(
+        "--accelerate_version",
+        default="latest",
+        help="The version of accelerate to install on the pod. If not specified, will use the latest pypi version. Specify 'dev' to install from GitHub.",
+    )
+    parser.add_argument(
+        "--debug", action="store_true", help="If set, will print the command that would be run instead of running it."
+    )
+
+    if subparsers is not None:
+        parser.set_defaults(func=pod_launcher)
+    return parser
+
+
+def pod_launcher(args):
+    defaults = None
+
+    # Get the default from the config file if it exists.
+    if args.config_file is not None or os.path.isfile(default_config_file):
+        defaults = load_config_from_file(args.config_file)
+        if not args.command_file and defaults.command_file is not None and not args.command:
+            args.command_file = defaults.command_file
+        if not args.command and defaults.command is not None:
+            args.command = defaults.command
+        if not args.tpu_name:
+            args.tpu_name = defaults.tpu_name
+        if not args.tpu_zone:
+            args.tpu_zone = defaults.tpu_zone
+    if args.accelerate_version == "dev":
+        args.accelerate_version = "git+https://github.com/huggingface/accelerate.git"
+    elif args.accelerate_version == "latest":
+        args.accelerate_version = "accelerate -U"
+    elif isinstance(parse(args.accelerate_version), Version):
+        args.accelerate_version = f"accelerate=={args.accelerate_version}"
+
+    if not args.command_file and not args.command:
+        raise ValueError("You must specify either a command file or a command to run on the pod.")
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/802
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index b9f6c1089..7c8534e93 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -592,6 +592,12 @@ def deepspeed_launcher(args):
     num_machines = getattr(args, "num_machines")
     main_process_ip = getattr(args, "main_process_ip")
     main_process_port = getattr(args, "main_process_port")
+
+    # make sure launcher is not None
+    if args.deepspeed_multinode_launcher is None:
+        # set to default pdsh
+        setattr(args, "deepspeed_multinode_launcher", DEEPSPEED_MULTINODE_LAUNCHERS[0])
+
     if num_machines > 1 and args.deepspeed_multinode_launcher != DEEPSPEED_MULTINODE_LAUNCHERS[1]:
         cmd = ["deepspeed", "--no_local_rank"]
         cmd.extend(["--hostfile", str(args.deepspeed_hostfile), "--launcher", str(args.deepspeed_multinode_launcher)])

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/800
diff --git a/src/accelerate/big_modeling.py b/src/accelerate/big_modeling.py
index 83c6738af..e277e89a2 100644
--- a/src/accelerate/big_modeling.py
+++ b/src/accelerate/big_modeling.py
@@ -245,11 +245,15 @@ def dispatch_model(
     check_device_map(model, device_map)
 
     if main_device is None:
-        main_device = [d for d in device_map.values() if d not in ["cpu", "disk"]][0]
+        if set(device_map.values()) == {"cpu"} or set(device_map.values()) == {"cpu", "disk"}:
+            main_device = "cpu"
+        else:
+            main_device = [d for d in device_map.values() if d not in ["cpu", "disk"]][0]
 
-    cpu_modules = [name for name, device in device_map.items() if device == "cpu"]
-    if state_dict is None and len(cpu_modules) > 0:
-        state_dict = extract_submodules_state_dict(model.state_dict(), cpu_modules)
+    if main_device != "cpu":
+        cpu_modules = [name for name, device in device_map.items() if device == "cpu"]
+        if state_dict is None and len(cpu_modules) > 0:
+            state_dict = extract_submodules_state_dict(model.state_dict(), cpu_modules)
 
     disk_modules = [name for name, device in device_map.items() if device == "disk"]
     if offload_dir is None and len(disk_modules) > 0:
@@ -266,7 +270,8 @@ def dispatch_model(
     execution_device = {
         name: main_device if device in ["cpu", "disk"] else device for name, device in device_map.items()
     }
-    offload = {name: device in ["cpu", "disk"] for name, device in device_map.items()}
+    offloaded_devices = ["disk"] if main_device == "cpu" else ["cpu", "disk"]
+    offload = {name: device in offloaded_devices for name, device in device_map.items()}
     save_folder = offload_dir if len(disk_modules) > 0 else None
     if state_dict is not None or save_folder is not None:
         weights_map = OffloadedWeightsLoader(state_dict=state_dict, save_folder=save_folder)

_The documentation is not available anymore as the PR was closed or merged._
Is this expected to work on the `mps` device too? I've checked #596, but the only device I see here is `cpu`:
https://github.com/huggingface/accelerate/blob/0de16441266ba9a799aa338aa75c73bf078429e9/src/accelerate/utils/modeling.py#L464

Furthermore, `get_max_memory` doesn't seem to take it into account: https://github.com/huggingface/accelerate/blob/b816e258a95904736c8e3a5a2b14a18d5f407f4a/src/accelerate/utils/modeling.py#L275-L284

Alternatively, is it possible to override using a custom `device_map`?

Reference: https://github.com/huggingface/diffusers/pull/1042
No we don't support the "mps" device at all in big model inference. `device_map="auto"` will load the model on the CPU and then the user can't move them to the MPS device if they want.
Ok understood, thanks! However, we get errors after moving to the MPS device, if we use `device_map="auto"`. I'm not sure what's causing it, but some tensors are being moved to the CPU when performing inference. I'll investigate further.
That's the side effect of using `device_map="auto"` @Narsil was warning you about. It does more than just load the model with empty weights but adds hooks to the model.
Yeah I saw the hooks while stepping with the debugger. One of the reasons why we want to use this is for easier communication (just use `accelerate` and `device_map="auto"`). We'll see what we can do :)
@pcuenca Loading weights efficiently is like 2 functions from `accelerate`.

https://github.com/huggingface/accelerate/blob/main/src/accelerate/big_modeling.py#L35-L110

And loading after the fact https://github.com/huggingface/transformers_bloom_parallel/blob/main/generate.py#L211-L270
(Couldn't fetch rapidly from `accelerate` but it has some more complete and robust function).

To be completely honest, I feel like all this should be in Pytorch proper. the `meta` device is really a bad workaround.

In the meantime maybe instead of pushing towards `device_map="auto"` you could instead use the 2 accelerate functions and wrapping them up nicely in `from_pretrained` ?
> In the meantime maybe instead of pushing towards `device_map="auto"` you could instead use the 2 accelerate functions and wrapping them up nicely in `from_pretrained` ?

That's what I was thinking, yes! What I'm not sure about is whether to apply those only for the `mps` device (and keep using `device_map` in the others), or something else. I'll do some experiments to see what it looks like in practice. Thanks a lot!
https://api.github.com/repos/huggingface/accelerate/pulls/797
diff --git a/src/accelerate/utils/offload.py b/src/accelerate/utils/offload.py
index 750ff9d57..bc0fcdef1 100644
--- a/src/accelerate/utils/offload.py
+++ b/src/accelerate/utils/offload.py
@@ -178,5 +178,13 @@ def extract_submodules_state_dict(state_dict: Dict[str, torch.Tensor], submodule
     """
     result = {}
     for module_name in submodule_names:
-        result.update({key: param for key, param in state_dict.items() if key.startswith(module_name)})
+        # We want to catch module_name parameter (module_name.xxx) or potentially module_name, but not any of the
+        # submodules that could being like module_name (transformers.h.1 and transformers.h.10 for instance)
+        result.update(
+            {
+                key: param
+                for key, param in state_dict.items()
+                if key == module_name or key.startswith(module_name + ".")
+            }
+        )
     return result
diff --git a/tests/test_offload.py b/tests/test_offload.py
index 765a9685d..a9ebd36ee 100644
--- a/tests/test_offload.py
+++ b/tests/test_offload.py
@@ -21,6 +21,7 @@
 
 from accelerate.utils import (
     OffloadedWeightsLoader,
+    extract_submodules_state_dict,
     is_torch_version,
     load_offloaded_weight,
     offload_state_dict,
@@ -105,3 +106,12 @@ def test_offload_weights_loader(self):
             self.assertEqual(sorted(weight_map), sorted(state_dict.keys()))
             for key, param in state_dict.items():
                 self.assertTrue(torch.allclose(param, weight_map[key]))
+
+    def test_extract_submodules_state_dict(self):
+        state_dict = {"a.1": 0, "a.10": 1, "a.2": 2}
+        extracted = extract_submodules_state_dict(state_dict, ["a.1", "a.2"])
+        self.assertDictEqual(extracted, {"a.1": 0, "a.2": 2})
+
+        state_dict = {"a.1.a": 0, "a.10.a": 1, "a.2.a": 2}
+        extracted = extract_submodules_state_dict(state_dict, ["a.1", "a.2"])
+        self.assertDictEqual(extracted, {"a.1.a": 0, "a.2.a": 2})

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/795
diff --git a/src/accelerate/tracking.py b/src/accelerate/tracking.py
index a438ab45b..3a31fed93 100644
--- a/src/accelerate/tracking.py
+++ b/src/accelerate/tracking.py
@@ -15,15 +15,23 @@
 # Expectation:
 # Provide a project dir name, then each type of logger gets stored in project/{`logging_dir`}
 
+import json
 import os
 import time
 from abc import ABCMeta, abstractmethod, abstractproperty
-from typing import List, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 import yaml
 
 from .logging import get_logger
-from .utils import LoggerType, is_aim_available, is_comet_ml_available, is_tensorboard_available, is_wandb_available
+from .utils import (
+    LoggerType,
+    is_aim_available,
+    is_comet_ml_available,
+    is_mlflow_available,
+    is_tensorboard_available,
+    is_wandb_available,
+)
 
 
 _available_trackers = []
@@ -48,6 +56,10 @@
 
     _available_trackers.append(LoggerType.AIM)
 
+if is_mlflow_available():
+    import mlflow
+
+    _available_trackers.append(LoggerType.MLFLOW)
 
 logger = get_logger(__name__)
 
@@ -398,9 +410,140 @@ def finish(self):
         self.writer.close()
 
 
+class MLflowTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `mlflow`. Should be initialized at the start of your script.
+
+    Args:
+        experiment_name (`str`, *optional*):
+            Name of the experiment. Environment variable MLFLOW_EXPERIMENT_NAME has priority over this argument.
+        logging_dir (`str` or `os.PathLike`, defaults to `"."`):
+            Location for mlflow logs to be stored.
+        run_id (`str`, *optional*):
+            If specified, get the run with the specified UUID and log parameters and metrics under that run. The run’s
+            end time is unset and its status is set to running, but the run’s other attributes (source_version,
+            source_type, etc.) are not changed. Environment variable MLFLOW_RUN_ID has priority over this argument.
+        tags (`Dict[str, str]`, *optional*):
+            An optional `dict` of `str` keys and values, or a `str` dump from a `dict`, to set as tags on the run. If a
+            run is being resumed, these tags are set on the resumed run. If a new run is being created, these tags are
+            set on the new run. Environment variable MLFLOW_TAGS has priority over this argument.
+        nested_run (`bool`, *optional*, defaults to `False`):
+            Controls whether run is nested in parent run. True creates a nested run. Environment variable
+            MLFLOW_NESTED_RUN has priority over this argument.
+        run_name (`str`, *optional*):
+            Name of new run (stored as a mlflow.runName tag). Used only when `run_id` is unspecified.
+        description (`str`, *optional*):
+            An optional string that populates the description box of the run. If a run is being resumed, the
+            description is set on the resumed run. If a new run is being created, the description is set on the new
+            run.
+    """
+
+    name = "mlflow"
+    requires_logging_directory = True
+
+    def __init__(
+        self,
+        experiment_name: str = None,
+        logging_dir: Optional[Union[str, os.PathLike]] = ".",
+        run_id: Optional[str] = None,
+        tags: Optional[Union[Dict[str, Any], str]] = None,
+        nested_run: Optional[bool] = False,
+        run_name: Optional[str] = None,
+        description: Optional[str] = None,
+    ):
+
+        experiment_name = os.getenv("MLFLOW_EXPERIMENT_NAME", experiment_name)
+        run_id = os.getenv("MLFLOW_RUN_ID", run_id)
+        tags = os.getenv("MLFLOW_TAGS", tags)
+        if isinstance(tags, str):
+            tags = json.loads(tags)
+
+        nested_run = os.getenv("MLFLOW_NESTED_RUN", nested_run)
+
+        experiment_id = mlflow.create_experiment(
+            name=experiment_name,
+            artifact_location=logging_dir,
+            tags=tags,
+        )
+
+        self.active_run = mlflow.start_run(
+            run_id=run_id,
+            experiment_id=experiment_id,
+            run_name=run_name,
+            nested=nested_run,
+            tags=tags,
+            description=description,
+        )
+
+        logger.debug(f"Initialized mlflow experiment {experiment_name}")
+        logger.debug(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    @property
+    def tracker(self):
+        return self.active_run
+
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment.
+
+        Args:
+            values (`dict`):
+                Values to be stored as initial hyperparameters as key-value pairs.
+        """
+
+        for name, value in list(values.items()):
+            # internally, all values are converted to str in MLflow
+            if len(str(value)) > mlflow.utils.validation.MAX_PARAM_VAL_LENGTH:
+                logger.warning(
+                    f'Trainer is attempting to log a value of "{value}" for key "{name}" as a parameter. MLflow\'s'
+                    f" log_param() only accepts values no longer than {mlflow.utils.validation.MAX_PARAM_VAL_LENGTH} characters so we dropped this attribute."
+                )
+                del values[name]
+
+        values_list = list(values.items())
+
+        # MLflow cannot log more than 100 values in one go, so we have to split it
+        for i in range(0, len(values_list), mlflow.utils.validation.MAX_PARAMS_TAGS_PER_BATCH):
+            mlflow.log_params(dict(values_list[i : i + mlflow.utils.validation.MAX_PARAMS_TAGS_PER_BATCH]))
+
+        logger.debug("Stored initial configuration hyperparameters to MLflow")
+
+    def log(self, values: dict, step: Optional[int]):
+        """
+        Logs `values` to the current run.
+
+        Args:
+            values (`dict`):
+                Values to be logged as key-value pairs.
+            step (`int`, *optional*):
+                The run step. If included, the log will be affiliated with this step.
+        """
+        metrics = {}
+        for k, v in values.items():
+            if isinstance(v, (int, float)):
+                metrics[k] = v
+            else:
+                logger.warning(
+                    f'MLflowTracker is attempting to log a value of "{v}" of type {type(v)} for key "{k}" as a metric. '
+                    "MLflow's log_metric() only accepts float and int types so we dropped this attribute."
+                )
+
+        mlflow.log_metrics(metrics, step=step)
+        logger.debug("Successfully logged to mlflow")
+
+    def finish(self):
+        """
+        End the active MLflow run.
+        """
+        mlflow.end_run()
+
+
 LOGGER_TYPE_TO_CLASS = {
     "aim": AimTracker,
     "comet_ml": CometMLTracker,
+    "mlflow": MLflowTracker,
     "tensorboard": TensorBoardTracker,
     "wandb": WandBTracker,
 }
@@ -424,6 +567,7 @@ def filter_trackers(
             - `"tensorboard"`
             - `"wandb"`
             - `"comet_ml"`
+            - `"mlflow"`
             If `"all"` is selected, will pick up all available trackers in the environment and initialize them. Can
             also accept implementations of `GeneralTracker` for custom trackers, and can be combined with `"all"`.
         logging_dir (`str`, `os.PathLike`, *optional*):
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
index b22518f9a..a017a2ad5 100644
--- a/src/accelerate/utils/__init__.py
+++ b/src/accelerate/utils/__init__.py
@@ -31,6 +31,7 @@
     is_datasets_available,
     is_deepspeed_available,
     is_megatron_lm_available,
+    is_mlflow_available,
     is_rich_available,
     is_sagemaker_available,
     is_tensorboard_available,
diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py
index 8227fa5bc..7dd8798e9 100644
--- a/src/accelerate/utils/dataclasses.py
+++ b/src/accelerate/utils/dataclasses.py
@@ -203,6 +203,7 @@ class LoggerType(BaseEnum):
     TENSORBOARD = "tensorboard"
     WANDB = "wandb"
     COMETML = "comet_ml"
+    MLFLOW = "mlflow"
 
 
 class PrecisionType(BaseEnum):
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index c33258a92..892ff63f8 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -136,3 +136,7 @@ def is_sagemaker_available():
 
 def is_tqdm_available():
     return importlib.util.find_spec("tqdm") is not None
+
+
+def is_mlflow_available():
+    return importlib.util.find_spec("mlflow") is not None

```suggestion
        tags (`Dict[str, str]`, *optional*):
```
@@ -398,9 +410,143 @@ def finish(self):
         self.writer.close()
 
 
+class MLflowTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `mlflow`. Should be initialized at the start of your script.
+
+    Args:
+        experiment_name (`str`):
+            Name of the experiment. Environment variable MLFLOW_EXPERIMENT_NAME has priority over this argument.
+        logging_dir (`str`, `os.PathLike`):
+            Location for mlflow logs to be stored.
+        run_id (`str`):
+            If specified, get the run with the specified UUID and log parameters and metrics under that run. The run’s
+            end time is unset and its status is set to running, but the run’s other attributes (source_version,
+            source_type, etc.) are not changed. Environment variable MLFLOW_RUN_ID has priority over this argument.
+        tags (`dict`, `str`):
```suggestion
        run_id (`str`, *optional*):
```
@@ -398,9 +410,143 @@ def finish(self):
         self.writer.close()
 
 
+class MLflowTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `mlflow`. Should be initialized at the start of your script.
+
+    Args:
+        experiment_name (`str`):
+            Name of the experiment. Environment variable MLFLOW_EXPERIMENT_NAME has priority over this argument.
+        logging_dir (`str`, `os.PathLike`):
+            Location for mlflow logs to be stored.
+        run_id (`str`):
```suggestion
        logging_dir (`str` or `os.PathLike`, defaults to `"."`):
```
@@ -398,9 +410,143 @@ def finish(self):
         self.writer.close()
 
 
+class MLflowTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `mlflow`. Should be initialized at the start of your script.
+
+    Args:
+        experiment_name (`str`):
+            Name of the experiment. Environment variable MLFLOW_EXPERIMENT_NAME has priority over this argument.
+        logging_dir (`str`, `os.PathLike`):
```suggestion
        experiment_name (`str`, *optional*):
```
@@ -398,9 +410,143 @@ def finish(self):
         self.writer.close()
 
 
+class MLflowTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `mlflow`. Should be initialized at the start of your script.
+
+    Args:
+        experiment_name (`str`):
```suggestion
        nested_run (`bool`, *optional*, defaults to `False`):
```
@@ -398,9 +410,143 @@ def finish(self):
         self.writer.close()
 
 
+class MLflowTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `mlflow`. Should be initialized at the start of your script.
+
+    Args:
+        experiment_name (`str`):
+            Name of the experiment. Environment variable MLFLOW_EXPERIMENT_NAME has priority over this argument.
+        logging_dir (`str`, `os.PathLike`):
+            Location for mlflow logs to be stored.
+        run_id (`str`):
+            If specified, get the run with the specified UUID and log parameters and metrics under that run. The run’s
+            end time is unset and its status is set to running, but the run’s other attributes (source_version,
+            source_type, etc.) are not changed. Environment variable MLFLOW_RUN_ID has priority over this argument.
+        tags (`dict`, `str`):
+            An optional `dict` of `str` keys and values, or a `str` dump from a `dict`, to set as tags on the run. If a
+            run is being resumed, these tags are set on the resumed run. If a new run is being created, these tags are
+            set on the new run. Environment variable MLFLOW_TAGS has priority over this argument.
+        nested_run (`bool`):
```suggestion
        run_name (`str`, *optional*):
```
@@ -398,9 +410,143 @@ def finish(self):
         self.writer.close()
 
 
+class MLflowTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `mlflow`. Should be initialized at the start of your script.
+
+    Args:
+        experiment_name (`str`):
+            Name of the experiment. Environment variable MLFLOW_EXPERIMENT_NAME has priority over this argument.
+        logging_dir (`str`, `os.PathLike`):
+            Location for mlflow logs to be stored.
+        run_id (`str`):
+            If specified, get the run with the specified UUID and log parameters and metrics under that run. The run’s
+            end time is unset and its status is set to running, but the run’s other attributes (source_version,
+            source_type, etc.) are not changed. Environment variable MLFLOW_RUN_ID has priority over this argument.
+        tags (`dict`, `str`):
+            An optional `dict` of `str` keys and values, or a `str` dump from a `dict`, to set as tags on the run. If a
+            run is being resumed, these tags are set on the resumed run. If a new run is being created, these tags are
+            set on the new run. Environment variable MLFLOW_TAGS has priority over this argument.
+        nested_run (`bool`):
+            Controls whether run is nested in parent run. True creates a nested run. Environment variable
+            MLFLOW_NESTED_RUN has priority over this argument.
+        run_name (`str`):
```suggestion
            Name of new run (stored as a mlflow.runName tag). Used only when `run_id` is unspecified.
```
@@ -398,9 +410,143 @@ def finish(self):
         self.writer.close()
 
 
+class MLflowTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `mlflow`. Should be initialized at the start of your script.
+
+    Args:
+        experiment_name (`str`):
+            Name of the experiment. Environment variable MLFLOW_EXPERIMENT_NAME has priority over this argument.
+        logging_dir (`str`, `os.PathLike`):
+            Location for mlflow logs to be stored.
+        run_id (`str`):
+            If specified, get the run with the specified UUID and log parameters and metrics under that run. The run’s
+            end time is unset and its status is set to running, but the run’s other attributes (source_version,
+            source_type, etc.) are not changed. Environment variable MLFLOW_RUN_ID has priority over this argument.
+        tags (`dict`, `str`):
+            An optional `dict` of `str` keys and values, or a `str` dump from a `dict`, to set as tags on the run. If a
+            run is being resumed, these tags are set on the resumed run. If a new run is being created, these tags are
+            set on the new run. Environment variable MLFLOW_TAGS has priority over this argument.
+        nested_run (`bool`):
+            Controls whether run is nested in parent run. True creates a nested run. Environment variable
+            MLFLOW_NESTED_RUN has priority over this argument.
+        run_name (`str`):
+            Name of new run (stored as a mlflow.runName tag). Used only when run_id is unspecified.
```suggestion
        description (`str`, *optional*):
```
@@ -398,9 +410,143 @@ def finish(self):
         self.writer.close()
 
 
+class MLflowTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `mlflow`. Should be initialized at the start of your script.
+
+    Args:
+        experiment_name (`str`):
+            Name of the experiment. Environment variable MLFLOW_EXPERIMENT_NAME has priority over this argument.
+        logging_dir (`str`, `os.PathLike`):
+            Location for mlflow logs to be stored.
+        run_id (`str`):
+            If specified, get the run with the specified UUID and log parameters and metrics under that run. The run’s
+            end time is unset and its status is set to running, but the run’s other attributes (source_version,
+            source_type, etc.) are not changed. Environment variable MLFLOW_RUN_ID has priority over this argument.
+        tags (`dict`, `str`):
+            An optional `dict` of `str` keys and values, or a `str` dump from a `dict`, to set as tags on the run. If a
+            run is being resumed, these tags are set on the resumed run. If a new run is being created, these tags are
+            set on the new run. Environment variable MLFLOW_TAGS has priority over this argument.
+        nested_run (`bool`):
+            Controls whether run is nested in parent run. True creates a nested run. Environment variable
+            MLFLOW_NESTED_RUN has priority over this argument.
+        run_name (`str`):
+            Name of new run (stored as a mlflow.runName tag). Used only when run_id is unspecified.
+        description (`str`):
Instead of storing those constants, let's use them directly when necessary (otherwise the reader has to figure out what those things are :-) 0
@@ -398,9 +410,143 @@ def finish(self):
         self.writer.close()
 
 
+class MLflowTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `mlflow`. Should be initialized at the start of your script.
+
+    Args:
+        experiment_name (`str`):
+            Name of the experiment. Environment variable MLFLOW_EXPERIMENT_NAME has priority over this argument.
+        logging_dir (`str`, `os.PathLike`):
+            Location for mlflow logs to be stored.
+        run_id (`str`):
+            If specified, get the run with the specified UUID and log parameters and metrics under that run. The run’s
+            end time is unset and its status is set to running, but the run’s other attributes (source_version,
+            source_type, etc.) are not changed. Environment variable MLFLOW_RUN_ID has priority over this argument.
+        tags (`dict`, `str`):
+            An optional `dict` of `str` keys and values, or a `str` dump from a `dict`, to set as tags on the run. If a
+            run is being resumed, these tags are set on the resumed run. If a new run is being created, these tags are
+            set on the new run. Environment variable MLFLOW_TAGS has priority over this argument.
+        nested_run (`bool`):
+            Controls whether run is nested in parent run. True creates a nested run. Environment variable
+            MLFLOW_NESTED_RUN has priority over this argument.
+        run_name (`str`):
+            Name of new run (stored as a mlflow.runName tag). Used only when run_id is unspecified.
+        description (`str`):
+            An optional string that populates the description box of the run. If a run is being resumed, the
+            description is set on the resumed run. If a new run is being created, the description is set on the new
+            run.
+    """
+
+    name = "mlflow"
+    requires_logging_directory = True
+
+    def __init__(
+        self,
+        experiment_name: str = None,
+        logging_dir: Optional[Union[str, os.PathLike]] = ".",
+        run_id: Optional[str] = None,
+        tags: Optional[Union[Dict[str, Any], str]] = None,
+        nested_run: Optional[bool] = False,
+        run_name: Optional[str] = None,
+        description: Optional[str] = None,
+    ):
+
+        self._MAX_PARAM_VAL_LENGTH = mlflow.utils.validation.MAX_PARAM_VAL_LENGTH
+        self._MAX_PARAMS_TAGS_PER_BATCH = mlflow.utils.validation.MAX_PARAMS_TAGS_PER_BATCH
_The documentation is not available anymore as the PR was closed or merged._
@sgugger, how is it now?
https://api.github.com/repos/huggingface/accelerate/pulls/794
diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
index c14833557..e471731a4 100644
--- a/src/accelerate/utils/modeling.py
+++ b/src/accelerate/utils/modeling.py
@@ -476,7 +476,11 @@ def infer_auto_device_map(
     current_memory_used = 0
 
     # Direct submodules and parameters
-    modules_to_treat = list(model.named_parameters(recurse=False)) + list(model.named_children())
+    modules_to_treat = (
+        list(model.named_parameters(recurse=False))
+        + list(model.named_children())
+        + list(model.named_buffers(recurse=False))
+    )
     # Initialize maximum largest layer, to know which space to keep in memory
     max_layer_size, max_layer_names = get_max_layer_size(modules_to_treat, module_sizes, no_split_module_classes)
 

_The documentation is not available anymore as the PR was closed or merged._
So if I understood it correctly, if you have some modules such as `nn.BatchNorm` in your model (as it is done in the `accelerate` CI test), the buffers `running_mean` and `running_var` will not be stored inside `model._buffers` but in `model.named_buffers()`. That is why I had to "filter out" the buffers by considering only the ones that are inside `model._buffers` and `model.named_buffers()` 

Here is an example that I have quickly tried: 
```
model = nn.Sequential(nn.Linear(1, 1), nn.BatchNorm1d(1), nn.Embedding(1, 1), nn.LayerNorm(1)) 
print(list(model.named_buffers()))
>>>[('1.running_mean', tensor([0.])), ('1.running_var', tensor([1.])), ('1.num_batches_tracked', tensor(0))]
print(list(model._buffers))
>>> []
model.register_buffer("position_bias", torch.ones(1))
print(list(model._buffers))
>>> ['position_bias']
print(list(model.named_buffers()))
>>> [('position_bias', tensor([1.])), ('1.running_mean', tensor([0.])), ('1.running_var', tensor([1.])), ('1.num_batches_tracked', tensor(0))]
```


I think in this case, it's just the difference between `named_buffers(recurse=True)` and `named_buffers(recurse=False)`. I'm not convinced this fix is the right fix, so would like to learn more what is failing.
Ah yes I see, you're probably right here! Let me dig a bit more and get back to you here 
@sgugger I might have more clue on what is failing 
I think that the problem comes from the fact that the `infer_auto_device_map` does take into account only modules and submodules, I have made a script below to better illustrate the problem

```
import torch.nn as nn
import torch
from accelerate.utils import infer_auto_device_map
from accelerate.big_modeling import dispatch_model

class SubModule(nn.Module):
    def __init__(self):
        super().__init__()

        self.register_buffer("position_bias", torch.ones(1, 1000))

class Model(nn.Module):
    def __init__(self, wrap_module=True):
        super().__init__()
        self.l1 = nn.Linear(1000, 1000)
        self.l2 = nn.Linear(1000, 1000)
        self.l3 = nn.Linear(1000, 1000)

        self.bn1 = nn.BatchNorm1d(1000)
        self.bn2 = nn.BatchNorm1d(1000)

        if wrap_module:
            self.position_bias = SubModule()
        else:
            self.register_buffer("position_bias", torch.ones(1, 1000))

# Test 1: wrapping with a module - this will pass
model = Model()
device_map = infer_auto_device_map(model, {0:"10MB", "cpu":"100MB"})
model = dispatch_model(model, device_map)

# Test 2: below will fail
model = Model(wrap_module=False)
device_map = infer_auto_device_map(model, {0:"10MB", "cpu":"100MB"})
model = dispatch_model(model, device_map)
```
Let me know what do you think! 

I guess this failed for `BartPreTrainedModel` since the `position_bias` buffer is on the parent module itself
Ah, in this case it looks very much like the problem #747 fixed for top-level parameters, so the fix should be pretty similar here too!
The whole testing suite (including slow tests) is green! 🟢 Merging ! 


https://api.github.com/repos/huggingface/accelerate/pulls/792
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index cdae54af9..8fbc8758f 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -352,16 +352,16 @@ class DataLoaderShard(DataLoader):
         - **total_dataset_length** (`int`) -- Total length of the inner dataset across all processes.
     """
 
-    def __init__(self, dataset, device=None, rng_types=None, generator=None, **kwargs):
+    def __init__(self, dataset, device=None, rng_types=None, synchronized_generator=None, **kwargs):
         super().__init__(dataset, **kwargs)
         self.device = device
         self.rng_types = rng_types
-        self.generator = generator
+        self.synchronized_generator = synchronized_generator
         self.gradient_state = GradientState()
 
     def __iter__(self):
         if self.rng_types is not None:
-            synchronize_rng_states(self.rng_types, self.generator)
+            synchronize_rng_states(self.rng_types, self.synchronized_generator)
         self.gradient_state._set_end_of_dataloader(False)
         try:
             length = getattr(self.dataset, "total_dataset_length", len(self.dataset))
@@ -650,12 +650,12 @@ def prepare_data_loader(
     # Iterable dataset doesn't like batch_sampler, but data_loader creates a default one for it
     new_batch_sampler = dataloader.batch_sampler if not isinstance(new_dataset, IterableDataset) else None
     sampler_is_batch_sampler = False
-    generator = getattr(dataloader, "generator", None)
+    synchronized_generator = None
     # No change if no multiprocess
     if (num_processes != 1 or state.distributed_type == DistributedType.MEGATRON_LM) and not dispatch_batches:
         if isinstance(new_dataset, IterableDataset):
             if getattr(dataloader.dataset, "generator", None) is not None:
-                generator = dataloader.dataset.generator
+                synchronized_generator = dataloader.dataset.generator
             new_dataset = IterableDatasetShard(
                 new_dataset,
                 batch_size=dataloader.batch_size,
@@ -674,8 +674,7 @@ def prepare_data_loader(
             if hasattr(sampler, "generator"):
                 if sampler.generator is None:
                     sampler.generator = torch.Generator()
-                generator = sampler.generator
-                generator.manual_seed(int(torch.empty((), dtype=torch.int64).random_().item()))
+                synchronized_generator = sampler.generator
 
             batch_sampler = dataloader.sampler if sampler_is_batch_sampler else dataloader.batch_sampler
             new_batch_sampler = BatchSamplerShard(
@@ -692,10 +691,9 @@ def prepare_data_loader(
         "sampler",
         "batch_sampler",
         "drop_last",
-        "generator",
     ]
 
-    if rng_types is not None and generator is None and "generator" in rng_types:
+    if rng_types is not None and synchronized_generator is None and "generator" in rng_types:
         rng_types.remove("generator")
 
     kwargs = {
@@ -710,6 +708,7 @@ def prepare_data_loader(
         kwargs["batch_size"] = dataloader.batch_size // num_processes if split_batches else dataloader.batch_size
 
     if dispatch_batches:
+        kwargs.pop("generator")
         dataloader = DataLoaderDispatcher(
             new_dataset,
             split_batches=split_batches,
@@ -722,9 +721,9 @@ def prepare_data_loader(
             new_dataset,
             device=device if put_on_device and state.distributed_type != DistributedType.TPU else None,
             sampler=new_batch_sampler,
-            batch_size=getattr(dataloader, "batch_size", _PYTORCH_DATALOADER_KWARGS["batch_size"]),
+            batch_size=dataloader.batch_size,
             rng_types=rng_types,
-            generator=generator,
+            synchronized_generator=synchronized_generator,
             **kwargs,
         )
     else:
@@ -733,7 +732,7 @@ def prepare_data_loader(
             device=device if put_on_device and state.distributed_type != DistributedType.TPU else None,
             batch_sampler=new_batch_sampler,
             rng_types=rng_types,
-            generator=generator,
+            synchronized_generator=synchronized_generator,
             **kwargs,
         )
 

We should have an else branch here in case `self.rng_types` contains something different than `"generator"`, by sending `None` for the generator argument (since we have removed `"generator"` from the rng_types at this stage, if there is none available.
@@ -361,7 +361,15 @@ def __init__(self, dataset, device=None, rng_types=None, generator=None, **kwarg
 
     def __iter__(self):
         if self.rng_types is not None:
-            synchronize_rng_states(self.rng_types, self.generator)
+            batch_sampler = self.sampler if isinstance(self.sampler, BatchSampler) else self.batch_sampler
+            sampler = (
+                batch_sampler.batch_sampler.sampler
+                if hasattr(batch_sampler, "batch_sampler")
+                else batch_sampler.sampler
+            )
+            if hasattr(sampler, "generator"):
+                generator = sampler.generator
+                synchronize_rng_states(self.rng_types, generator)
         self.gradient_state._set_end_of_dataloader(False)
I wonder why `dataloader.sampler` can be `BatchSampler`?
@@ -361,7 +361,15 @@ def __init__(self, dataset, device=None, rng_types=None, generator=None, **kwarg
 
     def __iter__(self):
         if self.rng_types is not None:
-            synchronize_rng_states(self.rng_types, self.generator)
+            batch_sampler = self.sampler if isinstance(self.sampler, BatchSampler) else self.batch_sampler
See the example script above in the PR, it has various scenarios. 
@@ -361,7 +361,15 @@ def __init__(self, dataset, device=None, rng_types=None, generator=None, **kwarg
 
     def __iter__(self):
         if self.rng_types is not None:
-            synchronize_rng_states(self.rng_types, self.generator)
+            batch_sampler = self.sampler if isinstance(self.sampler, BatchSampler) else self.batch_sampler
see https://github.com/huggingface/accelerate/issues/679, it helps reduce the number of queries

@@ -361,7 +361,15 @@ def __init__(self, dataset, device=None, rng_types=None, generator=None, **kwarg
 
     def __iter__(self):
         if self.rng_types is not None:
-            synchronize_rng_states(self.rng_types, self.generator)
+            batch_sampler = self.sampler if isinstance(self.sampler, BatchSampler) else self.batch_sampler
https://huggingface.co/docs/datasets/v2.4.0/en/use_with_pytorch#use-a-batchsampler
@@ -361,7 +361,15 @@ def __init__(self, dataset, device=None, rng_types=None, generator=None, **kwarg
 
     def __iter__(self):
         if self.rng_types is not None:
-            synchronize_rng_states(self.rng_types, self.generator)
+            batch_sampler = self.sampler if isinstance(self.sampler, BatchSampler) else self.batch_sampler
Wow! Thanks a lot!
I spent hours finding a way to query my dataset with a list of indices instead of collating data outside the dataset(so that I can write batching logic in the dataset/reduce the number of queries), and failed(thus I was forced to use DataLoader as a sampler(sharded by accelerate), as mentioned in #624 ).
Now I know I can just use batch sampler as sampler without passing batch size to `DataLoader`.
Thanks!!!😁👍
@@ -361,7 +361,15 @@ def __init__(self, dataset, device=None, rng_types=None, generator=None, **kwarg
 
     def __iter__(self):
         if self.rng_types is not None:
-            synchronize_rng_states(self.rng_types, self.generator)
+            batch_sampler = self.sampler if isinstance(self.sampler, BatchSampler) else self.batch_sampler
We should probably leave this one out and pass it as the generator (as initially suggested in your PR).
@@ -650,12 +650,12 @@ def prepare_data_loader(
     # Iterable dataset doesn't like batch_sampler, but data_loader creates a default one for it
     new_batch_sampler = dataloader.batch_sampler if not isinstance(new_dataset, IterableDataset) else None
     sampler_is_batch_sampler = False
-    generator = getattr(dataloader, "generator", None)
+    synchronized_generator = getattr(dataloader, "generator", None)
Here we should pass the main generator of the original dataloader if there was one (see comment above).
@@ -722,9 +721,9 @@ def prepare_data_loader(
             new_dataset,
             device=device if put_on_device and state.distributed_type != DistributedType.TPU else None,
             sampler=new_batch_sampler,
-            batch_size=getattr(dataloader, "batch_size", _PYTORCH_DATALOADER_KWARGS["batch_size"]),
+            batch_size=dataloader.batch_size,
             rng_types=rng_types,
-            generator=generator,
Same there.
@@ -733,7 +732,7 @@ def prepare_data_loader(
             device=device if put_on_device and state.distributed_type != DistributedType.TPU else None,
             batch_sampler=new_batch_sampler,
             rng_types=rng_types,
-            generator=generator,
I remove `"generator"` from `ignore_kwargs`. It should be enough.
@@ -722,9 +721,9 @@ def prepare_data_loader(
             new_dataset,
             device=device if put_on_device and state.distributed_type != DistributedType.TPU else None,
             sampler=new_batch_sampler,
-            batch_size=getattr(dataloader, "batch_size", _PYTORCH_DATALOADER_KWARGS["batch_size"]),
+            batch_size=dataloader.batch_size,
             rng_types=rng_types,
-            generator=generator,
"generator" is removed from `ignore_kwargs`.
@@ -650,12 +650,12 @@ def prepare_data_loader(
     # Iterable dataset doesn't like batch_sampler, but data_loader creates a default one for it
     new_batch_sampler = dataloader.batch_sampler if not isinstance(new_dataset, IterableDataset) else None
     sampler_is_batch_sampler = False
-    generator = getattr(dataloader, "generator", None)
+    synchronized_generator = getattr(dataloader, "generator", None)
Yes, but this one should not be synchronized I think. Or do we want to default to the dataloader generator if tehre is no generator in the sampler?
@@ -650,12 +650,12 @@ def prepare_data_loader(
     # Iterable dataset doesn't like batch_sampler, but data_loader creates a default one for it
     new_batch_sampler = dataloader.batch_sampler if not isinstance(new_dataset, IterableDataset) else None
     sampler_is_batch_sampler = False
-    generator = getattr(dataloader, "generator", None)
+    synchronized_generator = getattr(dataloader, "generator", None)
Indeed, resolving comment.
@@ -722,9 +721,9 @@ def prepare_data_loader(
             new_dataset,
             device=device if put_on_device and state.distributed_type != DistributedType.TPU else None,
             sampler=new_batch_sampler,
-            batch_size=getattr(dataloader, "batch_size", _PYTORCH_DATALOADER_KWARGS["batch_size"]),
+            batch_size=dataloader.batch_size,
             rng_types=rng_types,
-            generator=generator,
Sorry I didn't get your point. Yes, this line should be `synchronized_generator = None`, as #790 .
@@ -650,12 +650,12 @@ def prepare_data_loader(
     # Iterable dataset doesn't like batch_sampler, but data_loader creates a default one for it
     new_batch_sampler = dataloader.batch_sampler if not isinstance(new_dataset, IterableDataset) else None
     sampler_is_batch_sampler = False
-    generator = getattr(dataloader, "generator", None)
+    synchronized_generator = getattr(dataloader, "generator", None)
Oops! Yup, resolved the comment
@@ -650,12 +650,12 @@ def prepare_data_loader(
     # Iterable dataset doesn't like batch_sampler, but data_loader creates a default one for it
     new_batch_sampler = dataloader.batch_sampler if not isinstance(new_dataset, IterableDataset) else None
     sampler_is_batch_sampler = False
-    generator = getattr(dataloader, "generator", None)
+    synchronized_generator = getattr(dataloader, "generator", None)
_The documentation is not available anymore as the PR was closed or merged._
This PR currently will break 
https://github.com/huggingface/accelerate/blob/37b2aa0173fe9b6f46e7df07b92bb55e54b77eb3/src/accelerate/data_loader.py#L653
https://github.com/huggingface/accelerate/blob/37b2aa0173fe9b6f46e7df07b92bb55e54b77eb3/src/accelerate/data_loader.py#L658
https://github.com/huggingface/accelerate/blob/37b2aa0173fe9b6f46e7df07b92bb55e54b77eb3/src/accelerate/data_loader.py#L677
(Of course https://github.com/huggingface/accelerate/blob/37b2aa0173fe9b6f46e7df07b92bb55e54b77eb3/src/accelerate/data_loader.py#L698 should be modified)

---

This code duplication (for finding the sampler) also impairs readability instead of improving readability.
```python
batch_sampler = self.sampler if isinstance(self.sampler, BatchSampler) else self.batch_sampler
sampler = (
    batch_sampler.batch_sampler.sampler
    if hasattr(batch_sampler, "batch_sampler")
    else batch_sampler.sampler
)
if hasattr(sampler, "generator"):
    generator = sampler.generator
    synchronize_rng_states(self.rng_types, generator)
```
I believe what needs to be synchronized is IterableDataset's generator(according to `IterableDatasetShard` implementation) OR sampler's generator. We can store the synchronized generator in a private attribute.

@YouJiacheng I am not seeing any breaks in all the lines you mention, which are completely orthogonal to the change suggested. As for the code duplication, let us worry about readability as maintainers :-) We cannot store the generator as a private attribute as it is fetched in two different functions.
By "break", I means that these lines will become useless. And for `generator = dataloader.dataset.generator`, it is intended to synchronize the generator of IterableDataset. So after proposed change, IterableDataset's generator won't be synchronized.
Ah!, I get what you mean, thanks for clarifying!

We should indeed make the difference between the `sampler_generator` (for batch samplers) or `dataset.generator` (in iterable dataset)  and the Dataloader generator. The latter should be passed along again when creating the final DataLoader (or use the default if it's none) while the formers should be the one synchronized.

The check on `rng_types` should probably be moved to the init of `BatchSamplerShard`/`IterableDatasetShard`.
It's hard to explain my proposal using text. So I open a parallel PR to show the code.(It is somewhat "pseudocode"/proof of concept, since I didn't test it).
@YouJiacheng Thanks for taking the time to draft a PR to show your points, it's much clearer this way! I think we need to merge the two PRs somehow as they both contain important things the other has not.
https://api.github.com/repos/huggingface/accelerate/pulls/789
diff --git a/src/accelerate/launchers.py b/src/accelerate/launchers.py
index 0387109b6..ea8fffa51 100644
--- a/src/accelerate/launchers.py
+++ b/src/accelerate/launchers.py
@@ -43,12 +43,12 @@ def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, mix
             The port to use to communicate between processes when launching a multi-GPU training.
     """
     # Are we in a google colab or a Kaggle Kernel?
+    in_colab = False
+    in_kaggle = False
     if any(key.startswith("KAGGLE") for key in os.environ.keys()):
-        in_colab_or_kaggle = True
+        in_kaggle = True
     elif "IPython" in sys.modules:
-        in_colab_or_kaggle = "google.colab" in str(sys.modules["IPython"].get_ipython())
-    else:
-        in_colab_or_kaggle = False
+        in_colab = "google.colab" in str(sys.modules["IPython"].get_ipython())
 
     try:
         mixed_precision = PrecisionType(mixed_precision.lower())
@@ -57,31 +57,29 @@ def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, mix
             f"Unknown mixed_precision mode: {args.mixed_precision.lower()}. Choose between {PrecisionType.list()}."
         )
 
-    if in_colab_or_kaggle:
-        if os.environ.get("TPU_NAME", None) is not None:
-            # TPU launch
-            import torch_xla.distributed.xla_multiprocessing as xmp
-
-            if len(AcceleratorState._shared_state) > 0:
-                raise ValueError(
-                    "To train on TPU in Colab or Kaggle Kernel, the `Accelerator` should only be initialized inside "
-                    "your training function. Restart your notebook and make sure no cells initializes an "
-                    "`Accelerator`."
-                )
-            if num_processes is None:
-                num_processes = 8
+    if (in_colab or in_kaggle) and (os.environ.get("TPU_NAME", None) is not None):
+        # TPU launch
+        import torch_xla.distributed.xla_multiprocessing as xmp
 
-            launcher = PrepareForLaunch(function, distributed_type="TPU")
-            print(f"Launching a training on {num_processes} TPU cores.")
-            xmp.spawn(launcher, args=args, nprocs=num_processes, start_method="fork")
+        if len(AcceleratorState._shared_state) > 0:
+            raise ValueError(
+                "To train on TPU in Colab or Kaggle Kernel, the `Accelerator` should only be initialized inside "
+                "your training function. Restart your notebook and make sure no cells initializes an "
+                "`Accelerator`."
+            )
+        if num_processes is None:
+            num_processes = 8
+
+        launcher = PrepareForLaunch(function, distributed_type="TPU")
+        print(f"Launching a training on {num_processes} TPU cores.")
+        xmp.spawn(launcher, args=args, nprocs=num_processes, start_method="fork")
+    elif in_colab:
+        # No need for a distributed launch otherwise as it's either CPU or one GPU.
+        if torch.cuda.is_available():
+            print("Launching training on one GPU.")
         else:
-            # No need for a distributed launch otherwise as it's either CPU or one GPU.
-            if torch.cuda.is_available():
-                print("Launching training on one GPU.")
-            else:
-                print("Launching training on one CPU.")
-            function(*args)
-
+            print("Launching training on one CPU.")
+        function(*args)
     else:
         if num_processes is None:
             raise ValueError(

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/783
diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json
new file mode 100644
index 000000000..9d44afde7
--- /dev/null
+++ b/.devcontainer/devcontainer.json
@@ -0,0 +1,25 @@
+// File only needed for VSCode users to have proper Docker based interpreters
+{
+    "name": "accelerate_dev_environment",
+    "build": {
+        // ACTION NEEDED: comment/uncomment the relevant line depending on whether you are in a CPU/GPU environment
+         "dockerfile": "../docker/accelerate-cpu/Dockerfile"
+//        "dockerfile": "../docker/accelerate-gpu/Dockerfile"
+    },
+    "runArgs": [
+        // ACTION NEEDED: uncomment the next line if your local machine has GPUs available
+//        "--gpus", "all",
+        // Enable the docker container to access system resources
+        "--ipc", "host"
+    ],
+    "remoteEnv": {
+        "PYTHONPATH": "${containerEnv:PATH}:${containerWorkspaceFolder}"
+    },
+    "extensions": [
+        // Ensure we have IntelliSense in VSCode when running inside container
+        "ms-python.python"
+    ],
+    "workspaceFolder": "/workspaces/accelerate",
+    // Need git for VSCode to color code modifications. Only runs when building environment.
+    "onCreateCommand": "apt-get update && apt-get install -y git && pip install -e '.[dev]'"
+}
\ No newline at end of file
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index fcc7d9b1c..d0e142b16 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -130,6 +130,9 @@ Follow these steps to start contributing:
    it with `pip uninstall accelerate` before reinstalling it in editable
    mode with the `-e` flag.)
 
+   Alternatively, if you are using [Visual Studio Code](https://code.visualstudio.com/Download), the fastest way to get set up is by using
+   the provided Dev Container. Documentation on how to get started with dev containers is available [here](https://code.visualstudio.com/docs/remote/containers).
+
 5. Develop the features on your branch.
 
    As you work on the features, you should make sure that the test suite

I think we should enable CPU by default and have the user work on the container to activate GPU.
@@ -0,0 +1,25 @@
+// File only needed for VSCode users to have proper Docker based interpreters
+{
+    "name": "accelerate_dev_environment",
+    "build": {
+        // ACTION NEEDED: comment/uncomment the relevant line depending on whether you are in a CPU/GPU environment
+        // "dockerfile": "../docker/accelerate-cpu/Dockerfile"
Same here.
@@ -0,0 +1,25 @@
+// File only needed for VSCode users to have proper Docker based interpreters
+{
+    "name": "accelerate_dev_environment",
+    "build": {
+        // ACTION NEEDED: comment/uncomment the relevant line depending on whether you are in a CPU/GPU environment
+        // "dockerfile": "../docker/accelerate-cpu/Dockerfile"
+        "dockerfile": "../docker/accelerate-gpu/Dockerfile"
+    },
+    "runArgs": [
+        // ACTION NEEDED: comment the next line if your local machine has no GPUs available
Sure, I was torn on this, so am happy to go the other way.
@@ -0,0 +1,25 @@
+// File only needed for VSCode users to have proper Docker based interpreters
+{
+    "name": "accelerate_dev_environment",
+    "build": {
+        // ACTION NEEDED: comment/uncomment the relevant line depending on whether you are in a CPU/GPU environment
+        // "dockerfile": "../docker/accelerate-cpu/Dockerfile"
_The documentation is not available anymore as the PR was closed or merged._
> Also, if the users change the devcontainer.json for their local use, won't those changes be pushed to GitHub automatically?

Unfortunately, VSCode doesn't (yet) support multiple dev container configurations to my knowledge, so we have to rely on people not committing dev container changes.
https://api.github.com/repos/huggingface/accelerate/pulls/782
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index bf17f57ce..81759962c 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -153,6 +153,10 @@ class Accelerator:
             If set to `True`, the dataloader prepared by the Accelerator is only iterated through on the main process
             and then the batches are split and broadcast to each process. Will default to `True` for `DataLoader` whose
             underlying dataset is an `IterableDataset`, `False` otherwise.
+        even_batches (`bool`, *optional*, defaults to `True`):
+            If set to `True`, in cases where the total batch size across all processes does not exactly divide the
+            dataset, samples at the start of the dataset will be duplicated so the batch can be divided equally among
+            all workers.
         step_scheduler_with_optimizer (`bool`, *optional`, defaults to `True`):
             Set `True` if the learning rate scheduler is stepped at the same time as the optimizer, `False` if only
             done under certain circumstances (at the end of each epoch, for instance).
@@ -191,6 +195,7 @@ def __init__(
         log_with: Optional[List[Union[str, LoggerType, GeneralTracker]]] = None,
         logging_dir: Optional[Union[str, os.PathLike]] = None,
         dispatch_batches: Optional[bool] = None,
+        even_batches: bool = True,
         step_scheduler_with_optimizer: bool = True,
         kwargs_handlers: Optional[List[KwargsHandler]] = None,
     ):
@@ -305,6 +310,7 @@ def __init__(
             raise ImportError(
                 "Using `DataLoaderDispatcher` requires PyTorch 1.8.0 minimum. You have {torch.__version__}."
             )
+        self.even_batches = even_batches
         self.step_scheduler_with_optimizer = step_scheduler_with_optimizer
 
         # Mixed precision attributes
@@ -1109,6 +1115,7 @@ def prepare_data_loader(self, data_loader: torch.utils.data.DataLoader, device_p
             put_on_device=device_placement,
             rng_types=self.rng_types.copy(),
             dispatch_batches=self.dispatch_batches,
+            even_batches=self.even_batches,
         )
 
     def prepare_optimizer(self, optimizer: torch.optim.Optimizer, device_placement=None):
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index cdae54af9..39f843430 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -118,7 +118,8 @@ class BatchSamplerShard(BatchSampler):
 
     <Tip warning={true}>
 
-    This does not support `BatchSampler` with varying batch size yet.
+    `BatchSampler`s with varying batch sizes are not enabled by default. To enable this behaviour, set `even_batches`
+    equal to `False`
 
     </Tip>"""
 
@@ -570,6 +571,7 @@ def prepare_data_loader(
     put_on_device: bool = False,
     rng_types: Optional[List[Union[str, RNGType]]] = None,
     dispatch_batches: Optional[bool] = None,
+    even_batches: bool = True,
 ) -> DataLoader:
     """
     Wraps a PyTorch `DataLoader` to generate batches for one of the processes only.
@@ -615,15 +617,21 @@ def prepare_data_loader(
             If set to `True`, the datalaoder prepared is only iterated through on the main process and then the batches
             are split and broadcast to each process. Will default to `True` when the underlying dataset is an
             `IterableDataset`, `False` otherwise.
+        even_batches (`bool`, *optional*, defaults to `True`):
+            If set to `True`, in cases where the total batch size across all processes does not exactly divide the
+            dataset, samples at the start of the dataset will be duplicated so the batch can be divided equally among
+            all workers.
 
     Returns:
         `torch.utils.data.dataloader.DataLoader`: A new data loader that will yield the portion of the batches
 
     <Tip warning={true}>
 
-    This does not support `BatchSampler` with varying batch size yet.
+    `BatchSampler`s with varying batch sizes are not enabled by default. To enable this behaviour, set `even_batches`
+    equal to `False`
 
-    </Tip>"""
+    </Tip>
+    """
     if dispatch_batches is None:
         if is_torch_version("<", "1.8.0") or not put_on_device:
             dispatch_batches = False
@@ -683,6 +691,7 @@ def prepare_data_loader(
                 num_processes=num_processes,
                 process_index=process_index,
                 split_batches=split_batches,
+                even_batches=even_batches,
             )
 
     # We ignore all of those since they are all dealt with by our new_batch_sampler
diff --git a/src/accelerate/test_utils/scripts/test_distributed_data_loop.py b/src/accelerate/test_utils/scripts/test_distributed_data_loop.py
new file mode 100644
index 000000000..eaf7c6a34
--- /dev/null
+++ b/src/accelerate/test_utils/scripts/test_distributed_data_loop.py
@@ -0,0 +1,113 @@
+#!/usr/bin/env python
+
+# 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.
+
+from typing import List
+
+import torch
+from torch.utils.data import DataLoader, TensorDataset
+
+from accelerate.accelerator import Accelerator
+
+
+def create_accelerator(even_batches=True):
+    accelerator = Accelerator(even_batches=even_batches)
+    assert accelerator.num_processes == 2, "this script expects that two GPUs are available"
+    return accelerator
+
+
+def create_dataloader(accelerator: Accelerator, dataset_size: int, batch_size: int):
+    """
+    Create a simple DataLoader to use during the test cases
+    """
+    dataset = TensorDataset(torch.as_tensor(range(dataset_size)))
+
+    dl = DataLoader(dataset, batch_size=batch_size)
+    dl = accelerator.prepare(dl)
+
+    return dl
+
+
+def verify_dataloader_batch_sizes(
+    accelerator: Accelerator,
+    dataset_size: int,
+    batch_size: int,
+    process_0_expected_batch_sizes: List[int],
+    process_1_expected_batch_sizes: List[int],
+):
+    """
+    A helper function for verifying the batch sizes coming from a prepared dataloader in each process
+    """
+    dl = create_dataloader(accelerator=accelerator, dataset_size=dataset_size, batch_size=batch_size)
+
+    batch_sizes = [len(batch[0]) for batch in dl]
+
+    if accelerator.process_index == 0:
+        assert batch_sizes == process_0_expected_batch_sizes
+    elif accelerator.process_index == 1:
+        assert batch_sizes == process_1_expected_batch_sizes
+
+
+def test_default_ensures_even_batch_sizes():
+
+    accelerator = create_accelerator()
+
+    # without padding, we would expect a different number of batches
+    verify_dataloader_batch_sizes(
+        accelerator,
+        dataset_size=3,
+        batch_size=1,
+        process_0_expected_batch_sizes=[1, 1],
+        process_1_expected_batch_sizes=[1, 1],
+    )
+
+    # without padding, we would expect the same number of batches, but different sizes
+    verify_dataloader_batch_sizes(
+        accelerator,
+        dataset_size=7,
+        batch_size=2,
+        process_0_expected_batch_sizes=[2, 2],
+        process_1_expected_batch_sizes=[2, 2],
+    )
+
+
+def test_can_disable_even_batches():
+    accelerator = create_accelerator(even_batches=False)
+
+    verify_dataloader_batch_sizes(
+        accelerator,
+        dataset_size=3,
+        batch_size=1,
+        process_0_expected_batch_sizes=[1, 1],
+        process_1_expected_batch_sizes=[1],
+    )
+
+    verify_dataloader_batch_sizes(
+        accelerator,
+        dataset_size=7,
+        batch_size=2,
+        process_0_expected_batch_sizes=[2, 2],
+        process_1_expected_batch_sizes=[2, 1],
+    )
+
+
+if __name__ == "__main__":
+    accelerator = create_accelerator()
+
+    accelerator.print("Test that even_batches variable ensures uniform batches across processes")
+    test_default_ensures_even_batch_sizes()
+
+    accelerator.print("Run tests with even_batches disabled")
+    test_can_disable_even_batches()
diff --git a/tests/test_multigpu.py b/tests/test_multigpu.py
index 2c0403ed5..54a5bca75 100644
--- a/tests/test_multigpu.py
+++ b/tests/test_multigpu.py
@@ -28,6 +28,9 @@ class MultiGPUTester(unittest.TestCase):
     def setUp(self):
         mod_file = inspect.getfile(accelerate.test_utils)
         self.test_file_path = os.path.sep.join(mod_file.split(os.path.sep)[:-1] + ["scripts", "test_script.py"])
+        self.data_loop_file_path = os.path.sep.join(
+            mod_file.split(os.path.sep)[:-1] + ["scripts", "test_distributed_data_loop.py"]
+        )
 
     @require_multi_gpu
     def test_multi_gpu(self):
@@ -42,6 +45,17 @@ def test_pad_across_processes(self):
         with patch_environment(omp_num_threads=1):
             execute_subprocess_async(cmd, env=os.environ.copy())
 
+    @require_multi_gpu
+    def test_distributed_data_loop(self):
+        """
+        This TestCase checks the behaviour that occurs during distributed training or evaluation,
+        when the batch size does not evenly divide the dataset size.
+        """
+        print(f"Found {torch.cuda.device_count()} devices, using 2 devices only")
+        cmd = get_launch_prefix() + [f"--nproc_per_node={torch.cuda.device_count()}", self.data_loop_file_path]
+        with patch_environment(omp_num_threads=1, cuda_visible_devices="0,1"):
+            execute_subprocess_async(cmd, env=os.environ.copy())
+
 
 if __name__ == "__main__":
     accelerator = Accelerator()

```suggestion
        even_batches (`bool`, *optional*, defaults to `True`):
```
@@ -153,6 +153,10 @@ class Accelerator:
             If set to `True`, the dataloader prepared by the Accelerator is only iterated through on the main process
             and then the batches are split and broadcast to each process. Will default to `True` for `DataLoader` whose
             underlying dataset is an `IterableDataset`, `False` otherwise.
+        even_batches ('bool', defaults to `True`):
```suggestion
        even_batches (`bool`, *optional*, defaults to `True`):
```
@@ -615,15 +616,15 @@ def prepare_data_loader(
             If set to `True`, the datalaoder prepared is only iterated through on the main process and then the batches
             are split and broadcast to each process. Will default to `True` when the underlying dataset is an
             `IterableDataset`, `False` otherwise.
+        even_batches ('bool', defaults to `True`):
You should leave the warning but rephrase to state a `BatchSampler` with varying batch size will require `even_batches=False`.
@@ -615,15 +616,15 @@ def prepare_data_loader(
             If set to `True`, the datalaoder prepared is only iterated through on the main process and then the batches
             are split and broadcast to each process. Will default to `True` when the underlying dataset is an
             `IterableDataset`, `False` otherwise.
+        even_batches ('bool', defaults to `True`):
+            If set to `True`, in cases where the total batch size across all processes does not exactly divide the
+            dataset, samples at the start of the dataset will be duplicated so the batch can be divided equally among
+            all workers.
 
     Returns:
         `torch.utils.data.dataloader.DataLoader`: A new data loader that will yield the portion of the batches
 
-    <Tip warning={true}>
-
-    This does not support `BatchSampler` with varying batch size yet.
Sure, I can update that
@@ -615,15 +616,15 @@ def prepare_data_loader(
             If set to `True`, the datalaoder prepared is only iterated through on the main process and then the batches
             are split and broadcast to each process. Will default to `True` when the underlying dataset is an
             `IterableDataset`, `False` otherwise.
+        even_batches ('bool', defaults to `True`):
+            If set to `True`, in cases where the total batch size across all processes does not exactly divide the
+            dataset, samples at the start of the dataset will be duplicated so the batch can be divided equally among
+            all workers.
 
     Returns:
         `torch.utils.data.dataloader.DataLoader`: A new data loader that will yield the portion of the batches
 
-    <Tip warning={true}>
-
-    This does not support `BatchSampler` with varying batch size yet.
You're currently not testing anything at all, it's never being ran. Please see here to learn how: https://github.com/huggingface/accelerate/blob/main/src/accelerate/test_utils/scripts/test_sync.py#L229-L274

I also highly recommend using print statements so we can see just what is being tested at a given moment, and we know if they've run 😄 
@@ -0,0 +1,103 @@
+#!/usr/bin/env python
+
+# 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.
+
+from typing import List
+
+import torch
+from torch.utils.data import DataLoader, TensorDataset
+
+from accelerate.accelerator import Accelerator
+
+
+def create_accelerator(even_batches=True):
+    accelerator = Accelerator(even_batches=even_batches)
+    assert accelerator.num_processes == 2, "this script expects that two GPUs are available"
+    return accelerator
+
+
+def create_dataloader(accelerator: Accelerator, dataset_size: int, batch_size: int):
+    """
+    Create a simple DataLoader to use during the test cases
+    """
+    dataset = TensorDataset(torch.as_tensor(range(dataset_size)))
+
+    dl = DataLoader(dataset, batch_size=batch_size)
+    dl = accelerator.prepare(dl)
+
+    return dl
+
+
+def verify_dataloader_batch_sizes(
+    accelerator: Accelerator,
+    dataset_size: int,
+    batch_size: int,
+    process_0_expected_batch_sizes: List[int],
+    process_1_expected_batch_sizes: List[int],
+):
+    """
+    A helper function for verifying the batch sizes coming from a prepared dataloader in each process
+    """
+    dl = create_dataloader(accelerator=accelerator, dataset_size=dataset_size, batch_size=batch_size)
+
+    batch_sizes = [len(batch[0]) for batch in dl]
+
+    if accelerator.process_index == 0:
+        assert batch_sizes == process_0_expected_batch_sizes
+    elif accelerator.process_index == 1:
+        assert batch_sizes == process_1_expected_batch_sizes
+
+
+def test_default_ensures_even_batch_sizes():
+
+    accelerator = create_accelerator()
+
+    # without padding, we would expect a different number of batches
+    verify_dataloader_batch_sizes(
+        accelerator,
+        dataset_size=3,
+        batch_size=1,
+        process_0_expected_batch_sizes=[1, 1],
+        process_1_expected_batch_sizes=[1, 1],
+    )
+
+    # without padding, we would expect the same number of batches, but different sizes
+    verify_dataloader_batch_sizes(
+        accelerator,
+        dataset_size=7,
+        batch_size=2,
+        process_0_expected_batch_sizes=[2, 2],
+        process_1_expected_batch_sizes=[2, 2],
+    )
+
+
+def test_can_disable_even_batches():
Gah! I think I misinterpreted what you meant when you said PyTest format. Looking at the other example, it looks like it still needs a main method to drive running the tests; just what I removed a few commits back :satisfied:! I'll reinsert the main method and some print statements.

That said, when I run it with PyTest locally, using `CUDA_VISIBLE_DEVICES="0,1" pytest -sv tests/test_multigpu.py` it does seem to pick them up.
@@ -0,0 +1,103 @@
+#!/usr/bin/env python
+
+# 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.
+
+from typing import List
+
+import torch
+from torch.utils.data import DataLoader, TensorDataset
+
+from accelerate.accelerator import Accelerator
+
+
+def create_accelerator(even_batches=True):
+    accelerator = Accelerator(even_batches=even_batches)
+    assert accelerator.num_processes == 2, "this script expects that two GPUs are available"
+    return accelerator
+
+
+def create_dataloader(accelerator: Accelerator, dataset_size: int, batch_size: int):
+    """
+    Create a simple DataLoader to use during the test cases
+    """
+    dataset = TensorDataset(torch.as_tensor(range(dataset_size)))
+
+    dl = DataLoader(dataset, batch_size=batch_size)
+    dl = accelerator.prepare(dl)
+
+    return dl
+
+
+def verify_dataloader_batch_sizes(
+    accelerator: Accelerator,
+    dataset_size: int,
+    batch_size: int,
+    process_0_expected_batch_sizes: List[int],
+    process_1_expected_batch_sizes: List[int],
+):
+    """
+    A helper function for verifying the batch sizes coming from a prepared dataloader in each process
+    """
+    dl = create_dataloader(accelerator=accelerator, dataset_size=dataset_size, batch_size=batch_size)
+
+    batch_sizes = [len(batch[0]) for batch in dl]
+
+    if accelerator.process_index == 0:
+        assert batch_sizes == process_0_expected_batch_sizes
+    elif accelerator.process_index == 1:
+        assert batch_sizes == process_1_expected_batch_sizes
+
+
+def test_default_ensures_even_batch_sizes():
+
+    accelerator = create_accelerator()
+
+    # without padding, we would expect a different number of batches
+    verify_dataloader_batch_sizes(
+        accelerator,
+        dataset_size=3,
+        batch_size=1,
+        process_0_expected_batch_sizes=[1, 1],
+        process_1_expected_batch_sizes=[1, 1],
+    )
+
+    # without padding, we would expect the same number of batches, but different sizes
+    verify_dataloader_batch_sizes(
+        accelerator,
+        dataset_size=7,
+        batch_size=2,
+        process_0_expected_batch_sizes=[2, 2],
+        process_1_expected_batch_sizes=[2, 2],
+    )
+
+
+def test_can_disable_even_batches():
I've just pushed some changes that hopefully resolve this. Just FYI, it is getting pretty late here in the UK, and I am heading off on vacation tomorrow until the 31st. If there are other changes required, and you would like this merged before I'm back, please feel free to update the branch with any further changes!
@@ -0,0 +1,103 @@
+#!/usr/bin/env python
+
+# 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.
+
+from typing import List
+
+import torch
+from torch.utils.data import DataLoader, TensorDataset
+
+from accelerate.accelerator import Accelerator
+
+
+def create_accelerator(even_batches=True):
+    accelerator = Accelerator(even_batches=even_batches)
+    assert accelerator.num_processes == 2, "this script expects that two GPUs are available"
+    return accelerator
+
+
+def create_dataloader(accelerator: Accelerator, dataset_size: int, batch_size: int):
+    """
+    Create a simple DataLoader to use during the test cases
+    """
+    dataset = TensorDataset(torch.as_tensor(range(dataset_size)))
+
+    dl = DataLoader(dataset, batch_size=batch_size)
+    dl = accelerator.prepare(dl)
+
+    return dl
+
+
+def verify_dataloader_batch_sizes(
+    accelerator: Accelerator,
+    dataset_size: int,
+    batch_size: int,
+    process_0_expected_batch_sizes: List[int],
+    process_1_expected_batch_sizes: List[int],
+):
+    """
+    A helper function for verifying the batch sizes coming from a prepared dataloader in each process
+    """
+    dl = create_dataloader(accelerator=accelerator, dataset_size=dataset_size, batch_size=batch_size)
+
+    batch_sizes = [len(batch[0]) for batch in dl]
+
+    if accelerator.process_index == 0:
+        assert batch_sizes == process_0_expected_batch_sizes
+    elif accelerator.process_index == 1:
+        assert batch_sizes == process_1_expected_batch_sizes
+
+
+def test_default_ensures_even_batch_sizes():
+
+    accelerator = create_accelerator()
+
+    # without padding, we would expect a different number of batches
+    verify_dataloader_batch_sizes(
+        accelerator,
+        dataset_size=3,
+        batch_size=1,
+        process_0_expected_batch_sizes=[1, 1],
+        process_1_expected_batch_sizes=[1, 1],
+    )
+
+    # without padding, we would expect the same number of batches, but different sizes
+    verify_dataloader_batch_sizes(
+        accelerator,
+        dataset_size=7,
+        batch_size=2,
+        process_0_expected_batch_sizes=[2, 2],
+        process_1_expected_batch_sizes=[2, 2],
+    )
+
+
+def test_can_disable_even_batches():
_The documentation is not available anymore as the PR was closed or merged._
Hi, I have made a start on the work discussed in https://github.com/huggingface/accelerate/issues/684, and added an argument to the accelerator. 

I am slightly unsure about your testing strategy though, and which tests you require for this; I couldn't seem to find any integration tests for accelerator arguments. So far, I have tried to follow similar ideas to those seen in the end-to-end multi-GPU tests. 

I am slightly confused by exactly how the Multi-GPU tests are triggered though; whilst I see that there is a decorator to skip certain tests when multiple GPUs are not available, when I try to run this locally (on a machine with 2 GPUs), those tests only seem to execute on a single process.

Any guidance here would be appreciated! cc: @muellerzr @sgugger 
Hey @Chris-hughes10! Thanks!

You're absolutely right, I should probably write this all up in a contribution doc and will add it to my list of tasks!

Accelerates tests are setup in a few ways, but the key idea to remember is:

If it's something to be tested across multiple devices/workers, such as CPU, TPU, and mutli-GPU (such as your tests you want to add), these all get thrown into a script located in `accelerate/test_utils/scripts`

For example, the major integration script that checks the dataloaders operate properly is located at `accelerate/test_utils/scripts/test_script.py` (I wonder if we should section this out a bit more @sgugger since we now have a whole folder dedicated to it and we can clean it up some),  which you can find here: https://github.com/huggingface/accelerate/blob/main/src/accelerate/test_utils/scripts/test_script.py

These tests are setup like how a `pytest` test would be, where you have `def some_test` and `assert` statements to check if things are working properly. 

From here your script should be added as a test case to `tests/test.py`, `tests/test_cpu.py` (if applicable for multi-cpu),  `tests/test_multigpu.py`, and `tests/test_tpu.py`.

It has to be done this way because each platform has different launching configurations :) 

Finally, when running pytest you need to set the cuda visible devices for it to be able to pick up on the multiple GPUs, such as:

```bash
CUDA_VISIBLE_DEVICES="0,1" pytest -sv tests/test_multigpu.py
```
Or:
```bash
CUDA_VISIBLE_DEVICES="" pytest -sv tests/test_multigpu.py
```
(the latter will have all their tests skipped). 

All skipping decorator conditionals can be found here: https://github.com/huggingface/accelerate/blob/main/src/accelerate/test_utils/testing.py

Let me know if there are any other questions you have on this @Chris-hughes10, afterwards I can throw this all up into some documentation :) 
Hey @muellerzr, Thanks for the info!

So, it looks like I was on the right track regarding the structure, but I was running with unittests and not Pytest. I have refactored it to pytest, and it seems to work alright using `CUDA_VISIBLE_DEVICES="0,1" pytest -sv tests/test_multigpu.py, so it is probably ready for an initial review. 

I noticed that the main integration script seems to be doing a lot, so I created a new script to be super clear on what it is testing. The tests for the context manager and joining can probably live here too when implemented.

As an aside, how do you guys set up your development environment? I have set up a dev container, using the accelerate-gpu image, if it is useful, I could push that in another PR.
@Chris-hughes10 I use conda for mine, and the CI's all run based off the docker images as there's hosted computes we use
> @Chris-hughes10 I use conda for mine, and the CI's all run based off the conda images as there's hosted computes we use

Ok, if you are using conda envs, I think that this could be useful. I'll create a small PR and you guys can see what you think.
Hey @muellerzr, are there any further changes required on this?
@Chris-hughes10 close, your tests still weren't passing, and this is because `nproc_per_node` needed to be specified in the test function. (Discussing with Sylvain if this should be *everywhere* or not) but basically without that `Accelerator.num_processes` was only ever equal to one. 

Went ahead and pushed the fix :)
Thanks for all your work on this!
https://api.github.com/repos/huggingface/accelerate/pulls/781
diff --git a/docs/source/package_reference/cli.mdx b/docs/source/package_reference/cli.mdx
index a9e955d43..52e8ecfb6 100644
--- a/docs/source/package_reference/cli.mdx
+++ b/docs/source/package_reference/cli.mdx
@@ -95,10 +95,13 @@ The following arguments are useful for customization of worker machines
 * `--num_machines NUM_MACHINES` (`int`) -- The total number of machines used in this training.
 * `--num_processes NUM_PROCESSES` (`int`) -- The total number of processes to be launched in parallel.
 * `--gpu_ids` (`str`) -- What GPUs (by id) should be used for training on this machine as a comma-seperated list
+* `--same_network` (`bool`) -- Whether all machines used for multinode training exist on the same local network.
 * `--main_process_ip MAIN_PROCESS_IP` (`str`) -- The IP address of the machine of rank 0.
 * `--main_process_port MAIN_PROCESS_PORT` (`int`) -- The port to use to communicate with the machine of rank 0.
+* `--rdzv_conf` (`str`) -- Additional rendezvous configuration (<key1>=<value1>,<key2>=<value2>,...).
 * `--num_cpu_threads_per_process NUM_CPU_THREADS_PER_PROCESS` (`int`) -- The number of CPU threads per process. Can be tuned for optimal performance.
-
+* `--max_restarts` (`int`) -- Maximum number of worker group restarts before failing.
+* `--monitor_interval` (`float`) -- Interval, in seconds, to monitor the state of workers.
 
 **DeepSpeed Arguments**:
 
@@ -125,8 +128,9 @@ The following arguments are only useful when `use_fdsp` is passed:
 **TPU Arguments**:
 
 The following arguments are only useful when `tpu` is passed:
-* `--tpu` (`bool`) - Whether or not this should launch a TPU training.
+* `--tpu` (`bool`) -- Whether or not this should launch a TPU training.
 * `--main_training_function MAIN_TRAINING_FUNCTION` (`str`) -- The name of the main function to be executed in your script.
+* `--downcast_bf16` (`bool`) -- Whether when using bf16 precision on TPUs if both float and double tensors are cast to bfloat16 or if double tensors remain as float32.
 
 **AWS SageMaker Arguments**:
 
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 2198bb3f1..b9f6c1089 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -325,6 +325,12 @@ def launch_command_parser(subparsers=None):
         default=None,
         help="What GPUs (by id) should be used for training on this machine as a comma-seperated list",
     )
+    parser.add_argument(
+        "--same_network",
+        default=False,
+        action="store_true",
+        help="Whether all machines used for multinode training exist on the same local network.",
+    )
     parser.add_argument(
         "--machine_rank", type=int, default=None, help="The rank of the machine on which this script is launched."
     )
@@ -363,7 +369,7 @@ def launch_command_parser(subparsers=None):
     parser.add_argument(
         "--downcast_bf16",
         action="store_true",
-        help="Whether when using bf16 precision on TPUs if both float and double tensors are cast to bfloat16 or if double tensors remain as float32",
+        help="Whether when using bf16 precision on TPUs if both float and double tensors are cast to bfloat16 or if double tensors remain as float32.",
     )
     parser.add_argument(
         "-m",

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/780
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index fa7588b0a..29e1b7bad 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -573,7 +573,7 @@ def multi_gpu_launcher(args):
         try:
             distrib_run.run(args)
         except:
-            if debug:
+            if is_rich_available() and debug:
                 console = get_console()
                 console.print("\n[bold red]Using --debug, `torch.distributed` Stack Trace:[/bold red]")
                 console.print_exception(suppress=[__file__], show_locals=False)
@@ -684,7 +684,7 @@ def deepspeed_launcher(args):
             try:
                 distrib_run.run(args)
             except:
-                if debug:
+                if is_rich_available() and debug:
                     console = get_console()
                     console.print("\n[bold red]Using --debug, `torch.distributed` Stack Trace:[/bold red]")
                     console.print_exception(suppress=[__file__], show_locals=False)
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index 82ea288f0..00ba3ab67 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -14,11 +14,19 @@
 
 import os
 import warnings
-from distutils.util import strtobool
 
 import torch
 
-from .utils import DistributedType, get_ccl_version, is_ccl_available, is_deepspeed_available, is_tpu_available
+from .utils import (
+    DistributedType,
+    get_ccl_version,
+    get_int_from_env,
+    is_ccl_available,
+    is_deepspeed_available,
+    is_tpu_available,
+    parse_choice_from_env,
+    parse_flag_from_env,
+)
 from .utils.dataclasses import SageMakerDistributedType
 
 
@@ -26,25 +34,6 @@
     import torch_xla.core.xla_model as xm
 
 
-def get_int_from_env(env_keys, default):
-    """Returns the first positive env value found in the `env_keys` list or the default."""
-    for e in env_keys:
-        val = int(os.environ.get(e, -1))
-        if val >= 0:
-            return val
-    return default
-
-
-def parse_flag_from_env(key, default=False):
-    value = os.environ.get(key, str(default))
-    return strtobool(value) == 1  # As its name indicates `strtobool` actually returns an int...
-
-
-def parse_choice_from_env(key, default="no"):
-    value = os.environ.get(key, str(default))
-    return value
-
-
 # Inspired by Alex Martelli's 'Borg'.
 class AcceleratorState:
     """
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
index b472ec556..b22518f9a 100644
--- a/src/accelerate/utils/__init__.py
+++ b/src/accelerate/utils/__init__.py
@@ -19,6 +19,7 @@
     SageMakerDistributedType,
     TensorInformation,
 )
+from .environment import get_int_from_env, parse_choice_from_env, parse_flag_from_env
 from .imports import (
     get_ccl_version,
     is_aim_available,
diff --git a/src/accelerate/utils/environment.py b/src/accelerate/utils/environment.py
new file mode 100644
index 000000000..9247bb4fe
--- /dev/null
+++ b/src/accelerate/utils/environment.py
@@ -0,0 +1,36 @@
+# 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
+from distutils.util import strtobool
+
+
+def get_int_from_env(env_keys, default):
+    """Returns the first positive env value found in the `env_keys` list or the default."""
+    for e in env_keys:
+        val = int(os.environ.get(e, -1))
+        if val >= 0:
+            return val
+    return default
+
+
+def parse_flag_from_env(key, default=False):
+    """Returns truthy value for `key` from the env if available else the default."""
+    value = os.environ.get(key, str(default))
+    return strtobool(value) == 1  # As its name indicates `strtobool` actually returns an int...
+
+
+def parse_choice_from_env(key, default="no"):
+    value = os.environ.get(key, str(default))
+    return value
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index bdb53988a..c33258a92 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -20,6 +20,7 @@
 
 from packaging.version import parse
 
+from .environment import parse_flag_from_env
 from .versions import compare_versions, is_torch_version
 
 
@@ -126,7 +127,7 @@ def is_boto3_available():
 
 
 def is_rich_available():
-    return importlib.util.find_spec("rich") is not None
+    return (importlib.util.find_spec("rich") is not None) and (not parse_flag_from_env("DISABLE_RICH"))
 
 
 def is_sagemaker_available():

Let's accept anything that's truthy here? Like yes/no, true/false, 0/1
@@ -126,7 +127,7 @@ def is_boto3_available():
 
 
 def is_rich_available():
-    return importlib.util.find_spec("rich") is not None
+    return (importlib.util.find_spec("rich") is not None) and (not os.environ.get("DISABLE_RICH", "0") == "1")
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/779
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index fa7588b0a..071c86a1d 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -470,7 +470,7 @@ def multi_gpu_launcher(args):
         setattr(args, "nproc_per_node", str(num_processes // num_machines))
         setattr(args, "nnodes", str(num_machines))
         setattr(args, "node_rank", int(args.machine_rank))
-        if getattr(args, "same_network"):
+        if getattr(args, "same_network", False):
             setattr(args, "master_addr", str(main_process_ip))
             setattr(args, "master_port", str(main_process_port))
         else:
@@ -488,7 +488,7 @@ def multi_gpu_launcher(args):
         setattr(args, "no_python", True)
 
     current_env = os.environ.copy()
-    gpu_ids = getattr(args, "gpu_ids")
+    gpu_ids = getattr(args, "gpu_ids", "all")
     if gpu_ids != "all" and args.gpu_ids is not None:
         current_env["CUDA_VISIBLE_DEVICES"] = gpu_ids
     mixed_precision = args.mixed_precision.lower()
@@ -618,7 +618,7 @@ def deepspeed_launcher(args):
         setattr(args, "nproc_per_node", str(num_processes // num_machines))
         setattr(args, "nnodes", str(num_machines))
         setattr(args, "node_rank", int(args.machine_rank))
-        if getattr(args, "same_network"):
+        if getattr(args, "same_network", False):
             setattr(args, "master_addr", str(main_process_ip))
             setattr(args, "master_port", str(main_process_port))
         else:
@@ -636,7 +636,7 @@ def deepspeed_launcher(args):
         setattr(args, "no_python", True)
 
     current_env = os.environ.copy()
-    gpu_ids = getattr(args, "gpu_ids")
+    gpu_ids = getattr(args, "gpu_ids", "all")
     if gpu_ids != "all" and args.gpu_ids is not None:
         current_env["CUDA_VISIBLE_DEVICES"] = gpu_ids
     try:

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/778
diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
index 1760e72bc..c14833557 100644
--- a/src/accelerate/utils/modeling.py
+++ b/src/accelerate/utils/modeling.py
@@ -408,8 +408,9 @@ def get_balanced_memory(
     per_gpu += buffer
 
     max_memory = get_max_memory(max_memory)
+    last_gpu = max(i for i in max_memory if isinstance(i, int) and max_memory[i] > 0)
     # The last device is left with max_memory just in case the buffer is not enough.
-    for i in range(len(max_memory) - 1):
+    for i in range(last_gpu):
         max_memory[i] = min(0 if low_zero and i == 0 else per_gpu, max_memory[i])
 
     if low_zero:
diff --git a/tests/test_modeling_utils.py b/tests/test_modeling_utils.py
index 02ab867d7..16243d5e1 100644
--- a/tests/test_modeling_utils.py
+++ b/tests/test_modeling_utils.py
@@ -376,6 +376,10 @@ def test_get_balanced_memory(self):
         max_memory = get_balanced_memory(model, max_memory={0: 300, 1: 500})
         self.assertDictEqual({0: 215, 1: 500}, max_memory)
 
+        # Last device always get max memory to give more buffer, even if CPU is provided
+        max_memory = get_balanced_memory(model, max_memory={0: 300, "cpu": 1000})
+        self.assertDictEqual({0: 300, "cpu": 1000}, max_memory)
+
         # If we set a device to 0, it's not counted.
         max_memory = get_balanced_memory(model, max_memory={0: 0, 1: 300, 2: 300})
         self.assertDictEqual({0: 0, 1: 215, 2: 300}, max_memory)

```suggestion
    offset = 2 if "cpu" in (list(max_memory.keys())) else 1
    for i in range(len(max_memory) - offset):
```
I think this way it seems that the slow tests of `accelerate` are not broken and we also keep the `transformers` tests pass! Can you confirm this pass the `transformers` tests that were broken @ArthurZucker ? 🙏 
@@ -409,7 +409,7 @@ def get_balanced_memory(
 
     max_memory = get_max_memory(max_memory)
     # The last device is left with max_memory just in case the buffer is not enough.
-    for i in range(len(max_memory) - 1):
+    for i in range(num_devices - 1):
Will have a look😄 I am also adding the following 
```python
max_memory = get_balanced_memory(model, max_memory={0: 300, 1: 500, "cpu":1000})
```
@@ -409,7 +409,7 @@ def get_balanced_memory(
 
     max_memory = get_max_memory(max_memory)
     # The last device is left with max_memory just in case the buffer is not enough.
-    for i in range(len(max_memory) - 1):
+    for i in range(num_devices - 1):
Ahhh great! I was missing this test! Thanks for clarifying 💪 
@@ -409,7 +409,7 @@ def get_balanced_memory(
 
     max_memory = get_max_memory(max_memory)
     # The last device is left with max_memory just in case the buffer is not enough.
-    for i in range(len(max_memory) - 1):
+    for i in range(num_devices - 1):
This won't work when the last device is GPU 3 but the user forced 0 at GPU 0, since `num_devices` will be 3 (instead of 4) and thus the max memory will be at GPU 2 instead of 3. (you can try running the multi-GPU tests and they will fail).

Agreed that `len(max_memory)` is not good either, so we should just compute the biggest integer in the keys of `max_memory`:
```py
last_gpu = max(i for i in max_memory if isintance(i, int))
```
then use this.


@@ -409,7 +409,7 @@ def get_balanced_memory(
 
     max_memory = get_max_memory(max_memory)
     # The last device is left with max_memory just in case the buffer is not enough.
-    for i in range(len(max_memory) - 1):
+    for i in range(num_devices - 1):
Okay will add this suggestion 💯 
@@ -409,7 +409,7 @@ def get_balanced_memory(
 
     max_memory = get_max_memory(max_memory)
     # The last device is left with max_memory just in case the buffer is not enough.
-    for i in range(len(max_memory) - 1):
+    for i in range(num_devices - 1):
Oh and we should avoid a GPU that the user set at 0, so actually
```
last_gpu = max(i for i in max_memory if isintance(i, int) and max_memory[i] > 0)
```
@@ -409,7 +409,7 @@ def get_balanced_memory(
 
     max_memory = get_max_memory(max_memory)
     # The last device is left with max_memory just in case the buffer is not enough.
-    for i in range(len(max_memory) - 1):
+    for i in range(num_devices - 1):
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/777
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index c4a015156..cdae54af9 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -112,6 +112,9 @@ class BatchSamplerShard(BatchSampler):
               this argument is set to `False`.
             - the sampler on process 0 to yield `[0, 1]` then `[4, 5]` and the sampler on process 1 to yield `[2, 3]`
               then `[6, 7]` if this argument is set to `True`.
+        even_batches (`bool`, *optional*, defaults to `True`):
+            Whether or not to loop back at the beginning of the sampler when the number of samples is not a round
+            multiple of (original batch size / number of processes).
 
     <Tip warning={true}>
 
@@ -125,6 +128,7 @@ def __init__(
         num_processes: int = 1,
         process_index: int = 0,
         split_batches: bool = False,
+        even_batches: bool = True,
     ):
         if split_batches and batch_sampler.batch_size % num_processes != 0:
             raise ValueError(
@@ -135,8 +139,11 @@ def __init__(
         self.num_processes = num_processes
         self.process_index = process_index
         self.split_batches = split_batches
-        self.batch_size = batch_sampler.batch_size
-        self.drop_last = batch_sampler.drop_last
+        self.even_batches = even_batches
+        self.batch_size = getattr(batch_sampler, "batch_size", None)
+        self.drop_last = getattr(batch_sampler, "drop_last", False)
+        if self.batch_size is None and self.even_batches:
+            raise ValueError("You need to use `even_batches=False` when the batch sampler has no batch size.")
 
     @property
     def total_length(self):
@@ -144,11 +151,21 @@ def total_length(self):
 
     def __len__(self):
         if self.split_batches:
+            # Split batches does not change the length of the batch sampler
             return len(self.batch_sampler)
         if len(self.batch_sampler) % self.num_processes == 0:
+            # If the length is a round multiple of the number of processes, it's easy.
             return len(self.batch_sampler) // self.num_processes
         length = len(self.batch_sampler) // self.num_processes
-        return length if self.drop_last else length + 1
+        if self.drop_last:
+            # Same if we drop the remainder.
+            return length
+        elif self.even_batches:
+            # When we even batches we always get +1
+            return length + 1
+        else:
+            # Otherwise it depends on the process index.
+            return length + 1 if self.process_index < len(self.batch_sampler) % self.num_processes else length
 
     def __iter__(self):
         return self._iter_with_split() if self.split_batches else self._iter_with_no_split()
@@ -165,11 +182,15 @@ def _iter_with_split(self):
 
         # If drop_last is True of the last batch was full, iteration is over, otherwise...
         if not self.drop_last and len(initial_data) > 0 and len(batch) < self.batch_size:
-            # For degenerate cases where the dataset has less than num_process * batch_size samples
-            while len(initial_data) < self.batch_size:
-                initial_data += initial_data
-            batch = batch + initial_data
-            yield batch[batch_length * self.process_index : batch_length * (self.process_index + 1)]
+            if not self.even_batches:
+                if len(batch) > batch_length * self.process_index:
+                    yield batch[batch_length * self.process_index : batch_length * (self.process_index + 1)]
+            else:
+                # For degenerate cases where the dataset has less than num_process * batch_size samples
+                while len(initial_data) < self.batch_size:
+                    initial_data += initial_data
+                batch = batch + initial_data
+                yield batch[batch_length * self.process_index : batch_length * (self.process_index + 1)]
 
     def _iter_with_no_split(self):
         initial_data = []
@@ -182,35 +203,41 @@ def _iter_with_no_split(self):
             # yielding it.
             if idx % self.num_processes == self.process_index:
                 batch_to_yield = batch
-            if idx % self.num_processes == self.num_processes - 1 and len(batch) == self.batch_size:
+            if idx % self.num_processes == self.num_processes - 1 and (
+                self.batch_size is None or len(batch) == self.batch_size
+            ):
                 yield batch_to_yield
                 batch_to_yield = []
 
         # If drop_last is True, iteration is over, otherwise...
         if not self.drop_last and len(initial_data) > 0:
-            # ... we yield the complete batch we had saved before if it has the proper length
-            if len(batch_to_yield) == self.batch_size:
-                yield batch_to_yield
-
-            # For degenerate cases where the dataset has less than num_process * batch_size samples
-            while len(initial_data) < self.num_processes * self.batch_size:
-                initial_data += initial_data
-
-            # If the last batch seen was of the proper size, it has been yielded by its process so we move to the next
-            if len(batch) == self.batch_size:
-                batch = []
-                idx += 1
-
-            # Make sure we yield a multiple of self.num_processes batches
-            cycle_index = 0
-            while idx % self.num_processes != 0 or len(batch) > 0:
-                end_index = cycle_index + self.batch_size - len(batch)
-                batch += initial_data[cycle_index:end_index]
-                if idx % self.num_processes == self.process_index:
-                    yield batch
-                cycle_index = end_index
-                batch = []
-                idx += 1
+            if not self.even_batches:
+                if len(batch_to_yield) > 0:
+                    yield batch_to_yield
+            else:
+                # ... we yield the complete batch we had saved before if it has the proper length
+                if len(batch_to_yield) == self.batch_size:
+                    yield batch_to_yield
+
+                # For degenerate cases where the dataset has less than num_process * batch_size samples
+                while len(initial_data) < self.num_processes * self.batch_size:
+                    initial_data += initial_data
+
+                # If the last batch seen was of the proper size, it has been yielded by its process so we move to the next
+                if len(batch) == self.batch_size:
+                    batch = []
+                    idx += 1
+
+                # Make sure we yield a multiple of self.num_processes batches
+                cycle_index = 0
+                while idx % self.num_processes != 0 or len(batch) > 0:
+                    end_index = cycle_index + self.batch_size - len(batch)
+                    batch += initial_data[cycle_index:end_index]
+                    if idx % self.num_processes == self.process_index:
+                        yield batch
+                    cycle_index = end_index
+                    batch = []
+                    idx += 1
 
 
 class IterableDatasetShard(IterableDataset):
diff --git a/tests/test_data_loader.py b/tests/test_data_loader.py
index 2d3b6a394..a18dae732 100644
--- a/tests/test_data_loader.py
+++ b/tests/test_data_loader.py
@@ -36,8 +36,11 @@ def __iter__(self):
 
 
 class DataLoaderTester(unittest.TestCase):
-    def check_batch_sampler_shards(self, batch_sampler, expected, split_batches=False):
-        batch_sampler_shards = [BatchSamplerShard(batch_sampler, 2, i, split_batches) for i in range(2)]
+    def check_batch_sampler_shards(self, batch_sampler, expected, split_batches=False, even_batches=True):
+        batch_sampler_shards = [
+            BatchSamplerShard(batch_sampler, 2, i, split_batches=split_batches, even_batches=even_batches)
+            for i in range(2)
+        ]
         batch_sampler_lists = [list(batch_sampler_shard) for batch_sampler_shard in batch_sampler_shards]
         if not split_batches:
             self.assertListEqual([len(shard) for shard in batch_sampler_shards], [len(e) for e in expected])
@@ -164,6 +167,137 @@ def test_batch_sampler_shards_with_splits(self):
         expected = [[], []]
         self.check_batch_sampler_shards(batch_sampler, expected, split_batches=True)
 
+    def test_batch_sampler_shards_with_no_splits_no_even(self):
+        # Check the shards when the dataset is a round multiple of total batch size.
+        batch_sampler = BatchSampler(range(24), batch_size=3, drop_last=False)
+        expected = [
+            [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
+            [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 22, 23]],
+        ]
+        self.check_batch_sampler_shards(batch_sampler, expected, even_batches=False)
+
+        batch_sampler = BatchSampler(range(24), batch_size=3, drop_last=True)
+        # Expected shouldn't change
+        self.check_batch_sampler_shards(batch_sampler, expected, even_batches=False)
+
+        # Check the shards when the dataset is a round multiple of batch size but not total batch size.
+        batch_sampler = BatchSampler(range(21), batch_size=3, drop_last=False)
+        expected = [
+            [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
+            [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
+        ]
+        self.check_batch_sampler_shards(batch_sampler, expected, even_batches=False)
+
+        batch_sampler = BatchSampler(range(21), batch_size=3, drop_last=True)
+        expected = [
+            [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
+            [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
+        ]
+        self.check_batch_sampler_shards(batch_sampler, expected, even_batches=False)
+
+        # Check the shards when the dataset is not a round multiple of batch size but has a multiple of
+        # num_processes batch.
+        batch_sampler = BatchSampler(range(22), batch_size=3, drop_last=False)
+        expected = [
+            [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
+            [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21]],
+        ]
+        self.check_batch_sampler_shards(batch_sampler, expected, even_batches=False)
+
+        batch_sampler = BatchSampler(range(22), batch_size=3, drop_last=True)
+        expected = [
+            [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
+            [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
+        ]
+        self.check_batch_sampler_shards(batch_sampler, expected, even_batches=False)
+
+        # Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
+        # num_processes batch.
+        batch_sampler = BatchSampler(range(20), batch_size=3, drop_last=False)
+        expected = [
+            [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19]],
+            [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
+        ]
+        self.check_batch_sampler_shards(batch_sampler, expected, even_batches=False)
+
+        batch_sampler = BatchSampler(range(20), batch_size=3, drop_last=True)
+        expected = [
+            [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
+            [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
+        ]
+        self.check_batch_sampler_shards(batch_sampler, expected, even_batches=False)
+
+        # Check the shards when the dataset is very small.
+        batch_sampler = BatchSampler(range(2), batch_size=3, drop_last=False)
+        expected = [[[0, 1]], []]
+        self.check_batch_sampler_shards(batch_sampler, expected, even_batches=False)
+
+        batch_sampler = BatchSampler(range(2), batch_size=3, drop_last=True)
+        expected = [[], []]
+        self.check_batch_sampler_shards(batch_sampler, expected, even_batches=False)
+
+    def test_batch_sampler_shards_with_splits_no_even(self):
+        # Check the shards when the dataset is a round multiple of batch size.
+        batch_sampler = BatchSampler(range(24), batch_size=4, drop_last=False)
+        expected = [
+            [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
+            [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [22, 23]],
+        ]
+        self.check_batch_sampler_shards(batch_sampler, expected, split_batches=True, even_batches=False)
+
+        batch_sampler = BatchSampler(range(24), batch_size=4, drop_last=True)
+        # Expected shouldn't change
+        self.check_batch_sampler_shards(batch_sampler, expected, split_batches=True, even_batches=False)
+
+        # Check the shards when the dataset is not a round multiple of batch size.
+        batch_sampler = BatchSampler(range(22), batch_size=4, drop_last=False)
+        expected = [
+            [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
+            [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
+        ]
+        self.check_batch_sampler_shards(batch_sampler, expected, split_batches=True, even_batches=False)
+
+        batch_sampler = BatchSampler(range(22), batch_size=4, drop_last=True)
+        expected = [
+            [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
+            [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
+        ]
+        self.check_batch_sampler_shards(batch_sampler, expected, split_batches=True, even_batches=False)
+
+        # Check the shards when the dataset is not a round multiple of batch size or num_processes.
+        batch_sampler = BatchSampler(range(21), batch_size=4, drop_last=False)
+        expected = [
+            [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20]],
+            [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
+        ]
+        self.check_batch_sampler_shards(batch_sampler, expected, split_batches=True, even_batches=False)
+
+        batch_sampler = BatchSampler(range(21), batch_size=4, drop_last=True)
+        expected = [
+            [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
+            [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
+        ]
+        self.check_batch_sampler_shards(batch_sampler, expected, split_batches=True, even_batches=False)
+
+        # Check the shards when the dataset is very small.
+        batch_sampler = BatchSampler(range(2), batch_size=4, drop_last=False)
+        expected = [[[0, 1]], []]
+        self.check_batch_sampler_shards(batch_sampler, expected, split_batches=True, even_batches=False)
+
+        batch_sampler = BatchSampler(range(2), batch_size=4, drop_last=True)
+        expected = [[], []]
+        self.check_batch_sampler_shards(batch_sampler, expected, split_batches=True, even_batches=False)
+
+    def test_batch_sampler_with_varying_batch_size(self):
+        batch_sampler = [[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
+        batch_sampler_shards = [BatchSamplerShard(batch_sampler, 2, i, even_batches=False) for i in range(2)]
+
+        self.assertEqual(len(batch_sampler_shards[0]), 3)
+        self.assertEqual(len(batch_sampler_shards[1]), 2)
+
+        self.assertListEqual(list(batch_sampler_shards[0]), [[0, 1, 2], [5, 6, 7, 8], [12, 13]])
+        self.assertListEqual(list(batch_sampler_shards[1]), [[3, 4], [9, 10, 11]])
+
     def check_iterable_dataset_shards(
         self, dataset, seed, batch_size, drop_last=False, num_processes=2, split_batches=False
     ):

```suggestion
            Whether or not to loop back at the beginning of the sampler when the number of samples is not a round
            multiple of (original batch size / number of processes).
```
@@ -112,6 +112,9 @@ class BatchSamplerShard(BatchSampler):
               this argument is set to `False`.
             - the sampler on process 0 to yield `[0, 1]` then `[4, 5]` and the sampler on process 1 to yield `[2, 3]`
               then `[6, 7]` if this argument is set to `True`.
+        even_batches (`bool`, *optional*, defaults to `True`):
+            Whether or not to loop back at the beginning of the sampler when the number of samples is not a round
+            multiple of the batch size/number of processes.
`self.batch_size` can't be `None` else this logic would result in errors. 
@@ -182,35 +201,41 @@ def _iter_with_no_split(self):
             # yielding it.
             if idx % self.num_processes == self.process_index:
                 batch_to_yield = batch
-            if idx % self.num_processes == self.num_processes - 1 and len(batch) == self.batch_size:
+            if idx % self.num_processes == self.num_processes - 1 and (
+                self.batch_size is None or len(batch) == self.batch_size
+            ):
                 yield batch_to_yield
                 batch_to_yield = []
 
         # If drop_last is True, iteration is over, otherwise...
         if not self.drop_last and len(initial_data) > 0:
-            # ... we yield the complete batch we had saved before if it has the proper length
-            if len(batch_to_yield) == self.batch_size:
-                yield batch_to_yield
-
-            # For degenerate cases where the dataset has less than num_process * batch_size samples
-            while len(initial_data) < self.num_processes * self.batch_size:
-                initial_data += initial_data
-
-            # If the last batch seen was of the proper size, it has been yielded by its process so we move to the next
-            if len(batch) == self.batch_size:
-                batch = []
-                idx += 1
-
-            # Make sure we yield a multiple of self.num_processes batches
-            cycle_index = 0
-            while idx % self.num_processes != 0 or len(batch) > 0:
-                end_index = cycle_index + self.batch_size - len(batch)
-                batch += initial_data[cycle_index:end_index]
-                if idx % self.num_processes == self.process_index:
-                    yield batch
-                cycle_index = end_index
-                batch = []
-                idx += 1
+            if not self.even_batches:
+                if len(batch_to_yield) > 0:
+                    yield batch_to_yield
+            else:
+                # ... we yield the complete batch we had saved before if it has the proper length
+                if len(batch_to_yield) == self.batch_size:
+                    yield batch_to_yield
+
+                # For degenerate cases where the dataset has less than num_process * batch_size samples
+                while len(initial_data) < self.num_processes * self.batch_size:
+                    initial_data += initial_data
+
+                # If the last batch seen was of the proper size, it has been yielded by its process so we move to the next
+                if len(batch) == self.batch_size:
+                    batch = []
+                    idx += 1
+
+                # Make sure we yield a multiple of self.num_processes batches
+                cycle_index = 0
+                while idx % self.num_processes != 0 or len(batch) > 0:
+                    end_index = cycle_index + self.batch_size - len(batch)
+                    batch += initial_data[cycle_index:end_index]
+                    if idx % self.num_processes == self.process_index:
+                        yield batch
+                    cycle_index = end_index
+                    batch = []
+                    idx += 1
Yes it can't be `None` if `even_batches` is not `False`, will add an early check for that.
@@ -182,35 +201,41 @@ def _iter_with_no_split(self):
             # yielding it.
             if idx % self.num_processes == self.process_index:
                 batch_to_yield = batch
-            if idx % self.num_processes == self.num_processes - 1 and len(batch) == self.batch_size:
+            if idx % self.num_processes == self.num_processes - 1 and (
+                self.batch_size is None or len(batch) == self.batch_size
+            ):
                 yield batch_to_yield
                 batch_to_yield = []
 
         # If drop_last is True, iteration is over, otherwise...
         if not self.drop_last and len(initial_data) > 0:
-            # ... we yield the complete batch we had saved before if it has the proper length
-            if len(batch_to_yield) == self.batch_size:
-                yield batch_to_yield
-
-            # For degenerate cases where the dataset has less than num_process * batch_size samples
-            while len(initial_data) < self.num_processes * self.batch_size:
-                initial_data += initial_data
-
-            # If the last batch seen was of the proper size, it has been yielded by its process so we move to the next
-            if len(batch) == self.batch_size:
-                batch = []
-                idx += 1
-
-            # Make sure we yield a multiple of self.num_processes batches
-            cycle_index = 0
-            while idx % self.num_processes != 0 or len(batch) > 0:
-                end_index = cycle_index + self.batch_size - len(batch)
-                batch += initial_data[cycle_index:end_index]
-                if idx % self.num_processes == self.process_index:
-                    yield batch
-                cycle_index = end_index
-                batch = []
-                idx += 1
+            if not self.even_batches:
+                if len(batch_to_yield) > 0:
+                    yield batch_to_yield
+            else:
+                # ... we yield the complete batch we had saved before if it has the proper length
+                if len(batch_to_yield) == self.batch_size:
+                    yield batch_to_yield
+
+                # For degenerate cases where the dataset has less than num_process * batch_size samples
+                while len(initial_data) < self.num_processes * self.batch_size:
+                    initial_data += initial_data
+
+                # If the last batch seen was of the proper size, it has been yielded by its process so we move to the next
+                if len(batch) == self.batch_size:
+                    batch = []
+                    idx += 1
+
+                # Make sure we yield a multiple of self.num_processes batches
+                cycle_index = 0
+                while idx % self.num_processes != 0 or len(batch) > 0:
+                    end_index = cycle_index + self.batch_size - len(batch)
+                    batch += initial_data[cycle_index:end_index]
+                    if idx % self.num_processes == self.process_index:
+                        yield batch
+                    cycle_index = end_index
+                    batch = []
+                    idx += 1
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/776
diff --git a/tests/test_tracking.py b/tests/test_tracking.py
index 5e26eb9b5..d6ff4908a 100644
--- a/tests/test_tracking.py
+++ b/tests/test_tracking.py
@@ -17,6 +17,7 @@
 import logging
 import os
 import re
+import subprocess
 import tempfile
 import unittest
 import zipfile
@@ -85,59 +86,61 @@ def setUp(self):
         self.add_mocks(mock.patch.dict(os.environ, {"WANDB_DIR": self.tmpdir}))
 
     @staticmethod
-    def get_value_from_log(key: str, log: str, key_occurrence: int = 0):
+    def parse_log(log: str, section: str, record: bool = True):
         """
-        Parses wandb log for `key` and returns the value.
-        If parsing through multiple calls to .log, pass in a `key_occurrence`
+        Parses wandb log for `section` and returns a dictionary of
+        all items in that section. Section names are based on the
+        output of `wandb sync --view --verbose` and items starting
+        with "Record" in that result
         """
-        res = re.findall(rf"(?<={key} )[^\s]+", log)[key_occurrence]
-        if '"' in res:
-            return re.findall(r'"([^"]*)"', res)[0]
+        # Big thanks to the W&B team for helping us parse their logs
+        pattern = rf"{section} ([\S\s]*?)\n\n"
+        if record:
+            pattern = rf"Record: {pattern}"
+        cleaned_record = re.findall(pattern, log)[0]
+        # A config
+        if section == "config" or section == "history":
+            cleaned_record = re.findall(r'"([a-zA-Z0-9_.,]+)', cleaned_record)
+            return {key: val for key, val in zip(cleaned_record[0::2], cleaned_record[1::2])}
+        # Everything else
         else:
-            return res
+            return dict(re.findall(r'(\w+): "([^\s]+)"', cleaned_record))
 
-    def test_init_trackers(self):
+    def test_wandb(self):
         project_name = "test_project_with_config"
         accelerator = Accelerator(log_with="wandb")
         config = {"num_iterations": 12, "learning_rate": 1e-2, "some_boolean": False, "some_string": "some_value"}
         kwargs = {"wandb": {"tags": ["my_tag"]}}
         accelerator.init_trackers(project_name, config, kwargs)
-        accelerator.end_training()
-        # The latest offline log is stored at wandb/latest-run/*.wandb
-        for child in Path(f"{self.tmpdir}/wandb/latest-run").glob("*"):
-            logger.info(child)
-            if child.is_file() and child.suffix == ".wandb":
-                with open(child, "rb") as f:
-                    content = f.read()
-                break
-
-        # Check HPS through careful parsing and cleaning
-        cleaned_log = re.sub(r"[\x00-\x1f]+", " ", content.decode("utf8", "ignore"))
-        self.assertEqual(self.get_value_from_log("num_iterations", cleaned_log), "12")
-        self.assertEqual(self.get_value_from_log("learning_rate", cleaned_log), "0.01")
-        self.assertEqual(self.get_value_from_log("some_boolean", cleaned_log), "false")
-        self.assertEqual(self.get_value_from_log("some_string", cleaned_log), "some_value")
-        self.assertIn("my_tag", cleaned_log)
-
-    def test_log(self):
-        project_name = "test_project_with_log"
-        accelerator = Accelerator(log_with="wandb")
-        accelerator.init_trackers(project_name)
         values = {"total_loss": 0.1, "iteration": 1, "my_text": "some_value"}
         accelerator.log(values, step=0)
         accelerator.end_training()
         # The latest offline log is stored at wandb/latest-run/*.wandb
         for child in Path(f"{self.tmpdir}/wandb/latest-run").glob("*"):
             if child.is_file() and child.suffix == ".wandb":
-                with open(child, "rb") as f:
-                    content = f.read()
+                content = subprocess.check_output(
+                    ["wandb", "sync", "--view", "--verbose", str(child)], env=os.environ.copy()
+                ).decode("utf8", "ignore")
                 break
+
         # Check HPS through careful parsing and cleaning
-        cleaned_log = re.sub(r"[\x00-\x1f]+", " ", content.decode("utf8", "ignore"))
-        self.assertTrue("0.1" in self.get_value_from_log("total_loss", cleaned_log))
-        self.assertTrue("1" in self.get_value_from_log("iteration", cleaned_log))
-        self.assertTrue("some_value" in self.get_value_from_log("my_text", cleaned_log))
-        self.assertTrue("0" in self.get_value_from_log("_step", cleaned_log))
+        logged_items = self.parse_log(content, "config")
+        self.assertEqual(logged_items["num_iterations"], "12")
+        self.assertEqual(logged_items["learning_rate"], "0.01")
+        self.assertEqual(logged_items["some_boolean"], "false")
+        self.assertEqual(logged_items["some_string"], "some_value")
+        self.assertEqual(logged_items["some_string"], "some_value")
+
+        # Run tags
+        logged_items = self.parse_log(content, "run", False)
+        self.assertEqual(logged_items["tags"], "my_tag")
+
+        # Actual logging
+        logged_items = self.parse_log(content, "history")
+        self.assertEqual(logged_items["total_loss"], "0.1")
+        self.assertEqual(logged_items["iteration"], "1")
+        self.assertEqual(logged_items["my_text"], "some_value")
+        self.assertEqual(logged_items["_step"], "0")
 
 
 # Comet has a special `OfflineExperiment` we need to use for testing

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/775
diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
index 08ee3a65a..1760e72bc 100644
--- a/src/accelerate/utils/modeling.py
+++ b/src/accelerate/utils/modeling.py
@@ -365,7 +365,7 @@ def get_balanced_memory(
     if not torch.cuda.is_available():
         return max_memory
 
-    num_devices = len([d for d in max_memory if torch.device(d).type == "cuda"])
+    num_devices = len([d for d in max_memory if torch.device(d).type == "cuda" and max_memory[d] > 0])
     module_sizes = compute_module_sizes(model, dtype=dtype)
     per_gpu = module_sizes[""] // (num_devices - 1 if low_zero else num_devices)
 
@@ -409,7 +409,7 @@ def get_balanced_memory(
 
     max_memory = get_max_memory(max_memory)
     # The last device is left with max_memory just in case the buffer is not enough.
-    for i in range(num_devices - 1):
+    for i in range(len(max_memory) - 1):
         max_memory[i] = min(0 if low_zero and i == 0 else per_gpu, max_memory[i])
 
     if low_zero:
diff --git a/tests/test_modeling_utils.py b/tests/test_modeling_utils.py
index 1c6f6088d..02ab867d7 100644
--- a/tests/test_modeling_utils.py
+++ b/tests/test_modeling_utils.py
@@ -375,3 +375,7 @@ def test_get_balanced_memory(self):
         # Last device always get max memory to give more buffer and avoid accidental CPU offload
         max_memory = get_balanced_memory(model, max_memory={0: 300, 1: 500})
         self.assertDictEqual({0: 215, 1: 500}, max_memory)
+
+        # If we set a device to 0, it's not counted.
+        max_memory = get_balanced_memory(model, max_memory={0: 0, 1: 300, 2: 300})
+        self.assertDictEqual({0: 0, 1: 215, 2: 300}, max_memory)

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/774
diff --git a/.github/workflows/build-docker-images-release.yml b/.github/workflows/build-docker-images-release.yml
index 654259f87..be0c76c48 100644
--- a/.github/workflows/build-docker-images-release.yml
+++ b/.github/workflows/build-docker-images-release.yml
@@ -17,7 +17,7 @@ jobs:
     steps:
       - uses: actions/checkout@v3
       - id: step1
-        run: echo "::set-output name=version::$(python setup.py --version)"
+        run: echo "version=$(python setup.py --version)" >> $GITHUB_OUTPUT
 
   version-cpu:
     name: "Latest Accelerate CPU [version]"
diff --git a/.github/workflows/build_and_run_tests.yml b/.github/workflows/build_and_run_tests.yml
index abf9f3883..68e11c44a 100644
--- a/.github/workflows/build_and_run_tests.yml
+++ b/.github/workflows/build_and_run_tests.yml
@@ -16,7 +16,7 @@ jobs:
     outputs:
       changed: ${{ steps.was_changed.outputs.changed }}
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v3.1.0
         with: 
           fetch-depth: "2"
       
@@ -29,7 +29,7 @@ jobs:
         run: |
           for file in ${{ steps.changed-files.outputs.all_changed_files }}; do
             if [ `basename "${file}"` == "setup.py" ]; then
-              echo ::set-output name=changed::"1"
+              echo "changed=1" >> $GITHUB_OUTPUT
             fi
           done
           
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 2959316d8..691397b0d 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -38,7 +38,7 @@ jobs:
           test_rest
         ]
     steps:
-    - uses: actions/checkout@v3
+    - uses: actions/checkout@v3.1.0
     - name: Set up python 3.7
       uses: actions/setup-python@v3
       with:
@@ -50,7 +50,7 @@ jobs:
         path: |
           ${{ env.pythonLocation }}
           ${{ env.HF_HOME }}
-        key: ${{ env.pythonLocation }}-${{ matrix.test-kind }}-${{ hashFiles('setup.py') }}
+        key: ${{ env.pythonLocation }}-${{ matrix.pytorch-version }}-${{ matrix.test-kind }}-${{ hashFiles('setup.py') }}
     
     - name: Install the library
       run: |
@@ -62,6 +62,8 @@ jobs:
         pip install pytest-reportlog
     
     - name: Run Tests
+      env: 
+        PYTORCH_VERSION: ${{ matrix.pytorch-version }}
       run: |
         make ${{ matrix.test-kind }}
 
diff --git a/Makefile b/Makefile
index 558e98347..9aaecb076 100644
--- a/Makefile
+++ b/Makefile
@@ -25,43 +25,43 @@ style:
 	
 # Run tests for the library
 test:
-	python -m pytest -s -v ./tests/ --ignore=./tests/test_examples.py $(if $(IS_GITHUB_CI),--report-log 'all.log',)
+	python -m pytest -s -v ./tests/ --ignore=./tests/test_examples.py $(if $(IS_GITHUB_CI),--report-log "$(PYTORCH_VERSION)_all.log",)
 
 test_big_modeling:
-	python -m pytest -s -v ./tests/test_big_modeling.py $(if $(IS_GITHUB_CI),--report-log 'big_modeling.log',)
+	python -m pytest -s -v ./tests/test_big_modeling.py $(if $(IS_GITHUB_CI),--report-log "$(PYTORCH_VERSION)_big_modeling.log",)
 
 test_core:
 	python -m pytest -s -v ./tests/ --ignore=./tests/test_examples.py --ignore=./tests/deepspeed --ignore=./tests/test_big_modeling.py \
-	--ignore=./tests/fsdp --ignore=./tests/test_cli.py $(if $(IS_GITHUB_CI),--report-log 'core.log',)
+	--ignore=./tests/fsdp --ignore=./tests/test_cli.py $(if $(IS_GITHUB_CI),--report-log "$(PYTORCH_VERSION)_core.log",)
 
 test_cli:
-	python -m pytest -s -v ./tests/test_cli.py $(if $(IS_GITHUB_CI),--report-log 'cli.log',)
+	python -m pytest -s -v ./tests/test_cli.py $(if $(IS_GITHUB_CI),--report-log "$(PYTORCH_VERSION)_cli.log",)
 
 test_deepspeed:
-	python -m pytest -s -v ./tests/deepspeed $(if $(IS_GITHUB_CI),--report-log 'deepspeed.log',)
+	python -m pytest -s -v ./tests/deepspeed $(if $(IS_GITHUB_CI),--report-log "$(PYTORCH_VERSION)_deepspeed.log",)
 
 test_fsdp:
-	python -m pytest -s -v ./tests/fsdp $(if $(IS_GITHUB_CI),--report-log 'fsdp.log',)
+	python -m pytest -s -v ./tests/fsdp $(if $(IS_GITHUB_CI),--report-log "$(PYTORCH_VERSION)_fsdp.log",)
 
 test_examples:
-	python -m pytest -s -v ./tests/test_examples.py $(if $(IS_GITHUB_CI),--report-log 'examples.log',)
+	python -m pytest -s -v ./tests/test_examples.py $(if $(IS_GITHUB_CI),--report-log "$(PYTORCH_VERSION)_examples.log",)
 
 # Broken down example tests for the CI runners
 test_integrations:
-	python -m pytest -s -v ./tests/deepspeed ./tests/fsdp $(if $(IS_GITHUB_CI),--report-log 'integrations.log',)
+	python -m pytest -s -v ./tests/deepspeed ./tests/fsdp $(if $(IS_GITHUB_CI),--report-log "$(PYTORCH_VERSION)_integrations.log",)
 
 test_example_differences:
-	python -m pytest -s -v ./tests/test_examples.py::ExampleDifferenceTests $(if $(IS_GITHUB_CI),--report-log 'example_diff.log',)
+	python -m pytest -s -v ./tests/test_examples.py::ExampleDifferenceTests $(if $(IS_GITHUB_CI),--report-log "$(PYTORCH_VERSION)_example_diff.log",)
 
 test_checkpoint_epoch:
-	python -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k "by_epoch" $(if $(IS_GITHUB_CI),--report-log 'checkpoint_epoch.log',)
+	python -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k "by_epoch" $(if $(IS_GITHUB_CI),--report-log "$(PYTORCH_VERSION)_checkpoint_epoch.log",)
 
 test_checkpoint_step:
-	python -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k "by_step" $(if $(IS_GITHUB_CI),--report-log 'checkpoint_step.log',)
+	python -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k "by_step" $(if $(IS_GITHUB_CI),--report-log "$(PYTORCH_VERSION)_checkpoint_step.log",)
 
 # Same as test but used to install only the base dependencies
 test_prod:
 	$(MAKE) test_core
 
 test_rest:
-	python -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k "not by_step and not by_epoch" $(if $(IS_GITHUB_CI),--report-log 'rest.log',)
+	python -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k "not by_step and not by_epoch" $(if $(IS_GITHUB_CI),--report-log "$(PYTORCH_VERSION)_rest.log",)
diff --git a/utils/log_reports.py b/utils/log_reports.py
index f701f08c2..7f790dad3 100644
--- a/utils/log_reports.py
+++ b/utils/log_reports.py
@@ -1,6 +1,5 @@
 import json
-from pathlib import Path 
-import subprocess
+from pathlib import Path
 
 failed = []
 passed = []
@@ -19,16 +18,16 @@
                     duration = f'{line["duration"]:.4f}'
                     if line.get("outcome", "") == "failed":
                         section_num_failed += 1
-                        failed.append([test, duration])
+                        failed.append([test, duration, log.name.split('_')[0]])
                     else:
-                        passed.append([test, duration])
+                        passed.append([test, duration, log.name.split('_')[0]])
     group_info.append([str(log), section_num_failed])
 
 if len(failed) > 0:
     result = "## Failed Tests:\n"
-    failed_table = '| Test Location | Test Class | Test Name |\n|---|---|---|\n| '
+    failed_table = '| Test Location | Test Class | Test Name | PyTorch Version |\n|---|---|---|---|\n| '
     for test in failed:
         failed_table += ' | '.join(test[0].split("::"))
-    failed_table += " |"
+    failed_table += f" | {test[2]} |"
     result += failed_table
     print(result)
\ No newline at end of file

_The documentation is not available anymore as the PR was closed or merged._
Thank you @muellerzr for updating my knowledge!
https://api.github.com/repos/huggingface/accelerate/pulls/773
diff --git a/.github/workflows/run_merge_tests.yml b/.github/workflows/run_merge_tests.yml
index a794cd72b..fcbb62f53 100644
--- a/.github/workflows/run_merge_tests.yml
+++ b/.github/workflows/run_merge_tests.yml
@@ -29,6 +29,11 @@ jobs:
           pip install -e .[testing,test_trackers]
           pip install pytest-reportlog
 
+      - name: Run CLI tests
+        run: |
+          source activate accelerate
+          make test_cli
+          
       - name: Run test on GPUs
         run: |
           source activate accelerate
@@ -62,6 +67,11 @@ jobs:
           pip install -e .[testing,test_trackers]
           pip install pytest-reportlog
 
+      - name: Run CLI tests
+        run: |
+          source activate accelerate
+          make test_cli
+
       - name: Run test on GPUs
         run: |
           source activate accelerate
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 593eec741..2959316d8 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -28,6 +28,7 @@ jobs:
         test-kind: [
           test_prod,
           test_core,
+          test_cli,
           test_big_modeling,
           test_deepspeed,
           test_fsdp,
diff --git a/Makefile b/Makefile
index 52c594ae7..558e98347 100644
--- a/Makefile
+++ b/Makefile
@@ -32,7 +32,10 @@ test_big_modeling:
 
 test_core:
 	python -m pytest -s -v ./tests/ --ignore=./tests/test_examples.py --ignore=./tests/deepspeed --ignore=./tests/test_big_modeling.py \
-	--ignore=./tests/fsdp $(if $(IS_GITHUB_CI),--report-log 'core.log',)
+	--ignore=./tests/fsdp --ignore=./tests/test_cli.py $(if $(IS_GITHUB_CI),--report-log 'core.log',)
+
+test_cli:
+	python -m pytest -s -v ./tests/test_cli.py $(if $(IS_GITHUB_CI),--report-log 'cli.log',)
 
 test_deepspeed:
 	python -m pytest -s -v ./tests/deepspeed $(if $(IS_GITHUB_CI),--report-log 'deepspeed.log',)
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 65e9be9c7..fa7588b0a 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -432,7 +432,7 @@ def simple_launcher(args):
     current_env["USE_MPS_DEVICE"] = str(args.use_mps_device)
     if args.use_mps_device:
         current_env["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
-    elif args.gpu_ids != "all":
+    elif args.gpu_ids != "all" and args.gpu_ids is not None:
         current_env["CUDA_VISIBLE_DEVICES"] = args.gpu_ids
     if args.num_machines > 1:
         current_env["MASTER_ADDR"] = args.main_process_ip
@@ -489,7 +489,7 @@ def multi_gpu_launcher(args):
 
     current_env = os.environ.copy()
     gpu_ids = getattr(args, "gpu_ids")
-    if gpu_ids != "all":
+    if gpu_ids != "all" and args.gpu_ids is not None:
         current_env["CUDA_VISIBLE_DEVICES"] = gpu_ids
     mixed_precision = args.mixed_precision.lower()
     try:
@@ -637,7 +637,7 @@ def deepspeed_launcher(args):
 
     current_env = os.environ.copy()
     gpu_ids = getattr(args, "gpu_ids")
-    if gpu_ids != "all":
+    if gpu_ids != "all" and args.gpu_ids is not None:
         current_env["CUDA_VISIBLE_DEVICES"] = gpu_ids
     try:
         mixed_precision = PrecisionType(args.mixed_precision.lower())
@@ -947,7 +947,7 @@ def launch_command(args):
     else:
         if args.num_processes is None:
             args.num_processes = torch.cuda.device_count() if args.multi_gpu else 1
-            warned.append("\t`--num_processes` was set to a value of `{args.num_processes}`")
+            warned.append(f"\t`--num_processes` was set to a value of `{args.num_processes}`")
         if args.num_machines is None:
             warned.append("\t`--num_machines` was set to a value of `1`")
             args.num_machines = 1
diff --git a/src/accelerate/test_utils/scripts/test_cli.py b/src/accelerate/test_utils/scripts/test_cli.py
new file mode 100644
index 000000000..491410e5f
--- /dev/null
+++ b/src/accelerate/test_utils/scripts/test_cli.py
@@ -0,0 +1,13 @@
+import torch
+
+
+def main():
+    if torch.cuda.is_available():
+        num_gpus = torch.cuda.device_count()
+    else:
+        num_gpus = 0
+    print(f"Successfully ran on {num_gpus} GPUs")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/test_cli.py b/tests/test_cli.py
new file mode 100644
index 000000000..ceed58290
--- /dev/null
+++ b/tests/test_cli.py
@@ -0,0 +1,65 @@
+# 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 inspect
+import os
+import unittest
+from pathlib import Path
+
+import torch
+
+import accelerate
+from accelerate.test_utils import execute_subprocess_async
+
+
+class AccelerateLauncherTester(unittest.TestCase):
+    """
+    Test case for verifying the `accelerate launch` CLI operates correctly.
+    If a `default_config.yaml` file is located in the cache it will temporarily move it
+    for the duration of the tests.
+    """
+
+    mod_file = inspect.getfile(accelerate.test_utils)
+    test_file_path = os.path.sep.join(mod_file.split(os.path.sep)[:-1] + ["scripts", "test_cli.py"])
+
+    base_cmd = ["accelerate", "launch"]
+    config_folder = Path.home() / ".cache/huggingface/accelerate"
+    config_file = "default_config.yaml"
+    config_path = config_folder / config_file
+    changed_path = config_folder / "_default_config.yaml"
+
+    test_config_path = Path("tests/test_configs")
+
+    @classmethod
+    def setUpClass(cls):
+        if cls.config_path.is_file():
+            cls.config_path.rename(cls.changed_path)
+
+    @classmethod
+    def tearDownClass(cls):
+        if cls.changed_path.is_file():
+            cls.changed_path.rename(cls.config_path)
+
+    def test_no_config(self):
+        cmd = self.base_cmd
+        if torch.cuda.is_available() and (torch.cuda.device_count() > 1):
+            cmd += ["--multi_gpu"]
+        execute_subprocess_async(cmd + [self.test_file_path], env=os.environ.copy())
+
+    def test_config_compatibility(self):
+        for config in sorted(self.test_config_path.glob("**/*.yaml")):
+            with self.subTest(config_file=config):
+                execute_subprocess_async(
+                    self.base_cmd + ["--config_file", str(config), self.test_file_path], env=os.environ.copy()
+                )
diff --git a/tests/test_configs/0_11_0.yaml b/tests/test_configs/0_11_0.yaml
new file mode 100644
index 000000000..9ef829e6b
--- /dev/null
+++ b/tests/test_configs/0_11_0.yaml
@@ -0,0 +1,12 @@
+compute_environment: LOCAL_MACHINE
+deepspeed_config: {}
+distributed_type: 'NO'
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: 'no'
+num_machines: 1
+num_processes: 1
+use_cpu: false
\ No newline at end of file
diff --git a/tests/test_configs/0_12_0.yaml b/tests/test_configs/0_12_0.yaml
new file mode 100644
index 000000000..00d06aab9
--- /dev/null
+++ b/tests/test_configs/0_12_0.yaml
@@ -0,0 +1,13 @@
+compute_environment: LOCAL_MACHINE
+deepspeed_config: {}
+distributed_type: 'NO'
+downcast_bf16: 'no'
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: 'no'
+num_machines: 1
+num_processes: 1
+use_cpu: false
\ No newline at end of file
diff --git a/tests/test_configs/README.md b/tests/test_configs/README.md
new file mode 100644
index 000000000..fd88d066c
--- /dev/null
+++ b/tests/test_configs/README.md
@@ -0,0 +1,2 @@
+This folder contains test configs for `accelerate config`. These should be generated for each major version
+and are written based on `accelerate config` and selecting the "No distributed training" option.
\ No newline at end of file
diff --git a/tests/test_configs/latest.yaml b/tests/test_configs/latest.yaml
new file mode 100644
index 000000000..87b294cf8
--- /dev/null
+++ b/tests/test_configs/latest.yaml
@@ -0,0 +1,17 @@
+compute_environment: LOCAL_MACHINE
+deepspeed_config: {}
+distributed_type: 'NO'
+downcast_bf16: 'no'
+fsdp_config: {}
+gpu_ids: all
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+megatron_lm_config: {}
+mixed_precision: 'no'
+num_machines: 1
+num_processes: 1
+rdzv_backend: static
+same_network: true
+use_cpu: false

This addition is needed for the test case of `CUDA_VISIBLE_DEVICES=""`
@@ -489,7 +489,7 @@ def multi_gpu_launcher(args):
 
     current_env = os.environ.copy()
     gpu_ids = getattr(args, "gpu_ids")
-    if gpu_ids != "all":
+    if gpu_ids != "all" and args.gpu_ids is not None:
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/772
diff --git a/src/accelerate/big_modeling.py b/src/accelerate/big_modeling.py
index 9d57c144f..83c6738af 100644
--- a/src/accelerate/big_modeling.py
+++ b/src/accelerate/big_modeling.py
@@ -143,6 +143,8 @@ def cpu_offload(
         execution_device = next(iter(model.parameters())).device
     if state_dict is None:
         state_dict = {n: p.to("cpu") for n, p in model.state_dict().items()}
+
+    add_hook_to_module(model, AlignDevicesHook(io_same_device=True), append=True)
     attach_align_device_hook(
         model,
         execution_device=execution_device,
@@ -151,7 +153,7 @@ def cpu_offload(
         weights_map=state_dict,
         preload_module_classes=preload_module_classes,
     )
-    add_hook_to_module(model, AlignDevicesHook(io_same_device=True))
+
     return model
 
 
@@ -189,6 +191,8 @@ def disk_offload(
     if execution_device is None:
         execution_device = next(iter(model.parameters())).device
     weights_map = OffloadedWeightsLoader(save_folder=offload_dir)
+
+    add_hook_to_module(model, AlignDevicesHook(io_same_device=True), append=True)
     attach_align_device_hook(
         model,
         execution_device=execution_device,
@@ -197,7 +201,7 @@ def disk_offload(
         weights_map=weights_map,
         preload_module_classes=preload_module_classes,
     )
-    add_hook_to_module(model, AlignDevicesHook(io_same_device=True))
+
     return model
 
 
diff --git a/src/accelerate/hooks.py b/src/accelerate/hooks.py
index 493444daa..120bb7f42 100644
--- a/src/accelerate/hooks.py
+++ b/src/accelerate/hooks.py
@@ -108,26 +108,34 @@ def detach_hook(self, module):
         return module
 
 
-def add_hook_to_module(module: nn.Module, hook: ModelHook):
+def add_hook_to_module(module: nn.Module, hook: ModelHook, append: bool = False):
     """
     Adds a hook to a given module. This will rewrite the `forward` method of the module to include the hook, to remove
     this behavior and restore the original `forward` method, use `remove_hook_from_module`.
 
     <Tip warning={true}>
 
-    If the module already contains a hook, this will replace it with the new hook passed. To chain two hooks together,
-    use the `SequentialHook` class.
+    If the module already contains a hook, this will replace it with the new hook passed by default. To chain two hooks
+    together, pass `append=True`, so it chains the current and new hook into an instance of the `SequentialHook` class.
 
     </Tip>
 
     Args:
         module (`torch.nn.Module`): The module to attach a hook to.
         hook (`ModelHook`): The hook to attach.
+        append (`bool`, *optional*, defaults to `False`):
+            Whether the hook should be chained with an existing one (if module already contains a hook) or not.
 
     Returns:
         `torch.nn.Module`: The same module, with the hook attached (the module is modified in place, so the result can
         be discarded).
     """
+
+    if append and (getattr(module, "_hf_hook", None) is not None):
+        old_hook = module._hf_hook
+        remove_hook_from_module(module)
+        hook = SequentialHook(old_hook, hook)
+
     if hasattr(module, "_hf_hook") and hasattr(module, "_old_forward"):
         # If we already put some hook on this module, we replace it with the new one.
         old_forward = module._old_forward
@@ -352,7 +360,7 @@ def attach_align_device_hook(
             offload_buffers=offload_buffers,
             place_submodules=full_offload,
         )
-        add_hook_to_module(module, hook)
+        add_hook_to_module(module, hook, append=True)
 
     # We stop the recursion in case we hit the full offload.
     if full_offload:
diff --git a/tests/test_hooks.py b/tests/test_hooks.py
index 9d48db9e1..e3824809c 100644
--- a/tests/test_hooks.py
+++ b/tests/test_hooks.py
@@ -69,6 +69,25 @@ def test_add_and_remove_hooks(self):
         self.assertFalse(hasattr(test_model, "_hf_hook"))
         self.assertFalse(hasattr(test_model, "_old_forward"))
 
+    def test_append_and_remove_hooks(self):
+        test_model = ModelForTest()
+        test_hook = ModelHook()
+
+        add_hook_to_module(test_model, test_hook)
+        add_hook_to_module(test_model, test_hook, append=True)
+
+        self.assertEqual(isinstance(test_model._hf_hook, SequentialHook), True)
+        self.assertEqual(len(test_model._hf_hook.hooks), 2)
+        self.assertTrue(hasattr(test_model, "_old_forward"))
+
+        # Check adding the hook did not change the name or the signature
+        self.assertEqual(test_model.forward.__name__, "forward")
+        self.assertListEqual(list(inspect.signature(test_model.forward).parameters), ["x"])
+
+        remove_hook_from_module(test_model)
+        self.assertFalse(hasattr(test_model, "_hf_hook"))
+        self.assertFalse(hasattr(test_model, "_old_forward"))
+
     def test_pre_forward_hook_is_executed(self):
         test_model = ModelForTest()
         x = torch.randn(2, 3)

```suggestion
        append (`bool`, *optional*, defaults to `False`):
            Whether the hook should be chained with an existing one (if module already contains a hook) or not.
```

Also can't comment on line 119 but this should also be mentioned in the Note.
@@ -123,11 +123,21 @@ def add_hook_to_module(module: nn.Module, hook: ModelHook):
     Args:
         module (`torch.nn.Module`): The module to attach a hook to.
         hook (`ModelHook`): The hook to attach.
+        append (`bool`): Whether, if module already contains a hook, should chain the new one with SequentialHook
Those two ifs can be grouped.
Also the second test can be simplified in `getattr(module, "_hf_hook", None) is not None`
@@ -123,11 +123,21 @@ def add_hook_to_module(module: nn.Module, hook: ModelHook):
     Args:
         module (`torch.nn.Module`): The module to attach a hook to.
         hook (`ModelHook`): The hook to attach.
+        append (`bool`): Whether, if module already contains a hook, should chain the new one with SequentialHook
 
     Returns:
         `torch.nn.Module`: The same module, with the hook attached (the module is modified in place, so the result can
         be discarded).
     """
+
+    if append:
+        if hasattr(module, "_hf_hook") and (module._hf_hook is not None):
```suggestion
            hook = SequentialHook(old_hook, hook)
```
@@ -123,11 +123,21 @@ def add_hook_to_module(module: nn.Module, hook: ModelHook):
     Args:
         module (`torch.nn.Module`): The module to attach a hook to.
         hook (`ModelHook`): The hook to attach.
+        append (`bool`): Whether, if module already contains a hook, should chain the new one with SequentialHook
 
     Returns:
         `torch.nn.Module`: The same module, with the hook attached (the module is modified in place, so the result can
         be discarded).
     """
+
+    if append:
+        if hasattr(module, "_hf_hook") and (module._hf_hook is not None):
+            old_hook = module._hf_hook
+            remove_hook_from_module(module)
+
+            hooks = [old_hook, hook]
+            hook = SequentialHook(*hooks)
Done.
@@ -123,11 +123,21 @@ def add_hook_to_module(module: nn.Module, hook: ModelHook):
     Args:
         module (`torch.nn.Module`): The module to attach a hook to.
         hook (`ModelHook`): The hook to attach.
+        append (`bool`): Whether, if module already contains a hook, should chain the new one with SequentialHook
Done, thanks for the great idea. 
@@ -123,11 +123,21 @@ def add_hook_to_module(module: nn.Module, hook: ModelHook):
     Args:
         module (`torch.nn.Module`): The module to attach a hook to.
         hook (`ModelHook`): The hook to attach.
+        append (`bool`): Whether, if module already contains a hook, should chain the new one with SequentialHook
 
     Returns:
         `torch.nn.Module`: The same module, with the hook attached (the module is modified in place, so the result can
         be discarded).
     """
+
+    if append:
+        if hasattr(module, "_hf_hook") and (module._hf_hook is not None):
Done. Looks much better now.
@@ -123,11 +123,21 @@ def add_hook_to_module(module: nn.Module, hook: ModelHook):
     Args:
         module (`torch.nn.Module`): The module to attach a hook to.
         hook (`ModelHook`): The hook to attach.
+        append (`bool`): Whether, if module already contains a hook, should chain the new one with SequentialHook
 
     Returns:
         `torch.nn.Module`: The same module, with the hook attached (the module is modified in place, so the result can
         be discarded).
     """
+
+    if append:
+        if hasattr(module, "_hf_hook") and (module._hf_hook is not None):
+            old_hook = module._hf_hook
+            remove_hook_from_module(module)
+
+            hooks = [old_hook, hook]
+            hook = SequentialHook(*hooks)
_The documentation is not available anymore as the PR was closed or merged._
The fix is not exactly right: by doing so, the hook that ensures the input and output of the model are on the same device is now erased. In your code sample in #767, since `x` is on the CPU, `net(x)` should also be on the CPU. This is not the case with your PR. The solution would be to write a util function that will:
- just add the hook if none is present
- extract the current hook if one is present and chain it with this hook using a `SequentialHook`.

This is slightly more advanced than the current PR, so let me know if you'd prefer for me to do it :-)
@sgugger I would like to try if that's ok to you. 

What do you think of creating an `append_if_needed` flag on `add_hook_to_module` that does what you just said?
That works for me, though the name of the argument could simply be `append` :-)
Thanks for diving into this!
@sgugger `append` it is.
@sgugger it is ready for review, I've also added the tests. 
> Very nice, thanks! Left a couple of nits, and I think you should still put the hook with io first: just tested locally and we still have the same issue of `net(x)` being on the wrong device since it runs second and the input was already moved.

@sgugger I've just addressed your nits and moved the io hook to the top. Thanks for the review! I've tested it locally on the snippet of the bug report and it brings the tensor back to CPU after the inference. Thanks!
@sgugger, there is a test step that failed due to an http error when installing a lib. I've created and empty commit to try running it again. 
https://api.github.com/repos/huggingface/accelerate/pulls/768
diff --git a/docs/source/usage_guides/megatron_lm.mdx b/docs/source/usage_guides/megatron_lm.mdx
index bc0b90701..d26ecd4b5 100644
--- a/docs/source/usage_guides/megatron_lm.mdx
+++ b/docs/source/usage_guides/megatron_lm.mdx
@@ -103,65 +103,9 @@ cd ..
 
 4. Installing Megatron-LM
 
-    a. Cloning the Megatron-LM repo
-    ```
-    git clone https://github.com/NVIDIA/Megatron-LM.git
-    cd Megatron-LM
-    ```
-
-    b. Create a file `setup.py`, paste the below code and put in the root folder
-    ```python
-    """Setup for pip package."""
-
-    import os
-    import sys
-    import setuptools
-
-    if sys.version_info < (3,):
-        raise Exception("Python 2 is not supported by Megatron.")
-
-    with open("README.md", "r") as fh:
-        long_description = fh.read()
-
-    setuptools.setup(
-        name="megatron-lm",
-        version="3.0.0",
-        description="Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism.",
-        long_description=long_description,
-        long_description_content_type="text/markdown",
-        # The project's main homepage.
-        url="https://github.com/NVIDIA/Megatron-LM",
-        author="NVIDIA INC",
-        maintainer="NVIDIA INC",
-        # The licence under which the project is released
-        license="See https://github.com/NVIDIA/Megatron-LM/blob/master/LICENSE",
-        classifiers=[
-            "Intended Audience :: Developers",
-            "Intended Audience :: Science/Research",
-            "Intended Audience :: Information Technology",
-            # Indicate what your project relates to
-            "Topic :: Scientific/Engineering :: Artificial Intelligence",
-            "Topic :: Software Development :: Libraries :: Python Modules",
-            # Additional Setting
-            "Environment :: Console",
-            "Natural Language :: English",
-            "Operating System :: OS Independent",
-        ],
-        python_requires=">=3.6",
-        packages=setuptools.find_packages(),
-        install_requires=["nltk", "six", "regex", "torch>=1.12.0", "pybind11"],
-        # Add in any packaged data.
-        include_package_data=True,
-        zip_safe=False,
-        # PyPI package information.
-        keywords="deep learning, Megatron, gpu, NLP, nvidia, pytorch, torch, language",
-    )
-    ```
-
-    c. installing via below command
-    ```
-    pip install "."
-    ```
+```
+pip install git+https://github.com/huggingface/Megatron-LM.git
+```
 
 ## Accelerate Megatron-LM Plugin
 
@@ -183,7 +127,7 @@ What is the number of micro-batches? [1]:2
 Do you want to enable selective activation recomputation? [YES/no]: 
 Do you want to use distributed optimizer which shards optimizer state and gradients across data pralellel ranks? [YES/no]: 
 What is the gradient clipping value based on global L2 Norm (0 to disable)? [1.0]: 
-How many GPU(s) should be used for distributed training? [1]:8
+How many GPU(s) should be used for distributed training? [1]:4
 Do you wish to use FP16 or BF16 (mixed precision)? [NO/fp16/bf16]: bf16
 ```
 
@@ -210,7 +154,7 @@ megatron_lm_config:
   megatron_lm_use_distributed_optimizer: true
 mixed_precision: bf16
 num_machines: 1
-num_processes: 8
+num_processes: 4
 rdzv_backend: static
 same_network: true
 use_cpu: false
@@ -289,8 +233,8 @@ examples/by_feature/megatron_lm_gpt_pretraining.py \
 --dataset_config_name wikitext-2-raw-v1 \
 --block_size 1024 \
 --learning_rate 5e-5 \
---per_device_train_batch_size 4 \
---per_device_eval_batch_size 4 \
+--per_device_train_batch_size 24 \
+--per_device_eval_batch_size 24 \
 --num_train_epochs 5 \
 --with_tracking \
 --report_to "wandb" \
diff --git a/examples/by_feature/megatron_lm_gpt_pretraining.py b/examples/by_feature/megatron_lm_gpt_pretraining.py
index 68a425e4d..3fcf10207 100644
--- a/examples/by_feature/megatron_lm_gpt_pretraining.py
+++ b/examples/by_feature/megatron_lm_gpt_pretraining.py
@@ -672,8 +672,10 @@ def group_texts(examples):
                 output_dir = os.path.join(args.output_dir, output_dir)
             accelerator.save_state(output_dir)
 
-    if args.with_tracking:
-        accelerator.end_training()
+    # this is causing some issue with Megatron-LM when using `wandb` at the end of the main function.
+    # Everything works fine inspite of commenting this out. (wandb finishes/closes the run without error)
+    # if args.with_tracking:
+    #     accelerator.end_training()
 
     if args.output_dir is not None:
         accelerator.wait_for_everyone()

Maybe add some comment here on why it's not working for now and plans to fix in the future?
@@ -672,8 +672,8 @@ def group_texts(examples):
                 output_dir = os.path.join(args.output_dir, output_dir)
             accelerator.save_state(output_dir)
 
-    if args.with_tracking:
-        accelerator.end_training()
+    # if args.with_tracking:
+    #     accelerator.end_training()
Done.
@@ -672,8 +672,8 @@ def group_texts(examples):
                 output_dir = os.path.join(args.output_dir, output_dir)
             accelerator.save_state(output_dir)
 
-    if args.with_tracking:
-        accelerator.end_training()
+    # if args.with_tracking:
+    #     accelerator.end_training()
_The documentation is not available anymore as the PR was closed or merged._
Hello @sgugger, wrt maintaining fork, sure, I can do that. This would help in making sure the functionality works with accelerate and we can extend it for more features (main repo is barely active).
https://api.github.com/repos/huggingface/accelerate/pulls/766
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index e1282c0bb..bf17f57ce 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -777,10 +777,11 @@ def prepare_model(self, model: torch.nn.Module, device_placement=None):
         if device_placement:
             model = model.to(self.device)
         if self.distributed_type == DistributedType.MULTI_GPU:
-            kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {}
-            model = torch.nn.parallel.DistributedDataParallel(
-                model, device_ids=[self.local_process_index], output_device=self.local_process_index, **kwargs
-            )
+            if any(p.requires_grad for p in model.parameters()):
+                kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {}
+                model = torch.nn.parallel.DistributedDataParallel(
+                    model, device_ids=[self.local_process_index], output_device=self.local_process_index, **kwargs
+                )
         elif self.distributed_type == DistributedType.FSDP:
             from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP
 

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/761
diff --git a/README.md b/README.md
index deddf2986..7b912dd80 100644
--- a/README.md
+++ b/README.md
@@ -243,6 +243,7 @@ pip install accelerate
 - FP16 with native AMP (apex on the roadmap)
 - DeepSpeed support (Experimental)
 - PyTorch Fully Sharded Data Parallel (FSDP) support (Experimental)
+- Megatron-LM support (Experimental)
 
 ## Citing 🤗 Accelerate
 
diff --git a/docs/source/usage_guides/megatron_lm.mdx b/docs/source/usage_guides/megatron_lm.mdx
index 188dddb32..bc0b90701 100644
--- a/docs/source/usage_guides/megatron_lm.mdx
+++ b/docs/source/usage_guides/megatron_lm.mdx
@@ -23,18 +23,45 @@ For detailed information and how things work behind the scene please refer the g
 Accelerate integrates following feature of Megatron-LM to enable large scale pre-training/finetuning
 of BERT (Encoder), GPT (Decoder) or T5 models (Encoder and Decoder):
 
-a. **Tensor Parallelism (TP)**: Reduces memory footprint without much additional communication on intra-node ranks
+a. **Tensor Parallelism (TP)**: Reduces memory footprint without much additional communication on intra-node ranks.
+Each tensor is split into multiple chunks with each shard residing on separate GPU. At each step, the same mini-batch of data is processed
+independently and in parallel by each shard followed by syncing across all GPUs (`all-reduce` operation). 
+In a simple transformer layer, this leads to 2 `all-reduces` in the forward path and 2 in the backward path.
+For more details, please refer research paper [Megatron-LM: Training Multi-Billion Parameter Language Models Using
+Model Parallelism](https://arxiv.org/pdf/1909.08053.pdf) and 
+this section of 🤗 blogpost [The Technology Behind BLOOM Training](https://huggingface.co/blog/bloom-megatron-deepspeed#tensor-parallelism).
+
 
 b. **Pipeline Parallelism (PP)**: Reduces memory footprint and enables large scale training via inter-node parallelization. 
 Reduces the bubble of naive PP via PipeDream-Flush schedule/1F1B schedule and Interleaved 1F1B schedule. 
-Layers are distributed uniformly across PP stages.
+Layers are distributed uniformly across PP stages. For example, if a model has `24` layers and we have `4` GPUs for
+pipeline parallelism, each GPU will have `6` layers (24/4). For more details on schedules to reduce the idle time of PP,
+please refer to the research paper [Efficient Large-Scale Language Model Training on GPU Clusters
+Using Megatron-LM](https://arxiv.org/pdf/2104.04473.pdf) and 
+this section of 🤗 blogpost [The Technology Behind BLOOM Training](https://huggingface.co/blog/bloom-megatron-deepspeed#pipeline-parallelism).
 
 c. **Sequence Parallelism (SP)**: Reduces memory footprint without any additional communication. Only applicable when using TP.
-
-d. **Data Parallelism (DP)** via Distributed Optimizer: Reduces memory footprint by sharding optimizer states and gradients across DP ranks.
+It reduces activation memory required as it prevents the same copies to be on the tensor parallel ranks 
+post `all-reduce` by replacing then with `reduce-scatter` and `no-op` operation would be replaced by `all-gather`. 
+As `all-reduce = reduce-scatter + all-gather`, this saves a ton of activation memory at no added communication cost. 
+To put it simply, it shards the outputs of each transformer layer along sequence dimension, e.g., 
+if the sequence length is `1024` and the TP size is `4`, each GPU will have `256` tokens (1024/4) for each sample. 
+This increases the batch size that can be supported for training. For more details, please refer to the research paper
+[Reducing Activation Recomputation in Large Transformer Models](https://arxiv.org/pdf/2205.05198.pdf). 
+
+d. **Data Parallelism (DP)** via Distributed Optimizer: Reduces the memory footprint by sharding optimizer states and gradients across DP ranks
+(versus the traditional method of replicating the optimizer state across data parallel ranks). 
+For example, when using Adam optimizer with mixed-precision training, each parameter accounts for 12 bytes of memory.
+This gets distributed equally across the GPUs, i.e., each parameter would account for 3 bytes (12/4) if we have 4 GPUs.
+For more details, please refer the research paper [ZeRO: Memory Optimizations Toward Training Trillion
+Parameter Models](https://arxiv.org/pdf/1910.02054.pdf) and following section of 🤗 blog 
+[The Technology Behind BLOOM Training](https://huggingface.co/blog/bloom-megatron-deepspeed#zero-data-parallelism).
 
 e. **Selective Activation Recomputation**: Reduces the memory footprint of activations significantly via smart activation checkpointing.
 It doesn't store activations occupying large memory while being fast to recompute thereby achieving great tradeoff between memory and recomputation.
+For example, for GPT-3, this leads to 70% reduction in required memory for activations at the expense of
+only 2.7% FLOPs overhead for recomputation of activations. For more details, please refer to the research paper 
+[Reducing Activation Recomputation in Large Transformer Models](https://arxiv.org/pdf/2205.05198.pdf).
 
 f. **Fused Kernels**: Fused Softmax, Mixed Precision Fused Layer Norm and  Fused gradient accumulation to weight gradient computation of linear layer.
 PyTorch JIT compiled Fused GeLU and Fused Bias+Dropout+Residual addition.
diff --git a/docs/source/usage_guides/training_zoo.mdx b/docs/source/usage_guides/training_zoo.mdx
index cc388acac..444784b49 100644
--- a/docs/source/usage_guides/training_zoo.mdx
+++ b/docs/source/usage_guides/training_zoo.mdx
@@ -36,6 +36,7 @@ These examples showcase specific features that the Accelerate framework offers
 - [Memory-aware batch size finder](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/memory.py)
 - [Metric Computation](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/multi_process_metrics.py)
 - [Using Trackers](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/tracking.py)
+- [Using Megatron-LM](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/megatron_lm_gpt_pretraining.py)
 
 ### Full Examples 
 
@@ -55,6 +56,8 @@ These examples showcase every feature in Accelerate at once that was shown in "F
 - [Named entity recognition fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/token-classification/run_ner_no_trainer.py)
 - [Image classification fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/image-classification/run_image_classification_no_trainer.py)
 - [Summarization fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/summarization/run_summarization_no_trainer.py)
+- [End-to-end examples on how to use AWS SageMaker integration of Accelerate](https://github.com/huggingface/notebooks/blob/main/sagemaker/22_accelerate_sagemaker_examples/README.md)
+- [Megatron-LM examples for various NLp tasks](https://github.com/pacman100/accelerate-megatron-test) 
 
 ## Integration Examples 
 

```suggestion
Each tensor is split into multiple chunks with each shard residing on separate GPU. At each step, the same mini-batch of data is processed
```
@@ -23,18 +23,45 @@ For detailed information and how things work behind the scene please refer the g
 Accelerate integrates following feature of Megatron-LM to enable large scale pre-training/finetuning
 of BERT (Encoder), GPT (Decoder) or T5 models (Encoder and Decoder):
 
-a. **Tensor Parallelism (TP)**: Reduces memory footprint without much additional communication on intra-node ranks
+a. **Tensor Parallelism (TP)**: Reduces memory footprint without much additional communication on intra-node ranks.
+Each tensor are split into multiple chunks with each shard residing on separate GPU. At each step, same minimbatch of data is processed
```suggestion
It reduces activation memory required as it prevents the same copies to be on the tensor parallel ranks 
```
@@ -23,18 +23,45 @@ For detailed information and how things work behind the scene please refer the g
 Accelerate integrates following feature of Megatron-LM to enable large scale pre-training/finetuning
 of BERT (Encoder), GPT (Decoder) or T5 models (Encoder and Decoder):
 
-a. **Tensor Parallelism (TP)**: Reduces memory footprint without much additional communication on intra-node ranks
+a. **Tensor Parallelism (TP)**: Reduces memory footprint without much additional communication on intra-node ranks.
+Each tensor are split into multiple chunks with each shard residing on separate GPU. At each step, same minimbatch of data is processed
+independently and in parallel by each shard followed by syncing across all GPUs (`all-reduce` operation). 
+In a simple transformer layer, this leads to 2 `all-reduces` in the forward path and 2 in the backward path.
+For more details, please refer research paper [Megatron-LM: Training Multi-Billion Parameter Language Models Using
+Model Parallelism](https://arxiv.org/pdf/1909.08053.pdf) and 
+this section of 🤗 blogpost [The Technology Behind BLOOM Training](https://huggingface.co/blog/bloom-megatron-deepspeed#tensor-parallelism).
+
 
 b. **Pipeline Parallelism (PP)**: Reduces memory footprint and enables large scale training via inter-node parallelization. 
 Reduces the bubble of naive PP via PipeDream-Flush schedule/1F1B schedule and Interleaved 1F1B schedule. 
-Layers are distributed uniformly across PP stages.
+Layers are distributed uniformly across PP stages. For example, if a model has `24` layers and we have `4` GPUs for
+pipeline parallelism, each GPU will have `6` layers (24/4). For more details on schedules to reduce the idle time of PP,
+please refer to the research paper [Efficient Large-Scale Language Model Training on GPU Clusters
+Using Megatron-LM](https://arxiv.org/pdf/2104.04473.pdf) and 
+this section of 🤗 blogpost [The Technology Behind BLOOM Training](https://huggingface.co/blog/bloom-megatron-deepspeed#pipeline-parallelism).
 
 c. **Sequence Parallelism (SP)**: Reduces memory footprint without any additional communication. Only applicable when using TP.
-
-d. **Data Parallelism (DP)** via Distributed Optimizer: Reduces memory footprint by sharding optimizer states and gradients across DP ranks.
+It reduced activation memory required as it prevents same copies to be on the tensor parallel ranks 
```suggestion
post `all-reduce` by replacing then with `reduce-scatter` and `no-op` operation would be replaced by `all-gather`. 
```
@@ -23,18 +23,45 @@ For detailed information and how things work behind the scene please refer the g
 Accelerate integrates following feature of Megatron-LM to enable large scale pre-training/finetuning
 of BERT (Encoder), GPT (Decoder) or T5 models (Encoder and Decoder):
 
-a. **Tensor Parallelism (TP)**: Reduces memory footprint without much additional communication on intra-node ranks
+a. **Tensor Parallelism (TP)**: Reduces memory footprint without much additional communication on intra-node ranks.
+Each tensor are split into multiple chunks with each shard residing on separate GPU. At each step, same minimbatch of data is processed
+independently and in parallel by each shard followed by syncing across all GPUs (`all-reduce` operation). 
+In a simple transformer layer, this leads to 2 `all-reduces` in the forward path and 2 in the backward path.
+For more details, please refer research paper [Megatron-LM: Training Multi-Billion Parameter Language Models Using
+Model Parallelism](https://arxiv.org/pdf/1909.08053.pdf) and 
+this section of 🤗 blogpost [The Technology Behind BLOOM Training](https://huggingface.co/blog/bloom-megatron-deepspeed#tensor-parallelism).
+
 
 b. **Pipeline Parallelism (PP)**: Reduces memory footprint and enables large scale training via inter-node parallelization. 
 Reduces the bubble of naive PP via PipeDream-Flush schedule/1F1B schedule and Interleaved 1F1B schedule. 
-Layers are distributed uniformly across PP stages.
+Layers are distributed uniformly across PP stages. For example, if a model has `24` layers and we have `4` GPUs for
+pipeline parallelism, each GPU will have `6` layers (24/4). For more details on schedules to reduce the idle time of PP,
+please refer to the research paper [Efficient Large-Scale Language Model Training on GPU Clusters
+Using Megatron-LM](https://arxiv.org/pdf/2104.04473.pdf) and 
+this section of 🤗 blogpost [The Technology Behind BLOOM Training](https://huggingface.co/blog/bloom-megatron-deepspeed#pipeline-parallelism).
 
 c. **Sequence Parallelism (SP)**: Reduces memory footprint without any additional communication. Only applicable when using TP.
-
-d. **Data Parallelism (DP)** via Distributed Optimizer: Reduces memory footprint by sharding optimizer states and gradients across DP ranks.
+It reduced activation memory required as it prevents same copies to be on the tensor parallel ranks 
+post `all-reduce` by replacing it with `reduce-scatter` and `no-op` operation would be replaced by `all-gather`. 
```suggestion
d. **Data Parallelism (DP)** via Distributed Optimizer: Reduces the memory footprint by sharding optimizer states and gradients across DP ranks
```
@@ -23,18 +23,45 @@ For detailed information and how things work behind the scene please refer the g
 Accelerate integrates following feature of Megatron-LM to enable large scale pre-training/finetuning
 of BERT (Encoder), GPT (Decoder) or T5 models (Encoder and Decoder):
 
-a. **Tensor Parallelism (TP)**: Reduces memory footprint without much additional communication on intra-node ranks
+a. **Tensor Parallelism (TP)**: Reduces memory footprint without much additional communication on intra-node ranks.
+Each tensor are split into multiple chunks with each shard residing on separate GPU. At each step, same minimbatch of data is processed
+independently and in parallel by each shard followed by syncing across all GPUs (`all-reduce` operation). 
+In a simple transformer layer, this leads to 2 `all-reduces` in the forward path and 2 in the backward path.
+For more details, please refer research paper [Megatron-LM: Training Multi-Billion Parameter Language Models Using
+Model Parallelism](https://arxiv.org/pdf/1909.08053.pdf) and 
+this section of 🤗 blogpost [The Technology Behind BLOOM Training](https://huggingface.co/blog/bloom-megatron-deepspeed#tensor-parallelism).
+
 
 b. **Pipeline Parallelism (PP)**: Reduces memory footprint and enables large scale training via inter-node parallelization. 
 Reduces the bubble of naive PP via PipeDream-Flush schedule/1F1B schedule and Interleaved 1F1B schedule. 
-Layers are distributed uniformly across PP stages.
+Layers are distributed uniformly across PP stages. For example, if a model has `24` layers and we have `4` GPUs for
+pipeline parallelism, each GPU will have `6` layers (24/4). For more details on schedules to reduce the idle time of PP,
+please refer to the research paper [Efficient Large-Scale Language Model Training on GPU Clusters
+Using Megatron-LM](https://arxiv.org/pdf/2104.04473.pdf) and 
+this section of 🤗 blogpost [The Technology Behind BLOOM Training](https://huggingface.co/blog/bloom-megatron-deepspeed#pipeline-parallelism).
 
 c. **Sequence Parallelism (SP)**: Reduces memory footprint without any additional communication. Only applicable when using TP.
-
-d. **Data Parallelism (DP)** via Distributed Optimizer: Reduces memory footprint by sharding optimizer states and gradients across DP ranks.
+It reduced activation memory required as it prevents same copies to be on the tensor parallel ranks 
+post `all-reduce` by replacing it with `reduce-scatter` and `no-op` operation would be replaced by `all-gather`. 
+As `all-reduce = reduce-scatter + all-gather`, this saves a ton of activation memory at no added communication cost. 
+To put it simply, it shards the outputs of each transformer layer along sequence dimension, e.g., 
+if the sequence length is `1024` and the TP size is `4`, each GPU will have `256` tokens (1024/4) for each sample. 
+This increases the batch size that can be supported for training. For more details, please refer to the research paper
+[Reducing Activation Recomputation in Large Transformer Models](https://arxiv.org/pdf/2205.05198.pdf). 
+
+d. **Data Parallelism (DP)** via Distributed Optimizer: Reduces memory footprint by sharding optimizer states and gradients across DP ranks
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/759
diff --git a/src/accelerate/utils/constants.py b/src/accelerate/utils/constants.py
index 934923bef..f31027ee1 100644
--- a/src/accelerate/utils/constants.py
+++ b/src/accelerate/utils/constants.py
@@ -58,3 +58,5 @@
     "master_addr",
     "master_port",
 ]
+
+CUDA_DISTRIBUTED_TYPES = ["DEEPSPEED", "MULTI_GPU", "FSDP", "MEGATRON_LM"]
diff --git a/src/accelerate/utils/operations.py b/src/accelerate/utils/operations.py
index aef1fd589..7ba4482cc 100644
--- a/src/accelerate/utils/operations.py
+++ b/src/accelerate/utils/operations.py
@@ -24,6 +24,7 @@
 from torch.distributed import ReduceOp
 
 from ..state import AcceleratorState
+from .constants import CUDA_DISTRIBUTED_TYPES
 from .dataclasses import DistributedType, TensorInformation
 from .imports import is_tpu_available
 from .versions import is_torch_version
@@ -220,12 +221,7 @@ def gather(tensor):
     """
     if AcceleratorState().distributed_type == DistributedType.TPU:
         return _tpu_gather(tensor, name="accelerate.utils.gather")
-    elif AcceleratorState().distributed_type in [
-        DistributedType.DEEPSPEED,
-        DistributedType.MULTI_GPU,
-        DistributedType.FSDP,
-        DistributedType.MEGATRON_LM,
-    ]:
+    elif AcceleratorState().distributed_type in CUDA_DISTRIBUTED_TYPES:
         return _gpu_gather(tensor)
     elif AcceleratorState().distributed_type == DistributedType.MULTI_CPU:
         return _cpu_gather(tensor)
@@ -258,12 +254,7 @@ def gather_object(object: Any):
     """
     if AcceleratorState().distributed_type == DistributedType.TPU:
         raise NotImplementedError("gather objects in TPU is not supported")
-    elif AcceleratorState().distributed_type in [
-        DistributedType.DEEPSPEED,
-        DistributedType.MULTI_GPU,
-        DistributedType.FSDP,
-        DistributedType.MEGATRON_LM,
-    ]:
+    elif AcceleratorState().distributed_type in CUDA_DISTRIBUTED_TYPES:
         return _gpu_gather_object(object)
     elif AcceleratorState().distributed_type == DistributedType.MULTI_CPU:
         return _cpu_gather_object(object)
@@ -302,12 +293,7 @@ def broadcast(tensor, from_process: int = 0):
     """
     if AcceleratorState().distributed_type == DistributedType.TPU:
         return _tpu_broadcast(tensor, src=from_process, name="accelerate.utils.broadcast")
-    elif AcceleratorState().distributed_type in [
-        DistributedType.DEEPSPEED,
-        DistributedType.MULTI_GPU,
-        DistributedType.FSDP,
-        DistributedType.MEGATRON_LM,
-    ]:
+    elif AcceleratorState().distributed_type in CUDA_DISTRIBUTED_TYPES:
         return _gpu_broadcast(tensor, src=from_process)
     elif AcceleratorState().distributed_type == DistributedType.MULTI_CPU:
         return _gpu_broadcast(tensor, src=from_process)
@@ -331,12 +317,7 @@ def broadcast_object_list(object_list, from_process: int = 0):
     if AcceleratorState().distributed_type == DistributedType.TPU:
         for i, obj in enumerate(object_list):
             object_list[i] = xm.mesh_reduce("accelerate.utils.broadcast_object_list", obj, lambda x: x[from_process])
-    elif AcceleratorState().distributed_type in [
-        DistributedType.DEEPSPEED,
-        DistributedType.MULTI_GPU,
-        DistributedType.FSDP,
-        DistributedType.MEGATRON_LM,
-    ]:
+    elif AcceleratorState().distributed_type in CUDA_DISTRIBUTED_TYPES:
         torch.distributed.broadcast_object_list(object_list, src=from_process)
     elif AcceleratorState().distributed_type == DistributedType.MULTI_CPU:
         torch.distributed.broadcast_object_list(object_list, src=from_process)
@@ -452,12 +433,7 @@ def _reduce_across_processes(tensor, reduction="mean"):
         if state.distributed_type == DistributedType.TPU:
             xm.all_reduce("sum", cloned_tensor)
             return cloned_tensor
-        elif state.distributed_type in [
-            DistributedType.DEEPSPEED,
-            DistributedType.MULTI_GPU,
-            DistributedType.FSDP,
-            DistributedType.MEGATRON_LM,
-        ]:
+        elif state.distributed_type.value in CUDA_DISTRIBUTED_TYPES:
             torch.distributed.all_reduce(cloned_tensor, ReduceOp.SUM)
             return cloned_tensor
         else:
diff --git a/src/accelerate/utils/random.py b/src/accelerate/utils/random.py
index 8dc149a02..01c4df2af 100644
--- a/src/accelerate/utils/random.py
+++ b/src/accelerate/utils/random.py
@@ -19,6 +19,7 @@
 import torch
 
 from ..state import AcceleratorState
+from .constants import CUDA_DISTRIBUTED_TYPES
 from .dataclasses import DistributedType, RNGType
 from .imports import is_tpu_available
 
@@ -64,12 +65,7 @@ def synchronize_rng_state(rng_type: Optional[RNGType] = None, generator: Optiona
     state = AcceleratorState()
     if state.distributed_type == DistributedType.TPU:
         rng_state = xm.mesh_reduce("random_seed", rng_state, lambda x: x[0])
-    elif state.distributed_type in [
-        DistributedType.DEEPSPEED,
-        DistributedType.MULTI_GPU,
-        DistributedType.FSDP,
-        DistributedType.MEGATRON_LM,
-    ]:
+    elif state.distributed_type in CUDA_DISTRIBUTED_TYPES:
         rng_state = rng_state.to(state.device)
         torch.distributed.broadcast(rng_state, 0)
         rng_state = rng_state.cpu()

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/758
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index e40797b39..e1282c0bb 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -1194,6 +1194,9 @@ def clip_grad_norm_(self, parameters, max_norm, norm_type=2):
         """
         Should be used in place of `torch.nn.utils.clip_grad_norm_`.
 
+        Returns:
+            `torch.Tensor`: Total norm of the parameter gradients (viewed as a single vector).
+
         Example:
 
         ```python
@@ -1217,13 +1220,13 @@ def clip_grad_norm_(self, parameters, max_norm, norm_type=2):
             parameters = [p for p in parameters]
             for model in self._models:
                 if parameters == [p for p in model.parameters()]:
-                    model.clip_grad_norm_(max_norm, norm_type)
-                    return
+                    return model.clip_grad_norm_(max_norm, norm_type)
         elif self.distributed_type == DistributedType.DEEPSPEED:
             # `accelerator.backward(loss)` is doing that automatically. Therefore, it's implementation is not needed
-            return
+            # We cannot return the gradient norm because DeepSpeed does it.
+            return None
         self.unscale_gradients()
-        torch.nn.utils.clip_grad_norm_(parameters, max_norm, norm_type=norm_type)
+        return torch.nn.utils.clip_grad_norm_(parameters, max_norm, norm_type=norm_type)
 
     def clip_grad_value_(self, parameters, clip_value):
         """

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/756
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index 58009d03b..bdb53988a 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -18,7 +18,9 @@
 
 import torch
 
-from .versions import is_torch_version
+from packaging.version import parse
+
+from .versions import compare_versions, is_torch_version
 
 
 # The package importlib_metadata is in a different place, depending on the Python version.
@@ -88,7 +90,11 @@ def is_bf16_available(ignore_tpu=False):
 
 
 def is_megatron_lm_available():
-    return importlib.util.find_spec("megatron") is not None
+    package_exists = importlib.util.find_spec("megatron") is not None
+    if package_exists:
+        megatron_version = parse(importlib_metadata.version("megatron-lm"))
+        return compare_versions(megatron_version, ">=", "2.2.0")
+    return False
 
 
 def is_transformers_available():

Redundant `else` :)
@@ -88,7 +90,12 @@ def is_bf16_available(ignore_tpu=False):
 
 
 def is_megatron_lm_available():
-    return importlib.util.find_spec("megatron") is not None
+    package_exists = importlib.util.find_spec("megatron") is not None
+    if package_exists:
+        megatron_version = parse(importlib_metadata.version("megatron-lm"))
+        return compare_versions(megatron_version, ">=", "2.2.0")
+    else:
Good point, thanks @tjruwase :)
@@ -88,7 +90,12 @@ def is_bf16_available(ignore_tpu=False):
 
 
 def is_megatron_lm_available():
-    return importlib.util.find_spec("megatron") is not None
+    package_exists = importlib.util.find_spec("megatron") is not None
+    if package_exists:
+        megatron_version = parse(importlib_metadata.version("megatron-lm"))
+        return compare_versions(megatron_version, ">=", "2.2.0")
+    else:
_The documentation is not available anymore as the PR was closed or merged._
Thank you for the fix ✨. LGTM!
https://api.github.com/repos/huggingface/accelerate/pulls/755
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index f348bb0a6..3e5eff45e 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -878,15 +878,17 @@ def launch_command(args):
             args.use_cpu = args.cpu
 
     if args.num_cpu_threads_per_process is None:
-        local_size = get_int_from_env(
-            ["MPI_LOCALNRANKS", "OMPI_COMM_WORLD_LOCAL_SIZE", "MV2_COMM_WORLD_LOCAL_SIZE"], 1
-        )
-        args.num_cpu_threads_per_process = int(psutil.cpu_count(logical=False) / local_size)
-        if args.num_cpu_threads_per_process == 0:
-            args.num_cpu_threads_per_process = 1
-        warned.append(
-            f"\t`--num_cpu_threads_per_process` was set to `{args.num_cpu_threads_per_process}` to improve out-of-box performance"
-        )
+        args.num_cpu_threads_per_process = 1
+        if args.use_cpu and args.num_processes > 1:
+            local_size = get_int_from_env(
+                ["MPI_LOCALNRANKS", "OMPI_COMM_WORLD_LOCAL_SIZE", "MV2_COMM_WORLD_LOCAL_SIZE"], 1
+            )
+            threads_per_process = int(psutil.cpu_count(logical=False) / local_size)
+            if args.num_cpu_threads_per_process > 1:
+                args.num_cpu_threads_per_process = threads_per_process
+                warned.append(
+                    f"\t`--num_cpu_threads_per_process` was set to `{args.num_cpu_threads_per_process}` to improve out-of-box performance when training on CPUs"
+                )
 
     if any(warned):
         message = "The following values were not passed to `accelerate launch` and had defaults used instead:\n"

This means a chosen value of 1 is erased. Let's leave the default at `None` and change it at 1 when not in distributed CPU, or the other default when in distributed CPU.
@@ -877,15 +877,15 @@ def launch_command(args):
         if not hasattr(args, "use_cpu"):
             args.use_cpu = args.cpu
 
-    if args.num_cpu_threads_per_process is None:
+    if args.num_cpu_threads_per_process == 1 and args.use_cpu:
Adjusted :) 
@@ -877,15 +877,15 @@ def launch_command(args):
         if not hasattr(args, "use_cpu"):
             args.use_cpu = args.cpu
 
-    if args.num_cpu_threads_per_process is None:
+    if args.num_cpu_threads_per_process == 1 and args.use_cpu:
```suggestion
    if args.num_cpu_threads_per_process is None:
```
Let's not rely on Python bool conversion magic please :-)
@@ -877,16 +877,18 @@ def launch_command(args):
         if not hasattr(args, "use_cpu"):
             args.use_cpu = args.cpu
 
-    if args.num_cpu_threads_per_process is None:
-        local_size = get_int_from_env(
-            ["MPI_LOCALNRANKS", "OMPI_COMM_WORLD_LOCAL_SIZE", "MV2_COMM_WORLD_LOCAL_SIZE"], 1
-        )
-        args.num_cpu_threads_per_process = int(psutil.cpu_count(logical=False) / local_size)
-        if args.num_cpu_threads_per_process == 0:
-            args.num_cpu_threads_per_process = 1
-        warned.append(
-            f"\t`--num_cpu_threads_per_process` was set to `{args.num_cpu_threads_per_process}` to improve out-of-box performance"
-        )
+    if not args.num_cpu_threads_per_process:
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/753
diff --git a/src/accelerate/commands/config/config_args.py b/src/accelerate/commands/config/config_args.py
index 43faf45ca..716d5f602 100644
--- a/src/accelerate/commands/config/config_args.py
+++ b/src/accelerate/commands/config/config_args.py
@@ -165,6 +165,7 @@ class SageMakerConfig(BaseConfig):
     profile: Optional[str] = None
     region: str = "us-east-1"
     num_machines: int = 1
+    gpu_ids: str = "all"
     base_job_name: str = f"accelerate-sagemaker-{num_machines}"
     pytorch_version: str = SAGEMAKER_PYTORCH_VERSION
     transformers_version: str = SAGEMAKER_TRANSFORMERS_VERSION

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/751
diff --git a/docs/source/index.mdx b/docs/source/index.mdx
index 1664fd7a2..a87bd6e32 100644
--- a/docs/source/index.mdx
+++ b/docs/source/index.mdx
@@ -51,19 +51,19 @@ accelerate launch {my_script.py}
 
 <div class="mt-10">
   <div class="w-full flex flex-col space-y-4 md:space-y-0 md:grid md:grid-cols-2 md:gap-y-4 md:gap-x-5">
-    <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="/docs/accelerate/basic_tutorials/overview"
+    <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="./basic_tutorials/overview"
       ><div class="w-full text-center bg-gradient-to-br from-blue-400 to-blue-500 rounded-lg py-1.5 font-semibold mb-5 text-white text-lg leading-relaxed">Tutorials</div>
       <p class="text-gray-700">Learn the basics and become familiar with using 🤗 Accelerate. Start here if you are using 🤗 Accelerate for the first time!</p>
     </a>
-    <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="/docs/accelerate/usage_guides/gradient_accumulation"
+    <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="./usage_guides/gradient_accumulation"
       ><div class="w-full text-center bg-gradient-to-br from-indigo-400 to-indigo-500 rounded-lg py-1.5 font-semibold mb-5 text-white text-lg leading-relaxed">How-to guides</div>
       <p class="text-gray-700">Practical guides to help you achieve a specific goal. Take a look at these guides to learn how to use 🤗 Accelerate to solve real-world problems.</p>
     </a>
-    <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="/docs/accelerate/concept_guides/gradient_synchronization"
+    <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="./concept_guides/gradient_synchronization"
       ><div class="w-full text-center bg-gradient-to-br from-pink-400 to-pink-500 rounded-lg py-1.5 font-semibold mb-5 text-white text-lg leading-relaxed">Conceptual guides</div>
       <p class="text-gray-700">High-level explanations for building a better understanding of important topics such as avoiding subtle nuances and pitfalls in distributed training and DeepSpeed.</p>
    </a>
-    <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="/docs/accelerate/package_reference/accelerator"
+    <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="./package_reference/accelerator"
       ><div class="w-full text-center bg-gradient-to-br from-purple-400 to-purple-500 rounded-lg py-1.5 font-semibold mb-5 text-white text-lg leading-relaxed">Reference</div>
       <p class="text-gray-700">Technical descriptions of how 🤗 Accelerate classes and methods work.</p>
     </a>

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/749
diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
index 0b57ae6d8..08ee3a65a 100644
--- a/src/accelerate/utils/modeling.py
+++ b/src/accelerate/utils/modeling.py
@@ -258,7 +258,7 @@ def get_max_layer_size(
     modules_to_treat = modules.copy()
     while len(modules_to_treat) > 0:
         module_name, module = modules_to_treat.pop(0)
-        modules_children = list(module.named_children())
+        modules_children = list(module.named_children()) if isinstance(module, torch.nn.Module) else []
         if len(modules_children) == 0 or module.__class__.__name__ in no_split_module_classes:
             # No splitting this one so we compare to the max_size
             size = module_sizes[module_name]

```suggestion
        modules_children = list(module.named_children()) if isinstance(module, torch.nn.Module) else []
```
@@ -258,7 +258,7 @@ def get_max_layer_size(
     modules_to_treat = modules.copy()
     while len(modules_to_treat) > 0:
         module_name, module = modules_to_treat.pop(0)
-        modules_children = list(module.named_children())
+        modules_children = list(module.named_children()) if not isinstance(module, torch.nn.Parameter) else []
_The documentation is not available anymore as the PR was closed or merged._
> What's a reproducer? I'm not sure how you get to a parameter here since we only apply the `children` method.

Haha that was fast. Here a repro:

```py
from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker

StableDiffusionSafetyChecker._no_split_modules = ["CLIPEncoderLayer"]
pipe = StableDiffusionSafetyChecker.from_pretrained('CompVis/stable-diffusion-safety-checker', device_map='auto')
```

(sorry, it's 1GB of download)
Which branch should I check? On main I just get an error telling me device_map auto is not supported yet
Ah sorry just updated it above, can you try:

```python
from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker

StableDiffusionSafetyChecker._no_split_modules = ["CLIPEncoderLayer"]
pipe = StableDiffusionSafetyChecker.from_pretrained('CompVis/stable-diffusion-safety-checker', device_map='auto')
```
Awesome! Can I merge after one :heavy_check_mark: here or wait for another reviewer? 
Nope, you can go ahead and merge!
https://api.github.com/repos/huggingface/accelerate/pulls/747
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index 01c419f4f..e86a0abff 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -35,6 +35,7 @@ def get_cluster_input():
 
     machine_rank = 0
     num_machines = 1
+    num_processes = 1
     gpu_ids = None
     main_process_ip = None
     main_process_port = None
@@ -294,12 +295,6 @@ def get_cluster_input():
             default=1,
             error_message="Please enter an integer.",
         )
-
-    if distributed_type in [DistributedType.MULTI_GPU, DistributedType.NO] and not use_cpu:
-        gpu_ids = _ask_field(
-            "What GPU(s) (by id) should be used for training on this machine as a comma-seperated list? [all]:",
-            default="all",
-        )
     elif distributed_type in [DistributedType.FSDP, DistributedType.DEEPSPEED]:
         num_processes = _ask_field(
             "How many GPU(s) should be used for distributed training? [1]:",
@@ -310,6 +305,12 @@ def get_cluster_input():
     else:
         num_processes = 1
 
+    if distributed_type in [DistributedType.MULTI_GPU, DistributedType.NO] and not use_cpu:
+        gpu_ids = _ask_field(
+            "What GPU(s) (by id) should be used for training on this machine as a comma-seperated list? [all]:",
+            default="all",
+        )
+
     if distributed_type != DistributedType.TPU:
         if distributed_type == DistributedType.DEEPSPEED and use_deepspeed_config:
             mixed_precision = "no"

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/746
diff --git a/docs/source/usage_guides/training_zoo.mdx b/docs/source/usage_guides/training_zoo.mdx
index b3fff25e3..cc388acac 100644
--- a/docs/source/usage_guides/training_zoo.mdx
+++ b/docs/source/usage_guides/training_zoo.mdx
@@ -27,6 +27,7 @@ These examples showcase the base features of Accelerate and are a great starting
 
 These examples showcase specific features that the Accelerate framework offers
 
+- [Automatic memory-aware gradient accumulation](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/automatic_gradient_accumulation.py)
 - [Checkpointing states](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/checkpointing.py)
 - [Cross validation](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/cross_validation.py)
 - [DeepSpeed](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/deepspeed_with_config_support.py)
diff --git a/examples/by_feature/automatic_gradient_accumulation.py b/examples/by_feature/automatic_gradient_accumulation.py
new file mode 100644
index 000000000..d6e0cf028
--- /dev/null
+++ b/examples/by_feature/automatic_gradient_accumulation.py
@@ -0,0 +1,232 @@
+# 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 torch
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+
+# New Code #
+import evaluate
+from accelerate import Accelerator, DistributedType
+from accelerate.utils import find_executable_batch_size
+from datasets import load_dataset
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate,
+# specifically showcasing how to combine both the gradient accumulation
+# and automatic batch size finder utilities of Accelerate to perfrom
+# automatic gradient accumulation
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#
+# New additions from the base script can be found quickly by
+# looking for the # New Code # tags
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+EVAL_BATCH_SIZE = 32
+
+
+def get_dataloaders(accelerator: Accelerator, batch_size: int = 16):
+    """
+    Creates a set of `DataLoader`s for the `glue` dataset,
+    using "bert-base-cased" as the tokenizer.
+
+    Args:
+        accelerator (`Accelerator`):
+            An `Accelerator` object
+        batch_size (`int`, *optional*):
+            The batch size for the train and validation DataLoaders.
+    """
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    datasets = load_dataset("glue", "mrpc")
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    # starting with the main process first:
+    with accelerator.main_process_first():
+        tokenized_datasets = datasets.map(
+            tokenize_function,
+            batched=True,
+            remove_columns=["idx", "sentence1", "sentence2"],
+        )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    return train_dataloader, eval_dataloader
+
+
+# For testing only
+if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+    from accelerate.test_utils.training import mocked_dataloaders
+
+    get_dataloaders = mocked_dataloaders  # noqa: F811
+
+
+def training_function(config, args):
+    # For testing only
+    if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+        config["num_epochs"] = 2
+    # Initialize accelerator
+    accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    seed = int(config["seed"])
+    observed_batch_size = int(config["batch_size"])
+
+    metric = evaluate.load("glue", "mrpc")
+
+    # New Code #
+    # We use the `find_executable_batch_size` decorator, passing in the desired observed batch size
+    # to train on. If a CUDA OOM error occurs, it will retry this loop cutting the batch size in
+    # half each time. From this, we can calculate the number of gradient accumulation steps needed
+    # and modify the Accelerator object as a result
+    @find_executable_batch_size(starting_batch_size=int(observed_batch_size))
+    def inner_training_loop(batch_size):
+        # Since we need to modify the outside accelerator object, we need to bring it
+        # to the local scope
+        nonlocal accelerator
+
+        # We can calculate the number of gradient accumulation steps based on the current
+        # batch size vs the starting batch size
+        num_gradient_accumulation_steps = observed_batch_size // batch_size
+
+        # And then set it in the Accelerator directly:
+        accelerator.gradient_accumulation_steps = num_gradient_accumulation_steps
+
+        # Next we need to free all of the stored model references in the Accelerator each time
+        accelerator.free_memory()
+
+        # And set the seed so our results are reproducable each reset
+        set_seed(seed)
+
+        # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+        model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
+
+        # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
+        # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
+        # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
+        model = model.to(accelerator.device)
+
+        # Instantiate optimizer
+        optimizer = AdamW(params=model.parameters(), lr=lr)
+        train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)
+
+        # Instantiate scheduler
+        lr_scheduler = get_linear_schedule_with_warmup(
+            optimizer=optimizer,
+            num_warmup_steps=100,
+            num_training_steps=(len(train_dataloader) * num_epochs),
+        )
+
+        # Prepare everything
+        # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
+        # prepare method.
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+            model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+        )
+
+        # Now we train the model
+        for epoch in range(num_epochs):
+            model.train()
+            for step, batch in enumerate(train_dataloader):
+                # And perform gradient accumulation
+                with accelerator.accumulate(model):
+                    # We could avoid this line since we set the accelerator with `device_placement=True`.
+                    batch.to(accelerator.device)
+                    outputs = model(**batch)
+                    loss = outputs.loss
+                    accelerator.backward(loss)
+                    optimizer.step()
+                    lr_scheduler.step()
+                    optimizer.zero_grad()
+
+            model.eval()
+            for step, batch in enumerate(eval_dataloader):
+                # We could avoid this line since we set the accelerator with `device_placement=True`.
+                batch.to(accelerator.device)
+                with torch.no_grad():
+                    outputs = model(**batch)
+                predictions = outputs.logits.argmax(dim=-1)
+                predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
+                metric.add_batch(
+                    predictions=predictions,
+                    references=references,
+                )
+
+            eval_metric = metric.compute()
+            # Use accelerator.print to print only on the main process.
+            accelerator.print(f"epoch {epoch}:", eval_metric)
+
+    # New Code #
+    # And call it at the end with no arguments
+    # Note: You could also refactor this outside of your training loop function
+    inner_training_loop()
+
+
+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()
+    # New Code #
+    # We modify the starting batch size to be an observed batch size of 256, to guarentee an initial CUDA OOM
+    config = {"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 256}
+    training_function(config, args)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/test_examples.py b/tests/test_examples.py
index ac8c4de52..9197070c8 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -36,6 +36,7 @@
     "gradient_accumulation.py",
     "multi_process_metrics.py",
     "memory.py",
+    "automatic_gradient_accumulation.py",
     "fsdp_with_peak_mem_tracking.py",
     "deepspeed_with_config_support.py",
 ]

```suggestion
    # half each time. From this, we can calculate the number of gradient accumulation steps needed
```
@@ -0,0 +1,232 @@
+# 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 torch
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+
+# New Code #
+import evaluate
+from accelerate import Accelerator, DistributedType
+from accelerate.utils import find_executable_batch_size
+from datasets import load_dataset
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate,
+# specifically showcasing how to combine both the gradient accumulation
+# and automatic batch size finder utilities of Accelerate to perfrom
+# automatic gradient accumulation
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#
+# New additions from the base script can be found quickly by
+# looking for the # New Code # tags
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+EVAL_BATCH_SIZE = 32
+
+
+def get_dataloaders(accelerator: Accelerator, batch_size: int = 16):
+    """
+    Creates a set of `DataLoader`s for the `glue` dataset,
+    using "bert-base-cased" as the tokenizer.
+
+    Args:
+        accelerator (`Accelerator`):
+            An `Accelerator` object
+        batch_size (`int`, *optional*):
+            The batch size for the train and validation DataLoaders.
+    """
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    datasets = load_dataset("glue", "mrpc")
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    # starting with the main process first:
+    with accelerator.main_process_first():
+        tokenized_datasets = datasets.map(
+            tokenize_function,
+            batched=True,
+            remove_columns=["idx", "sentence1", "sentence2"],
+        )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    return train_dataloader, eval_dataloader
+
+
+# For testing only
+if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+    from accelerate.test_utils.training import mocked_dataloaders
+
+    get_dataloaders = mocked_dataloaders  # noqa: F811
+
+
+def training_function(config, args):
+    # For testing only
+    if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+        config["num_epochs"] = 2
+    # Initialize accelerator
+    accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    seed = int(config["seed"])
+    observed_batch_size = int(config["batch_size"])
+
+    metric = evaluate.load("glue", "mrpc")
+
+    # New Code #
+    # We use the `find_executable_batch_size` decorator, passing in the desired observed batch size
+    # to train on. If a CUDA OOM error occurs, it will retry this loop cutting the batch size in
+    # half each time. From this, we can calcualte the number of gradient accumulation steps needed
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_742). All of your documentation changes will be reflected on that endpoint.
https://api.github.com/repos/huggingface/accelerate/pulls/742
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 982293a33..f348bb0a6 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -45,6 +45,7 @@
 )
 from accelerate.utils.constants import DEEPSPEED_MULTINODE_LAUNCHERS
 from accelerate.utils.dataclasses import SageMakerDistributedType
+from accelerate.utils.launch import env_var_path_add
 
 
 if is_rich_available():
@@ -573,7 +574,7 @@ def deepspeed_launcher(args):
         warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', DeprecationWarning)
         mixed_precision = "fp16"
 
-    current_env["PYTHONPATH"] = sys.executable
+    current_env["PYTHONPATH"] = env_var_path_add("PYTHONPATH", os.path.abspath("."))
     current_env["MIXED_PRECISION"] = str(mixed_precision)
     current_env["USE_DEEPSPEED"] = "true"
     current_env["DEEPSPEED_ZERO_STAGE"] = str(args.zero_stage)
diff --git a/src/accelerate/utils/launch.py b/src/accelerate/utils/launch.py
index 31d71ad21..8642a441a 100644
--- a/src/accelerate/utils/launch.py
+++ b/src/accelerate/utils/launch.py
@@ -52,6 +52,16 @@ def _filter_args(args):
     return new_args
 
 
+def env_var_path_add(env_var_name, path_to_add):
+    """
+    Extends a path-based environment variable's value with a new path and returns the updated value. It's up to the
+    caller to set it in os.environ.
+    """
+    paths = [p for p in os.environ.get(env_var_name, "").split(":") if len(p) > 0]
+    paths.append(str(path_to_add))
+    return ":".join(paths)
+
+
 class PrepareForLaunch:
     """
     Prepare a function that will launched in a distributed setup.

Let's just write the three lines of code instead of adding a nested function? It will be shorter and as easy to read.
@@ -573,7 +573,13 @@ def deepspeed_launcher(args):
         warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', DeprecationWarning)
         mixed_precision = "fp16"
 
-    current_env["PYTHONPATH"] = sys.executable
+    def env_path_add(env_var_name, new_path):
+        """extend path-based env variable with a new path"""
+        paths = [p for p in os.environ.get(env_var_name, "").split(":") if len(p) > 0]
+        paths.append(str(new_path))
+        return ":".join(paths)
+
+    current_env["PYTHONPATH"] = env_path_add("PYTHONPATH", sys.executable)
as it's likely that at some point you might want to extend other path-based env vars like `PATH` or `LD_LIBRARY_PATH` - one more option is to move it to utils. 

but if you prefer inlined and @pacman100 is an agreement I will rewrite.
@@ -573,7 +573,13 @@ def deepspeed_launcher(args):
         warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', DeprecationWarning)
         mixed_precision = "fp16"
 
-    current_env["PYTHONPATH"] = sys.executable
+    def env_path_add(env_var_name, new_path):
+        """extend path-based env variable with a new path"""
+        paths = [p for p in os.environ.get(env_var_name, "").split(":") if len(p) > 0]
+        paths.append(str(new_path))
+        return ":".join(paths)
+
+    current_env["PYTHONPATH"] = env_path_add("PYTHONPATH", sys.executable)
I'm fine having it as a util too, just don't like the nested function approach where it can't be reused.
@@ -573,7 +573,13 @@ def deepspeed_launcher(args):
         warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', DeprecationWarning)
         mixed_precision = "fp16"
 
-    current_env["PYTHONPATH"] = sys.executable
+    def env_path_add(env_var_name, new_path):
+        """extend path-based env variable with a new path"""
+        paths = [p for p in os.environ.get(env_var_name, "").split(":") if len(p) > 0]
+        paths.append(str(new_path))
+        return ":".join(paths)
+
+    current_env["PYTHONPATH"] = env_path_add("PYTHONPATH", sys.executable)
pushed the change to move to utils, renamed the helper for better clarity now that it's far away from definition.
@@ -573,7 +573,13 @@ def deepspeed_launcher(args):
         warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', DeprecationWarning)
         mixed_precision = "fp16"
 
-    current_env["PYTHONPATH"] = sys.executable
+    def env_path_add(env_var_name, new_path):
+        """extend path-based env variable with a new path"""
+        paths = [p for p in os.environ.get(env_var_name, "").split(":") if len(p) > 0]
+        paths.append(str(new_path))
+        return ":".join(paths)
+
+    current_env["PYTHONPATH"] = env_path_add("PYTHONPATH", sys.executable)
```suggestion
    current_env["PYTHONPATH"] = env_var_path_add("PYTHONPATH", os.path.abspath('.'))
```
@@ -573,7 +574,7 @@ def deepspeed_launcher(args):
         warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', DeprecationWarning)
         mixed_precision = "fp16"
 
-    current_env["PYTHONPATH"] = sys.executable
+    current_env["PYTHONPATH"] = env_var_path_add("PYTHONPATH", sys.executable)
_The documentation is not available anymore as the PR was closed or merged._
Oops, Thanks a lot for pointing this out! Post the above PR, a sample run is shown below:
```
>>> export PYTHONPATH="sample_dir/fancy_library"
...

>>> accelerate launch --config_file ds_mn_hf.yaml accelerate/examples/complete_nlp_example.py --mixed_precision "fp16"
[2022-10-06 12:00:16,739] [INFO] [runner.py:504:main] cmd = pdsh -S -f 1024 -w localhost,sourab_huggingface_co@sourab-vm-image-1 
export PYTHONPATH=sample_dir/fancy_library:/home/sourab_huggingface_co;

...

export MIXED_PRECISION=fp16; export USE_DEEPSPEED=true; export DEEPSPEED_ZERO_STAGE=2; export GRADIENT_ACCUMULATION_STEPS=1; export GRADIENT_CLIPPING=1.0; export DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE=none; export DEEPSPEED_OFFLOAD_PARAM_DEVICE=none; export DEEPSPEED_ZERO3_INIT=false; export DEEPSPEED_ZERO3_SAVE_16BIT_MODEL=none;  cd /home/sourab_huggingface_co; /opt/conda/envs/ml/bin/python3.9 -u -m deepspeed.launcher.launch --world_info=eyJsb2NhbGhvc3QiOiBbMF0sICJzb3VyYWJfaHVnZ2luZ2ZhY2VfY29Ac291cmFiLXZtLWltYWdlLTEiOiBbMF19 --node_rank=%n --master_addr=xx.xxx.x.x --master_port=29500 --no_local_rank accelerate/examples/complete_nlp_example.py --mixed_precision 'fp16'
localhost: [2022-10-06 12:00:22,720] [INFO] [launch.py:136:main] WORLD INFO DICT: {'localhost': [0], 'sourab_huggingface_co@sourab-vm-image-1': [0]}
localhost: [2022-10-06 12:00:22,720] [INFO] [launch.py:142:main] nnodes=2, num_local_procs=1, node_rank=0
localhost: [2022-10-06 12:00:22,720] [INFO] [launch.py:155:main] global_rank_mapping=defaultdict(<class 'list'>, {'localhost': [0], 'sourab_huggingface_co@sourab-vm-image-1': [1]})
localhost: [2022-10-06 12:00:22,720] [INFO] [launch.py:156:main] dist_world_size=2
localhost: [2022-10-06 12:00:22,720] [INFO] [launch.py:158:main] Setting CUDA_VISIBLE_DEVICES=0
sourab-vm-image-1: [2022-10-06 12:00:23,404] [INFO] [launch.py:136:main] WORLD INFO DICT: {'localhost': [0], 'sourab_huggingface_co@sourab-vm-image-1': [0]}
sourab-vm-image-1: [2022-10-06 12:00:23,404] [INFO] [launch.py:142:main] nnodes=2, num_local_procs=1, node_rank=1
sourab-vm-image-1: [2022-10-06 12:00:23,404] [INFO] [launch.py:155:main] global_rank_mapping=defaultdict(<class 'list'>, {'localhost': [0], 'sourab_huggingface_co@sourab-vm-image-1': [1]})
sourab-vm-image-1: [2022-10-06 12:00:23,404] [INFO] [launch.py:156:main] dist_world_size=2
sourab-vm-image-1: [2022-10-06 12:00:23,404] [INFO] [launch.py:158:main] Setting CUDA_VISIBLE_DEVICES=0
localhost: [2022-10-06 12:00:27,932] [INFO] [comm.py:633:init_distributed] Initializing TorchBackend in DeepSpeed with backend nccl

...

localhost: [2022-10-06 12:01:10,571] [INFO] [config.py:976:print_user_config]   json = {
localhost:     "train_batch_size": 32, 
localhost:     "train_micro_batch_size_per_gpu": 16, 
localhost:     "gradient_accumulation_steps": 1, 
localhost:     "zero_optimization": {
localhost:         "stage": 2, 
localhost:         "offload_optimizer": {
localhost:             "device": "none"
localhost:         }, 
localhost:         "offload_param": {
localhost:             "device": "none"
localhost:         }, 
localhost:         "stage3_gather_16bit_weights_on_model_save": false
localhost:     }, 
localhost:     "gradient_clipping": 1.0, 
localhost:     "steps_per_print": inf, 
localhost:     "fp16": {
localhost:         "enabled": true, 
localhost:         "auto_cast": true
localhost:     }, 
localhost:     "zero_allow_untested_optimizer": true
localhost: }
localhost: Using /home/sourab_huggingface_co/.cache/torch_extensions/py39_cu113 as PyTorch extensions root...

...

localhost: epoch 0: {'accuracy': 0.7303921568627451, 'f1': 0.8318042813455658}
sourab-vm-image-1: epoch 0: {'accuracy': 0.7303921568627451, 'f1': 0.8318042813455658}
localhost: epoch 1: {'accuracy': 0.8308823529411765, 'f1': 0.880415944540728}
sourab-vm-image-1: epoch 1: {'accuracy': 0.8308823529411765, 'f1': 0.880415944540728}
localhost: epoch 2: {'accuracy': 0.8333333333333334, 'f1': 0.8811188811188811}
sourab-vm-image-1: epoch 2: {'accuracy': 0.8333333333333334, 'f1': 0.8811188811188811}

```
@pacman100, please feel free to merge when you feel it's ready. Thanks.
https://api.github.com/repos/huggingface/accelerate/pulls/741
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 5d8ff0c1b..593eec741 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -6,6 +6,7 @@ on:
       - "src/**"
       - "tests/**"
       - ".github/**"
+      - "examples/**"
       - "setup.py"
     types: [opened, synchronize, reopened]
 

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/740
diff --git a/docs/source/usage_guides/memory.mdx b/docs/source/usage_guides/memory.mdx
index 8841cd584..213a2f670 100644
--- a/docs/source/usage_guides/memory.mdx
+++ b/docs/source/usage_guides/memory.mdx
@@ -25,16 +25,20 @@ training script. To use it, restructure your training function to include an inn
 and build your dataloaders inside it. At a minimum, this could look like 4 new lines of code. 
 > Note: The inner function *must* take in the batch size as the first parameter, but we do not pass one to it when called. The wrapper handles this for us
 
+It should also be noted that anything which will consume CUDA memory and passed to the `accelerator` **must** be declared inside the inner function,
+such as models and optimizers.
+
 ```diff
 def training_function(args):
     accelerator = Accelerator()
-    model = get_model()
-    model.to(accelerator.device)
-    optimizer = get_optimizer()
 
 +   @find_executable_batch_size(starting_batch_size=args.batch_size)
 +   def inner_training_loop(batch_size):
-+       nonlocal model, optimizer # Ensure they can be used in our context
++       nonlocal accelerator # Ensure they can be used in our context
++       accelerator.free_memory() # Free all lingering references
+        model = get_model()
+        model.to(accelerator.device)
+        optimizer = get_optimizer()
         train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)
         lr_scheduler = get_scheduler(
             optimizer, 
diff --git a/examples/by_feature/memory.py b/examples/by_feature/memory.py
index 82dc34b8a..684a32853 100644
--- a/examples/by_feature/memory.py
+++ b/examples/by_feature/memory.py
@@ -122,24 +122,6 @@ def training_function(config, args):
 
     metric = evaluate.load("glue", "mrpc")
 
-    # If the batch size is too big we use gradient accumulation
-    gradient_accumulation_steps = 1
-    if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
-        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
-        batch_size = MAX_GPU_BATCH_SIZE
-
-    set_seed(seed)
-    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
-    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
-
-    # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
-    # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
-    # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
-    model = model.to(accelerator.device)
-
-    # Instantiate optimizer
-    optimizer = AdamW(params=model.parameters(), lr=lr)
-
     # New Code #
     # We now can define an inner training loop function. It should take a batch size as the only parameter,
     # and build the dataloaders in there.
@@ -147,16 +129,31 @@ def training_function(config, args):
     @find_executable_batch_size(starting_batch_size=batch_size)
     def inner_training_loop(batch_size):
         # And now just move everything below under this function
-        # Ensure that anything declared outside this function is set as `nonlocal`
-        # so it is in scope
-        nonlocal model, optimizer
+        # We need to bring in the Accelerator object from earlier
+        nonlocal accelerator
+        # And reset all of its attributes that could hold onto any memory:
+        accelerator.free_memory()
+
+        # Then we can declare the model, optimizer, and everything else:
+        set_seed(seed)
+
+        # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+        model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
+
+        # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
+        # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
+        # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
+        model = model.to(accelerator.device)
+
+        # Instantiate optimizer
+        optimizer = AdamW(params=model.parameters(), lr=lr)
         train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)
 
         # Instantiate scheduler
         lr_scheduler = get_linear_schedule_with_warmup(
             optimizer=optimizer,
             num_warmup_steps=100,
-            num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,
+            num_training_steps=(len(train_dataloader) * num_epochs),
         )
 
         # Prepare everything
@@ -174,12 +171,10 @@ def inner_training_loop(batch_size):
                 batch.to(accelerator.device)
                 outputs = model(**batch)
                 loss = outputs.loss
-                loss = loss / gradient_accumulation_steps
                 accelerator.backward(loss)
-                if step % gradient_accumulation_steps == 0:
-                    optimizer.step()
-                    lr_scheduler.step()
-                    optimizer.zero_grad()
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
 
             model.eval()
             for step, batch in enumerate(eval_dataloader):

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/739
diff --git a/docs/source/usage_guides/training_zoo.mdx b/docs/source/usage_guides/training_zoo.mdx
index d0510bf93..b3fff25e3 100644
--- a/docs/source/usage_guides/training_zoo.mdx
+++ b/docs/source/usage_guides/training_zoo.mdx
@@ -100,3 +100,7 @@ These are tutorials from libraries that integrate with 🤗 Accelerate:
 ### Tez 
 
 - [Leaf disease detection with Tez and Accelerate](https://www.kaggle.com/code/abhishek/tez-faster-and-easier-training-for-leaf-detection/notebook)
+
+### trlx 
+
+- [How to implement a sentiment learning task with trlx](https://github.com/CarperAI/trlx#example-how-to-add-a-task)
\ No newline at end of file

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/738
diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml
index 84c4e03f6..559c798c5 100644
--- a/docs/source/_toctree.yml
+++ b/docs/source/_toctree.yml
@@ -32,9 +32,11 @@
   - local: usage_guides/memory
     title: How to avoid CUDA Out-of-Memory
   - local: usage_guides/sagemaker
-    title: Using Accelerate on SageMaker
+    title: Using 🤗 Accelerate on SageMaker
   - local: usage_guides/mps
     title: How to use Apple Silicon M1 GPUs
+  - local: usage_guides/training_zoo
+    title: 🤗 Accelerate Example Zoo
   title: How-To Guides
 - sections:
   - local: concept_guides/gradient_synchronization
diff --git a/docs/source/usage_guides/training_zoo.mdx b/docs/source/usage_guides/training_zoo.mdx
new file mode 100644
index 000000000..d0510bf93
--- /dev/null
+++ b/docs/source/usage_guides/training_zoo.mdx
@@ -0,0 +1,102 @@
+<!--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.
+-->
+
+# Example Zoo
+
+Below contains a non-exhuastive list of tutorials and scripts showcasing Accelerate
+
+## Official Accelerate Examples:
+
+### Basic Examples
+
+These examples showcase the base features of Accelerate and are a great starting point
+
+- [Barebones NLP example](https://github.com/huggingface/accelerate/blob/main/examples/nlp_example.py)
+- [Barebones computer vision example](https://github.com/huggingface/accelerate/blob/main/examples/cv_example.py)
+
+### Feature Specific Examples
+
+These examples showcase specific features that the Accelerate framework offers
+
+- [Checkpointing states](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/checkpointing.py)
+- [Cross validation](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/cross_validation.py)
+- [DeepSpeed](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/deepspeed_with_config_support.py)
+- [Fully Sharded Data Parallelism](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/fsdp_with_peak_mem_tracking.py)
+- [Gradient accumulation](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/gradient_accumulation.py)
+- [Memory-aware batch size finder](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/memory.py)
+- [Metric Computation](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/multi_process_metrics.py)
+- [Using Trackers](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/tracking.py)
+
+### Full Examples 
+
+These examples showcase every feature in Accelerate at once that was shown in "Feature Specific Examples"
+
+- [Complete NLP example](https://github.com/huggingface/accelerate/blob/main/examples/complete_nlp_example.py)
+- [Complete computer vision example](https://github.com/huggingface/accelerate/blob/main/examples/complete_cv_example.py)
+- [Causal language model fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/language-modeling/run_clm_no_trainer.py)
+- [Masked language model fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/language-modeling/run_mlm_no_trainer.py)
+- [Speech pretraining example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/speech-pretraining/run_wav2vec2_pretraining_no_trainer.py)
+- [Translation fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/translation/run_translation_no_trainer.py)
+- [Text classification fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/text-classification/run_glue_no_trainer.py)
+- [Semantic segmentation fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/semantic-segmentation/run_semantic_segmentation_no_trainer.py)
+- [Question answering fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/question-answering/run_qa_no_trainer.py)
+- [Beam search question answering fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/question-answering/run_qa_beam_search_no_trainer.py)
+- [Multiple choice question answering fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/multiple-choice/run_swag_no_trainer.py)
+- [Named entity recognition fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/token-classification/run_ner_no_trainer.py)
+- [Image classification fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/image-classification/run_image_classification_no_trainer.py)
+- [Summarization fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/summarization/run_summarization_no_trainer.py)
+
+## Integration Examples 
+
+These are tutorials from libraries that integrate with 🤗 Accelerate: 
+
+### Catalyst
+
+- [Distributed training tutorial with Catalyst](https://catalyst-team.github.io/catalyst/tutorials/ddp.html)
+
+### DALLE2-pytorch 
+
+- [Fine-tuning DALLE2](https://github.com/lucidrains/DALLE2-pytorch#usage)
+
+### 🤗 diffusers
+
+- [Performing textual inversion with diffusers](https://github.com/huggingface/diffusers/tree/main/examples/textual_inversion)
+- [Training DreamBooth with diffusers](https://github.com/huggingface/diffusers/tree/main/examples/dreambooth)
+
+### fastai 
+
+- [Distributed training from Jupyter Notebooks with fastai](https://docs.fast.ai/tutorial.distributed.html)
+- [Basic distributed training examples with fastai](https://docs.fast.ai/examples/distributed_app_examples.html)
+
+### GradsFlow
+
+- [Auto Image Classification with GradsFlow](https://docs.gradsflow.com/en/latest/examples/nbs/01-ImageClassification/)
+
+### imagen-pytorch 
+
+- [Fine-tuning Imagen](https://github.com/lucidrains/imagen-pytorch#usage)
+
+### Kornia
+
+- [Fine-tuning vision models with Kornia's Trainer](https://kornia.readthedocs.io/en/latest/get-started/training.html)
+
+### PyTorch Accelerated 
+
+- [Quickstart distributed training tutorial with PyTorch Accelerated](https://pytorch-accelerated.readthedocs.io/en/latest/quickstart.html)
+
+### PyTorch3D
+
+- [Perform Deep Learning with 3D data](https://pytorch3d.org/tutorials/)
+
+### Tez 
+
+- [Leaf disease detection with Tez and Accelerate](https://www.kaggle.com/code/abhishek/tez-faster-and-easier-training-for-leaf-detection/notebook)

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/737
diff --git a/.github/ISSUE_TEMPLATE/bug-report.yml b/.github/ISSUE_TEMPLATE/bug-report.yml
index 69a101307..3ed39ac4b 100644
--- a/.github/ISSUE_TEMPLATE/bug-report.yml
+++ b/.github/ISSUE_TEMPLATE/bug-report.yml
@@ -1,6 +1,5 @@
 name: "\U0001F41B Bug Report"
 description: Submit a bug report to help us improve Accelerate
-labels: [ "bug" ]
 body:
   - type: textarea
     id: system-info

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/735
diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml
index 84c4e03f6..ab24fe8e5 100644
--- a/docs/source/_toctree.yml
+++ b/docs/source/_toctree.yml
@@ -37,6 +37,8 @@
     title: How to use Apple Silicon M1 GPUs
   title: How-To Guides
 - sections:
+  - local: concept_guides/performance
+    title: Comparing performance across distributed setups
   - local: concept_guides/gradient_synchronization
     title: Gradient synchronization
   - local: concept_guides/deferring_execution
diff --git a/docs/source/concept_guides/deferring_execution.mdx b/docs/source/concept_guides/deferring_execution.mdx
index 4297e2567..cb80ee0da 100644
--- a/docs/source/concept_guides/deferring_execution.mdx
+++ b/docs/source/concept_guides/deferring_execution.mdx
@@ -1,3 +1,15 @@
+<!--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.
+-->
+
 # Deferring Executions
 
 When you run your usual script, instructions are executed in order. Using 🤗 Accelerate to deploy your script on several
diff --git a/docs/source/concept_guides/performance.mdx b/docs/source/concept_guides/performance.mdx
new file mode 100644
index 000000000..c974b322f
--- /dev/null
+++ b/docs/source/concept_guides/performance.mdx
@@ -0,0 +1,91 @@
+<!--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.
+-->
+
+# Comparing performance between different device setups
+
+Evaluating and comparing the performance from different setups can be quite tricky if you don't know what to look for.
+For example, you cannot run the same script with the same batch size across TPU, multi-GPU, and single-GPU with Accelerate 
+and expect your results to line up. 
+
+But why?
+
+There's three reasons for this that this tutorial will cover: 
+
+1. **Setting the right seeds**
+2. **Observed Batch Sizes**
+3. **Learning Rates**
+
+## Setting the Seed 
+
+While this issue has not come up as much, make sure to use [`utils.set_seed`] to fully set the seed in all distributed cases so training will be reproducable:
+
+```python
+from accelerate import set_seed
+
+set_seed(42)
+```
+
+Why is this important? Under the hood this will set **5** different seed settings:
+
+```python
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    # ^^ safe to call this function even if cuda is not available
+    if is_tpu_available():
+        xm.set_rng_state(seed)
+```
+
+The random state, numpy's state, torch, torch's cuda state, and if TPUs are available torch_xla's cuda state.
+
+## Observed Batch Sizes 
+
+When training with Accelerate, the batch size passed to the dataloader is the **batch size per GPU**. What this entails is 
+a batch size of 64 on two GPUs is truly a batch size of 128. As a result, when testing on a single GPU this needs to be accounted for,
+as well as similarly for TPUs. 
+
+The below table can be used as a quick reference to try out different batch sizes:
+
+<Tip>
+
+In this example there are two GPUs for "Multi-GPU" and a TPU pod with 8 workers
+
+</Tip>
+
+| Single GPU Batch Size | Multi-GPU Equivalent Batch Size | TPU Equivalent Batch Size |
+|-----------------------|---------------------------------|---------------------------|
+| 256                   | 128                             | 32                        |
+| 128                   | 64                              | 16                        |
+| 64                    | 32                              | 8                         |
+| 32                    | 16                              | 4                         |
+
+## Learning Rates 
+
+As noted in multiple sources[[1](https://aws.amazon.com/blogs/machine-learning/scalable-multi-node-deep-learning-training-using-gpus-in-the-aws-cloud/)][[2](https://docs.nvidia.com/clara/tlt-mi_archive/clara-train-sdk-v2.0/nvmidl/appendix/training_with_multiple_gpus.html)], the learning rate should be scaled *linearly* based on the number of devices present. The below 
+snippet shows doing so with Accelerate:
+
+<Tip>
+
+Since users can have their own learning rate schedulers defined, we leave this up to the user to decide if they wish to scale their 
+learning rate or not.
+ 
+</Tip>
+
+```python
+learning_rate = 1e-3
+accelerator = Accelerator()
+learning_rate *= accelerator.num_processes
+
+optimizer = AdamW(params=model.parameters(), lr=learning_rate)
+```
+

```suggestion
<Tip>

In this example there are two GPUs for "Multi-GPU" and a TPU pod with 8 workers

</Tip>
```
@@ -0,0 +1,87 @@
+<!--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.
+-->
+
+# Comparing performance between different device setups
+
+Evaluating and comparing the performance from different setups can be quite tricky if you don't know what to look for.
+For example, you cannot run the same script with the same batch size across TPU, multi-GPU, and single-GPU with Accelerate 
+and expect your results to line up. 
+
+But why?
+
+There's three reasons for this that this tutorial will cover: 
+
+1. **Setting the right seeds**
+2. **Observed Batch Sizes**
+3. **Learning Rates**
+
+## Setting the Seed 
+
+While this issue has not come up as much, make sure to use [`utils.set_seed`] to fully set the seed in all distributed cases so training will be reproducable:
+
+```python
+from accelerate import set_seed
+
+set_seed(42)
+```
+
+Why is this important? Under the hood this will set **5** different seed settings:
+
+```python
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    # ^^ safe to call this function even if cuda is not available
+    if is_tpu_available():
+        xm.set_rng_state(seed)
+```
+
+The random state, numpy's state, torch, torch's cuda state, and if TPUs are available torch_xla's cuda state.
+
+## Observed Batch Sizes 
+
+When training with Accelerate, the batch size passed to the dataloader is the **batch size per GPU**. What this entails is 
+a batch size of 64 on two GPUs is truly a batch size of 128. As a result, when testing on a single GPU this needs to be accounted for,
+as well as similarly for TPUs. 
+
+The below table can be used as a quick reference to try out different batch sizes:
+
+<Tip>
+    In this example there are two GPUs for "Multi-GPU" and a TPU pod with 8 workers
+</Tip>
```suggestion
<Tip>

Since users can have their own learning rate schedulers defined, we leave this up to the user to decide if they wish to scale their 
learning rate or not.
 
</Tip>
```
@@ -0,0 +1,87 @@
+<!--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.
+-->
+
+# Comparing performance between different device setups
+
+Evaluating and comparing the performance from different setups can be quite tricky if you don't know what to look for.
+For example, you cannot run the same script with the same batch size across TPU, multi-GPU, and single-GPU with Accelerate 
+and expect your results to line up. 
+
+But why?
+
+There's three reasons for this that this tutorial will cover: 
+
+1. **Setting the right seeds**
+2. **Observed Batch Sizes**
+3. **Learning Rates**
+
+## Setting the Seed 
+
+While this issue has not come up as much, make sure to use [`utils.set_seed`] to fully set the seed in all distributed cases so training will be reproducable:
+
+```python
+from accelerate import set_seed
+
+set_seed(42)
+```
+
+Why is this important? Under the hood this will set **5** different seed settings:
+
+```python
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    # ^^ safe to call this function even if cuda is not available
+    if is_tpu_available():
+        xm.set_rng_state(seed)
+```
+
+The random state, numpy's state, torch, torch's cuda state, and if TPUs are available torch_xla's cuda state.
+
+## Observed Batch Sizes 
+
+When training with Accelerate, the batch size passed to the dataloader is the **batch size per GPU**. What this entails is 
+a batch size of 64 on two GPUs is truly a batch size of 128. As a result, when testing on a single GPU this needs to be accounted for,
+as well as similarly for TPUs. 
+
+The below table can be used as a quick reference to try out different batch sizes:
+
+<Tip>
+    In this example there are two GPUs for "Multi-GPU" and a TPU pod with 8 workers
+</Tip>
+
+| Single GPU Batch Size | Multi-GPU Equivalent Batch Size | TPU Equivalent Batch Size |
+|-----------------------|---------------------------------|---------------------------|
+| 256                   | 128                             | 32                        |
+| 128                   | 64                              | 16                        |
+| 64                    | 32                              | 8                         |
+| 32                    | 16                              | 4                         |
+
+## Learning Rates 
+
+As noted in multiple sources[[1](https://aws.amazon.com/blogs/machine-learning/scalable-multi-node-deep-learning-training-using-gpus-in-the-aws-cloud/)][[2](https://docs.nvidia.com/clara/tlt-mi_archive/clara-train-sdk-v2.0/nvmidl/appendix/training_with_multiple_gpus.html)], the learning rate should be scaled *linearly* based on the number of devices present. The below 
+snippet shows doing so with Accelerate:
+
+<Tip>
+    Since users can have their own learning rate schedulers defined, we leave this up to the user to decide if they wish to scale their 
+    learning rate or not.
+</Tip>
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/734
diff --git a/docs/source/package_reference/cli.mdx b/docs/source/package_reference/cli.mdx
index 092b72850..a9e955d43 100644
--- a/docs/source/package_reference/cli.mdx
+++ b/docs/source/package_reference/cli.mdx
@@ -94,6 +94,7 @@ The following arguments are useful for customization of worker machines
 * `--machine_rank MACHINE_RANK` (`int`) -- The rank of the machine on which this script is launched.
 * `--num_machines NUM_MACHINES` (`int`) -- The total number of machines used in this training.
 * `--num_processes NUM_PROCESSES` (`int`) -- The total number of processes to be launched in parallel.
+* `--gpu_ids` (`str`) -- What GPUs (by id) should be used for training on this machine as a comma-seperated list
 * `--main_process_ip MAIN_PROCESS_IP` (`str`) -- The IP address of the machine of rank 0.
 * `--main_process_port MAIN_PROCESS_PORT` (`int`) -- The port to use to communicate with the machine of rank 0.
 * `--num_cpu_threads_per_process NUM_CPU_THREADS_PER_PROCESS` (`int`) -- The number of CPU threads per process. Can be tuned for optimal performance.
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index a174115b1..01c419f4f 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -35,6 +35,7 @@ def get_cluster_input():
 
     machine_rank = 0
     num_machines = 1
+    gpu_ids = None
     main_process_ip = None
     main_process_port = None
     rdzv_backend = "static"
@@ -293,6 +294,12 @@ def get_cluster_input():
             default=1,
             error_message="Please enter an integer.",
         )
+
+    if distributed_type in [DistributedType.MULTI_GPU, DistributedType.NO] and not use_cpu:
+        gpu_ids = _ask_field(
+            "What GPU(s) (by id) should be used for training on this machine as a comma-seperated list? [all]:",
+            default="all",
+        )
     elif distributed_type in [DistributedType.FSDP, DistributedType.DEEPSPEED]:
         num_processes = _ask_field(
             "How many GPU(s) should be used for distributed training? [1]:",
@@ -325,6 +332,7 @@ def get_cluster_input():
         compute_environment=ComputeEnvironment.LOCAL_MACHINE,
         distributed_type=distributed_type,
         num_processes=num_processes,
+        gpu_ids=gpu_ids,
         mixed_precision=mixed_precision,
         downcast_bf16=downcast_bf16,
         machine_rank=machine_rank,
diff --git a/src/accelerate/commands/config/config_args.py b/src/accelerate/commands/config/config_args.py
index 68da64963..43faf45ca 100644
--- a/src/accelerate/commands/config/config_args.py
+++ b/src/accelerate/commands/config/config_args.py
@@ -135,6 +135,7 @@ class ClusterConfig(BaseConfig):
     num_processes: int
     machine_rank: int = 0
     num_machines: int = 1
+    gpu_ids: Optional[str] = None
     main_process_ip: Optional[str] = None
     main_process_port: Optional[int] = None
     rdzv_backend: Optional[str] = "static"
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index d1905d1c0..982293a33 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -259,6 +259,11 @@ def launch_command_parser(subparsers=None):
     parser.add_argument(
         "--num_machines", type=int, default=None, help="The total number of machines used in this training."
     )
+    parser.add_argument(
+        "--gpu_ids",
+        default=None,
+        help="What GPUs (by id) should be used for training on this machine as a comma-seperated list",
+    )
     parser.add_argument(
         "--machine_rank", type=int, default=None, help="The rank of the machine on which this script is launched."
     )
@@ -366,6 +371,8 @@ def simple_launcher(args):
     current_env["USE_MPS_DEVICE"] = str(args.use_mps_device)
     if args.use_mps_device:
         current_env["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+    elif args.gpu_ids != "all":
+        current_env["CUDA_VISIBLE_DEVICES"] = args.gpu_ids
     if args.num_machines > 1:
         current_env["MASTER_ADDR"] = args.main_process_ip
         current_env["MASTER_PORT"] = str(args.main_process_port)
@@ -420,6 +427,9 @@ def multi_gpu_launcher(args):
         setattr(args, "no_python", True)
 
     current_env = os.environ.copy()
+    gpu_ids = getattr(args, "gpu_ids")
+    if gpu_ids != "all":
+        current_env["CUDA_VISIBLE_DEVICES"] = gpu_ids
     mixed_precision = args.mixed_precision.lower()
     try:
         mixed_precision = PrecisionType(mixed_precision)
@@ -549,6 +559,9 @@ def deepspeed_launcher(args):
         setattr(args, "no_python", True)
 
     current_env = os.environ.copy()
+    gpu_ids = getattr(args, "gpu_ids")
+    if gpu_ids != "all":
+        current_env["CUDA_VISIBLE_DEVICES"] = gpu_ids
     try:
         mixed_precision = PrecisionType(args.mixed_precision.lower())
     except ValueError:
@@ -817,6 +830,14 @@ def launch_command(args):
             args.tpu = defaults.distributed_type == DistributedType.TPU
             args.use_fsdp = defaults.distributed_type == DistributedType.FSDP
             args.use_mps_device = defaults.distributed_type == DistributedType.MPS
+        if not args.use_mps_device:
+            if args.gpu_ids is None:
+                if defaults.gpu_ids is not None:
+                    args.gpu_ids = defaults.gpu_ids
+                else:
+                    args.gpu_ids = "all"
+            if len(args.gpu_ids.split(",")) < 2 and args.multi_gpu and (args.gpu_ids != "all"):
+                args.multi_gpu = False
         if defaults.compute_environment == ComputeEnvironment.LOCAL_MACHINE:
             # Update args with the defaults
             for name, attr in defaults.__dict__.items():

I meant with respect to these changes

@@ -817,6 +830,13 @@ def launch_command(args):
             args.tpu = defaults.distributed_type == DistributedType.TPU
             args.use_fsdp = defaults.distributed_type == DistributedType.FSDP
             args.use_mps_device = defaults.distributed_type == DistributedType.MPS
+        if args.gpu_ids is None:
+            if defaults.gpu_ids is not None:
+                args.gpu_ids = defaults.gpu_ids
+            else:
+                args.gpu_ids = "all"
+        if len(args.gpu_ids.split(",")) < 2 and args.multi_gpu and (args.gpu_ids != "all"):
+            args.multi_gpu = False
Let me know what you think about my solution here: https://github.com/huggingface/accelerate/pull/732/commits/0467d4ff0bacda7ffb0d6087a7ef636cfbbb4960
@@ -817,6 +830,13 @@ def launch_command(args):
             args.tpu = defaults.distributed_type == DistributedType.TPU
             args.use_fsdp = defaults.distributed_type == DistributedType.FSDP
             args.use_mps_device = defaults.distributed_type == DistributedType.MPS
+        if args.gpu_ids is None:
+            if defaults.gpu_ids is not None:
+                args.gpu_ids = defaults.gpu_ids
+            else:
+                args.gpu_ids = "all"
+        if len(args.gpu_ids.split(",")) < 2 and args.multi_gpu and (args.gpu_ids != "all"):
+            args.multi_gpu = False
Nice! This should have expected behaviour. 🤗
@@ -817,6 +830,13 @@ def launch_command(args):
             args.tpu = defaults.distributed_type == DistributedType.TPU
             args.use_fsdp = defaults.distributed_type == DistributedType.FSDP
             args.use_mps_device = defaults.distributed_type == DistributedType.MPS
+        if args.gpu_ids is None:
+            if defaults.gpu_ids is not None:
+                args.gpu_ids = defaults.gpu_ids
+            else:
+                args.gpu_ids = "all"
+        if len(args.gpu_ids.split(",")) < 2 and args.multi_gpu and (args.gpu_ids != "all"):
+            args.multi_gpu = False
_The documentation is not available anymore as the PR was closed or merged._
@pacman100 correct, this is solely for non-M1 GPUs, I'll specify that in the info description for it
https://api.github.com/repos/huggingface/accelerate/pulls/732
diff --git a/src/accelerate/scheduler.py b/src/accelerate/scheduler.py
index 4ae91a657..835d4e0d9 100644
--- a/src/accelerate/scheduler.py
+++ b/src/accelerate/scheduler.py
@@ -69,7 +69,9 @@ def step(self, *args, **kwargs):
             num_processes = AcceleratorState().num_processes
             for _ in range(num_processes):
                 # Special case when using OneCycle and `drop_last` was not used
-                if getattr(self.scheduler, "total_steps", 0) <= self.scheduler.last_epoch:
+                if hasattr(self.scheduler, "total_steps") and self.scheduler._step_count <= self.scheduler.total_steps:
+                    self.scheduler.step(*args, **kwargs)
+                else:
                     self.scheduler.step(*args, **kwargs)
 
     # Passthroughs
diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py
index be4f975fb..c1ef18f1e 100644
--- a/tests/test_scheduler.py
+++ b/tests/test_scheduler.py
@@ -21,12 +21,30 @@
 from accelerate.test_utils import require_cpu
 
 
-def scheduler_test(num_processes=2, step_scheduler_with_optimizer=True, split_batches=False):
+def one_cycle_test(num_processes=2, step_scheduler_with_optimizer=True, split_batches=False):
     accelerator = Accelerator(step_scheduler_with_optimizer=step_scheduler_with_optimizer, split_batches=split_batches)
     model = torch.nn.Linear(2, 4)
     optimizer = torch.optim.AdamW(model.parameters(), lr=1.0)
-    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda n: 1 - n / 10)
+    scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=0.01, steps_per_epoch=2, epochs=1)
+    model, optimizer, scheduler = accelerator.prepare(model, optimizer, scheduler)
 
+    # Optimizer has stepped
+    scheduler.step()
+    if step_scheduler_with_optimizer or (num_processes == 1):
+        assert (
+            scheduler.scheduler.last_epoch == num_processes
+        ), f"Last Epoch ({scheduler.scheduler.last_epoch}) != Num Processes ({num_processes})"
+    else:
+        assert (
+            scheduler.scheduler.last_epoch != num_processes
+        ), f"Last Epoch ({scheduler.scheduler.last_epoch}) == Num Processes ({num_processes})"
+
+
+def lambda_test(num_processes=2, step_scheduler_with_optimizer=True, split_batches=False):
+    accelerator = Accelerator(step_scheduler_with_optimizer=step_scheduler_with_optimizer, split_batches=split_batches)
+    model = torch.nn.Linear(2, 4)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=1.0)
+    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda n: 1 - n / 10)
     model, optimizer, scheduler = accelerator.prepare(model, optimizer, scheduler)
 
     # Optimizer has stepped
@@ -49,16 +67,30 @@ def scheduler_test(num_processes=2, step_scheduler_with_optimizer=True, split_ba
 
 @require_cpu
 class SchedulerTester(unittest.TestCase):
-    def test_scheduler_steps_with_optimizer_single_process(self):
-        debug_launcher(partial(scheduler_test, num_processes=1), num_processes=1)
-        debug_launcher(partial(scheduler_test, num_processes=1, split_batches=True), num_processes=1)
+    def test_lambda_scheduler_steps_with_optimizer_single_process(self):
+        debug_launcher(partial(lambda_test, num_processes=1), num_processes=1)
+        debug_launcher(partial(lambda_test, num_processes=1, split_batches=True), num_processes=1)
+
+    def test_one_cycle_scheduler_steps_with_optimizer_single_process(self):
+        debug_launcher(partial(one_cycle_test, num_processes=1), num_processes=1)
+        debug_launcher(partial(one_cycle_test, num_processes=1, split_batches=True), num_processes=1)
+
+    def test_lambda_scheduler_not_step_with_optimizer_single_process(self):
+        debug_launcher(partial(lambda_test, num_processes=1, step_scheduler_with_optimizer=False), num_processes=1)
+
+    def test_one_cycle_scheduler_not_step_with_optimizer_single_process(self):
+        debug_launcher(partial(one_cycle_test, num_processes=1, step_scheduler_with_optimizer=False), num_processes=1)
+
+    def test_lambda_scheduler_steps_with_optimizer_multiprocess(self):
+        debug_launcher(lambda_test)
+        debug_launcher(partial(lambda_test, num_processes=1, split_batches=True), num_processes=1)
 
-    def test_scheduler_not_step_with_optimizer_single_process(self):
-        debug_launcher(partial(scheduler_test, num_processes=1, step_scheduler_with_optimizer=False), num_processes=1)
+    def test_one_cycle_scheduler_steps_with_optimizer_multiprocess(self):
+        debug_launcher(one_cycle_test)
+        debug_launcher(partial(one_cycle_test, num_processes=1, split_batches=True), num_processes=1)
 
-    def test_scheduler_steps_with_optimizer_multiprocess(self):
-        debug_launcher(scheduler_test)
-        debug_launcher(partial(scheduler_test, num_processes=1, split_batches=True), num_processes=1)
+    def test_lambda_scheduler_not_step_with_optimizer_multiprocess(self):
+        debug_launcher(partial(lambda_test, step_scheduler_with_optimizer=False))
 
-    def test_scheduler_not_step_with_optimizer_multiprocess(self):
-        debug_launcher(partial(scheduler_test, step_scheduler_with_optimizer=False))
+    def test_one_cycle_scheduler_not_step_with_optimizer_multiprocess(self):
+        debug_launcher(partial(one_cycle_test, step_scheduler_with_optimizer=False))

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/728
diff --git a/.github/workflows/build_and_run_tests.yml b/.github/workflows/build_and_run_tests.yml
index a5378353d..abf9f3883 100644
--- a/.github/workflows/build_and_run_tests.yml
+++ b/.github/workflows/build_and_run_tests.yml
@@ -34,8 +34,8 @@ jobs:
           done
           
   build-docker-containers:
-    needs: check-for-setup
-    if: (github.event_name == 'push') && (needs.check-for-setup.outputs.changed == '1')
+    needs: check-for-source
+    if: (github.event_name == 'push') && (needs.check-for-source.outputs.changed == '1')
     uses: ./.github/workflows/build_docker_images.yml
     secrets: inherit
 

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/727
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index bf57d11a4..d1905d1c0 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -844,8 +844,8 @@ def launch_command(args):
                 args.mixed_precision = defaults.mixed_precision
     else:
         if args.num_processes is None:
-            warned.append("\t`--num_processes` was set to a value of `1`")
-            args.num_processes = 1
+            args.num_processes = torch.cuda.device_count() if args.multi_gpu else 1
+            warned.append("\t`--num_processes` was set to a value of `{args.num_processes}`")
         if args.num_machines is None:
             warned.append("\t`--num_machines` was set to a value of `1`")
             args.num_machines = 1
@@ -854,14 +854,6 @@ def launch_command(args):
             args.mixed_precision = "no"
         if not hasattr(args, "use_cpu"):
             args.use_cpu = args.cpu
-    if args.multi_gpu and args.num_processes == 1:
-        args.num_processes = torch.cuda.device_count()
-        if not any("--num_processes" in warn for warn in warned):
-            warned.append(f"\t`--num_processes` was set to `{args.num_processes}`")
-        else:
-            for i, warn in enumerate(warned):
-                if "--num_processes" in warn:
-                    warned[i] = warn.replace("`1`", f"`{args.num_processes}`")
 
     if args.num_cpu_threads_per_process is None:
         local_size = get_int_from_env(

```suggestion
            args.num_processes = torch.cuda.device_count() if args.multi_gpu else 1
            warned.append("\t`--num_processes` was set to a value of `{args.num_processes}`")
```
@@ -845,7 +845,7 @@ def launch_command(args):
     else:
         if args.num_processes is None:
             warned.append("\t`--num_processes` was set to a value of `1`")
```suggestion
```
@@ -845,7 +845,7 @@ def launch_command(args):
     else:
         if args.num_processes is None:
             warned.append("\t`--num_processes` was set to a value of `1`")
-            args.num_processes = 1
+            args.num_processes = torch.cuda.device_count() if args.multi_gpu else 1
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/726
diff --git a/.github/workflows/build-docker-images-release.yml b/.github/workflows/build-docker-images-release.yml
new file mode 100644
index 000000000..654259f87
--- /dev/null
+++ b/.github/workflows/build-docker-images-release.yml
@@ -0,0 +1,64 @@
+name: Build Docker images (releases)
+
+on:
+  workflow_dispatch:
+  release:
+    types: [published]
+
+concurrency:
+  group: docker-image-builds
+  cancel-in-progress: false
+
+jobs:
+  get-version:
+    runs-on: ubuntu-latest
+    outputs:
+      version: ${{ steps.step1.outputs.version }}
+    steps:
+      - uses: actions/checkout@v3
+      - id: step1
+        run: echo "::set-output name=version::$(python setup.py --version)"
+
+  version-cpu:
+    name: "Latest Accelerate CPU [version]"
+    runs-on: ubuntu-latest
+    needs: get-version
+    steps:
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v1
+      - name: Check out code
+        uses: actions/checkout@v2
+      - name: Login to DockerHub
+        uses: docker/login-action@v1
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_PASSWORD }}
+
+      - name: Build and Push CPU
+        uses: docker/build-push-action@v2
+        with:
+          context: ./docker/accelerate-cpu
+          push: true
+          tags: huggingface/accelerate-cpu:${{needs.get-version.outputs.version}}
+
+  version-cuda:
+    name: "Latest Accelerate GPU [version]"
+    runs-on: ubuntu-latest
+    needs: get-version
+    steps:
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v1
+      - name: Check out code
+        uses: actions/checkout@v2
+      - name: Login to DockerHub
+        uses: docker/login-action@v1
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_PASSWORD }}
+
+      - name: Build and Push GPU
+        uses: docker/build-push-action@v2
+        with:
+          context: ./docker/accelerate-gpu
+          push: true
+          tags: huggingface/accelerate-gpu:${{needs.get-version.outputs.version}}
\ No newline at end of file

_The documentation is not available anymore as the PR was closed or merged._
> FFT, should we just merge these two release workflows together and check if 'dev' is in the version to automatically tag the `latest` tag? Or maybe just publish the prerelease tags? WDYT?

I guess you are talking about the existing `build_docker_images.yml` here. I don't have particular opinion about if merging these 2 workflow files or not.
@ydshieh by other folks, not for testing. We've had people use the latest images already, so made sense to also do releases 

Okay, will think on this and do a follow up most likely with the merge 
https://api.github.com/repos/huggingface/accelerate/pulls/725
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index 1eb8ff79d..fac1b50e2 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -347,7 +347,6 @@ def __iter__(self):
         try:
             current_batch = next(dataloader_iter)
         except StopIteration:
-            self.gradient_state._iterate_samples_seen(find_batch_size(current_batch))
             yield
         while True:
             try:

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/724
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index ab4b48660..2b6558944 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -619,6 +619,12 @@ def prepare(self, *args, device_placement=None):
             device_placement (`List[bool]`, *optional*):
                 Used to customize whether automatic device placement should be performed for each object passed. Needs
                 to be a list of the same length as `args`.
+
+        <Tip>
+
+          You don't need to prepare a model if you only use it for inference without any kind of mixed precision
+
+        </Tip>
         """
         if device_placement is None:
             device_placement = [None for _ in args]
@@ -699,7 +705,8 @@ def prepare_model(self, model: torch.nn.Module, device_placement=None):
 
         Args:
             model (`torch.nn.Module`):
-                A PyTorch model to prepare
+                A PyTorch model to prepare. You don't need to prepare a model if it is used only for inference without
+                any kind of mixed precision
             device_placement (`bool`, *optional*):
                 Whether or not to place the model on the proper device. Will default to `self.device_placement`.
         """

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/722
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index cd74b68e2..ab4b48660 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -547,15 +547,15 @@ def print(self, *args, **kwargs):
         if self.is_local_main_process:
             print(*args, **kwargs)
 
-    def _prepare_one(self, obj, first_pass=False):
+    def _prepare_one(self, obj, first_pass=False, device_placement=None):
         # First pass of preparation: DataLoader, model, optimizer
         if first_pass:
             if isinstance(obj, torch.utils.data.DataLoader):
-                return self.prepare_data_loader(obj)
+                return self.prepare_data_loader(obj, device_placement=device_placement)
             elif isinstance(obj, torch.nn.Module):
-                return self.prepare_model(obj)
+                return self.prepare_model(obj, device_placement=device_placement)
             elif isinstance(obj, torch.optim.Optimizer):
-                optimizer = self.prepare_optimizer(obj)
+                optimizer = self.prepare_optimizer(obj, device_placement=device_placement)
                 return optimizer
         # Second pass of preparation: LR scheduler (which need the full list of optimizers)
         elif isinstance(obj, torch.optim.lr_scheduler._LRScheduler):
@@ -602,17 +602,33 @@ def _prepare_fsdp(self, *args):
         self._optimizers = optimizers
         return tuple(result)
 
-    def prepare(self, *args):
+    def prepare(self, *args, device_placement=None):
         """
         Prepare all objects passed in `args` for distributed training and mixed precision, then return them in the same
         order.
 
-        Accepts the following type of objects:
+        Args:
+            *args (list of objects):
+                Any of the following type of objects:
+
+                - `torch.utils.data.DataLoader`: PyTorch Dataloader
+                - `torch.nn.Module`: PyTorch Module
+                - `torch.optim.Optimizer`: PyTorch Optimizer
+                - `torch.optim.lr_scheduler._LRScheduler`: PyTorch LR Scheduler
 
-            - `torch.utils.data.DataLoader`: PyTorch Dataloader
-            - `torch.nn.Module`: PyTorch Module
-            - `torch.optim.Optimizer`: PyTorch Optimizer
+            device_placement (`List[bool]`, *optional*):
+                Used to customize whether automatic device placement should be performed for each object passed. Needs
+                to be a list of the same length as `args`.
         """
+        if device_placement is None:
+            device_placement = [None for _ in args]
+        elif self.distributed_type == DistributedType.DEEPSPEED:
+            raise ValueError("You can't customize device placements with DeepSpeed.")
+        elif len(device_placement) != len(args):
+            raise ValueError(
+                f"`device_placement` should be a list with {len(args)} elements (the number of objects passed)."
+            )
+
         if self.distributed_type == DistributedType.FSDP:
             model_count = 0
             optimizer_present = False
@@ -656,8 +672,10 @@ def prepare(self, *args):
         if self.distributed_type == DistributedType.DEEPSPEED:
             result = self._prepare_deepspeed(*args)
         else:
-            result = tuple(self._prepare_one(obj, first_pass=True) for obj in args)
-            result = tuple(self._prepare_one(obj) for obj in result)
+            result = tuple(
+                self._prepare_one(obj, first_pass=True, device_placement=d) for obj, d in zip(args, device_placement)
+            )
+            result = tuple(self._prepare_one(obj, device_placement=d) for obj, d in zip(result, device_placement))
 
         if tpu_should_fix_optimizer:
             # 2. grabbing new model parameters
@@ -674,7 +692,7 @@ def prepare(self, *args):
 
         return result if len(result) > 1 else result[0]
 
-    def prepare_model(self, model: torch.nn.Module):
+    def prepare_model(self, model: torch.nn.Module, device_placement=None):
         """
         Prepares a PyTorch model for training in any distributed setup. It is recommended to use
         [`Accelerator.prepare`] instead.
@@ -682,9 +700,13 @@ def prepare_model(self, model: torch.nn.Module):
         Args:
             model (`torch.nn.Module`):
                 A PyTorch model to prepare
+            device_placement (`bool`, *optional*):
+                Whether or not to place the model on the proper device. Will default to `self.device_placement`.
         """
+        if device_placement is None:
+            device_placement = self.device_placement and self.distributed_type != DistributedType.FSDP
         self._models.append(model)
-        if self.device_placement and self.distributed_type != DistributedType.FSDP:
+        if device_placement:
             model = model.to(self.device)
         if self.distributed_type == DistributedType.MULTI_GPU:
             kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {}
@@ -894,7 +916,7 @@ def _prepare_deepspeed(self, *args):
                 )
         return tuple(result)
 
-    def prepare_data_loader(self, data_loader: torch.utils.data.DataLoader):
+    def prepare_data_loader(self, data_loader: torch.utils.data.DataLoader, device_placement=None):
         """
         Prepares a PyTorch DataLoader for training in any distributed setup. It is recommended to use
         [`Accelerator.prepare`] instead.
@@ -902,19 +924,24 @@ def prepare_data_loader(self, data_loader: torch.utils.data.DataLoader):
         Args:
             data_loader (`torch.utils.data.DataLoader`):
                 A vanilla PyTorch DataLoader to prepare
+            device_placement (`bool`, *optional*):
+                Whether or not to place the batches on the proper device in the prepared dataloader. Will default to
+                `self.device_placement`.
         """
+        if device_placement is None:
+            device_placement = self.device_placement if self.distributed_type != DistributedType.TPU else False
         return prepare_data_loader(
             data_loader,
             self.device,
             num_processes=self.num_processes,
             process_index=self.process_index,
             split_batches=self.split_batches,
-            put_on_device=self.device_placement if self.distributed_type != DistributedType.TPU else False,
+            put_on_device=device_placement,
             rng_types=self.rng_types.copy(),
             dispatch_batches=self.dispatch_batches,
         )
 
-    def prepare_optimizer(self, optimizer: torch.optim.Optimizer):
+    def prepare_optimizer(self, optimizer: torch.optim.Optimizer, device_placement=None):
         """
         Prepares a PyTorch Optimizer for training in any distributed setup. It is recommended to use
         [`Accelerator.prepare`] instead.
@@ -922,8 +949,12 @@ def prepare_optimizer(self, optimizer: torch.optim.Optimizer):
         Args:
             optimizer (`torch.optim.Optimizer`):
                 A vanilla PyTorch optimizer to prepare
+            device_placement (`bool`, *optional*):
+                Whether or not to place the optimizer on the proper device. Will default to `self.device_placement`.
         """
-        optimizer = AcceleratedOptimizer(optimizer, device_placement=self.device_placement, scaler=self.scaler)
+        if device_placement is None:
+            device_placement = self.device_placement
+        optimizer = AcceleratedOptimizer(optimizer, device_placement=device_placement, scaler=self.scaler)
         self._optimizers.append(optimizer)
         return optimizer
 

Perhaps we can be a bit more specific here:

```suggestion
            device_placement (`List[bool]`, *optional*):
                Used to customize whether automatic device placement should be performed for each object passed. Needs to be a list of the same length as
                `args`.
```
@@ -602,17 +602,32 @@ def _prepare_fsdp(self, *args):
         self._optimizers = optimizers
         return tuple(result)
 
-    def prepare(self, *args):
+    def prepare(self, *args, device_placement=None):
         """
         Prepare all objects passed in `args` for distributed training and mixed precision, then return them in the same
         order.
 
-        Accepts the following type of objects:
+        Args:
+            *args:
+                Any of the following type of objects:
+
+                - `torch.utils.data.DataLoader`: PyTorch Dataloader
+                - `torch.nn.Module`: PyTorch Module
+                - `torch.optim.Optimizer`: PyTorch Optimizer
+                - `torch.optim.lr_scheduler._LRScheduler`: PyTorch LR Scheduler
+            device_placement (`List[bool]`, *optional*):
+                Used to customize the device placement for each object passed. Needs to be a list of the same length as
+                `args`.
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/716
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index bb58ca57f..cd74b68e2 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -674,7 +674,15 @@ def prepare(self, *args):
 
         return result if len(result) > 1 else result[0]
 
-    def prepare_model(self, model):
+    def prepare_model(self, model: torch.nn.Module):
+        """
+        Prepares a PyTorch model for training in any distributed setup. It is recommended to use
+        [`Accelerator.prepare`] instead.
+
+        Args:
+            model (`torch.nn.Module`):
+                A PyTorch model to prepare
+        """
         self._models.append(model)
         if self.device_placement and self.distributed_type != DistributedType.FSDP:
             model = model.to(self.device)
@@ -886,7 +894,15 @@ def _prepare_deepspeed(self, *args):
                 )
         return tuple(result)
 
-    def prepare_data_loader(self, data_loader):
+    def prepare_data_loader(self, data_loader: torch.utils.data.DataLoader):
+        """
+        Prepares a PyTorch DataLoader for training in any distributed setup. It is recommended to use
+        [`Accelerator.prepare`] instead.
+
+        Args:
+            data_loader (`torch.utils.data.DataLoader`):
+                A vanilla PyTorch DataLoader to prepare
+        """
         return prepare_data_loader(
             data_loader,
             self.device,
@@ -898,12 +914,28 @@ def prepare_data_loader(self, data_loader):
             dispatch_batches=self.dispatch_batches,
         )
 
-    def prepare_optimizer(self, optimizer):
+    def prepare_optimizer(self, optimizer: torch.optim.Optimizer):
+        """
+        Prepares a PyTorch Optimizer for training in any distributed setup. It is recommended to use
+        [`Accelerator.prepare`] instead.
+
+        Args:
+            optimizer (`torch.optim.Optimizer`):
+                A vanilla PyTorch optimizer to prepare
+        """
         optimizer = AcceleratedOptimizer(optimizer, device_placement=self.device_placement, scaler=self.scaler)
         self._optimizers.append(optimizer)
         return optimizer
 
-    def prepare_scheduler(self, scheduler):
+    def prepare_scheduler(self, scheduler: torch.optim.lr_scheduler._LRScheduler):
+        """
+        Prepares a PyTorch Scheduler for training in any distributed setup. It is recommended to use
+        [`Accelerator.prepare`] instead.
+
+        Args:
+            scheduler (`torch.optim.lr_scheduler._LRScheduler`):
+                A vanilla PyTorch scheduler to prepare
+        """
         # We try to find the optimizer associated with `scheduler`, the default is the full list.
         optimizer = self._optimizers
         for opt in self._optimizers:
@@ -1133,7 +1165,7 @@ def init_trackers(self, project_name: str, config: Optional[dict] = None, init_k
                 Optional starting configuration to be logged.
             init_kwargs (`dict`, *optional*):
                 A nested dictionary of kwargs to be passed to a specific tracker's `__init__` function. Should be
-                formatted like this:
+                formatted like so:
                 ```python
                 {"wandb": {"tags": ["tag_a", "tag_b"]}}
                 ```
@@ -1182,7 +1214,7 @@ def log(self, values: dict, step: Optional[int] = None, log_kwargs: Optional[dic
                 The run step. If included, the log will be affiliated with this step.
             log_kwargs (`dict`, *optional*):
                 A nested dictionary of kwargs to be passed to a specific tracker's `log` function. Should be formatted
-                like this:
+                like so:
                 ```python
                 {"wandb": {"tags": ["tag_a", "tag_b"]}}
                 ```
@@ -1193,7 +1225,8 @@ def log(self, values: dict, step: Optional[int] = None, log_kwargs: Optional[dic
     @on_main_process
     def end_training(self):
         """
-        Runs any special end training behaviors, such as stopping trackers on the main process only.
+        Runs any special end training behaviors, such as stopping trackers on the main process only. Should always be
+        called at the end of your script if using experiment tracking.
         """
         for tracker in self.trackers:
             tracker.finish()
@@ -1383,6 +1416,15 @@ def _get_devices(self, *args):
         return (model_device, optimizer_device)
 
     def get_state_dict(self, model, unwrap=True):
+        """
+        Returns the state dictionary of a model sent through [`Accelerator.prepare`] in full precision
+
+        Args:
+            model (`torch.nn.Module`):
+                A PyTorch model sent through [`Accelerator.prepare`]
+            unwrap (`bool`, *optional*, defaults to True):
+                Whether to return the original underlying state_dict of `model` or to return the wrapped state_dict
+        """
         is_zero_3 = False
         if self.distributed_type == DistributedType.DEEPSPEED:
             is_zero_3 = self.deepspeed_config["zero_optimization"]["stage"] == 3

I think the other types are all subclasses of `Module`
```suggestion
            model (`torch.nn.Module`):
```
@@ -1383,6 +1416,15 @@ def _get_devices(self, *args):
         return (model_device, optimizer_device)
 
     def get_state_dict(self, model, unwrap=True):
+        """
+        Returns the state dictionary of a model sent through [`Accelerator.prepare`] in full precision
+
+        Args:
+            model (`Module`, `DistributedDataParallel`, `FullyShardedDataParallel`, `MpModelWrapper`):
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/715
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 552c8392c..bb58ca57f 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -125,8 +125,8 @@ class Accelerator:
             - `"tensorboard"`
             - `"wandb"`
             - `"comet_ml"`
-            If `"all`" is selected, will pick up all available trackers in the environment and intialize them. Can also
-            accept implementations of `GeneralTracker` for custom trackers, and can be combined with `"all"`.
+            If `"all"` is selected, will pick up all available trackers in the environment and initialize them. Can
+            also accept implementations of `GeneralTracker` for custom trackers, and can be combined with `"all"`.
         logging_dir (`str`, `os.PathLike`, *optional*):
             A path to a directory for storing logs of locally-compatible loggers.
         dispatch_batches (`bool`, *optional*):
@@ -643,7 +643,7 @@ def prepare(self, *args):
                     "The model and the optimizer parameters are not on the same device, which probably means you "
                     "created an optimizer around your model **before** putting on the device. Make sure the line "
                     "model.to(device) is before the optimizer creation in your script or remove it entirely and use "
-                    "the flag default value for `devicement_placement` in your `Accelerator` to let it handle that "
+                    "the flag default value for `device_placement` in your `Accelerator` to let it handle that "
                     "part for you."
                 )
 
diff --git a/src/accelerate/big_modeling.py b/src/accelerate/big_modeling.py
index 3c67b6ee5..9d57c144f 100644
--- a/src/accelerate/big_modeling.py
+++ b/src/accelerate/big_modeling.py
@@ -44,7 +44,7 @@ def init_empty_weights(include_buffers: bool = False):
 
     Example:
 
-    ```pyton
+    ```python
     import torch.nn as nn
     from accelerate import init_empty_weights
 
@@ -324,7 +324,7 @@ def load_checkpoint_and_dispatch(
         dtype (`str` or `torch.dtype`, *optional*):
             If provided, the weights will be converted to that type when loaded.
         offload_state_dict (`bool`, *optional*):
-            If `True`, will temporarily offload the CPU state dict on the hard drive to avoig getting out of CPU RAM if
+            If `True`, will temporarily offload the CPU state dict on the hard drive to avoid getting out of CPU RAM if
             the weight of the CPU state dict + the biggest shard does not fit. Will default to `True` if the device map
             picked contains `"disk"` values.
         preload_module_classes (`List[str]`, *optional*):
diff --git a/src/accelerate/tracking.py b/src/accelerate/tracking.py
index ee6d194f1..a438ab45b 100644
--- a/src/accelerate/tracking.py
+++ b/src/accelerate/tracking.py
@@ -424,8 +424,8 @@ def filter_trackers(
             - `"tensorboard"`
             - `"wandb"`
             - `"comet_ml"`
-            If `"all`" is selected, will pick up all available trackers in the environment and intialize them. Can also
-            accept implementations of `GeneralTracker` for custom trackers, and can be combined with `"all"`.
+            If `"all"` is selected, will pick up all available trackers in the environment and initialize them. Can
+            also accept implementations of `GeneralTracker` for custom trackers, and can be combined with `"all"`.
         logging_dir (`str`, `os.PathLike`, *optional*):
             A path to a directory for storing logs of locally-compatible loggers.
     """
diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
index 467a757b4..0b57ae6d8 100644
--- a/src/accelerate/utils/modeling.py
+++ b/src/accelerate/utils/modeling.py
@@ -617,7 +617,7 @@ def load_checkpoint_in_model(
         dtype (`str` or `torch.dtype`, *optional*):
             If provided, the weights will be converted to that type when loaded.
         offload_state_dict (`bool`, *optional*, defaults to `False`):
-            If `True`, will temporarily offload the CPU state dict on the hard drive to avoig getting out of CPU RAM if
+            If `True`, will temporarily offload the CPU state dict on the hard drive to avoid getting out of CPU RAM if
             the weight of the CPU state dict + the biggest shard does not fit.
     """
     if offload_folder is None and device_map is not None and "disk" in device_map.values():

_The documentation is not available anymore as the PR was closed or merged._
Hey @sgugger ,

Trying to run your Makefile has proven really tough for me:
1. `python` doesn't exist on ubuntu. Setup a symlink for `python` -> `python3` ...

2. `python3.9` running `make style` tries to change  30 files
`30 files reformatted, 55 files left unchanged.`

3. Tried to match the runner and did a `docker run ... 3.7.... -v .:/accelerate` ...
This still gave different results than the CI runner:
`5 files reformatted, 80 files left unchanged.`

I'm unable to actually get my black/isort to match what the test unhappy about.

It appears the test is failing for a line length style issue from a file that wasn't actually touched in this commit... 

It feels like it should have tripped on #711  -- not here. 

```
python utils/style_doc.py src/accelerate docs/source --max_len 119 --check_only
Traceback (most recent call last):
  File "utils/style_doc.py", line 556, in <module>
    main(*args.files, max_len=args.max_len, check_only=args.check_only)
  File "utils/style_doc.py", line 544, in main
    raise ValueError(f"{len(changed)} files should be restyled!")
```

Can we just merge this in, and fix that style next time you lint?

Or could you jump in on this one and help run the `make style` to get that file to format correctly based on your black/isort spec?

Is there a docker image or something setup to run your `Makefile` in a reproducible environment? I'm probably missing something...
Hey @sgugger  --

This is jumping the shark, but I do have a (arguably) nice pre-commit template that handles autoflake, isort, black and works well cross-platform.

Here's a screenshot. If you want to re-do the makefile, I can start a new ticket and switch it to pre-commit.

![image](https://user-images.githubusercontent.com/523300/191733560-c07f973d-5c9c-4087-a8b0-6303ba028169.png)

This also checks for conventional commits format(optional).

This would be a bit bigger in scope and more opinionated than what I thought I was starting with some readability PR's :)

Let me know if this is something that's of interest to swapping out for the huggingface team... if so, we can throw it in a new issue.

Hi @ryanrussell, I think the issue here is the wrong black/flake8 may be installed? Can you try running `make style` again after doing `pip install -e .[quality]` (or `pip install accelerate[quality]` depending on how you have it installed)?
Hi @ryanrussell. 90% on the team is on Ubuntu and have no problem with running `make style` without doing any aliases. Not sure why you don't have the `python` command available after installing Python.

Then as @muellerzr said, you need to have the pinned versions we use for the libraries formatting the code, which is why we recommend using a virtual environment. I've pushed a style commit to move forward with this PR.

I don't want to add pre-commit hooks as I prefer having the ability to do the styling in a separate commit when I'm afraid it might break something.
Hey @sgugger -- thanks for helping fix this. Sounds like it's a user error on my side.

I'll wait to send any more PR's until I setup a venv with correct version pins and can replicate the expected linting from the CI...

Just found the discord, so joining there... if I run into any issues, much easier to iterate there than through comments :+1: 
Failures are unrelated (something in the dev of evaluate messed up) so merging!
https://api.github.com/repos/huggingface/accelerate/pulls/713
diff --git a/src/accelerate/hooks.py b/src/accelerate/hooks.py
index 5c04b8621..493444daa 100644
--- a/src/accelerate/hooks.py
+++ b/src/accelerate/hooks.py
@@ -71,7 +71,7 @@ def post_forward(self, module, output):
 
     def detach_hook(self, module):
         """
-        To be executed when the hook is deached from a module.
+        To be executed when the hook is detached from a module.
 
         Args:
             module (`torch.nn.Module`): The module detached from this hook.
@@ -182,7 +182,7 @@ class AlignDevicesHook(ModelHook):
     Args:
         execution_device (`torch.device`, *optional*):
             The device on which inputs and model weights should be placed before the forward pass.
-        offload (`bool`, *optional*, defauts to `False`):
+        offload (`bool`, *optional*, defaults to `False`):
             Whether or not the weights should be offloaded after the forward pass.
         io_same_device (`bool`, *optional*, defaults to `False`):
             Whether or not the output should be placed on the same device as the input was.
@@ -319,7 +319,7 @@ def attach_align_device_hook(
             The module where we want to attach the hooks.
         execution_device (`torch.device`, *optional*):
             The device on which inputs and model weights should be placed before the forward pass.
-        offload (`bool`, *optional*, defauts to `False`):
+        offload (`bool`, *optional*, defaults to `False`):
             Whether or not the weights should be offloaded after the forward pass.
         weights_map (`Mapping[str, torch.Tensor]`, *optional*):
             When the model weights are offloaded, a (potentially lazy) map from param names to the tensor values.
@@ -402,7 +402,7 @@ def attach_align_device_hook_on_blocks(
         execution_device (`torch.device` or `Dict[str, torch.device]`, *optional*):
             The device on which inputs and model weights should be placed before the forward pass. It can be one device
             for the whole module, or a dictionary mapping module name to device.
-        offload (`bool`, *optional*, defauts to `False`):
+        offload (`bool`, *optional*, defaults to `False`):
             Whether or not the weights should be offloaded after the forward pass. It can be one boolean for the whole
             module, or a dictionary mapping module name to boolean.
         weights_map (`Mapping[str, torch.Tensor]`, *optional*):

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/712
diff --git a/utils/style_doc.py b/utils/style_doc.py
index 2a325805f..0422ebeb4 100644
--- a/utils/style_doc.py
+++ b/utils/style_doc.py
@@ -28,7 +28,7 @@
 # Regexes
 # Re pattern that catches list introduction (with potential indent)
 _re_list = re.compile(r"^(\s*-\s+|\s*\*\s+|\s*\d+\.\s+)")
-# Re pattern that catches code block introduction (with potentinal indent)
+# Re pattern that catches code block introduction (with potential indent)
 _re_code = re.compile(r"^(\s*)```(.*)$")
 # Re pattern that catches rst args blocks of the form `Parameters:`.
 _re_args = re.compile("^\s*(Args?|Arguments?|Params?|Parameters?):\s*$")
@@ -62,7 +62,7 @@ def parse_code_example(code_lines):
 
     Args:
         code_lines (`List[str]`): The code lines to parse.
-        max_len (`int`): The maximum lengh per line.
+        max_len (`int`): The maximum length per line.
 
     Returns:
         (List[`str`], List[`str`]): The list of code samples and the list of outputs.
@@ -109,7 +109,7 @@ def format_code_example(code: str, max_len: int, in_docstring: bool = False):
 
     Args:
         code (`str`): The code example to format.
-        max_len (`int`): The maximum lengh per line.
+        max_len (`int`): The maximum length per line.
         in_docstring (`bool`, *optional*, defaults to `False`): Whether or not the code example is inside a docstring.
 
     Returns:

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/711
diff --git a/tests/test_tracking.py b/tests/test_tracking.py
index 917ee329d..5e26eb9b5 100644
--- a/tests/test_tracking.py
+++ b/tests/test_tracking.py
@@ -85,12 +85,12 @@ def setUp(self):
         self.add_mocks(mock.patch.dict(os.environ, {"WANDB_DIR": self.tmpdir}))
 
     @staticmethod
-    def get_value_from_log(key: str, log: str, key_occurance: int = 0):
+    def get_value_from_log(key: str, log: str, key_occurrence: int = 0):
         """
         Parses wandb log for `key` and returns the value.
-        If parsing through multiple calls to .log, pass in a `key_occurance`
+        If parsing through multiple calls to .log, pass in a `key_occurrence`
         """
-        res = re.findall(rf"(?<={key} )[^\s]+", log)[key_occurance]
+        res = re.findall(rf"(?<={key} )[^\s]+", log)[key_occurrence]
         if '"' in res:
             return re.findall(r'"([^"]*)"', res)[0]
         else:

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/710
diff --git a/examples/by_feature/cross_validation.py b/examples/by_feature/cross_validation.py
index 707c93c4c..87f804ce6 100644
--- a/examples/by_feature/cross_validation.py
+++ b/examples/by_feature/cross_validation.py
@@ -58,7 +58,7 @@
 EVAL_BATCH_SIZE = 32
 
 # New Code #
-# We need a different `get_dataloaders` function that will build dataloaders by indexs
+# We need a different `get_dataloaders` function that will build dataloaders by index
 
 
 def get_fold_dataloaders(
@@ -71,9 +71,9 @@ def get_fold_dataloaders(
         accelerator (`Accelerator`):
             The main `Accelerator` object
         train_idxs (list of `int`):
-            The split indicies for the training dataset
+            The split indices for the training dataset
         valid_idxs (list of `int`):
-            The split indicies for the validation dataset
+            The split indices for the validation dataset
         batch_size (`int`):
             The size of the minibatch. Default is 16
     """
diff --git a/examples/by_feature/deepspeed_with_config_support.py b/examples/by_feature/deepspeed_with_config_support.py
index 89234a911..36ace84b9 100755
--- a/examples/by_feature/deepspeed_with_config_support.py
+++ b/examples/by_feature/deepspeed_with_config_support.py
@@ -525,7 +525,7 @@ def group_texts(examples):
         },
     ]
     # New Code #
-    # Creates Dummy Optimizer if `optimizer` was spcified in the config file else creates Adam Optimizer
+    # Creates Dummy Optimizer if `optimizer` was specified in the config file else creates Adam Optimizer
     optimizer_cls = (
         torch.optim.AdamW
         if accelerator.state.deepspeed_plugin is None
@@ -554,7 +554,7 @@ def group_texts(examples):
         args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
 
     # New Code #
-    # Creates Dummy Scheduler if `scheduler` was spcified in the config file else creates `args.lr_scheduler_type` Scheduler
+    # Creates Dummy Scheduler if `scheduler` was specified in the config file else creates `args.lr_scheduler_type` Scheduler
     if (
         accelerator.state.deepspeed_plugin is None
         or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
diff --git a/examples/by_feature/memory.py b/examples/by_feature/memory.py
index 9f0d55c69..82dc34b8a 100644
--- a/examples/by_feature/memory.py
+++ b/examples/by_feature/memory.py
@@ -29,7 +29,7 @@
 ########################################################################
 # This is a fully working simple example to use Accelerate,
 # specifically showcasing how to ensure out-of-memory errors never
-# iterrupt training, and builds off the `nlp_example.py` script.
+# interrupt training, and builds off the `nlp_example.py` script.
 #
 # This example trains a Bert base model on GLUE MRPC
 # in any of the following settings (with the same script):

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/709
diff --git a/docs/source/usage_guides/sagemaker.mdx b/docs/source/usage_guides/sagemaker.mdx
index 6195783e9..0afe52e29 100644
--- a/docs/source/usage_guides/sagemaker.mdx
+++ b/docs/source/usage_guides/sagemaker.mdx
@@ -129,7 +129,26 @@ You can find your model data at: s3://your-bucket/accelerate-sagemaker-1-2021-04
 
 ### Distributed Training: Data Parallelism
 
-*currently in development, will be supported soon.*
+Set up the accelerate config by running `accelerate config` and answer the SageMaker questions and set it up.
+To use SageMaker DDP, select it when asked 
+`What is the distributed mode? ([0] No distributed training, [1] data parallelism):`.
+Example config below:
+```yaml
+base_job_name: accelerate-sagemaker-1
+compute_environment: AMAZON_SAGEMAKER
+distributed_type: DATA_PARALLEL
+ec2_instance_type: ml.p3.16xlarge
+iam_role_name: xxxxx
+image_uri: null
+mixed_precision: fp16
+num_machines: 1
+profile: xxxxx
+py_version: py38
+pytorch_version: 1.10.2
+region: us-east-1
+transformers_version: 4.17.0
+use_cpu: false
+```
 
 ### Distributed Training: Model Parallelism
 
diff --git a/examples/README.md b/examples/README.md
index c4f050b3f..6a3c0a11a 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -187,6 +187,9 @@ To run it in each of these various modes, use the following commands:
 ### Simple vision example (GANs)
 
 - [huggan project](https://github.com/huggingface/community-events/tree/main/huggan)
+
+### Using AWS SageMaker integration
+- [Examples showcasing AWS SageMaker integration of 🤗 Accelerate.](https://github.com/pacman100/accelerate-aws-sagemaker)
     
 ## Finer Examples
 
diff --git a/src/accelerate/commands/config/sagemaker.py b/src/accelerate/commands/config/sagemaker.py
index 786881edf..b3a45c9e4 100644
--- a/src/accelerate/commands/config/sagemaker.py
+++ b/src/accelerate/commands/config/sagemaker.py
@@ -157,7 +157,7 @@ def get_sagemaker_input():
         )
 
     distributed_type = _ask_field(
-        "Which type of machine are you using? ([0] No distributed training, [1] data parallelism): ",
+        "What is the distributed mode? ([0] No distributed training, [1] data parallelism): ",
         _convert_sagemaker_distributed_mode,
         error_message="Please enter 0 or 1",
     )

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/708
diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml
index cb0a7ee42..277a81df9 100644
--- a/.github/workflows/nightly.yml
+++ b/.github/workflows/nightly.yml
@@ -7,6 +7,7 @@ on:
 
 env:
   RUN_SLOW: "yes"
+  IS_GITHUB_CI: "1"
 
 jobs:
   run_all_tests_single_gpu:
@@ -27,6 +28,7 @@ jobs:
           git config --global --add safe.directory '*'
           git fetch && git checkout ${{ github.sha }} 
           pip install -e . --no-deps
+          pip install pytest-reportlog
 
       - name: Run test on GPUs
         run: |
@@ -37,6 +39,11 @@ jobs:
           source activate accelerate
           pip uninstall comet_ml -y
           make test_examples
+          
+      - name: Generate Report
+        if: always()
+        run: |
+          python utils/log_reports.py >> $GITHUB_STEP_SUMMARY
 
   run_all_tests_multi_gpu:
     runs-on: [self-hosted, docker-gpu, multi-gpu]
@@ -56,6 +63,7 @@ jobs:
           git config --global --add safe.directory '*'
           git fetch && git checkout ${{ github.sha }}
           pip install -e . --no-deps
+          pip install pytest-reportlog
 
       - name: Run core and big modeling tests on GPUs
         run: |
@@ -72,4 +80,9 @@ jobs:
         run: |
           source activate accelerate
           pip uninstall comet_ml -y
-          make test_examples
\ No newline at end of file
+          make test_examples
+
+      - name: Generate Report
+        if: always()
+        run: |
+          python utils/log_reports.py >> $GITHUB_STEP_SUMMARY
\ No newline at end of file
diff --git a/.github/workflows/on-merge.yml b/.github/workflows/on-merge.yml
index 028b99d9a..a794cd72b 100644
--- a/.github/workflows/on-merge.yml
+++ b/.github/workflows/on-merge.yml
@@ -6,6 +6,7 @@ on:
 
 env:
   TESTING_MOCKED_DATALOADERS: "1"
+  IS_GITHUB_CI: "1"
 
 jobs:
   run_all_tests_single_gpu:
@@ -26,6 +27,7 @@ jobs:
           git config --global --add safe.directory '*'
           git fetch && git checkout ${{ github.sha }}
           pip install -e .[testing,test_trackers]
+          pip install pytest-reportlog
 
       - name: Run test on GPUs
         run: |
@@ -37,6 +39,11 @@ jobs:
           pip uninstall comet_ml -y
           make test_examples
 
+      - name: Generate Report
+        if: always()
+        run: |
+          python utils/log_reports.py >> $GITHUB_STEP_SUMMARY
+
   run_all_tests_multi_gpu:
     runs-on: [self-hosted, docker-gpu, multi-gpu]
     container:
@@ -53,6 +60,7 @@ jobs:
           git config --global --add safe.directory '*'
           git fetch && git checkout ${{ github.sha }}
           pip install -e .[testing,test_trackers]
+          pip install pytest-reportlog
 
       - name: Run test on GPUs
         run: |
@@ -63,4 +71,9 @@ jobs:
         run: |
           source activate accelerate
           pip uninstall comet_ml -y
-          make test_examples
\ No newline at end of file
+          make test_examples
+
+      - name: Generate Report
+        if: always()
+        run: |
+          python utils/log_reports.py >> $GITHUB_STEP_SUMMARY
\ No newline at end of file
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 2cf8e34a7..d39d36951 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -5,7 +5,7 @@ on: [pull_request]
 env:
   HF_HOME: ~/hf_cache
   TESTING_MOCKED_DATALOADERS: "1"
-  CI: 1
+  IS_GITHUB_CI: "1"
 
 jobs:
   run-tests:
diff --git a/Makefile b/Makefile
index 0c9ad0ca9..52c594ae7 100644
--- a/Makefile
+++ b/Makefile
@@ -25,40 +25,40 @@ style:
 	
 # Run tests for the library
 test:
-	python -m pytest -s -v ./tests/ --ignore=./tests/test_examples.py $(if $(CI),--report-log 'all.log',)
+	python -m pytest -s -v ./tests/ --ignore=./tests/test_examples.py $(if $(IS_GITHUB_CI),--report-log 'all.log',)
 
 test_big_modeling:
-	python -m pytest -s -v ./tests/test_big_modeling.py $(if $(CI),--report-log 'big_modeling.log',)
+	python -m pytest -s -v ./tests/test_big_modeling.py $(if $(IS_GITHUB_CI),--report-log 'big_modeling.log',)
 
 test_core:
 	python -m pytest -s -v ./tests/ --ignore=./tests/test_examples.py --ignore=./tests/deepspeed --ignore=./tests/test_big_modeling.py \
-	--ignore=./tests/fsdp $(if $(CI),--report-log 'core.log',)
+	--ignore=./tests/fsdp $(if $(IS_GITHUB_CI),--report-log 'core.log',)
 
 test_deepspeed:
-	python -m pytest -s -v ./tests/deepspeed $(if $(CI),--report-log 'deepspeed.log',)
+	python -m pytest -s -v ./tests/deepspeed $(if $(IS_GITHUB_CI),--report-log 'deepspeed.log',)
 
 test_fsdp:
-	python -m pytest -s -v ./tests/fsdp $(if $(CI),--report-log 'fsdp.log',)
+	python -m pytest -s -v ./tests/fsdp $(if $(IS_GITHUB_CI),--report-log 'fsdp.log',)
 
 test_examples:
-	python -m pytest -s -v ./tests/test_examples.py $(if $(CI),--report-log 'examples.log',)
+	python -m pytest -s -v ./tests/test_examples.py $(if $(IS_GITHUB_CI),--report-log 'examples.log',)
 
 # Broken down example tests for the CI runners
 test_integrations:
-	python -m pytest -s -v ./tests/deepspeed ./tests/fsdp $(if $(CI),--report-log 'integrations.log',)
+	python -m pytest -s -v ./tests/deepspeed ./tests/fsdp $(if $(IS_GITHUB_CI),--report-log 'integrations.log',)
 
 test_example_differences:
-	python -m pytest -s -v ./tests/test_examples.py::ExampleDifferenceTests $(if $(CI),--report-log 'example_diff.log',)
+	python -m pytest -s -v ./tests/test_examples.py::ExampleDifferenceTests $(if $(IS_GITHUB_CI),--report-log 'example_diff.log',)
 
 test_checkpoint_epoch:
-	python -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k "by_epoch" $(if $(CI),--report-log 'checkpoint_epoch.log',)
+	python -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k "by_epoch" $(if $(IS_GITHUB_CI),--report-log 'checkpoint_epoch.log',)
 
 test_checkpoint_step:
-	python -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k "by_step" $(if $(CI),--report-log 'checkpoint_step.log',)
+	python -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k "by_step" $(if $(IS_GITHUB_CI),--report-log 'checkpoint_step.log',)
 
 # Same as test but used to install only the base dependencies
 test_prod:
 	$(MAKE) test_core
 
 test_rest:
-	python -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k "not by_step and not by_epoch" $(if $(CI),--report-log 'rest.log',)
+	python -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k "not by_step and not by_epoch" $(if $(IS_GITHUB_CI),--report-log 'rest.log',)

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/704
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 3fec513d2..2cf8e34a7 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -5,11 +5,13 @@ on: [pull_request]
 env:
   HF_HOME: ~/hf_cache
   TESTING_MOCKED_DATALOADERS: "1"
+  CI: 1
 
 jobs:
   run-tests:
     runs-on: ubuntu-latest
     strategy:
+      fail-fast: false
       matrix:
         pytorch-version: [
           latest,
@@ -48,7 +50,13 @@ jobs:
         if [[ ${{ matrix.test-kind }} != test_prod ]]; then pip install -e .[testing,test_trackers]; fi
         if [[ ${{ matrix.test-kind }} = test_rest ]]; then pip uninstall comet_ml -y; fi
         if [[ ${{ matrix.pytorch-version }} = minimum ]]; then pip install torch==1.6.0; fi
+        pip install pytest-reportlog
     
     - name: Run Tests
       run: |
-        make ${{ matrix.test-kind }}
\ No newline at end of file
+        make ${{ matrix.test-kind }}
+
+    - name: Generate Report
+      if: always()
+      run: |
+        python utils/log_reports.py >> $GITHUB_STEP_SUMMARY
\ No newline at end of file
diff --git a/Makefile b/Makefile
index c1c79d6e8..0c9ad0ca9 100644
--- a/Makefile
+++ b/Makefile
@@ -25,39 +25,40 @@ style:
 	
 # Run tests for the library
 test:
-	python -m pytest -s -v ./tests/ --ignore=./tests/test_examples.py
+	python -m pytest -s -v ./tests/ --ignore=./tests/test_examples.py $(if $(CI),--report-log 'all.log',)
 
 test_big_modeling:
-	python -m pytest -s -v ./tests/test_big_modeling.py
+	python -m pytest -s -v ./tests/test_big_modeling.py $(if $(CI),--report-log 'big_modeling.log',)
 
 test_core:
 	python -m pytest -s -v ./tests/ --ignore=./tests/test_examples.py --ignore=./tests/deepspeed --ignore=./tests/test_big_modeling.py \
-	--ignore=./tests/fsdp
+	--ignore=./tests/fsdp $(if $(CI),--report-log 'core.log',)
 
 test_deepspeed:
-	python -m pytest -s -v ./tests/deepspeed
+	python -m pytest -s -v ./tests/deepspeed $(if $(CI),--report-log 'deepspeed.log',)
 
 test_fsdp:
-	python -m pytest -s -v ./tests/fsdp
+	python -m pytest -s -v ./tests/fsdp $(if $(CI),--report-log 'fsdp.log',)
 
 test_examples:
-	python -m pytest -s -v ./tests/test_examples.py
+	python -m pytest -s -v ./tests/test_examples.py $(if $(CI),--report-log 'examples.log',)
 
 # Broken down example tests for the CI runners
 test_integrations:
-	python -m pytest -s -v ./tests/deepspeed ./tests/fsdp
+	python -m pytest -s -v ./tests/deepspeed ./tests/fsdp $(if $(CI),--report-log 'integrations.log',)
+
 test_example_differences:
-	python -m pytest -s -v ./tests/test_examples.py::ExampleDifferenceTests
+	python -m pytest -s -v ./tests/test_examples.py::ExampleDifferenceTests $(if $(CI),--report-log 'example_diff.log',)
 
 test_checkpoint_epoch:
-	python -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k "by_epoch"
+	python -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k "by_epoch" $(if $(CI),--report-log 'checkpoint_epoch.log',)
 
 test_checkpoint_step:
-	python -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k "by_step"
+	python -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k "by_step" $(if $(CI),--report-log 'checkpoint_step.log',)
 
 # Same as test but used to install only the base dependencies
 test_prod:
 	$(MAKE) test_core
 
 test_rest:
-	python -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k "not by_step and not by_epoch"
+	python -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k "not by_step and not by_epoch" $(if $(CI),--report-log 'rest.log',)
diff --git a/utils/log_reports.py b/utils/log_reports.py
new file mode 100644
index 000000000..f701f08c2
--- /dev/null
+++ b/utils/log_reports.py
@@ -0,0 +1,34 @@
+import json
+from pathlib import Path 
+import subprocess
+
+failed = []
+passed = []
+
+group_info = []
+
+total_num_failed = 0
+for log in Path().glob("*.log"):
+    section_num_failed = 0
+    with open(log, "r") as f:
+        for line in f:
+            line = json.loads(line)
+            if line.get("nodeid", "") != "":
+                test = line["nodeid"]
+                if line.get("duration", None) is not None:
+                    duration = f'{line["duration"]:.4f}'
+                    if line.get("outcome", "") == "failed":
+                        section_num_failed += 1
+                        failed.append([test, duration])
+                    else:
+                        passed.append([test, duration])
+    group_info.append([str(log), section_num_failed])
+
+if len(failed) > 0:
+    result = "## Failed Tests:\n"
+    failed_table = '| Test Location | Test Class | Test Name |\n|---|---|---|\n| '
+    for test in failed:
+        failed_table += ' | '.join(test[0].split("::"))
+    failed_table += " |"
+    result += failed_table
+    print(result)
\ No newline at end of file

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/703
diff --git a/docs/source/usage_guides/memory.mdx b/docs/source/usage_guides/memory.mdx
index 9c5674977..8841cd584 100644
--- a/docs/source/usage_guides/memory.mdx
+++ b/docs/source/usage_guides/memory.mdx
@@ -48,4 +48,4 @@ def training_function(args):
 +   inner_training_loop()
 ```
 
-To find out more, check the documentation [here](package_reference/utilities#accelerate.find_executable_batch_size)
\ No newline at end of file
+To find out more, check the documentation [here](../package_reference/utilities#accelerate.find_executable_batch_size).

Hello @tomaarsen, looks good, thanks for your contribution 😄
_The documentation is not available anymore as the PR was closed or merged._
Thanks a lot!
https://api.github.com/repos/huggingface/accelerate/pulls/702
diff --git a/src/accelerate/big_modeling.py b/src/accelerate/big_modeling.py
index 311111ad4..0750248de 100644
--- a/src/accelerate/big_modeling.py
+++ b/src/accelerate/big_modeling.py
@@ -75,15 +75,35 @@ def register_empty_buffer(module, name, buffer):
         if buffer is not None:
             module._buffers[name] = module._buffers[name].to(torch.device("meta"))
 
+    # Patch tensor creation
+    if include_buffers:
+        tensor_constructors_to_patch = {
+            torch_function_name: getattr(torch, torch_function_name)
+            for torch_function_name in ["empty", "zeros", "ones", "full"]
+        }
+    else:
+        tensor_constructors_to_patch = {}
+
+    def patch_tensor_constructor(fn):
+        def wrapper(*args, **kwargs):
+            kwargs["device"] = torch.device("meta")
+            return fn(*args, **kwargs)
+
+        return wrapper
+
     try:
         nn.Module.register_parameter = register_empty_parameter
         if include_buffers:
             nn.Module.register_buffer = register_empty_buffer
+        for torch_function_name in tensor_constructors_to_patch.keys():
+            setattr(torch, torch_function_name, patch_tensor_constructor(getattr(torch, torch_function_name)))
         yield
     finally:
         nn.Module.register_parameter = old_register_parameter
         if include_buffers:
             nn.Module.register_buffer = old_register_buffer
+        for torch_function_name, old_torch_function in tensor_constructors_to_patch.items():
+            setattr(torch, torch_function_name, old_torch_function)
 
 
 def cpu_offload(

_The documentation is not available anymore as the PR was closed or merged._
I don't really understand why this is needed: while load a pretrained model inside the context manager and complain it takes time?
```py
from accelerate import init_empty_weights
from transformers import AutoConfig, AutoModel

config = AutoConfig.from_pretrained("gpt2")

with init_empty_weights():
    model = AutoModel.from_config(config)
```
is way faster than 6s
Hum it's doing it 10 times, so 0.6 sec per load. Benchmarking your solution displays the same order of magnitude: `14.567796499999996` (I think the difference is that the use of config doesn't require to read from disk the checkpoint anymore)

Though your workaround removed the need to override `torch.load` and especially the hack that I introduced.
```
Using config + this PR: 4.1660165419999995
```
Also I'm not sure why this wasn't detected, but the test `tests/test_big_modeling.py::BigModelingTester::test_init_empty_weights` pass on my MAC contrary to CI ... Moving back to draft as I need to figure this one out.
Actually if we activate this feature only when `include_buffers=True` (I'm guessing the assumption is that all pytorch tensors are expected to be `meta`) then that should be fine.
https://api.github.com/repos/huggingface/accelerate/pulls/699
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index b17ea387d..552c8392c 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -747,6 +747,12 @@ def _prepare_deepspeed(self, *args):
             batch_size_per_device = deepspeed_plugin.deepspeed_config["train_micro_batch_size_per_gpu"]
             result = [obj for obj in args]
 
+        if self.gradient_accumulation_steps != deepspeed_plugin.deepspeed_config["gradient_accumulation_steps"]:
+            logger.info(
+                f"Updating DeepSpeed's gradient accumulation steps to {self.gradient_accumulation_steps} from "
+                f"{deepspeed_plugin.deepspeed_config['gradient_accumulation_steps']}."
+            )
+            deepspeed_plugin.deepspeed_config["gradient_accumulation_steps"] = self.gradient_accumulation_steps
         config_kwargs = {
             "train_micro_batch_size_per_gpu": batch_size_per_device,
             "train_batch_size": batch_size_per_device

No need for a warning here, I think a logging info is sufficient (the fact that args to the `Accelerator` supercede the config is a common behavior, you can do it to force CPU or a given mixed precision already).
@@ -747,6 +747,12 @@ def _prepare_deepspeed(self, *args):
             batch_size_per_device = deepspeed_plugin.deepspeed_config["train_micro_batch_size_per_gpu"]
             result = [obj for obj in args]
 
+        if self.gradient_accumulation_steps != deepspeed_plugin.deepspeed_config["gradient_accumulation_steps"]:
+            warnings.warn(
Done.
@@ -747,6 +747,12 @@ def _prepare_deepspeed(self, *args):
             batch_size_per_device = deepspeed_plugin.deepspeed_config["train_micro_batch_size_per_gpu"]
             result = [obj for obj in args]
 
+        if self.gradient_accumulation_steps != deepspeed_plugin.deepspeed_config["gradient_accumulation_steps"]:
+            warnings.warn(
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/698
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index 84c2ebd67..1eb8ff79d 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -364,10 +364,11 @@ def __iter__(self):
 
     @property
     def total_batch_size(self):
+        batch_sampler = self.sampler if isinstance(self.sampler, BatchSampler) else self.batch_sampler
         return (
-            self.batch_sampler.batch_size
-            if self.batch_sampler.split_batches
-            else (self.batch_sampler.batch_size * self.batch_sampler.num_processes)
+            batch_sampler.batch_size
+            if batch_sampler.split_batches
+            else (batch_sampler.batch_size * batch_sampler.num_processes)
         )
 
     @property
@@ -639,14 +640,17 @@ def prepare_data_loader(
             )
         else:
             # New batch sampler for the current process.
-            if hasattr(dataloader.sampler, "generator"):
-                if dataloader.sampler.generator is None:
-                    dataloader.sampler.generator = torch.Generator()
-                    generator = dataloader.sampler.generator
-                    generator.manual_seed(int(torch.empty((), dtype=torch.int64).random_().item()))
-            elif getattr(dataloader.batch_sampler, "generator", None) is not None:
-                generator = dataloader.batch_sampler.generator
             sampler_is_batch_sampler = isinstance(dataloader.sampler, BatchSampler)
+            if sampler_is_batch_sampler:
+                sampler = dataloader.sampler.sampler
+            else:
+                sampler = dataloader.batch_sampler.sampler
+            if hasattr(sampler, "generator"):
+                if sampler.generator is None:
+                    sampler.generator = torch.Generator()
+                generator = sampler.generator
+                generator.manual_seed(int(torch.empty((), dtype=torch.int64).random_().item()))
+
             batch_sampler = dataloader.sampler if sampler_is_batch_sampler else dataloader.batch_sampler
             new_batch_sampler = BatchSamplerShard(
                 batch_sampler,
@@ -692,6 +696,7 @@ def prepare_data_loader(
             new_dataset,
             device=device if put_on_device and state.distributed_type != DistributedType.TPU else None,
             sampler=new_batch_sampler,
+            batch_size=getattr(dataloader, "batch_size", _PYTORCH_DATALOADER_KWARGS["batch_size"]),
             rng_types=rng_types,
             generator=generator,
             **kwargs,

I think we should use a instance check on the sampler here, as in my tests, there was a `batch_sampler` when using a `BatchSampler` for `sampler`:
```py
batch_sampler = self.sampler if isinstance(self.sampler, BatchSampler) else self.batch_sampler
```
@@ -364,10 +364,14 @@ def __iter__(self):
 
     @property
     def total_batch_size(self):
+        if self.batch_sampler is None:
+            batch_sampler = self.sampler
+        else:
+            batch_sampler = self.batch_sampler
Would reuse the same `batch_size` as for the original `datalaoder` here.
@@ -692,6 +699,7 @@ def prepare_data_loader(
             new_dataset,
             device=device if put_on_device and state.distributed_type != DistributedType.TPU else None,
             sampler=new_batch_sampler,
+            batch_size=None,
This is necessary else `batch_sampler` would get created with `batch_size=1` when it shouldn't do that as the `sampler` is already an instance of `batch_sampler`. This line here: https://github.com/pytorch/pytorch/blob/master/torch/utils/data/dataloader.py#L357

This leads to the issue `1.a` specified in the PR description

@@ -692,6 +699,7 @@ def prepare_data_loader(
             new_dataset,
             device=device if put_on_device and state.distributed_type != DistributedType.TPU else None,
             sampler=new_batch_sampler,
+            batch_size=None,
But it's not our place to fix users mistakes and they might have wanted to use `batch_size=1` with this for some reason instead of `batch_size=None`.
@@ -692,6 +699,7 @@ def prepare_data_loader(
             new_dataset,
             device=device if put_on_device and state.distributed_type != DistributedType.TPU else None,
             sampler=new_batch_sampler,
+            batch_size=None,
Here, the Dataloader default is batch_size=1 and not the one fixed by user. The one fixed by user is used which is part of sampler which is an instance of BatchSampler. When sampler is an object of BatchSampler, batch_sampler shouldn't be created and for that batch_size to dataloader has to be None so that dataloader can use sampler directly to create batches. 
@@ -692,6 +699,7 @@ def prepare_data_loader(
             new_dataset,
             device=device if put_on_device and state.distributed_type != DistributedType.TPU else None,
             sampler=new_batch_sampler,
+            batch_size=None,
Yes, `batch_size` should be None, but in the examples of the issue linked above, the user did not do that. We should honor the same argument as was passed to the original dataloader which is saved [here](https://github.com/pytorch/pytorch/blob/d2d145a40001d1e1f815a144160bd0b8d0f60ea0/torch/utils/data/dataloader.py#L361).
@@ -692,6 +699,7 @@ def prepare_data_loader(
             new_dataset,
             device=device if put_on_device and state.distributed_type != DistributedType.TPU else None,
             sampler=new_batch_sampler,
+            batch_size=None,
Got it. However, even if the user passes `None` to dataloader, it isn't being passed to the final `DataLoaderShard` as batch_size is part of `ignore_kwargs` and that leads to default value of 1 being used. https://github.com/huggingface/accelerate/blob/main/src/accelerate/data_loader.py#L659
@@ -692,6 +699,7 @@ def prepare_data_loader(
             new_dataset,
             device=device if put_on_device and state.distributed_type != DistributedType.TPU else None,
             sampler=new_batch_sampler,
+            batch_size=None,
Yes when there is a `batch_sampler`, because PyTorch raises an error and does not let us do that. But here we only pass a sampler ;-)
@@ -692,6 +699,7 @@ def prepare_data_loader(
             new_dataset,
             device=device if put_on_device and state.distributed_type != DistributedType.TPU else None,
             sampler=new_batch_sampler,
+            batch_size=None,
Done.
@@ -364,10 +364,14 @@ def __iter__(self):
 
     @property
     def total_batch_size(self):
+        if self.batch_sampler is None:
+            batch_sampler = self.sampler
+        else:
+            batch_sampler = self.batch_sampler
Done.
@@ -692,6 +699,7 @@ def prepare_data_loader(
             new_dataset,
             device=device if put_on_device and state.distributed_type != DistributedType.TPU else None,
             sampler=new_batch_sampler,
+            batch_size=None,
If `(batch_size != 1 and batch_size is not None) and sampler_is_batch_sampler`, is there any "default" way to shard batches without explicit instruction from user? I think at least a warning should be raised.
@@ -692,6 +696,7 @@ def prepare_data_loader(
             new_dataset,
             device=device if put_on_device and state.distributed_type != DistributedType.TPU else None,
             sampler=new_batch_sampler,
+            batch_size=getattr(dataloader, "batch_size", _PYTORCH_DATALOADER_KWARGS["batch_size"]),
BTW, `hasattr(dataloader, "batch_size")` is always `True` for `DataLoader`.
@@ -692,6 +696,7 @@ def prepare_data_loader(
             new_dataset,
             device=device if put_on_device and state.distributed_type != DistributedType.TPU else None,
             sampler=new_batch_sampler,
+            batch_size=getattr(dataloader, "batch_size", _PYTORCH_DATALOADER_KWARGS["batch_size"]),
> BTW, `hasattr(dataloader, "batch_size")` is always `True` for `DataLoader`.

resolving this in the other PR. wrt warning, current logic follows PyTorch API and as such warning isn't warranted.
@@ -692,6 +696,7 @@ def prepare_data_loader(
             new_dataset,
             device=device if put_on_device and state.distributed_type != DistributedType.TPU else None,
             sampler=new_batch_sampler,
+            batch_size=getattr(dataloader, "batch_size", _PYTORCH_DATALOADER_KWARGS["batch_size"]),
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/696
diff --git a/docs/source/usage_guides/mps.mdx b/docs/source/usage_guides/mps.mdx
index 3648bdce8..7a7f8a628 100644
--- a/docs/source/usage_guides/mps.mdx
+++ b/docs/source/usage_guides/mps.mdx
@@ -19,7 +19,7 @@ This will map computational graphs and primitives on the MPS Graph framework and
 For more information please refer official documents [Introducing Accelerated PyTorch Training on Mac](https://pytorch.org/blog/introducing-accelerated-pytorch-training-on-mac/)
 and [MPS BACKEND](https://pytorch.org/docs/stable/notes/mps.html).
 
-### Benefits of Training and Inference using Apple M1 Chips
+### Benefits of Training and Inference using Apple Silicon Chips
 
 1. Enables users to train larger networks or batch sizes locally
 2. Reduces data retrieval latency and provides the GPU with direct access to the full memory store due to unified memory architecture. 
diff --git a/src/accelerate/launchers.py b/src/accelerate/launchers.py
index dcdead956..0387109b6 100644
--- a/src/accelerate/launchers.py
+++ b/src/accelerate/launchers.py
@@ -121,12 +121,17 @@ def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, mix
                 start_processes(launcher, args=args, nprocs=num_processes, start_method="fork")
 
         else:
-            # No need for a distributed launch otherwise as it's either CPU or one GPU.
-            if torch.cuda.is_available():
+            # No need for a distributed launch otherwise as it's either CPU, GPU or MPS.
+            use_mps_device = "false"
+            if torch.backends.mps.is_available():
+                print("Launching training on MPS.")
+                use_mps_device = "true"
+            elif torch.cuda.is_available():
                 print("Launching training on one GPU.")
             else:
                 print("Launching training on CPU.")
-            function(*args)
+            with patch_environment(use_mps_device=use_mps_device):
+                function(*args)
 
 
 def debug_launcher(function, args=(), num_processes=2):

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_695). All of your documentation changes will be reflected on that endpoint.
https://api.github.com/repos/huggingface/accelerate/pulls/695
diff --git a/src/accelerate/utils/random.py b/src/accelerate/utils/random.py
index 51b7c3719..e95ed03bf 100644
--- a/src/accelerate/utils/random.py
+++ b/src/accelerate/utils/random.py
@@ -64,7 +64,7 @@ def synchronize_rng_state(rng_type: Optional[RNGType] = None, generator: Optiona
     state = AcceleratorState()
     if state.distributed_type == DistributedType.TPU:
         rng_state = xm.mesh_reduce("random_seed", rng_state, lambda x: x[0])
-    elif state.distributed_type in [DistributedType.DEEPSPEED, DistributedType.MULTI_GPU]:
+    elif state.distributed_type in [DistributedType.DEEPSPEED, DistributedType.MULTI_GPU, DistributedType.FSDP]:
         rng_state = rng_state.to(state.device)
         torch.distributed.broadcast(rng_state, 0)
         rng_state = rng_state.cpu()

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/688
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index fb5393d34..84c2ebd67 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -622,6 +622,7 @@ def prepare_data_loader(
     new_dataset = dataloader.dataset
     # Iterable dataset doesn't like batch_sampler, but data_loader creates a default one for it
     new_batch_sampler = dataloader.batch_sampler if not isinstance(new_dataset, IterableDataset) else None
+    sampler_is_batch_sampler = False
     generator = getattr(dataloader, "generator", None)
     # No change if no multiprocess
     if num_processes != 1 and not dispatch_batches:
@@ -645,8 +646,10 @@ def prepare_data_loader(
                     generator.manual_seed(int(torch.empty((), dtype=torch.int64).random_().item()))
             elif getattr(dataloader.batch_sampler, "generator", None) is not None:
                 generator = dataloader.batch_sampler.generator
+            sampler_is_batch_sampler = isinstance(dataloader.sampler, BatchSampler)
+            batch_sampler = dataloader.sampler if sampler_is_batch_sampler else dataloader.batch_sampler
             new_batch_sampler = BatchSamplerShard(
-                dataloader.batch_sampler,
+                batch_sampler,
                 num_processes=num_processes,
                 process_index=process_index,
                 split_batches=split_batches,
@@ -684,6 +687,15 @@ def prepare_data_loader(
             _drop_last=dataloader.drop_last,
             **kwargs,
         )
+    elif sampler_is_batch_sampler:
+        dataloader = DataLoaderShard(
+            new_dataset,
+            device=device if put_on_device and state.distributed_type != DistributedType.TPU else None,
+            sampler=new_batch_sampler,
+            rng_types=rng_types,
+            generator=generator,
+            **kwargs,
+        )
     else:
         dataloader = DataLoaderShard(
             new_dataset,

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/687
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index cbe39efa3..fb5393d34 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -462,10 +462,6 @@ def _fetch_batches(self, iterator):
                 else:
                     batch_info = [None, True]
                 broadcast_object_list(batch_info)
-                if batch_info[1]:
-                    return batch, batch_info
-            else:
-                return batch, batch_info
         return batch, batch_info
 
     def __iter__(self):

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/686
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 989774148..b17ea387d 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -745,6 +745,7 @@ def _prepare_deepspeed(self, *args):
                 )
         else:
             batch_size_per_device = deepspeed_plugin.deepspeed_config["train_micro_batch_size_per_gpu"]
+            result = [obj for obj in args]
 
         config_kwargs = {
             "train_micro_batch_size_per_gpu": batch_size_per_device,

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/683
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index 9236ee310..cbe39efa3 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -463,10 +463,10 @@ def _fetch_batches(self, iterator):
                     batch_info = [None, True]
                 broadcast_object_list(batch_info)
                 if batch_info[1]:
-                    return batch, batch_info, True
+                    return batch, batch_info
             else:
-                return batch, batch_info, True
-        return batch, batch_info, False
+                return batch, batch_info
+        return batch, batch_info
 
     def __iter__(self):
         self.gradient_state._set_end_of_dataloader(False)
@@ -477,11 +477,10 @@ def __iter__(self):
         stop_iteration = False
         self._stop_iteration = False
         first_batch = None
-        next_batch, next_batch_info, next_skip = self._fetch_batches(main_iterator)
+        next_batch, next_batch_info = self._fetch_batches(main_iterator)
         while not stop_iteration:
-            batch, batch_info, skip = next_batch, next_batch_info, next_skip
-            if skip:
-                continue
+            batch, batch_info = next_batch, next_batch_info
+
             if self.state.process_index != 0:
                 # Initialize tensors on other processes than process 0.
                 batch = initialize_tensors(batch_info[0])
@@ -500,7 +499,7 @@ def __iter__(self):
             if not stop_iteration:
                 # We may still be at the end of the dataloader without knowing it yet: if there is nothing left in
                 # the dataloader since the number of batches is a round multiple of the number of processes.
-                next_batch, next_batch_info, next_skip = self._fetch_batches(main_iterator)
+                next_batch, next_batch_info = self._fetch_batches(main_iterator)
                 # next_batch_info[0] is None when there are no more batches, otherwise we still need to process them.
                 if self._stop_iteration and next_batch_info[0] is None:
                     stop_iteration = True

_The documentation is not available anymore as the PR was closed or merged._
Hello, for some reason I wasn't notified by github about review request, looks good. I was thinking if the below part of `_fetch_batches` could be simplified 

```diff
def _fetch_batches(self, iterator):
        batches, batch = None, None
        # On process 0, we gather the batch to dispatch.
        if self.state.process_index == 0:
            try:
                if self.split_batches:
                    # One batch of the main iterator is dispatched and split.
                    batch = next(iterator)
                else:
                    # num_processes batches of the main iterator are concatenated then dispatched and split.
                    # We add the batches one by one so we have the remainder available when drop_last=False.
                    batches = []
                    for _ in range(self.state.num_processes):
                        batches.append(next(iterator))
                    batch = concatenate(batches, dim=0)
                # In both cases, we need to get the structure of the batch that we will broadcast on other
                # processes to initialize the tensors with the right shape.
                # data_structure, stop_iteration
                batch_info = [get_data_structure(batch), False]
            except StopIteration:
                batch_info = [None, True]
        else:
            batch_info = [None, self._stop_iteration]
        # This is inplace, so after this instruction, every process has the same `batch_info` as process 0.
        broadcast_object_list(batch_info)
        self._stop_iteration = batch_info[1]
        if self._stop_iteration:
            # If drop_last is False and split_batches is False, we may have a remainder to take care of.
            if not self.split_batches and not self._drop_last:
                if self.state.process_index == 0 and len(batches) > 0:
                    batch = concatenate(batches, dim=0)
                    batch_info = [get_data_structure(batch), False]
                else:
                    batch_info = [None, True]
                broadcast_object_list(batch_info)
-                if batch_info[1]:
-                    return batch, batch_info
-            else:
-                return batch, batch_info
        return batch, batch_info
```
Yes, you're right thanks! Will make a follow-up PR.
https://api.github.com/repos/huggingface/accelerate/pulls/682
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index c7f235843..fad42792d 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -270,6 +270,25 @@ def launch_command_parser(subparsers=None):
         default=None,
         help="The port to use to communicate with the machine of rank 0.",
     )
+    # Rendezvous related arguments
+    parser.add_argument(
+        "--rdzv_conf",
+        type=str,
+        default="",
+        help="Additional rendezvous configuration (<key1>=<value1>,<key2>=<value2>,...).",
+    )
+    parser.add_argument(
+        "--max_restarts",
+        type=int,
+        default=0,
+        help="Maximum number of worker group restarts before failing.",
+    )
+    parser.add_argument(
+        "--monitor_interval",
+        type=float,
+        default=5,
+        help="Interval, in seconds, to monitor the state of workers.",
+    )
     parser.add_argument(
         "--main_training_function",
         type=str,

_The documentation is not available anymore as the PR was closed or merged._
> It looks like those new fields are not used anywhere afterwards?

Hello, it is being used via below lines:

https://github.com/huggingface/accelerate/blob/4400eb90b29638deec117f95fecdcb79bd7684cb/src/accelerate/commands/launch.py#L465

https://github.com/huggingface/accelerate/blob/4400eb90b29638deec117f95fecdcb79bd7684cb/src/accelerate/utils/launch.py#L51
https://api.github.com/repos/huggingface/accelerate/pulls/680
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 27db19b9b..989774148 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -723,25 +723,28 @@ def _prepare_deepspeed(self, *args):
 
         deepspeed_plugin = self.state.deepspeed_plugin
 
-        result = [
-            self._prepare_one(obj, first_pass=True) if isinstance(obj, torch.utils.data.DataLoader) else obj
-            for obj in args
-        ]
-
-        batch_sizes = [obj.batch_size for obj in args if hasattr(obj, "batch_size")]
-        if self.split_batches:
-            batch_sizes = [batch_size // self.num_processes for batch_size in batch_sizes]
-        if len(batch_sizes) == 0:
-            raise ValueError(
-                "You must specify a training or evaluation dataloader in `accelerate.prepare()` when using DeepSpeed."
-            )
+        if deepspeed_plugin.deepspeed_config["train_micro_batch_size_per_gpu"] == "auto":
+            result = [
+                self._prepare_one(obj, first_pass=True) if isinstance(obj, torch.utils.data.DataLoader) else obj
+                for obj in args
+            ]
+
+            batch_sizes = [obj.batch_size for obj in args if hasattr(obj, "batch_size")]
+            if self.split_batches:
+                batch_sizes = [batch_size // self.num_processes for batch_size in batch_sizes]
+            if len(batch_sizes) == 0:
+                raise ValueError(
+                    "You must specify a training or evaluation dataloader in `accelerate.prepare()` when using DeepSpeed."
+                )
 
-        batch_size_per_device = min(batch_sizes) if deepspeed_plugin.is_train_batch_min else max(batch_sizes)
-        if len(batch_sizes) > 1:
-            logger.info(
-                "Since you passed both train and evaluation dataloader, `is_train_batch_min` (here "
-                f"{deepspeed_plugin.is_train_batch_min} will decide the `train_batch_size` ({batch_size_per_device})."
-            )
+            batch_size_per_device = min(batch_sizes) if deepspeed_plugin.is_train_batch_min else max(batch_sizes)
+            if len(batch_sizes) > 1:
+                logger.info(
+                    "Since you passed both train and evaluation dataloader, `is_train_batch_min` (here "
+                    f"{deepspeed_plugin.is_train_batch_min} will decide the `train_batch_size` ({batch_size_per_device})."
+                )
+        else:
+            batch_size_per_device = deepspeed_plugin.deepspeed_config["train_micro_batch_size_per_gpu"]
 
         config_kwargs = {
             "train_micro_batch_size_per_gpu": batch_size_per_device,
@@ -916,7 +919,9 @@ def backward(self, loss, **kwargs):
 
         Should be used in lieu of `loss.backward()`.
         """
-        loss /= self.gradient_accumulation_steps
+        if self.distributed_type != DistributedType.DEEPSPEED:
+            # deepspeed handles loss scaling by gradient_accumulation_steps in its `backward`
+            loss = loss / self.gradient_accumulation_steps
         if self.distributed_type == DistributedType.DEEPSPEED:
             self.deepspeed_engine_wrapped.backward(loss, **kwargs)
         elif self.scaler is not None:
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index c7f235843..30e497e84 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -549,7 +549,8 @@ def deepspeed_launcher(args):
     current_env["DEEPSPEED_OFFLOAD_PARAM_DEVICE"] = str(args.offload_param_device).lower()
     current_env["DEEPSPEED_ZERO3_INIT"] = str(args.zero3_init_flag).lower()
     current_env["DEEPSPEED_ZERO3_SAVE_16BIT_MODEL"] = str(args.zero3_save_16bit_model).lower()
-    current_env["DEEPSPEED_CONFIG_FILE"] = str(args.deepspeed_config_file).lower()
+    if args.deepspeed_config_file is not None:
+        current_env["DEEPSPEED_CONFIG_FILE"] = str(args.deepspeed_config_file)
 
     if args.num_machines > 1 and args.deepspeed_multinode_launcher != DEEPSPEED_MULTINODE_LAUNCHERS[1]:
         with open(".deepspeed_env", "a") as f:

_The documentation is not available anymore as the PR was closed or merged._
Hmm, I tried to remove the fake dataloader workaround that was discussed at https://discuss.huggingface.co/t/when-using-deepspeed-why-do-i-need-to-pass-dataloaders-to-the-accelerator-prepare/22432

and it's super cumbersome. It appears that the only way to get to the batch size is from dataloader? why can't it be derived from a `batch_size` argument? or is it by design that `batch_size`  is derived from dataloader?

Specifically to this PR:

1. The original idea behind auto was to try to make the ds config file as simple as possible wrt to hardcoding values and have command line args set these things once. So this doesn't work:

```
    "train_micro_batch_size_per_gpu": "auto",
```

To remind, the intention of creating `auto`-values in DS config was to avoid mismatch of the same configuration value used by different sub-systems. So the way I was thinking when integrating DS in HF Trainer, is let's have one place where `batch_size` is set and then all other sub-systems would inherit that value, rather than relying on the user to remember to change the same value in many places. I hope this makes sense. I don't understand the design here so I'm surely missing something important.

I was trying to remove the originally used workaround 
https://github.com/huggingface/m4/pull/610
but I think the workaround with the dummy dataloader is a way user-friendlier than the hardcoded `train_micro_batch_size_per_gpu` config value.

2. Also while this feature that supports ` train_micro_batch_size_per_gpu` config was added, the end user has no idea it's supported w/o reading the source code - the error message is still:

> You must specify a training or evaluation dataloader in `accelerate.prepare()` when using DeepSpeed.

Should it perhaps say:

> when using DeepSpeed `accelerate.prepare()` requires you to pass at least one of training or evaluation dataloaders or alternatively set an integer value in `train_micro_batch_size_per_gpu` in the deepspeed config file.

or something of a kind? otherwise the feature is there but nobody knows about it.

The API doc could also say that, except it's private so there is no documentation.

Thank you for reading.
Hello Stas,

1. Inline with the intention of using `auto` values, we use the dataloader to infer the `batch_size`, one reason reason being that `accelerator.prepare` function or `accelerator` object doesn't have access to command line arguments, i.e., `batch_size` argument from the main code.  Another reason being that the single place where `batch_size` gets set is while the user creates dataloaders as part of conventional training, e.g., the user can have `batch_size` argument which they may modify because of certain custom logic and the modified batch_size now gets used to create dataloaders. As we have no control over training loop/code unlike Trainer, it makes sense to infer `batch_size` directly from dataloaders. 

Would the alternative suggestion work?
`AcceleratorState().deepspeed_plugin.deepspeed_config["train_micro_batch_size_per_gpu"]=my_batch_size` before calling `accelerator.prepare()`?

2. Thank you for the pointer on the error message, I will update it.

Hi Sourab!

It appears that `accelerator.prepare` relies on the `deepspeed_plugin.*` settings - if those were already parsed so that any `auto`s get replaced with the proper values then it'd have access to the correct setting, no? or is it still too early in the game - perhaps it's a matter of an order of execution?

Perhaps there should be another wrapper that a user should call explicitly for deepspeed with args like bs early in the code, so that no `auto` values remain and then it'd be easy to rely on the actual values later on. All the values should be available early on. e.g. in HF Trainer we only had to wait for later for  `num_training_steps`:
https://github.com/huggingface/transformers/blob/0a144b8c6bd16f7b08119a233f8e7cbd33fe5bfd/src/transformers/deepspeed.py#L167

> Would the alternative suggestion work?
> AcceleratorState().deepspeed_plugin.deepspeed_config["train_micro_batch_size_per_gpu"]=my_batch_size before calling accelerator.prepare()?

That would definitely work.

My first reaction is that suggestion could potentially be much more problematic should the user set the value in the ds config file and it might be an unexpected override (even though if written correctly it should be the same value). Somehow this feels like replacing one hack with another hack.

I think the dummy dataset wrapped dataloader is a much cleaner way over the above, especially if the code isn't necessarily always using the deepspeed backend.

If this is the best that can be done, and there is no simpler way, let's just leave it as is.

> Perhaps there should be another wrapper that a user should call explicitly for deepspeed with args like bs early in the code, so that no `auto` values remain and then it'd be easy to rely on the actual values later on

There is already a way to do this `deepspeed_plugin.deepspeed_config_process(**kwargs)` (Process the DeepSpeed config with the values from the kwargs.): https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L417
Example below:

```
if AcceleratorState().deepspeed_plugin is not None:
  kwargs = {
              "fp16.enabled": True,
              "optimizer.params.lr": 5e-5,
              "optimizer.params.weight_decay": 0.0,
              "scheduler.params.warmup_min_lr": 0.0,
              "scheduler.params.warmup_max_lr": 5e-5,
              "scheduler.params.warmup_num_steps": 0,
              "train_micro_batch_size_per_gpu": my_batch_size,
              "gradient_clipping": 1.0,
          } 
  AcceleratorState().deepspeed_plugin.deepspeed_config_process(must_match=True, **kwargs)

# call `accelerator.prepare` below `_` are just dummy placeholders 
_, _, _ = accelerator.prepare(_, _, _)

...

```
should `deepspeed_config_process` call that last `accelerator.prepare` command internally? This looks super-weird with `_` input and `_` output args, unless you meant something else there.
Hello @stas00 , in the example above, it is the user code, I was just mentioning/showcasing that `deepspeed_config_process` should be called before `accelerator.prepare`. As mentioned, `deepspeed_config_process` only sets `auto` values from the given kwargs dict, doesn't internally call `accelerator.prepare` at all and those `_` were placeholders 😅. 
Thank you for clarifying, @pacman100! It's crystal clear now.
https://api.github.com/repos/huggingface/accelerate/pulls/676
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index 9b14aeaed..a174115b1 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -59,7 +59,7 @@ def get_cluster_input():
                 lambda x: int(x),
             )
             same_network = _ask_field(
-                "Are all the machines on the same network? [YES/no]: ",
+                "Are all the machines on the same local network? Answer `no` if nodes are on the cloud and/or on different network hosts [YES/no]: ",
                 _convert_yes_no_to_bool,
                 default=True,
                 error_message="Please enter yes or no.",

```suggestion
                "Are all the machines on the same local network? Answer `no` if nodes are on the cloud and/or on the different network hosts [YES/no]: ",
```

@@ -59,7 +59,7 @@ def get_cluster_input():
                 lambda x: int(x),
             )
             same_network = _ask_field(
-                "Are all the machines on the same network? [YES/no]: ",
+                "Are all the machines on the same local network? [YES/no]: ",
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/674
diff --git a/setup.py b/setup.py
index 4bc977029..409399b17 100644
--- a/setup.py
+++ b/setup.py
@@ -21,9 +21,10 @@
 extras["test_prod"] = ["pytest", "pytest-xdist", "pytest-subtests", "parameterized"]
 extras["test_dev"] = ["datasets", "evaluate", "transformers", "scipy", "sklearn", "deepspeed<0.7.0", "tqdm"]
 extras["testing"] = extras["test_prod"] + extras["test_dev"]
+extras["rich"] = ["rich"]
 
 extras["test_trackers"] = ["wandb", "comet-ml", "tensorboard"]
-extras["dev"] = extras["quality"] + extras["testing"]
+extras["dev"] = extras["quality"] + extras["testing"] + extras["rich"]
 
 extras["sagemaker"] = [
     "sagemaker",  # boto3 is a required package in sagemaker
@@ -50,7 +51,7 @@
         ]
     },
     python_requires=">=3.7.0",
-    install_requires=["numpy>=1.17", "packaging>=20.0", "psutil", "pyyaml", "torch>=1.4.0", "rich"],
+    install_requires=["numpy>=1.17", "packaging>=20.0", "psutil", "pyyaml", "torch>=1.4.0"],
     extras_require=extras,
     classifiers=[
         "Development Status :: 5 - Production/Stable",
diff --git a/src/accelerate/__init__.py b/src/accelerate/__init__.py
index 03a7198ac..10c258efa 100644
--- a/src/accelerate/__init__.py
+++ b/src/accelerate/__init__.py
@@ -16,7 +16,11 @@
     InitProcessGroupKwargs,
     find_executable_batch_size,
     infer_auto_device_map,
+    is_rich_available,
     load_checkpoint_in_model,
-    rich,
     synchronize_rng_states,
 )
+
+
+if is_rich_available():
+    from .utils import rich
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 491ea5ad3..c7f235843 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -39,22 +39,25 @@
     _filter_args,
     get_launch_prefix,
     is_deepspeed_available,
+    is_rich_available,
     is_sagemaker_available,
     is_torch_version,
     patch_environment,
 )
 from accelerate.utils.constants import DEEPSPEED_MULTINODE_LAUNCHERS
 from accelerate.utils.dataclasses import SageMakerDistributedType
-from rich import get_console
-from rich.logging import RichHandler
 
 
-if is_torch_version(">=", "1.9.0"):
-    import torch.distributed.run as distrib_run
+if is_rich_available():
+    from rich import get_console
+    from rich.logging import RichHandler
+
+    FORMAT = "%(message)s"
+    logging.basicConfig(format=FORMAT, datefmt="[%X]", handlers=[RichHandler()])
 
 
-FORMAT = "%(message)s"
-logging.basicConfig(format=FORMAT, datefmt="[%X]", handlers=[RichHandler()])
+if is_torch_version(">=", "1.9.0"):
+    import torch.distributed.run as distrib_run
 
 logger = logging.getLogger(__name__)
 
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
index 2b8f566f0..6ec2f4ac4 100644
--- a/src/accelerate/utils/__init__.py
+++ b/src/accelerate/utils/__init__.py
@@ -27,6 +27,7 @@
     is_comet_ml_available,
     is_datasets_available,
     is_deepspeed_available,
+    is_rich_available,
     is_sagemaker_available,
     is_tensorboard_available,
     is_tpu_available,
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index 3efb61034..6015f2c14 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -111,6 +111,10 @@ def is_boto3_available():
     return importlib.util.find_spec("boto3") is not None
 
 
+def is_rich_available():
+    return importlib.util.find_spec("rich") is not None
+
+
 def is_sagemaker_available():
     return importlib.util.find_spec("sagemaker") is not None
 
diff --git a/src/accelerate/utils/rich.py b/src/accelerate/utils/rich.py
index 7a54c8b5a..2d48661b7 100644
--- a/src/accelerate/utils/rich.py
+++ b/src/accelerate/utils/rich.py
@@ -12,7 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from rich.traceback import install
+from .imports import is_rich_available
 
 
-install(show_locals=False)
+if is_rich_available():
+    from rich.traceback import install
+
+    install(show_locals=False)
+
+else:
+    raise ModuleNotFoundError("To use the rich extension, install rich with `pip install rich`")

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/673
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index 65c0e56b2..9b14aeaed 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -38,6 +38,7 @@ def get_cluster_input():
     main_process_ip = None
     main_process_port = None
     rdzv_backend = "static"
+    same_network = True
     if distributed_type in [DistributedType.MULTI_GPU, DistributedType.MULTI_CPU]:
         num_machines = _ask_field(
             "How many different machines will you use (use more than 1 for multi-node training)? [1]: ",
@@ -57,9 +58,16 @@ def get_cluster_input():
                 "What is the port you will use to communicate with the main process? ",
                 lambda x: int(x),
             )
-            rdzv_backend = _ask_field(
-                "What rendezvous backend will you use? ('static', 'c10d', ...)", default="static"
+            same_network = _ask_field(
+                "Are all the machines on the same network? [YES/no]: ",
+                _convert_yes_no_to_bool,
+                default=True,
+                error_message="Please enter yes or no.",
             )
+            if not same_network:
+                rdzv_backend = _ask_field(
+                    "What rendezvous backend will you use? ('static', 'c10d', ...): ", default="static"
+                )
 
     if distributed_type == DistributedType.NO:
         use_cpu = _ask_field(
@@ -328,4 +336,5 @@ def get_cluster_input():
         fsdp_config=fsdp_config,
         use_cpu=use_cpu,
         rdzv_backend=rdzv_backend,
+        same_network=same_network,
     )
diff --git a/src/accelerate/commands/config/config_args.py b/src/accelerate/commands/config/config_args.py
index e78aad181..68da64963 100644
--- a/src/accelerate/commands/config/config_args.py
+++ b/src/accelerate/commands/config/config_args.py
@@ -138,6 +138,7 @@ class ClusterConfig(BaseConfig):
     main_process_ip: Optional[str] = None
     main_process_port: Optional[int] = None
     rdzv_backend: Optional[str] = "static"
+    same_network: Optional[bool] = False
     main_training_function: str = "main"
 
     # args for deepspeed_plugin
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index fcc80bb7f..491ea5ad3 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -375,15 +375,21 @@ def simple_launcher(args):
 def multi_gpu_launcher(args):
     num_processes = getattr(args, "num_processes")
     num_machines = getattr(args, "num_machines")
+    main_process_ip = getattr(args, "main_process_ip")
+    main_process_port = getattr(args, "main_process_port")
     if num_machines > 1:
         setattr(args, "nproc_per_node", str(num_processes // num_machines))
         setattr(args, "nnodes", str(num_machines))
-        setattr(args, "node_rank", str(args.machine_rank))
-        setattr(args, "rdzv_endpoint", f"{args.main_process_ip}:{args.main_process_port}")
+        setattr(args, "node_rank", int(args.machine_rank))
+        if getattr(args, "same_network"):
+            setattr(args, "master_addr", str(main_process_ip))
+            setattr(args, "master_port", str(main_process_port))
+        else:
+            setattr(args, "rdzv_endpoint", f"{main_process_ip}:{main_process_port}")
     else:
         setattr(args, "nproc_per_node", str(num_processes))
-        if args.main_process_port is not None:
-            setattr(args, "master_port", str(args.main_process_port))
+        if main_process_port is not None:
+            setattr(args, "master_port", str(main_process_port))
 
     if args.module and args.no_python:
         raise ValueError("--module and --no_python cannot be used together")

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/672
diff --git a/docs/source/usage_guides/big_modeling.mdx b/docs/source/usage_guides/big_modeling.mdx
index 9b3252809..1e13849c6 100644
--- a/docs/source/usage_guides/big_modeling.mdx
+++ b/docs/source/usage_guides/big_modeling.mdx
@@ -35,7 +35,13 @@ While this works very well for regularly sized models, this workflow has some cl
 
 </Tip>
 
-## Instantiating an empty model
+## How the Process Works: A Quick Overview
+
+<Youtube id="MWCSGj9jEAo" />
+
+## How the Process Works: Working with Code
+
+### Instantiating an empty model
 
 The first tool 🤗 Accelerate introduces to help with big models is a context manager [`init_empty_weights`] that helps you initialize a model without using any RAM, so that step 1 can be done on models of any size. Here is how it works:
 
@@ -61,7 +67,7 @@ initializes an empty model with a bit more than 100B parameters. Behind the scen
 
 </Tip>
 
-## Sharded checkpoints
+### Sharded checkpoints
 
 It's possible your model is so big that even a single copy won't fit in RAM. That doesn't mean it can't be loaded: if you have one or several GPUs, this is more memory available to store your model. In this case, it's better if your checkpoint is split in several smaller files that we call checkpoint shards.
 
@@ -86,7 +92,7 @@ with index.json being the following file:
 
 and `first_state_dict.bin` containing the weights for `"linear1.weight"` and `"linear1.bias"`, `second_state_dict.bin` the ones for `"linear2.weight"` and `"linear2.bias"`
 
-## Loading weights
+### Loading weights
 
 The second tool 🤗 Accelerate introduces is a function [`load_checkpoint_and_dispatch`], that will allow you to load a checkpoint inside your empty model. This supports full checkpoints (a single file containing the whole state dict) as well as sharded checkpoints. It will also automatically dispatch those weights across the devices you have available (GPUs, CPU RAM), so if you are loading a sharded checkpoint, the maximum RAM usage will be the size of the biggest shard.
 
@@ -176,7 +182,7 @@ You can also design your `device_map` yourself, if you prefer to explicitly deci
 model = load_checkpoint_and_dispatch(model, "sharded-gpt-j-6B", device_map=my_device_map)
 ```
 
-## Run the model
+### Run the model
 
 Now that we have done this, our model lies across several devices, and maybe the hard drive. But it can still be used as a regular PyTorch model:
 
@@ -203,7 +209,7 @@ This way, you model can run for inference even if it doesn't fit on one of the G
 
 </Tip>
 
-## Designing a device map
+### Designing a device map
 
 You can let 🤗 Accelerate handle the device map computation by setting `device_map` to one of the supported options (`"auto"`, `"balanced"`, `"balanced_low_0"`, `"sequential"`) or create one yourself, if you want more control over where each layer should go.
 
diff --git a/manim_animations/big_model_inference/stage_1.py b/manim_animations/big_model_inference/stage_1.py
new file mode 100644
index 000000000..81ec0c965
--- /dev/null
+++ b/manim_animations/big_model_inference/stage_1.py
@@ -0,0 +1,108 @@
+# 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.
+
+from manim import *
+
+
+class Stage1(Scene):
+    def construct(self):
+        mem = Rectangle(height=0.5,width=0.5)
+        fill = Rectangle(height=0.46,width=0.46).set_stroke(width=0)
+
+        cpu_left_col_base = [mem.copy() for i in range(6)]
+        cpu_right_col_base = [mem.copy() for i in range(6)]
+        cpu_left_col = VGroup(*cpu_left_col_base).arrange(UP, buff=0)
+        cpu_right_col = VGroup(*cpu_right_col_base).arrange(UP, buff=0)
+        cpu_rects = VGroup(cpu_left_col,cpu_right_col).arrange(RIGHT, buff=0)
+        cpu_text = Text("CPU", font_size=24)
+        cpu = Group(cpu_rects,cpu_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        cpu.move_to([-2.5,-.5,0])
+        self.add(cpu)
+
+        gpu_base = [mem.copy() for i in range(1)]
+        gpu_rect = VGroup(*gpu_base).arrange(UP,buff=0)
+        gpu_text = Text("GPU", font_size=24)
+        gpu = Group(gpu_rect,gpu_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        gpu.align_to(cpu, DOWN)
+        gpu.set_x(gpu.get_x() - 1)
+        
+        self.add(gpu)
+
+        model_base = [mem.copy() for i in range(6)]
+        model_rect = VGroup(*model_base).arrange(RIGHT,buff=0)
+
+        model_text = Text("Model", font_size=24)
+        model = Group(model_rect,model_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        model.move_to([3, -1., 0])
+        
+        self.play(
+            Create(cpu_left_col, run_time=1),
+            Create(cpu_right_col, run_time=1),
+            Create(gpu_rect, run_time=1),
+        )
+
+        step_1 = MarkupText(
+            f"First, an empty model skeleton is loaded\ninto <span fgcolor='{YELLOW}'>memory</span> without using much RAM.", 
+            font_size=24
+        )
+
+        key = Square(side_length=2.2)
+        key.move_to([-5, 2, 0])
+
+        key_text = MarkupText(
+            f"<b>Key:</b>\n\n<span fgcolor='{YELLOW}'>●</span> Empty Model",
+            font_size=18,
+        )
+
+        key_text.move_to([-5, 2.4, 0])
+
+
+        step_1.move_to([2, 2, 0])
+        self.play(
+            Write(step_1, run_time=2.5),
+            Write(key_text),
+            Write(key)
+        )
+
+        self.add(model)
+        
+
+        cpu_targs = []
+        first_animations = []
+        second_animations = []
+        for i,rect in enumerate(model_base):
+
+            cpu_target = Rectangle(height=0.46,width=0.46).set_stroke(width=0.).set_fill(YELLOW, opacity=0.7)
+            cpu_target.move_to(rect)
+            cpu_target.generate_target()
+            cpu_target.target.height = 0.46/4
+            cpu_target.target.width = 0.46/3
+            
+            if i == 0:
+                cpu_target.target.next_to(cpu_left_col_base[0].get_corner(DOWN+LEFT), buff=0.02, direction=UP)
+                cpu_target.target.set_x(cpu_target.target.get_x()+0.1)
+            elif i == 3:
+                cpu_target.target.next_to(cpu_targs[0].target, direction=UP, buff=0.)
+            else:
+                cpu_target.target.next_to(cpu_targs[i-1].target, direction=RIGHT, buff=0.)
+            cpu_targs.append(cpu_target)
+
+            first_animations.append(rect.animate(run_time=0.5).set_stroke(YELLOW))
+            second_animations.append(MoveToTarget(cpu_target, run_time=1.5))
+
+        self.play(*first_animations)
+        self.play(*second_animations)
+                 
+
+        self.wait()
\ No newline at end of file
diff --git a/manim_animations/big_model_inference/stage_2.py b/manim_animations/big_model_inference/stage_2.py
new file mode 100644
index 000000000..a30e9593b
--- /dev/null
+++ b/manim_animations/big_model_inference/stage_2.py
@@ -0,0 +1,126 @@
+# 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.
+
+from manim import *
+
+class Stage2(Scene):
+    def construct(self):
+        mem = Rectangle(height=0.5,width=0.5)
+        fill = Rectangle(height=0.46,width=0.46).set_stroke(width=0)
+
+        cpu_left_col_base = [mem.copy() for i in range(6)]
+        cpu_right_col_base = [mem.copy() for i in range(6)]
+        cpu_left_col = VGroup(*cpu_left_col_base).arrange(UP, buff=0)
+        cpu_right_col = VGroup(*cpu_right_col_base).arrange(UP, buff=0)
+        cpu_rects = VGroup(cpu_left_col,cpu_right_col).arrange(RIGHT, buff=0)
+        cpu_text = Text("CPU", font_size=24)
+        cpu = Group(cpu_rects,cpu_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        cpu.move_to([-2.5,-.5,0])
+        self.add(cpu)
+
+        gpu_base = [mem.copy() for i in range(4)]
+        gpu_rect = VGroup(*gpu_base).arrange(UP,buff=0)
+        gpu_text = Text("GPU", font_size=24)
+        gpu = Group(gpu_rect,gpu_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        gpu.move_to([-1,-1,0])
+        self.add(gpu)
+
+        model_base = [mem.copy() for i in range(6)]
+        model_rect = VGroup(*model_base).arrange(RIGHT,buff=0)
+
+        model_text = Text("Model", font_size=24)
+        model = Group(model_rect,model_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        model.move_to([3, -1., 0])
+        self.add(model)
+        
+        cpu_targs = []
+        for i,rect in enumerate(model_base):
+            rect.set_stroke(YELLOW)
+            # target = fill.copy().set_fill(YELLOW, opacity=0.7)
+            # target.move_to(rect)
+            # self.add(target)
+
+            cpu_target = Rectangle(height=0.46/4,width=0.46/3).set_stroke(width=0.).set_fill(YELLOW, opacity=0.7)
+            
+            if i == 0:
+                cpu_target.next_to(cpu_left_col_base[0].get_corner(DOWN+LEFT), buff=0.02, direction=UP)
+                cpu_target.set_x(cpu_target.get_x()+0.1)
+            elif i == 3:
+                cpu_target.next_to(cpu_targs[0], direction=UP, buff=0.)
+            else:
+                cpu_target.next_to(cpu_targs[i-1], direction=RIGHT, buff=0.)
+            self.add(cpu_target)
+            cpu_targs.append(cpu_target)
+
+              
+
+        checkpoint_base = [mem.copy() for i in range(6)]
+        checkpoint_rect = VGroup(*checkpoint_base).arrange(RIGHT,buff=0)
+
+        checkpoint_text = Text("Loaded Checkpoint", font_size=24)
+        checkpoint = Group(checkpoint_rect,checkpoint_text).arrange(DOWN, aligned_edge=DOWN, buff=0.4)
+        checkpoint.move_to([3, .5, 0])
+            
+        key = Square(side_length=2.2)
+        key.move_to([-5, 2, 0])
+
+        key_text = MarkupText(
+            f"<b>Key:</b>\n\n<span fgcolor='{YELLOW}'>●</span> Empty Model",
+            font_size=18,
+        )
+
+        key_text.move_to([-5, 2.4, 0])
+
+        self.add(key_text, key)
+
+        blue_text = MarkupText(
+            f"<span fgcolor='{BLUE}'>●</span> Checkpoint",
+            font_size=18,
+        )
+
+        blue_text.next_to(key_text, DOWN*2.4, aligned_edge=key_text.get_left())
+
+        step_2 = MarkupText(
+            f'Next, a <i><span fgcolor="{BLUE}">second</span></i> model is loaded into memory,\nwith the weights of a <span fgcolor="{BLUE}">single shard</span>.', 
+            font_size=24
+        )
+        step_2.move_to([2, 2, 0])
+        self.play(
+            Write(step_2),
+            Write(blue_text)
+        )
+
+        self.play(
+            Write(checkpoint_text, run_time=1),
+            Create(checkpoint_rect, run_time=1)
+        )
+
+        first_animations = []
+        second_animations = []
+        for i,rect in enumerate(checkpoint_base):
+            target = fill.copy().set_fill(BLUE, opacity=0.7)
+            target.move_to(rect)
+            first_animations.append(GrowFromCenter(target, run_time=1))
+
+            cpu_target = target.copy()
+            cpu_target.generate_target()
+            if i < 5:
+                cpu_target.target.move_to(cpu_left_col_base[i+1])
+            else:
+                cpu_target.target.move_to(cpu_right_col_base[i-5])
+            second_animations.append(MoveToTarget(cpu_target, run_time=1.5))
+            
+        self.play(*first_animations)
+        self.play(*second_animations)
+        self.wait()
\ No newline at end of file
diff --git a/manim_animations/big_model_inference/stage_3.py b/manim_animations/big_model_inference/stage_3.py
new file mode 100644
index 000000000..4ba20c4b5
--- /dev/null
+++ b/manim_animations/big_model_inference/stage_3.py
@@ -0,0 +1,158 @@
+# 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.
+
+from manim import *
+
+class Stage3(Scene):
+    def construct(self):
+        mem = Rectangle(height=0.5,width=0.5)
+        meta_mem = Rectangle(height=0.25,width=0.25)
+        fill = Rectangle(height=0.46,width=0.46).set_stroke(width=0)
+
+        cpu_left_col_base = [mem.copy() for i in range(6)]
+        cpu_right_col_base = [mem.copy() for i in range(6)]
+        cpu_left_col = VGroup(*cpu_left_col_base).arrange(UP, buff=0)
+        cpu_right_col = VGroup(*cpu_right_col_base).arrange(UP, buff=0)
+        cpu_rects = VGroup(cpu_left_col,cpu_right_col).arrange(RIGHT, buff=0)
+        cpu_text = Text("CPU", font_size=24)
+        cpu = Group(cpu_rects,cpu_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        cpu.move_to([-2.5,-.5,0])
+        self.add(cpu)
+
+        gpu_base = [mem.copy() for i in range(4)]
+        gpu_rect = VGroup(*gpu_base).arrange(UP,buff=0)
+        gpu_text = Text("GPU", font_size=24)
+        gpu = Group(gpu_rect,gpu_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        gpu.move_to([-1,-1,0])
+        self.add(gpu)
+
+        model_base = [mem.copy() for i in range(6)]
+        model_rect = VGroup(*model_base).arrange(RIGHT,buff=0)
+
+        model_text = Text("Model", font_size=24)
+        model = Group(model_rect,model_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        model.move_to([3, -1., 0])
+        self.add(model)
+
+        model_arr = []
+        model_cpu_arr = []
+        model_meta_arr = []
+        
+        for i,rect in enumerate(model_base):
+            rect.set_stroke(YELLOW)
+
+            cpu_target = Rectangle(height=0.46/4,width=0.46/3).set_stroke(width=0.).set_fill(YELLOW, opacity=0.7)
+            
+            if i == 0:
+                cpu_target.next_to(cpu_left_col_base[0].get_corner(DOWN+LEFT), buff=0.02, direction=UP)
+                cpu_target.set_x(cpu_target.get_x()+0.1)
+            elif i == 3:
+                cpu_target.next_to(model_cpu_arr[0], direction=UP, buff=0.)
+            else:
+                cpu_target.next_to(model_cpu_arr[i-1], direction=RIGHT, buff=0.)
+            self.add(cpu_target)
+            model_cpu_arr.append(cpu_target)
+
+        self.add(*model_arr, *model_cpu_arr, *model_meta_arr)
+
+        checkpoint_base = [mem.copy() for i in range(6)]
+        checkpoint_rect = VGroup(*checkpoint_base).arrange(RIGHT,buff=0)
+
+        checkpoint_text = Text("Loaded Checkpoint", font_size=24)
+        checkpoint = Group(checkpoint_rect,checkpoint_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        checkpoint.move_to([3, .5, 0])
+            
+        self.add(checkpoint)
+
+        ckpt_arr = []
+        ckpt_cpu_arr = []
+
+        for i,rect in enumerate(checkpoint_base):
+            target = fill.copy().set_fill(BLUE, opacity=0.7)
+            target.move_to(rect)
+            ckpt_arr.append(target)
+
+            cpu_target = target.copy()
+            if i < 5:
+                cpu_target.move_to(cpu_left_col_base[i+1])
+            else:
+                cpu_target.move_to(cpu_right_col_base[i-5])
+            ckpt_cpu_arr.append(cpu_target)
+        self.add(*ckpt_arr, *ckpt_cpu_arr)
+
+        key = Square(side_length=2.2)
+        key.move_to([-5, 2, 0])
+
+        key_text = MarkupText(
+            f"<b>Key:</b>\n\n<span fgcolor='{YELLOW}'>●</span> Empty Model",
+            font_size=18,
+        )
+
+        key_text.move_to([-5, 2.4, 0])
+
+        self.add(key_text, key)
+
+        blue_text = MarkupText(
+            f"<span fgcolor='{BLUE}'>●</span> Checkpoint",
+            font_size=18,
+        )
+
+        blue_text.next_to(key_text, DOWN*2.4, aligned_edge=key_text.get_left())
+        self.add(blue_text)
+
+        step_3 = MarkupText(
+            f'Based on the passed in configuration, weights are stored in\na variety of np.memmaps on disk or to a particular device.', 
+            font_size=24
+        )
+        step_3.move_to([2, 2, 0])
+
+        disk_left_col_base = [meta_mem.copy() for i in range(6)]
+        disk_right_col_base = [meta_mem.copy() for i in range(6)]
+        disk_left_col = VGroup(*disk_left_col_base).arrange(UP, buff=0)
+        disk_right_col = VGroup(*disk_right_col_base).arrange(UP, buff=0)
+        disk_rects = VGroup(disk_left_col,disk_right_col).arrange(RIGHT, buff=0)
+        disk_text = Text("Disk", font_size=24)
+        disk = Group(disk_rects,disk_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        disk.move_to([-4.,-1.25,0])
+        self.play(
+            Write(step_3, run_time=3),
+            Write(disk_text, run_time=1),
+            Create(disk_rects, run_time=1)
+        )
+
+        animations = []
+        for i,rect in enumerate(ckpt_cpu_arr):
+            target = rect.copy()
+            target.generate_target()
+            target.target.move_to(disk_left_col_base[i]).scale(0.5)
+            animations.append(MoveToTarget(target, run_time=1.5))
+        self.play(*animations)
+
+        self.play(FadeOut(step_3))
+
+        step_4 = MarkupText(
+            f'Then, the checkpoint is removed from memory\nthrough garbage collection.', 
+            font_size=24
+        )
+        step_4.move_to([2, 2, 0])
+
+        self.play(
+            Write(step_4, run_time=3)
+        )
+
+        self.play(
+            FadeOut(checkpoint_rect, checkpoint_text, *ckpt_arr, *ckpt_cpu_arr),
+        )
+
+        self.wait()      
\ No newline at end of file
diff --git a/manim_animations/big_model_inference/stage_4.py b/manim_animations/big_model_inference/stage_4.py
new file mode 100644
index 000000000..3a79ad97c
--- /dev/null
+++ b/manim_animations/big_model_inference/stage_4.py
@@ -0,0 +1,156 @@
+# 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.
+
+from manim import *
+
+class Stage4(Scene):
+    def construct(self):
+        mem = Rectangle(height=0.5,width=0.5)
+        fill = Rectangle(height=0.46,width=0.46).set_stroke(width=0)
+        meta_mem = Rectangle(height=0.25,width=0.25)
+
+        cpu_left_col_base = [mem.copy() for i in range(6)]
+        cpu_right_col_base = [mem.copy() for i in range(6)]
+        cpu_left_col = VGroup(*cpu_left_col_base).arrange(UP, buff=0)
+        cpu_right_col = VGroup(*cpu_right_col_base).arrange(UP, buff=0)
+        cpu_rects = VGroup(cpu_left_col,cpu_right_col).arrange(RIGHT, buff=0)
+        cpu_text = Text("CPU", font_size=24)
+        cpu = Group(cpu_rects,cpu_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        cpu.move_to([-2.5,-.5,0])
+        self.add(cpu)
+
+        gpu_base = [mem.copy() for i in range(4)]
+        gpu_rect = VGroup(*gpu_base).arrange(UP,buff=0)
+        gpu_text = Text("GPU", font_size=24)
+        gpu = Group(gpu_rect,gpu_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        gpu.move_to([-1,-1,0])
+        self.add(gpu)
+
+        model_base = [mem.copy() for i in range(6)]
+        model_rect = VGroup(*model_base).arrange(RIGHT,buff=0)
+
+        model_text = Text("Model", font_size=24)
+        model = Group(model_rect,model_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        model.move_to([3, -1., 0])
+        self.add(model)
+
+        model_cpu_arr = []
+        model_meta_arr = []
+        
+        for i,rect in enumerate(model_base):
+            rect.set_stroke(YELLOW)
+
+            cpu_target = Rectangle(height=0.46/4,width=0.46/3).set_stroke(width=0.).set_fill(YELLOW, opacity=0.7)
+            
+            if i == 0:
+                cpu_target.next_to(cpu_left_col_base[0].get_corner(DOWN+LEFT), buff=0.02, direction=UP)
+                cpu_target.set_x(cpu_target.get_x()+0.1)
+            elif i == 3:
+                cpu_target.next_to(model_cpu_arr[0], direction=UP, buff=0.)
+            else:
+                cpu_target.next_to(model_cpu_arr[i-1], direction=RIGHT, buff=0.)
+            self.add(cpu_target)
+            model_cpu_arr.append(cpu_target)
+
+        self.add(*model_cpu_arr, *model_meta_arr)
+
+        disk_left_col_base = [meta_mem.copy() for i in range(6)]
+        disk_right_col_base = [meta_mem.copy() for i in range(6)]
+        disk_left_col = VGroup(*disk_left_col_base).arrange(UP, buff=0)
+        disk_right_col = VGroup(*disk_right_col_base).arrange(UP, buff=0)
+        disk_rects = VGroup(disk_left_col,disk_right_col).arrange(RIGHT, buff=0)
+        disk_text = Text("Disk", font_size=24)
+        disk = Group(disk_rects,disk_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        disk.move_to([-4.,-1.25,0])
+        self.add(disk_text, disk_rects)
+
+        cpu_disk_arr = []
+
+        for i in range(6):
+            target = fill.copy().set_fill(BLUE, opacity=0.8)
+            target.move_to(disk_left_col_base[i]).scale(0.5)
+            cpu_disk_arr.append(target)
+
+        self.add(*cpu_disk_arr)
+
+        key = Square(side_length=2.2)
+        key.move_to([-5, 2, 0])
+
+        key_text = MarkupText(
+            f"<b>Key:</b>\n\n<span fgcolor='{YELLOW}'>●</span> Empty Model",
+            font_size=18,
+        )
+
+        key_text.move_to([-5, 2.4, 0])
+
+        self.add(key_text, key)
+
+        blue_text = MarkupText(
+            f"<span fgcolor='{BLUE}'>●</span> Checkpoint",
+            font_size=18,
+        )
+
+        blue_text.next_to(key_text, DOWN*2.4, aligned_edge=key_text.get_left())
+        self.add(blue_text)
+
+        step_5 = MarkupText(
+            f'The offloaded weights are all sent to the CPU.', 
+            font_size=24
+        )
+        step_5.move_to([2, 2, 0])
+
+        self.play(Write(step_5, run_time=3))
+
+        for i in range(6):
+            rect = cpu_disk_arr[i]
+            cp2 = rect.copy().set_fill(BLUE, opacity=0.8).scale(2.0)
+            cp2.generate_target()
+            cp2.target.move_to(model_base[i])
+
+            if i == 0:
+                rect.set_fill(BLUE, opacity=0.8)
+                rect.generate_target()
+                rect.target.move_to(cpu_left_col_base[0]).scale(2.0)
+                
+                self.remove(*model_meta_arr, 
+                    *model_cpu_arr,
+                )
+
+            else:
+                rect.generate_target()
+                rect.target.move_to(cpu_left_col_base[i]).scale(2.0)
+            self.play(
+                MoveToTarget(rect),
+                MoveToTarget(cp2),
+                model_base[i].animate.set_stroke(WHITE)
+            )
+        self.play(FadeOut(step_5))
+
+        step_5 = MarkupText(
+            f'Finally, hooks are added to each weight in the model\nto transfer the weights from CPU to GPU\n\t\tand back when needed.', 
+            font_size=24
+        )
+        step_5.move_to([2, 2, 0])
+
+        self.play(Write(step_5, run_time=3))
+
+        arrows = []
+        animations = []
+        for i in range(6):
+            a = Arrow(start=UP, end=DOWN, color=RED, buff=.5)
+            a.next_to(model_base[i].get_left(), UP, buff=0.2)
+            arrows.append(a)
+            animations.append(Write(a))
+        self.play(*animations)
+        self.wait()  
\ No newline at end of file
diff --git a/manim_animations/big_model_inference/stage_5.py b/manim_animations/big_model_inference/stage_5.py
new file mode 100644
index 000000000..8b2ff3357
--- /dev/null
+++ b/manim_animations/big_model_inference/stage_5.py
@@ -0,0 +1,221 @@
+# 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.
+
+from manim import *
+
+class Stage5(Scene):
+    def construct(self):
+        mem = Rectangle(height=0.5,width=0.5)
+        fill = Rectangle(height=0.46,width=0.46).set_stroke(width=0)
+
+        meta_mem = Rectangle(height=0.25,width=0.25)
+
+        cpu_left_col_base = [mem.copy() for i in range(6)]
+        cpu_right_col_base = [mem.copy() for i in range(6)]
+        cpu_left_col = VGroup(*cpu_left_col_base).arrange(UP, buff=0)
+        cpu_right_col = VGroup(*cpu_right_col_base).arrange(UP, buff=0)
+        cpu_rects = VGroup(cpu_left_col,cpu_right_col).arrange(RIGHT, buff=0)
+        cpu_text = Text("CPU", font_size=24)
+        cpu = Group(cpu_rects,cpu_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        cpu.move_to([-2.5,-.5,0])
+        self.add(cpu)
+
+        gpu_base = [mem.copy() for i in range(4)]
+        gpu_rect = VGroup(*gpu_base).arrange(UP,buff=0)
+        gpu_text = Text("GPU", font_size=24)
+        gpu = Group(gpu_rect,gpu_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        gpu.move_to([-1,-1,0])
+        self.add(gpu)
+
+        model_base = [mem.copy() for i in range(6)]
+        model_rect = VGroup(*model_base).arrange(RIGHT,buff=0)
+
+        model_text = Text("Model", font_size=24)
+        model = Group(model_rect,model_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        model.move_to([3, -1., 0])
+        self.add(model)
+
+        model_arr = []
+        model_cpu_arr = []
+        
+        for i,rect in enumerate(model_base):
+            target = fill.copy().set_fill(BLUE, opacity=0.8)
+            target.move_to(rect)
+            model_arr.append(target)
+
+            cpu_target = Rectangle(height=0.46,width=0.46).set_stroke(width=0.).set_fill(BLUE, opacity=0.8)
+            cpu_target.move_to(cpu_left_col_base[i])
+            model_cpu_arr.append(cpu_target)
+
+        self.add(*model_arr, *model_cpu_arr)
+
+        disk_left_col_base = [meta_mem.copy() for i in range(6)]
+        disk_right_col_base = [meta_mem.copy() for i in range(6)]
+        disk_left_col = VGroup(*disk_left_col_base).arrange(UP, buff=0)
+        disk_right_col = VGroup(*disk_right_col_base).arrange(UP, buff=0)
+        disk_rects = VGroup(disk_left_col,disk_right_col).arrange(RIGHT, buff=0)
+        disk_text = Text("Disk", font_size=24)
+        disk = Group(disk_rects,disk_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        disk.move_to([-4,-1.25,0])
+        self.add(disk_text, disk_rects)
+
+        key = Square(side_length=2.2)
+        key.move_to([-5, 2, 0])
+
+        key_text = MarkupText(
+            f"<b>Key:</b>\n\n<span fgcolor='{YELLOW}'>●</span> Empty Model",
+            font_size=18,
+        )
+
+        key_text.move_to([-5, 2.4, 0])
+
+        self.add(key_text, key)
+
+        blue_text = MarkupText(
+            f"<span fgcolor='{BLUE}'>●</span> Checkpoint",
+            font_size=18,
+        )
+
+        blue_text.next_to(key_text, DOWN*2.4, aligned_edge=key_text.get_left())
+        self.add(blue_text)
+
+        step_6 = MarkupText(
+            f'Now watch as an input is passed through the model\nand how the memory is utilized and handled.', 
+            font_size=24
+        )
+        step_6.move_to([2, 2, 0])
+
+        self.play(Write(step_6))
+
+        input = Square(0.3)
+        input.set_fill(RED, opacity=1.)
+        input.set_stroke(width=0.)
+        input.next_to(model_base[0], LEFT, buff=.5)
+
+        self.play(Write(input))
+
+        input.generate_target()
+        input.target.next_to(model_arr[0], direction=LEFT, buff=0.02)
+        self.play(MoveToTarget(input))
+
+        self.play(FadeOut(step_6))
+
+
+        a = Arrow(start=UP, end=DOWN, color=RED, buff=.5)
+        a.next_to(model_arr[0].get_left(), UP, buff=0.2)
+
+        model_cpu_arr[0].generate_target()
+        model_cpu_arr[0].target.move_to(gpu_rect[0])
+
+        step_7 = MarkupText(
+            f'As the input reaches a layer, the hook triggers\nand weights are moved from the CPU\nto the GPU and back.', 
+            font_size=24
+        )
+        step_7.move_to([2, 2, 0])
+
+        self.play(Write(step_7, run_time=3))
+
+        circ_kwargs = {"run_time":1, "fade_in":True, "fade_out":True, "buff":0.02}
+
+        self.play(
+            Write(a), 
+            Circumscribe(model_arr[0], color=ORANGE, **circ_kwargs),
+            Circumscribe(model_cpu_arr[0], color=ORANGE, **circ_kwargs),
+            Circumscribe(gpu_rect[0], color=ORANGE, **circ_kwargs),
+        )
+        self.play(
+            MoveToTarget(model_cpu_arr[0])
+        )
+
+        a_c = a.copy()
+        for i in range(6):
+            a_c.next_to(model_arr[i].get_right()+0.02, UP, buff=0.2)
+
+            input.generate_target()
+            input.target.move_to(model_arr[i].get_right()+0.02)
+
+            grp = AnimationGroup(
+                FadeOut(a, run_time=.5), 
+                MoveToTarget(input, run_time=.5), 
+                FadeIn(a_c, run_time=.5),
+                lag_ratio=0.2
+            )
+
+            self.play(grp)
+
+
+            model_cpu_arr[i].generate_target()
+            model_cpu_arr[i].target.move_to(cpu_left_col_base[i])
+
+
+            if i < 5:
+                model_cpu_arr[i+1].generate_target()
+                model_cpu_arr[i+1].target.move_to(gpu_rect[0])
+                if i >= 1:
+                    circ_kwargs["run_time"] = .7
+
+                self.play(
+                    Circumscribe(model_arr[i], **circ_kwargs),
+                    Circumscribe(cpu_left_col_base[i], **circ_kwargs),
+                    Circumscribe(cpu_left_col_base[i+1], color=ORANGE, **circ_kwargs),                    
+                    Circumscribe(gpu_rect[0], color=ORANGE, **circ_kwargs),
+                    Circumscribe(model_arr[i+1], color=ORANGE, **circ_kwargs),
+                )
+                if i < 1:
+                    self.play(
+                        MoveToTarget(model_cpu_arr[i]), 
+                        MoveToTarget(model_cpu_arr[i+1]),
+                    )
+                else:
+                    self.play(
+                        MoveToTarget(model_cpu_arr[i], run_time=.7), 
+                        MoveToTarget(model_cpu_arr[i+1], run_time=.7),
+                    )
+            else:
+                model_cpu_arr[i].generate_target()
+                model_cpu_arr[i].target.move_to(cpu_left_col_base[-1])
+                input.generate_target()
+                input.target.next_to(model_arr[-1].get_right(), RIGHT+0.02, buff=0.2)
+
+                self.play(
+                    Circumscribe(model_arr[-1], color=ORANGE, **circ_kwargs),
+                    Circumscribe(cpu_left_col_base[-1], color=ORANGE, **circ_kwargs),
+                    Circumscribe(gpu_rect[0], color=ORANGE, **circ_kwargs),
+                )
+
+                self.play(
+                    MoveToTarget(model_cpu_arr[i])
+                )
+
+            a = a_c
+            a_c = a_c.copy()
+
+        input.generate_target()
+        input.target.next_to(model_base[-1], RIGHT+0.02, buff=.5)
+        self.play(
+            FadeOut(step_7),
+            FadeOut(a, run_time=.5), 
+        )
+
+        step_8 = MarkupText(
+            f'Inference on a model too large for GPU memory\nis successfully completed.', font_size=24
+        )
+        step_8.move_to([2, 2, 0])
+
+        self.play(
+            Write(step_8, run_time=3),
+            MoveToTarget(input)
+        )
+
+        self.wait()
\ No newline at end of file

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/671
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index 074d02e4a..3efb61034 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -14,6 +14,7 @@
 
 import importlib
 import sys
+from functools import lru_cache
 
 import torch
 
@@ -50,6 +51,7 @@ def is_apex_available():
     return importlib.util.find_spec("apex") is not None
 
 
+@lru_cache()
 def is_tpu_available(check_device=True):
     "Checks if `torch_xla` is installed and potentially if a TPU is in the environment"
     if _tpu_available and check_device:

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/670
diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml
index 8160cb4ee..6b9825523 100644
--- a/docs/source/_toctree.yml
+++ b/docs/source/_toctree.yml
@@ -35,6 +35,8 @@
     title: Using 🤗 Accelerate on SageMaker
   - local: usage_guides/mps
     title: How to use Apple Silicon M1 GPUs
+  - local: usage_guides/megatron_lm
+    title: How to use Megatron-LM
   - local: usage_guides/training_zoo
     title: 🤗 Accelerate Example Zoo
   title: How-To Guides
@@ -71,4 +73,6 @@
     title: Kwargs handlers
   - local: package_reference/utilities
     title: Utility functions and classes
+  - local: package_reference/megatron_lm
+    title: Megatron-LM Utilities
   title: "Reference"
\ No newline at end of file
diff --git a/docs/source/package_reference/megatron_lm.mdx b/docs/source/package_reference/megatron_lm.mdx
new file mode 100644
index 000000000..b59768bec
--- /dev/null
+++ b/docs/source/package_reference/megatron_lm.mdx
@@ -0,0 +1,29 @@
+<!--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.
+-->
+
+# Utilities for Megatron-LM
+
+[[autodoc]] utils.MegatronLMPlugin
+
+[[autodoc]] utils.MegatronLMDummyScheduler
+
+[[autodoc]] utils.MegatronLMDummyDataLoader
+
+[[autodoc]] utils.AbstractTrainStep
+
+[[autodoc]] utils.GPTTrainStep
+
+[[autodoc]] utils.BertTrainStep
+
+[[autodoc]] utils.T5TrainStep
+
+[[autodoc]] utils.avg_losses_across_data_parallel_group
diff --git a/docs/source/usage_guides/megatron_lm.mdx b/docs/source/usage_guides/megatron_lm.mdx
new file mode 100644
index 000000000..188dddb32
--- /dev/null
+++ b/docs/source/usage_guides/megatron_lm.mdx
@@ -0,0 +1,507 @@
+<!--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.
+-->
+
+
+# Megatron-LM
+
+[Megatron-LM](https://github.com/NVIDIA/Megatron-LM) enables training large transformer language models at scale.
+It provides efficient tensor, pipeline and sequence based model parallelism for pre-training transformer based
+Language Models such as [GPT](https://arxiv.org/abs/2005.14165) (Decoder Only), [BERT](https://arxiv.org/pdf/1810.04805.pdf) (Encoder Only) and [T5](https://arxiv.org/abs/1910.10683) (Encoder-Decoder).
+For detailed information and how things work behind the scene please refer the github [repo](https://github.com/NVIDIA/Megatron-LM).
+
+## What is integrated?
+
+Accelerate integrates following feature of Megatron-LM to enable large scale pre-training/finetuning
+of BERT (Encoder), GPT (Decoder) or T5 models (Encoder and Decoder):
+
+a. **Tensor Parallelism (TP)**: Reduces memory footprint without much additional communication on intra-node ranks
+
+b. **Pipeline Parallelism (PP)**: Reduces memory footprint and enables large scale training via inter-node parallelization. 
+Reduces the bubble of naive PP via PipeDream-Flush schedule/1F1B schedule and Interleaved 1F1B schedule. 
+Layers are distributed uniformly across PP stages.
+
+c. **Sequence Parallelism (SP)**: Reduces memory footprint without any additional communication. Only applicable when using TP.
+
+d. **Data Parallelism (DP)** via Distributed Optimizer: Reduces memory footprint by sharding optimizer states and gradients across DP ranks.
+
+e. **Selective Activation Recomputation**: Reduces the memory footprint of activations significantly via smart activation checkpointing.
+It doesn't store activations occupying large memory while being fast to recompute thereby achieving great tradeoff between memory and recomputation.
+
+f. **Fused Kernels**: Fused Softmax, Mixed Precision Fused Layer Norm and  Fused gradient accumulation to weight gradient computation of linear layer.
+PyTorch JIT compiled Fused GeLU and Fused Bias+Dropout+Residual addition.
+
+g. **Support for Indexed datasets**: Efficient binary format of datasets for large scale training. Support for the `mmap`, `cached` index file and the `lazy` loader format.
+
+h. **Checkpoint reshaping and interoperability**: Utility for reshaping Megatron-LM checkpoints of variable 
+tensor and pipeline parallel sizes to the beloved 🤗 Transformers sharded checkpoints as it has great support with plethora of tools
+such as 🤗 Accelerate Big Model Inference, Megatron-DeepSpeed Inference etc. 
+Support is also available for converting 🤗 Transformers sharded checkpoints to Megatron-LM checkpoint of variable tensor and pipeline parallel sizes
+for large scale training.  
+
+
+## Pre-Requisites 
+
+You will need to install the latest pytorch, cuda, nccl, and NVIDIA [APEX](https://github.com/NVIDIA/apex#quick-start) releases and the nltk library.
+See [documentation](https://github.com/NVIDIA/Megatron-LM#setup) for more details. 
+Another way to setup the environment is to pull an NVIDIA PyTorch Container that comes with all the required installations from NGC.
+
+Below is a step-by-step method to set up the conda environment:
+
+1. Create a virtual environment
+```
+conda create --name ml
+```
+
+2. Assuming that the machine has CUDA 11.3 installed, installing the corresponding PyTorch GPU Version
+```
+conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch
+```
+
+3. Install Nvidia APEX
+```
+git clone https://github.com/NVIDIA/apex
+cd apex
+pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
+cd ..
+```
+
+4. Installing Megatron-LM
+
+    a. Cloning the Megatron-LM repo
+    ```
+    git clone https://github.com/NVIDIA/Megatron-LM.git
+    cd Megatron-LM
+    ```
+
+    b. Create a file `setup.py`, paste the below code and put in the root folder
+    ```python
+    """Setup for pip package."""
+
+    import os
+    import sys
+    import setuptools
+
+    if sys.version_info < (3,):
+        raise Exception("Python 2 is not supported by Megatron.")
+
+    with open("README.md", "r") as fh:
+        long_description = fh.read()
+
+    setuptools.setup(
+        name="megatron-lm",
+        version="3.0.0",
+        description="Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism.",
+        long_description=long_description,
+        long_description_content_type="text/markdown",
+        # The project's main homepage.
+        url="https://github.com/NVIDIA/Megatron-LM",
+        author="NVIDIA INC",
+        maintainer="NVIDIA INC",
+        # The licence under which the project is released
+        license="See https://github.com/NVIDIA/Megatron-LM/blob/master/LICENSE",
+        classifiers=[
+            "Intended Audience :: Developers",
+            "Intended Audience :: Science/Research",
+            "Intended Audience :: Information Technology",
+            # Indicate what your project relates to
+            "Topic :: Scientific/Engineering :: Artificial Intelligence",
+            "Topic :: Software Development :: Libraries :: Python Modules",
+            # Additional Setting
+            "Environment :: Console",
+            "Natural Language :: English",
+            "Operating System :: OS Independent",
+        ],
+        python_requires=">=3.6",
+        packages=setuptools.find_packages(),
+        install_requires=["nltk", "six", "regex", "torch>=1.12.0", "pybind11"],
+        # Add in any packaged data.
+        include_package_data=True,
+        zip_safe=False,
+        # PyPI package information.
+        keywords="deep learning, Megatron, gpu, NLP, nvidia, pytorch, torch, language",
+    )
+    ```
+
+    c. installing via below command
+    ```
+    pip install "."
+    ```
+
+## Accelerate Megatron-LM Plugin
+
+Important features are directly supported via the `accelerate config` command. 
+An example of thr corresponding questions for using Megatron-LM features is shown below:
+
+```bash
+:~$ accelerate config --config_file "megatron_gpt_config.yaml"
+In which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)): 0
+Which type of machine are you using? ([0] No distributed training, [1] multi-CPU, [2] multi-GPU, [3] TPU [4] MPS): 2
+How many different machines will you use (use more than 1 for multi-node training)? [1]: 
+Do you want to use DeepSpeed? [yes/NO]: 
+Do you want to use FullyShardedDataParallel? [yes/NO]: 
+Do you want to use Megatron-LM ? [yes/NO]: yes
+What is the Tensor Parallelism degree/size? [1]:2
+Do you want to enable Sequence Parallelism? [YES/no]: 
+What is the Pipeline Parallelism degree/size? [1]:2
+What is the number of micro-batches? [1]:2
+Do you want to enable selective activation recomputation? [YES/no]: 
+Do you want to use distributed optimizer which shards optimizer state and gradients across data pralellel ranks? [YES/no]: 
+What is the gradient clipping value based on global L2 Norm (0 to disable)? [1.0]: 
+How many GPU(s) should be used for distributed training? [1]:8
+Do you wish to use FP16 or BF16 (mixed precision)? [NO/fp16/bf16]: bf16
+```
+
+The resulting config is shown below:
+
+```
+~$ cat megatron_gpt_config.yaml 
+compute_environment: LOCAL_MACHINE
+deepspeed_config: {}
+distributed_type: MEGATRON_LM
+downcast_bf16: 'no'
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+megatron_lm_config:
+  megatron_lm_gradient_clipping: 1.0
+  megatron_lm_num_micro_batches: 2
+  megatron_lm_pp_degree: 2
+  megatron_lm_recompute_activations: true
+  megatron_lm_sequence_parallelism: true
+  megatron_lm_tp_degree: 2
+  megatron_lm_use_distributed_optimizer: true
+mixed_precision: bf16
+num_machines: 1
+num_processes: 8
+rdzv_backend: static
+same_network: true
+use_cpu: false
+```
+
+We will take the example of GPT pre-training. The minimal changes required to the official `run_clm_no_trainer.py` 
+to use Megatron-LM are as follows:
+
+1. As Megatron-LM uses its own implementation of Optimizer, the corresponding scheduler compatible with it needs to be used.
+As such, support for only the Megatron-LM's scheduler is present. User will need to create `accelerate.utils.MegatronLMDummyScheduler`.
+Example is given below:
+
+```python
+from accelerate.utils import MegatronLMDummyScheduler
+
+if accelerator.distributed_type == DistributedType.MEGATRON_LM:
+    lr_scheduler = MegatronLMDummyScheduler(
+        optimizer=optimizer,
+        total_num_steps=args.max_train_steps,
+        warmup_num_steps=args.num_warmup_steps,
+    )
+else:
+    lr_scheduler = get_scheduler(
+        name=args.lr_scheduler_type,
+        optimizer=optimizer,
+        num_warmup_steps=args.num_warmup_steps * args.gradient_accumulation_steps,
+        num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
+    )
+```
+
+2. Getting the details of the total batch size now needs to be cognization of tensor and pipeline parallel sizes.
+Example of getting the effective total batch size is shown below:
+
+```python
+if accelerator.distributed_type == DistributedType.MEGATRON_LM:
+    total_batch_size = accelerator.state.megatron_lm_plugin.global_batch_size
+else:
+    total_batch_size = args.per_device_train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+```
+
+3. When using Megatron-LM, the losses are already averaged across the data parallel group
+
+```python
+if accelerator.distributed_type == DistributedType.MEGATRON_LM:
+    losses.append(loss)
+else:
+    losses.append(accelerator.gather_for_metrics(loss.repeat(args.per_device_eval_batch_size)))
+
+if accelerator.distributed_type == DistributedType.MEGATRON_LM:
+    losses = torch.tensor(losses)
+else:
+    losses = torch.cat(losses)
+```
+
+4. For Megatron-LM, we need to save the model using `accelerator.save_state`
+
+```python
+if accelerator.distributed_type == DistributedType.MEGATRON_LM:
+    accelerator.save_state(args.output_dir)
+else:
+    unwrapped_model = accelerator.unwrap_model(model)
+    unwrapped_model.save_pretrained(
+        args.output_dir, is_main_process=accelerator.is_main_process, save_function=accelerator.save
+    )
+```
+
+That's it! We are good to go 🚀. Please find the example script in the examples folder at the path `accelerate/examples/by_feature/megatron_lm_gpt_pretraining.py`.
+Let's run it for `gpt-large` model architecture using 4 A100-80GB GPUs.
+
+```bash
+accelerate launch --config_file megatron_gpt_config.yaml \
+examples/by_feature/megatron_lm_gpt_pretraining.py \
+--config_name "gpt2-large" \
+--tokenizer_name "gpt2-large" \
+--dataset_name wikitext \
+--dataset_config_name wikitext-2-raw-v1 \
+--block_size 1024 \
+--learning_rate 5e-5 \
+--per_device_train_batch_size 4 \
+--per_device_eval_batch_size 4 \
+--num_train_epochs 5 \
+--with_tracking \
+--report_to "wandb" \
+--output_dir "awesome_model"
+```
+
+Below are some important excerpts from the output logs:
+
+```bash
+Loading extension module fused_dense_cuda...
+>>> done with compiling and loading fused kernels. Compilation time: 3.569 seconds
+ > padded vocab (size: 50257) with 175 dummy tokens (new size: 50432)
+Building gpt model in the pre-training mode.
+The Megatron LM model weights are initialized at random in `accelerator.prepare`. Please use `accelerator.load_checkpoint` to load a pre-trained checkpoint matching the distributed setup.
+Preparing dataloader
+Preparing dataloader
+Preparing model
+ > number of parameters on (tensor, pipeline) model parallel rank (1, 0): 210753280
+ > number of parameters on (tensor, pipeline) model parallel rank (1, 1): 209445120
+ > number of parameters on (tensor, pipeline) model parallel rank (0, 0): 210753280
+ > number of parameters on (tensor, pipeline) model parallel rank (0, 1): 209445120
+Preparing optimizer
+Preparing scheduler
+> learning rate decay style: linear
+10/10/2022 22:57:22 - INFO - __main__ - ***** Running training *****
+10/10/2022 22:57:22 - INFO - __main__ -   Num examples = 2318
+10/10/2022 22:57:22 - INFO - __main__ -   Num Epochs = 5
+10/10/2022 22:57:22 - INFO - __main__ -   Instantaneous batch size per device = 24
+10/10/2022 22:57:22 - INFO - __main__ -   Total train batch size (w. parallel, distributed & accumulation) = 48
+10/10/2022 22:57:22 - INFO - __main__ -   Gradient Accumulation steps = 1
+10/10/2022 22:57:22 - INFO - __main__ -   Total optimization steps = 245
+ 20%|████████████▍                                                 | 49/245 [01:04<04:09,  1.27s/it]
+ 10/10/2022 22:58:29 - INFO - __main__ - epoch 0: perplexity: 1222.1594275215962 eval_loss: 7.10837459564209
+ 40%|████████████████████████▊                                     | 98/245 [02:10<03:07,  1.28s/it]
+ 10/10/2022 22:59:35 - INFO - __main__ - epoch 1: perplexity: 894.5236583794557 eval_loss: 6.796291351318359
+ 60%|████████████████████████████████████▌                        | 147/245 [03:16<02:05,  1.28s/it]
+ 10/10/2022 23:00:40 - INFO - __main__ - epoch 2: perplexity: 702.8458788508042 eval_loss: 6.555137634277344
+ 80%|████████████████████████████████████████████████▊            | 196/245 [04:22<01:02,  1.28s/it]
+ 10/10/2022 23:01:46 - INFO - __main__ - epoch 3: perplexity: 600.3220028695281 eval_loss: 6.39746618270874
+100%|█████████████████████████████████████████████████████████████| 245/245 [05:27<00:00,  1.28s/it]
+```
+
+There are a large number of other options/features that one can set using `accelerate.utils.MegatronLMPlugin`.
+
+## Advanced features to leverage writing custom train step and Megatron-LM Indexed Datasets
+
+For leveraging more features, please go through below details.
+
+1. Below is an example of changes required to customize the Train Step while using Megatron-LM. 
+You will implement the `accelerate.utils.AbstractTrainStep` or inherit from their corresponding children 
+`accelerate.utils.GPTTrainStep`, `accelerate.utils.BertTrainStep` or `accelerate.utils.T5TrainStep`.
+
+```python
+from accelerate.utils import MegatronLMDummyScheduler, GPTTrainStep, avg_losses_across_data_parallel_group
+
+# Custom loss function for the Megatron model
+class GPTTrainStepWithCustomLoss(GPTTrainStep):
+    def __init__(self, megatron_args, **kwargs):
+        super().__init__(megatron_args)
+        self.kwargs = kwargs
+
+    def get_loss_func(self):
+        def loss_func(inputs, loss_mask, output_tensor):
+            batch_size, seq_length = output_tensor.shape
+            losses = output_tensor.float()
+            loss_mask = loss_mask.view(-1).float()
+            loss = losses.view(-1) * loss_mask
+
+            # Resize and average loss per sample
+            loss_per_sample = loss.view(batch_size, seq_length).sum(axis=1)
+            loss_mask_per_sample = loss_mask.view(batch_size, seq_length).sum(axis=1)
+            loss_per_sample = loss_per_sample / loss_mask_per_sample
+
+            # Calculate and scale weighting
+            weights = torch.stack([(inputs == kt).float() for kt in self.kwargs["keytoken_ids"]]).sum(axis=[0, 2])
+            weights = 1.0 + self.kwargs["alpha"] * weights
+            # Calculate weighted average
+            weighted_loss = (loss_per_sample * weights).mean()
+
+            # Reduce loss across data parallel groups
+            averaged_loss = avg_losses_across_data_parallel_group([weighted_loss])
+
+            return weighted_loss, {"lm loss": averaged_loss[0]}
+
+        return loss_func
+
+    def get_forward_step_func(self):
+        def forward_step(data_iterator, model):
+            """Forward step."""
+            # Get the batch.
+            tokens, labels, loss_mask, attention_mask, position_ids = self.get_batch(data_iterator)
+            output_tensor = model(tokens, position_ids, attention_mask, labels=labels)
+
+            return output_tensor, partial(self.loss_func, tokens, loss_mask)
+
+        return forward_step
+
+
+def main():
+    # Custom loss function for the Megatron model
+    keytoken_ids = []
+    keywords = ["plt", "pd", "sk", "fit", "predict", " plt", " pd", " sk", " fit", " predict"]
+    for keyword in keywords:
+        ids = tokenizer([keyword]).input_ids[0]
+        if len(ids) == 1:
+            keytoken_ids.append(ids[0])
+    accelerator.print(f"Keytoken ids: {keytoken_ids}")
+    accelerator.state.megatron_lm_plugin.custom_train_step_class = GPTTrainStepWithCustomLoss
+    accelerator.state.megatron_lm_plugin.custom_train_step_kwargs = {
+        "keytoken_ids": keytoken_ids,
+        "alpha": 0.25,
+    }
+```
+
+2. For using the Megatron-LM datasets, a few more changes are required. Dataloaders for these datasets
+are available only on rank 0 of each tensor parallel group. As such, there are rank where dataloader won't be
+avaiable and this requires tweaks to the training loop. Being able to do all this shows how
+felixble and extensible 🤗 Accelerate is. The changes required are as follows.
+
+a. For Megatron-LM indexed datasets, we need to use `MegatronLMDummyDataLoader` 
+and pass the required dataset args to it such as `data_path`, `seq_length` etc. 
+See [here](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/arguments.py#L804) for the list of available args. 
+    
+```python
+from accelerate.utils import MegatronLMDummyDataLoader
+
+megatron_dataloader_config = {
+    "data_path": args.data_path,
+    "splits_string": args.splits_string,
+    "seq_length": args.block_size,
+    "micro_batch_size": args.per_device_train_batch_size,
+}
+megatron_dataloader = MegatronLMDummyDataLoader(**megatron_dataloader_config)
+accelerator.state.megatron_lm_plugin.megatron_dataset_flag = True
+```
+
+b. `megatron_dataloader` is repeated 3 times to get training, validation and test dataloaders
+as per the `args.splits_string` proportions
+    
+```python
+model, optimizer, lr_scheduler, train_dataloader, eval_dataloader, _ = accelerator.prepare(
+    model, optimizer, lr_scheduler, megatron_dataloader, megatron_dataloader, megatron_dataloader
+)
+```
+
+c. Changes to training and evaluation loops as dataloader is only available on tensor parallel ranks 0
+So, we need to iterate only if the dataloader isn't `None` else provide empty dict
+As such, we loop using `while` loop and break when `completed_steps` is equal to `args.max_train_steps`
+This is similar to the Megatron-LM setup wherein user has to provide `max_train_steps` when using Megaton-LM indexed datasets.
+This displays how flexible and extensible 🤗 Accelerate is.
+
+```python
+while completed_steps < args.max_train_steps:
+    model.train()
+    batch = next(train_dataloader) if train_dataloader is not None else {}
+    outputs = model(**batch)
+    loss = outputs.loss
+    ...
+
+    if completed_steps % eval_interval == 0:
+        eval_completed_steps = 0
+        losses = []
+        while eval_completed_steps < eval_iters:
+            model.eval()
+            with torch.no_grad():
+                batch = next(eval_dataloader) if eval_dataloader is not None else {}
+                outputs = model(**batch)
+```
+
+    
+## Utility for Checkpoint reshaping and interoperability
+
+1. The scripts for these are present in 🤗 Transformers library under respective models. 
+Currently, it is available for GPT model [checkpoint_reshaping_and_interoperability.py](https://github.com/huggingface/transformers/blob/main/src/transformers/models/megatron_gpt2/checkpoint_reshaping_and_interoperability.py)
+
+2. Below is an example of conversion of checkpoint from Megatron-LM to universal 🤗 Transformers sharded checkpoint.
+```bash
+python checkpoint_reshaping_and_interoperability.py \
+--convert_checkpoint_from_megatron_to_transformers \
+--load_path "gpt/iter_0005000" \
+--save_path "gpt/trfs_checkpoint" \
+--max_shard_size "200MB" \
+--tokenizer_name "gpt2" \
+--print-checkpoint-structure
+```
+
+3. Conversion of checkpoint from transformers to megatron with `tp_size=2`, `pp_size=2` and `dp_size=2`.
+```bash
+python checkpoint_utils/megatgron_gpt2/checkpoint_reshaping_and_interoperability.py \
+--load_path "gpt/trfs_checkpoint" \
+--save_path "gpt/megatron_lm_checkpoint" \
+--target_tensor_model_parallel_size 2 \
+--target_pipeline_model_parallel_size 2 \
+--target_data_parallel_size 2 \
+--target_params_dtype "bf16" \
+--make_vocab_size_divisible_by 128 \
+--use_distributed_optimizer \
+--print-checkpoint-structure
+```
+
+## Caveats
+
+1. Supports Transformers GPT2, Megatron-BERT and T5 models.
+This covers Decoder only, Encode only and Encoder-Decoder model classes.
+
+2. Only loss is returned from model forward pass as 
+there is quite complex interplay of pipeline, tensor and data parallelsim behind the scenes.
+The `model(**batch_data)` call return loss(es) averaged across the data parallel ranks.
+This is fine for most cases wherein pre-training jobs are run using Megatron-LM features and
+you can easily compute the `perplexity` using the loss. 
+
+3. The main process is the last rank as the losses are available in the last stage of pipeline.
+`accelerator.is_main_process` and `accelerator.is_local_main_process` return `True` for last rank when using 
+Megatron-LM integration.
+
+4. In `accelerator.prepare` call, a Megatron-LM model corresponding to a given Transformers model is created
+with random weights. Please use `accelerator.load_state` to load the Megatron-LM checkpoint with matching TP, PP and DP partitions.
+
+5. Currently, checkpoint reshaping and interoperability support is only available for GPT. 
+Soon it will be extended to BERT and T5.
+
+6. `gradient_accumulation_steps` needs to be 1. When using Megatron-LM, micro batches in pipeline parallelism 
+setting is synonymous with gradient accumulation. 
+
+7. When using Megatron-LM, use `accelerator.save_state` and `accelerator.load_state` for saving and loading checkpoints.
+
+8. Below are the mapping from Megatron-LM model architectures to the the equivalent 🤗 transformers model architectures.
+Only these 🤗 transformers model architectures are supported.
+
+a. Megatron-LM [BertModel](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/bert_model.py) : 
+🤗 transformers models with `megatron-bert` in config's model type, e.g., 
+[MegatronBERT](https://huggingface.co/docs/transformers/model_doc/megatron-bert)
+    
+b. Megatron-LM [GPTModel](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py) : 
+🤗 transformers models with `gpt2` in config's model type, e.g., 
+[OpenAI GPT2](https://huggingface.co/docs/transformers/model_doc/gpt2)
+   
+c. Megatron-LM [T5Model](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/t5_model.py) : 
+🤗 transformers models with `t5` in  config's model type, e.g., 
+[T5](https://huggingface.co/docs/transformers/model_doc/t5) and 
+[MT5](https://huggingface.co/docs/transformers/model_doc/mt5)
\ No newline at end of file
diff --git a/examples/by_feature/megatron_lm_gpt_pretraining.py b/examples/by_feature/megatron_lm_gpt_pretraining.py
new file mode 100644
index 000000000..68a425e4d
--- /dev/null
+++ b/examples/by_feature/megatron_lm_gpt_pretraining.py
@@ -0,0 +1,699 @@
+#!/usr/bin/env python
+# coding=utf-8
+# Copyright 2021 The HuggingFace Inc. 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.
+"""
+Fine-tuning the library models for causal language modeling (GPT, GPT-2, CTRL, ...)
+on a text file or a dataset without using HuggingFace Trainer.
+
+Here is the full list of checkpoints on the hub that can be fine-tuned by this script:
+https://huggingface.co/models?filter=text-generation
+"""
+# You can also adapt this script on your own causal language modeling task. Pointers for this are left as comments.
+
+import argparse
+import json
+import logging
+import math
+import os
+import random
+from itertools import chain
+from pathlib import Path
+
+import torch
+from torch.utils.data import DataLoader
+
+import datasets
+import transformers
+from accelerate import Accelerator, DistributedType
+from accelerate.logging import get_logger
+from accelerate.utils import MegatronLMDummyScheduler, set_seed
+from datasets import load_dataset
+from huggingface_hub import Repository
+from tqdm.auto import tqdm
+from transformers import (
+    CONFIG_MAPPING,
+    MODEL_MAPPING,
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    SchedulerType,
+    default_data_collator,
+    get_scheduler,
+)
+from transformers.utils import check_min_version, get_full_repo_name, send_example_telemetry
+from transformers.utils.versions import require_version
+
+
+# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
+check_min_version("4.23.0.dev0")
+
+logger = get_logger(__name__)
+
+require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt")
+
+MODEL_CONFIG_CLASSES = list(MODEL_MAPPING.keys())
+MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Finetune a transformers model on a causal language modeling task")
+    parser.add_argument(
+        "--dataset_name",
+        type=str,
+        default=None,
+        help="The name of the dataset to use (via the datasets library).",
+    )
+    parser.add_argument(
+        "--dataset_config_name",
+        type=str,
+        default=None,
+        help="The configuration name of the dataset to use (via the datasets library).",
+    )
+    parser.add_argument(
+        "--train_file", type=str, default=None, help="A csv or a json file containing the training data."
+    )
+    parser.add_argument(
+        "--validation_file", type=str, default=None, help="A csv or a json file containing the validation data."
+    )
+    parser.add_argument(
+        "--validation_split_percentage",
+        default=5,
+        help="The percentage of the train set used as validation set in case there's no validation split",
+    )
+    parser.add_argument(
+        "--model_name_or_path",
+        type=str,
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+        required=False,
+    )
+    parser.add_argument(
+        "--config_name",
+        type=str,
+        default=None,
+        help="Pretrained config name or path if not the same as model_name",
+    )
+    parser.add_argument(
+        "--tokenizer_name",
+        type=str,
+        default=None,
+        help="Pretrained tokenizer name or path if not the same as model_name",
+    )
+    parser.add_argument(
+        "--use_slow_tokenizer",
+        action="store_true",
+        help="If passed, will use a slow tokenizer (not backed by the 🤗 Tokenizers library).",
+    )
+    parser.add_argument(
+        "--per_device_train_batch_size",
+        type=int,
+        default=8,
+        help="Batch size (per device) for the training dataloader.",
+    )
+    parser.add_argument(
+        "--per_device_eval_batch_size",
+        type=int,
+        default=8,
+        help="Batch size (per device) for the evaluation dataloader.",
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=5e-5,
+        help="Initial learning rate (after the potential warmup period) to use.",
+    )
+    parser.add_argument("--weight_decay", type=float, default=0.0, help="Weight decay to use.")
+    parser.add_argument("--num_train_epochs", type=int, default=3, help="Total number of training epochs to perform.")
+    parser.add_argument(
+        "--max_train_steps",
+        type=int,
+        default=None,
+        help="Total number of training steps to perform. If provided, overrides num_train_epochs.",
+    )
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help="Number of updates steps to accumulate before performing a backward/update pass.",
+    )
+    parser.add_argument(
+        "--lr_scheduler_type",
+        type=SchedulerType,
+        default="linear",
+        help="The scheduler type to use.",
+        choices=["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"],
+    )
+    parser.add_argument(
+        "--num_warmup_steps", type=int, default=0, help="Number of steps for the warmup in the lr scheduler."
+    )
+    parser.add_argument("--output_dir", type=str, default=None, help="Where to store the final model.")
+    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
+    parser.add_argument(
+        "--model_type",
+        type=str,
+        default=None,
+        help="Model type to use if training from scratch.",
+        choices=MODEL_TYPES,
+    )
+    parser.add_argument(
+        "--block_size",
+        type=int,
+        default=None,
+        help=(
+            "Optional input sequence length after tokenization. The training dataset will be truncated in block of"
+            " this size for training. Default to the model max input length for single sentence inputs (take into"
+            " account special tokens)."
+        ),
+    )
+    parser.add_argument(
+        "--preprocessing_num_workers",
+        type=int,
+        default=None,
+        help="The number of processes to use for the preprocessing.",
+    )
+    parser.add_argument(
+        "--overwrite_cache", action="store_true", help="Overwrite the cached training and evaluation sets"
+    )
+    parser.add_argument(
+        "--no_keep_linebreaks", action="store_true", help="Do not keep line breaks when using TXT files."
+    )
+    parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")
+    parser.add_argument(
+        "--hub_model_id", type=str, help="The name of the repository to keep in sync with the local `output_dir`."
+    )
+    parser.add_argument("--hub_token", type=str, help="The token to use to push to the Model Hub.")
+    parser.add_argument(
+        "--checkpointing_steps",
+        type=str,
+        default=None,
+        help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.",
+    )
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        default=None,
+        help="If the training should continue from a checkpoint folder.",
+    )
+    parser.add_argument(
+        "--with_tracking",
+        action="store_true",
+        help="Whether to enable experiment trackers for logging.",
+    )
+    parser.add_argument(
+        "--report_to",
+        type=str,
+        default="all",
+        help=(
+            'The integration to report the results and logs to. Supported platforms are `"tensorboard"`,'
+            ' `"wandb"` and `"comet_ml"`. Use `"all"` (default) to report to all integrations.'
+            "Only applicable when `--with_tracking` is passed."
+        ),
+    )
+    args = parser.parse_args()
+
+    # Sanity checks
+    if args.dataset_name is None and args.train_file is None and args.validation_file is None:
+        raise ValueError("Need either a dataset name or a training/validation file.")
+    else:
+        if args.train_file is not None:
+            extension = args.train_file.split(".")[-1]
+            assert extension in ["csv", "json", "txt"], "`train_file` should be a csv, json or txt file."
+        if args.validation_file is not None:
+            extension = args.validation_file.split(".")[-1]
+            assert extension in ["csv", "json", "txt"], "`validation_file` should be a csv, json or txt file."
+
+    if args.push_to_hub:
+        assert args.output_dir is not None, "Need an `output_dir` to create a repo when `--push_to_hub` is passed."
+
+    return args
+
+
+def main():
+    args = parse_args()
+
+    # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
+    # information sent is the one passed as arguments along with your Python/PyTorch versions.
+    send_example_telemetry("run_clm_no_trainer", args)
+
+    # Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
+    # If we're using tracking, we also need to initialize it here and it will by default pick up all supported trackers
+    # in the environment
+    accelerator_log_kwargs = {}
+
+    if args.with_tracking:
+        accelerator_log_kwargs["log_with"] = args.report_to
+        accelerator_log_kwargs["logging_dir"] = args.output_dir
+
+    accelerator = Accelerator(gradient_accumulation_steps=args.gradient_accumulation_steps, **accelerator_log_kwargs)
+
+    # Make one log on every process with the configuration for debugging.
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO,
+    )
+    logger.info(accelerator.state, main_process_only=False)
+    if accelerator.is_local_main_process:
+        datasets.utils.logging.set_verbosity_warning()
+        transformers.utils.logging.set_verbosity_info()
+    else:
+        datasets.utils.logging.set_verbosity_error()
+        transformers.utils.logging.set_verbosity_error()
+
+    # If passed along, set the training seed now.
+    if args.seed is not None:
+        set_seed(args.seed)
+
+    # Handle the repository creation
+    if accelerator.is_main_process:
+        if args.push_to_hub:
+            if args.hub_model_id is None:
+                repo_name = get_full_repo_name(Path(args.output_dir).name, token=args.hub_token)
+            else:
+                repo_name = args.hub_model_id
+            repo = Repository(args.output_dir, clone_from=repo_name)
+
+            with open(os.path.join(args.output_dir, ".gitignore"), "w+") as gitignore:
+                if "step_*" not in gitignore:
+                    gitignore.write("step_*\n")
+                if "epoch_*" not in gitignore:
+                    gitignore.write("epoch_*\n")
+        elif args.output_dir is not None:
+            os.makedirs(args.output_dir, exist_ok=True)
+    accelerator.wait_for_everyone()
+
+    # Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)
+    # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
+    # (the dataset will be downloaded automatically from the datasets Hub).
+    #
+    # For CSV/JSON files, this script will use the column called 'text' or the first column if no column called
+    # 'text' is found. You can easily tweak this behavior (see below).
+    #
+    # In distributed training, the load_dataset function guarantee that only one local process can concurrently
+    # download the dataset.
+    if args.dataset_name is not None:
+        # Downloading and loading a dataset from the hub.
+        raw_datasets = load_dataset(args.dataset_name, args.dataset_config_name)
+        if "validation" not in raw_datasets.keys():
+            raw_datasets["validation"] = load_dataset(
+                args.dataset_name,
+                args.dataset_config_name,
+                split=f"train[:{args.validation_split_percentage}%]",
+            )
+            raw_datasets["train"] = load_dataset(
+                args.dataset_name,
+                args.dataset_config_name,
+                split=f"train[{args.validation_split_percentage}%:]",
+            )
+    else:
+        data_files = {}
+        dataset_args = {}
+        if args.train_file is not None:
+            data_files["train"] = args.train_file
+        if args.validation_file is not None:
+            data_files["validation"] = args.validation_file
+        extension = args.train_file.split(".")[-1]
+        if extension == "txt":
+            extension = "text"
+            dataset_args["keep_linebreaks"] = not args.no_keep_linebreaks
+        raw_datasets = load_dataset(extension, data_files=data_files, **dataset_args)
+        # If no validation data is there, validation_split_percentage will be used to divide the dataset.
+        if "validation" not in raw_datasets.keys():
+            raw_datasets["validation"] = load_dataset(
+                extension,
+                data_files=data_files,
+                split=f"train[:{args.validation_split_percentage}%]",
+                **dataset_args,
+            )
+            raw_datasets["train"] = load_dataset(
+                extension,
+                data_files=data_files,
+                split=f"train[{args.validation_split_percentage}%:]",
+                **dataset_args,
+            )
+
+    # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
+    # https://huggingface.co/docs/datasets/loading_datasets.html.
+
+    # Load pretrained model and tokenizer
+    #
+    # In distributed training, the .from_pretrained methods guarantee that only one local process can concurrently
+    # download model & vocab.
+    if args.config_name:
+        config = AutoConfig.from_pretrained(args.config_name)
+    elif args.model_name_or_path:
+        config = AutoConfig.from_pretrained(args.model_name_or_path)
+    else:
+        config = CONFIG_MAPPING[args.model_type]()
+        logger.warning("You are instantiating a new config instance from scratch.")
+
+    if args.tokenizer_name:
+        tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name, use_fast=not args.use_slow_tokenizer)
+    elif args.model_name_or_path:
+        tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, use_fast=not args.use_slow_tokenizer)
+    else:
+        raise ValueError(
+            "You are instantiating a new tokenizer from scratch. This is not supported by this script."
+            "You can do it from another script, save it, and load it from here, using --tokenizer_name."
+        )
+
+    if args.model_name_or_path:
+        model = AutoModelForCausalLM.from_pretrained(
+            args.model_name_or_path,
+            from_tf=bool(".ckpt" in args.model_name_or_path),
+            config=config,
+        )
+    else:
+        logger.info("Training new model from scratch")
+        model = AutoModelForCausalLM.from_config(config)
+
+    model.resize_token_embeddings(len(tokenizer))
+
+    # Preprocessing the datasets.
+    # First we tokenize all the texts.
+    column_names = raw_datasets["train"].column_names
+    text_column_name = "text" if "text" in column_names else column_names[0]
+
+    def tokenize_function(examples):
+        return tokenizer(examples[text_column_name])
+
+    with accelerator.main_process_first():
+        tokenized_datasets = raw_datasets.map(
+            tokenize_function,
+            batched=True,
+            num_proc=args.preprocessing_num_workers,
+            remove_columns=column_names,
+            load_from_cache_file=not args.overwrite_cache,
+            desc="Running tokenizer on dataset",
+        )
+
+    if args.block_size is None:
+        block_size = tokenizer.model_max_length
+        if block_size > 1024:
+            logger.warning(
+                f"The tokenizer picked seems to have a very large `model_max_length` ({tokenizer.model_max_length}). "
+                "Picking 1024 instead. You can change that default value by passing --block_size xxx."
+            )
+        block_size = 1024
+    else:
+        if args.block_size > tokenizer.model_max_length:
+            logger.warning(
+                f"The block_size passed ({args.block_size}) is larger than the maximum length for the model"
+                f"({tokenizer.model_max_length}). Using block_size={tokenizer.model_max_length}."
+            )
+        block_size = min(args.block_size, tokenizer.model_max_length)
+
+    # Main data processing function that will concatenate all texts from our dataset and generate chunks of block_size.
+    def group_texts(examples):
+        # Concatenate all texts.
+        concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}
+        total_length = len(concatenated_examples[list(examples.keys())[0]])
+        # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
+        # customize this part to your needs.
+        if total_length >= block_size:
+            total_length = (total_length // block_size) * block_size
+        # Split by chunks of max_len.
+        result = {
+            k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
+            for k, t in concatenated_examples.items()
+        }
+        result["labels"] = result["input_ids"].copy()
+        return result
+
+    # Note that with `batched=True`, this map processes 1,000 texts together, so group_texts throws away a remainder
+    # for each of those groups of 1,000 texts. You can adjust that batch_size here but a higher value might be slower
+    # to preprocess.
+    #
+    # To speed up this part, we use multiprocessing. See the documentation of the map method for more information:
+    # https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasets.Dataset.map
+
+    with accelerator.main_process_first():
+        lm_datasets = tokenized_datasets.map(
+            group_texts,
+            batched=True,
+            num_proc=args.preprocessing_num_workers,
+            load_from_cache_file=not args.overwrite_cache,
+            desc=f"Grouping texts in chunks of {block_size}",
+        )
+
+    train_dataset = lm_datasets["train"]
+    eval_dataset = lm_datasets["validation"]
+
+    # Log a few random samples from the training set:
+    for index in random.sample(range(len(train_dataset)), 3):
+        logger.info(f"Sample {index} of the training set: {train_dataset[index]}.")
+
+    # DataLoaders creation:
+    train_dataloader = DataLoader(
+        train_dataset, shuffle=True, collate_fn=default_data_collator, batch_size=args.per_device_train_batch_size
+    )
+    eval_dataloader = DataLoader(
+        eval_dataset, collate_fn=default_data_collator, batch_size=args.per_device_eval_batch_size
+    )
+
+    # Optimizer
+    # Split weights in two groups, one with weight decay and the other not.
+    no_decay = ["bias", "layer_norm.weight"]
+    optimizer_grouped_parameters = [
+        {
+            "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
+            "weight_decay": args.weight_decay,
+        },
+        {
+            "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
+            "weight_decay": 0.0,
+        },
+    ]
+    optimizer = torch.optim.AdamW(optimizer_grouped_parameters, lr=args.learning_rate)
+
+    # Scheduler and math around the number of training steps.
+    overrode_max_train_steps = False
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    if args.max_train_steps is None:
+        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+        overrode_max_train_steps = True
+
+    # New Code
+    # For Megatron-LM, we need to use `MegatronLMDummyScheduler` instead of regular schedulers
+    if accelerator.distributed_type == DistributedType.MEGATRON_LM:
+        lr_scheduler = MegatronLMDummyScheduler(
+            optimizer=optimizer,
+            total_num_steps=args.max_train_steps,
+            warmup_num_steps=args.num_warmup_steps,
+        )
+    else:
+        lr_scheduler = get_scheduler(
+            name=args.lr_scheduler_type,
+            optimizer=optimizer,
+            num_warmup_steps=args.num_warmup_steps * args.gradient_accumulation_steps,
+            num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
+        )
+
+    # Prepare everything with our `accelerator`.
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    # On TPU, the tie weights in our model have been disconnected, so we need to restore the ties.
+    if accelerator.distributed_type == DistributedType.TPU:
+        model.tie_weights()
+
+    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    if overrode_max_train_steps:
+        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+    # Afterwards we recalculate our number of training epochs
+    args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+    # Figure out how many steps we should save the Accelerator states
+    checkpointing_steps = args.checkpointing_steps
+    if checkpointing_steps is not None and checkpointing_steps.isdigit():
+        checkpointing_steps = int(checkpointing_steps)
+
+    # We need to initialize the trackers we use, and also store our configuration.
+    # The trackers initializes automatically on the main process.
+    if args.with_tracking:
+        experiment_config = vars(args)
+        # TensorBoard cannot log Enums, need the raw value
+        experiment_config["lr_scheduler_type"] = experiment_config["lr_scheduler_type"].value
+        accelerator.init_trackers("clm_no_trainer", experiment_config)
+
+    # Train!
+    # New Code
+    # For Megatron-LM, we need to get `global_batch_size` from megatron_lm_plugin
+    # as it handles the specifics related to data parallelism, tensor model parallelism and pipeline parallelism
+    if accelerator.distributed_type == DistributedType.MEGATRON_LM:
+        total_batch_size = accelerator.state.megatron_lm_plugin.global_batch_size
+    else:
+        total_batch_size = (
+            args.per_device_train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+        )
+
+    logger.info("***** Running training *****")
+    logger.info(f"  Num examples = {len(train_dataset)}")
+    logger.info(f"  Num Epochs = {args.num_train_epochs}")
+    logger.info(f"  Instantaneous batch size per device = {args.per_device_train_batch_size}")
+    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
+    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
+    logger.info(f"  Total optimization steps = {args.max_train_steps}")
+    # Only show the progress bar once on each machine.
+    progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process)
+    completed_steps = 0
+    starting_epoch = 0
+
+    # Potentially load in the weights and states from a previous save
+    if args.resume_from_checkpoint:
+        if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
+            accelerator.print(f"Resumed from checkpoint: {args.resume_from_checkpoint}")
+            accelerator.load_state(args.resume_from_checkpoint)
+            path = os.path.basename(args.resume_from_checkpoint)
+        else:
+            # Get the most recent checkpoint
+            dirs = [f.name for f in os.scandir(os.getcwd()) if f.is_dir()]
+            dirs.sort(key=os.path.getctime)
+            path = dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
+        # Extract `epoch_{i}` or `step_{i}`
+        training_difference = os.path.splitext(path)[0]
+
+        if "epoch" in training_difference:
+            starting_epoch = int(training_difference.replace("epoch_", "")) + 1
+            resume_step = None
+        else:
+            # need to multiply `gradient_accumulation_steps` to reflect real steps
+            resume_step = int(training_difference.replace("step_", "")) * args.gradient_accumulation_steps
+            starting_epoch = resume_step // len(train_dataloader)
+            resume_step -= starting_epoch * len(train_dataloader)
+
+    # update the progress_bar if load from checkpoint
+    progress_bar.update(starting_epoch * num_update_steps_per_epoch)
+    completed_steps = starting_epoch * num_update_steps_per_epoch
+
+    for epoch in range(starting_epoch, args.num_train_epochs):
+        model.train()
+        if args.with_tracking:
+            total_loss = 0
+        for step, batch in enumerate(train_dataloader):
+            # We need to skip steps until we reach the resumed step
+            if args.resume_from_checkpoint and epoch == starting_epoch:
+                if resume_step is not None and step < resume_step:
+                    if step % args.gradient_accumulation_steps == 0:
+                        progress_bar.update(1)
+                        completed_steps += 1
+                    continue
+
+            with accelerator.accumulate(model):
+                outputs = model(**batch)
+                loss = outputs.loss
+                # We keep track of the loss at each epoch
+                if args.with_tracking:
+                    total_loss += loss.detach().float()
+                accelerator.backward(loss)
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+
+            # Checks if the accelerator has performed an optimization step behind the scenes
+            if accelerator.sync_gradients:
+                progress_bar.update(1)
+                completed_steps += 1
+
+            if isinstance(checkpointing_steps, int):
+                if completed_steps % checkpointing_steps == 0:
+                    output_dir = f"step_{completed_steps }"
+                    if args.output_dir is not None:
+                        output_dir = os.path.join(args.output_dir, output_dir)
+                    accelerator.save_state(output_dir)
+            if completed_steps >= args.max_train_steps:
+                break
+
+        model.eval()
+        losses = []
+        for step, batch in enumerate(eval_dataloader):
+            with torch.no_grad():
+                outputs = model(**batch)
+
+            loss = outputs.loss
+            # New Code
+            # For Megatron-LM, the losses are already averaged across the data parallel group
+            if accelerator.distributed_type == DistributedType.MEGATRON_LM:
+                losses.append(loss)
+            else:
+                losses.append(accelerator.gather_for_metrics(loss.repeat(args.per_device_eval_batch_size)))
+        try:
+            if accelerator.distributed_type == DistributedType.MEGATRON_LM:
+                losses = torch.tensor(losses)
+            else:
+                losses = torch.cat(losses)
+            eval_loss = torch.mean(losses)
+            perplexity = math.exp(eval_loss)
+        except OverflowError:
+            perplexity = float("inf")
+
+        logger.info(f"epoch {epoch}: perplexity: {perplexity} eval_loss: {eval_loss}")
+
+        if args.with_tracking:
+            accelerator.log(
+                {
+                    "perplexity": perplexity,
+                    "eval_loss": eval_loss,
+                    "train_loss": total_loss.item() / len(train_dataloader),
+                    "epoch": epoch,
+                    "step": completed_steps,
+                },
+                step=completed_steps,
+            )
+
+        if args.push_to_hub and epoch < args.num_train_epochs - 1:
+            accelerator.wait_for_everyone()
+            unwrapped_model = accelerator.unwrap_model(model)
+            unwrapped_model.save_pretrained(
+                args.output_dir, is_main_process=accelerator.is_main_process, save_function=accelerator.save
+            )
+            if accelerator.is_main_process:
+                tokenizer.save_pretrained(args.output_dir)
+                repo.push_to_hub(
+                    commit_message=f"Training in progress epoch {epoch}", blocking=False, auto_lfs_prune=True
+                )
+
+        if args.checkpointing_steps == "epoch":
+            output_dir = f"epoch_{epoch}"
+            if args.output_dir is not None:
+                output_dir = os.path.join(args.output_dir, output_dir)
+            accelerator.save_state(output_dir)
+
+    if args.with_tracking:
+        accelerator.end_training()
+
+    if args.output_dir is not None:
+        accelerator.wait_for_everyone()
+        # New Code
+        # For Megatron-LM, we need to save the model using `accelerator.save_state`
+        if accelerator.distributed_type == DistributedType.MEGATRON_LM:
+            accelerator.save_state(args.output_dir)
+        else:
+            unwrapped_model = accelerator.unwrap_model(model)
+            unwrapped_model.save_pretrained(
+                args.output_dir, is_main_process=accelerator.is_main_process, save_function=accelerator.save
+            )
+        if accelerator.is_main_process:
+            tokenizer.save_pretrained(args.output_dir)
+            if args.push_to_hub:
+                repo.push_to_hub(commit_message="End of training", auto_lfs_prune=True)
+
+        with open(os.path.join(args.output_dir, "all_results.json"), "w") as f:
+            json.dump({"perplexity": perplexity}, f)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 2b6558944..e40797b39 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -41,6 +41,7 @@
     InitProcessGroupKwargs,
     KwargsHandler,
     LoggerType,
+    MegatronLMPlugin,
     PrecisionType,
     RNGType,
     compare_versions,
@@ -50,6 +51,7 @@
     get_pretty_name,
     is_bf16_available,
     is_deepspeed_available,
+    is_megatron_lm_available,
     is_torch_version,
     is_tpu_available,
     pad_across_processes,
@@ -71,6 +73,21 @@
         DummyScheduler,
     )
 
+if is_megatron_lm_available():
+    from .utils import (
+        MegatronEngine,
+        MegatronLMDummyDataLoader,
+        MegatronLMDummyScheduler,
+        MegatronLMOptimizerWrapper,
+        MegatronLMSchedulerWrapper,
+        megatron_lm_initialize,
+        megatron_lm_prepare_data_loader,
+        megatron_lm_prepare_model,
+        megatron_lm_prepare_optimizer,
+        megatron_lm_prepare_scheduler,
+    )
+
+
 if is_tpu_available(check_device=False):
     import torch_xla.distributed.xla_multiprocessing as xmp
 
@@ -107,6 +124,9 @@ class Accelerator:
         fsdp_plugin (`FullyShardedDataParallelPlugin`, *optional*):
             Tweak your FSDP related args using this argument. This argument is optional and can be configured directly
             using *accelerate config*
+        megatron_lm_plugin (`MegatronLMPlugin`, *optional*):
+            Tweak your MegatronLM related args using this argument. This argument is optional and can be configured
+            directly using *accelerate config*
         rng_types (list of `str` or [`~utils.RNGType`]):
             The list of random number generators to synchronize at the beginning of each iteration in your prepared
             dataloaders. Should be one or several of:
@@ -166,6 +186,7 @@ def __init__(
         cpu: bool = False,
         deepspeed_plugin: DeepSpeedPlugin = None,
         fsdp_plugin: FullyShardedDataParallelPlugin = None,
+        megatron_lm_plugin: MegatronLMPlugin = None,
         rng_types: Optional[List[Union[str, RNGType]]] = None,
         log_with: Optional[List[Union[str, LoggerType, GeneralTracker]]] = None,
         logging_dir: Optional[Union[str, os.PathLike]] = None,
@@ -218,6 +239,17 @@ def __init__(
                 raise TypeError("`fsdp_plugin` must be a FullyShardedDataParallelPlugin object.")
             os.environ["USE_FSDP"] = "true"  # use FSDP if plugin is provided
 
+        if megatron_lm_plugin is None:  # init from env variables
+            megatron_lm_plugin = MegatronLMPlugin() if os.environ.get("USE_MEGATRON_LM", "false") == "true" else None
+        else:
+            if not isinstance(megatron_lm_plugin, MegatronLMPlugin):
+                raise TypeError("`megatron_lm_plugin` must be a MegatronLMPlugin object.")
+            os.environ["USE_MEGATRON_LM"] = "true"  # use MegatronLM if plugin is provided
+
+        if megatron_lm_plugin:
+            if not is_megatron_lm_available():
+                raise ImportError("Megatron is not installed. please build it from source.")
+
         # Kwargs handlers
         self.ddp_handler = None
         self.scaler_handler = None
@@ -247,6 +279,7 @@ def __init__(
             cpu=cpu,
             deepspeed_plugin=deepspeed_plugin,
             fsdp_plugin=fsdp_plugin,
+            megatron_lm_plugin=megatron_lm_plugin,
             _from_accelerator=True,
             **kwargs,
         )
@@ -278,7 +311,7 @@ def __init__(
         self.scaler = None
         self.native_amp = False
         err = "{mode} mixed precision requires {requirement}"
-        if self.state.mixed_precision == "fp16":
+        if self.state.mixed_precision == "fp16" and self.distributed_type != DistributedType.MEGATRON_LM:
             self.native_amp = True
             if not torch.cuda.is_available() and not parse_flag_from_env("USE_MPS_DEVICE"):
                 raise ValueError(err.format(mode="fp16", requirement="a GPU"))
@@ -289,7 +322,11 @@ def __init__(
                 self.scaler = ShardedGradScaler(**kwargs)
             else:
                 self.scaler = torch.cuda.amp.GradScaler(**kwargs)
-        elif self.state.mixed_precision == "bf16" and self.distributed_type != DistributedType.FSDP:
+        elif (
+            self.state.mixed_precision == "bf16"
+            and self.distributed_type != DistributedType.FSDP
+            and self.distributed_type != DistributedType.MEGATRON_LM
+        ):
             self.native_amp = is_bf16_available(True)
             if mixed_precision == "bf16" and not self.native_amp and not is_tpu_available():
                 raise ValueError(err.format(mode="bf16", requirement="PyTorch >= 1.10 and a supported device."))
@@ -344,17 +381,27 @@ def device(self):
     @property
     def is_main_process(self):
         """True for one process only."""
-        return self.process_index == 0
+        return (
+            self.process_index == 0 if self.distributed_type != DistributedType.MEGATRON_LM else self.is_last_process
+        )
 
     @property
     def is_local_main_process(self):
         """True for one process per server."""
-        return self.local_process_index == 0
+        return (
+            self.local_process_index == 0
+            if self.distributed_type != DistributedType.MEGATRON_LM
+            else self.is_last_process
+        )
 
     @property
     def use_fp16(self):
         return self.mixed_precision != "no"
 
+    @property
+    def is_last_process(self):
+        return self.process_index == self.num_processes - 1
+
     @property
     def mixed_precision(self):
         if self.distributed_type == DistributedType.DEEPSPEED:
@@ -393,6 +440,18 @@ def wrapper(self, *args, **kwargs):
 
         return wrapper
 
+    def on_last_process(func):
+        """
+        A decorator that will run the decorated function on the last process only.
+        """
+
+        @wraps(func)
+        def wrapper(self, *args, **kwargs):
+            if self.is_last_process or not self.use_distributed:
+                return func(self, *args, **kwargs)
+
+        return wrapper
+
     def on_process(process_idx):
         """
         A decorator that will run the decorated function on a given process index only.
@@ -628,8 +687,8 @@ def prepare(self, *args, device_placement=None):
         """
         if device_placement is None:
             device_placement = [None for _ in args]
-        elif self.distributed_type == DistributedType.DEEPSPEED:
-            raise ValueError("You can't customize device placements with DeepSpeed.")
+        elif self.distributed_type in (DistributedType.DEEPSPEED, DistributedType.MEGATRON_LM):
+            raise ValueError("You can't customize device placements with DeepSpeed or Megatron-LM.")
         elif len(device_placement) != len(args):
             raise ValueError(
                 f"`device_placement` should be a list with {len(args)} elements (the number of objects passed)."
@@ -677,6 +736,8 @@ def prepare(self, *args, device_placement=None):
 
         if self.distributed_type == DistributedType.DEEPSPEED:
             result = self._prepare_deepspeed(*args)
+        elif self.distributed_type == DistributedType.MEGATRON_LM:
+            result = self._prepare_megatron_lm(*args)
         else:
             result = tuple(
                 self._prepare_one(obj, first_pass=True, device_placement=d) for obj, d in zip(args, device_placement)
@@ -923,6 +984,107 @@ def _prepare_deepspeed(self, *args):
                 )
         return tuple(result)
 
+    def _prepare_megatron_lm(self, *args):
+        megatron_lm_plugin = self.state.megatron_lm_plugin
+        if not megatron_lm_plugin.megatron_dataset_flag:
+            batch_sizes = [obj.batch_size for obj in args if hasattr(obj, "batch_size")]
+            if len(batch_sizes) == 0:
+                raise ValueError(
+                    "You must specify a training or evaluation dataloader in `accelerate.prepare()` when using Megatron-LM."
+                )
+
+            micro_batch_size = min(batch_sizes) if megatron_lm_plugin.is_train_batch_min else max(batch_sizes)
+            if len(batch_sizes) > 1:
+                logger.info(
+                    "Since you passed both train and evaluation dataloader, `is_train_batch_min` (here "
+                    f"{megatron_lm_plugin.is_train_batch_min} will decide the `train_batch_size` ({micro_batch_size})."
+                )
+        else:
+            for obj in args:
+                if isinstance(obj, MegatronLMDummyDataLoader):
+                    micro_batch_size = obj.dataset_args["micro_batch_size"]
+                    break
+
+        dp_degree = self.num_processes // (megatron_lm_plugin.tp_degree * megatron_lm_plugin.pp_degree)
+        megatron_lm_plugin.set_training_args(micro_batch_size, dp_degree)
+
+        model = None
+        optimizer = None
+        scheduler = None
+        is_dummy_scheduler = False
+        batch_data = None
+        for obj in args:
+            if isinstance(obj, torch.utils.data.DataLoader) and batch_data is None:
+                batch_data = next(iter(obj))
+            if isinstance(obj, torch.nn.Module):
+                model = obj
+            elif isinstance(obj, (torch.optim.Optimizer)):
+                optimizer = obj
+            elif isinstance(obj, (torch.optim.lr_scheduler._LRScheduler, MegatronLMDummyScheduler)):
+                scheduler = obj
+
+        if model is not None:
+            megatron_lm_plugin.set_network_size_args(model, batch_data)
+        if optimizer is not None:
+            megatron_lm_plugin.set_optimizer_type(optimizer)
+        if scheduler is not None:
+            is_dummy_scheduler = isinstance(scheduler, MegatronLMDummyScheduler)
+            if not is_dummy_scheduler:
+                raise ValueError(
+                    "You can't use a custom scheduler with Megatron-LM. Please use the `accelerate.utils.MegatronLMDummyScheduler` instead."
+                )
+            megatron_lm_plugin.set_scheduler_args(scheduler)
+
+        # initialize megatron-lm
+        megatron_lm_initialize(self, args_defaults=megatron_lm_plugin.megatron_lm_default_args)
+        counter = 0
+        result = []
+        for obj in args:
+            if isinstance(obj, torch.utils.data.DataLoader):
+                result.append(megatron_lm_prepare_data_loader(self, obj))
+                counter += 1
+            elif isinstance(obj, MegatronLMDummyDataLoader):
+                if counter == 0:
+                    obj.set_megatron_data_args()
+                    dataloaders = megatron_lm_prepare_data_loader(self, obj)
+                result.append(dataloaders[counter])
+                counter += 1
+            else:
+                result.append(obj)
+
+        if model is not None:
+            model = megatron_lm_prepare_model(self)
+        if optimizer is not None:
+            optimizer = megatron_lm_prepare_optimizer(self, model)
+        if scheduler is not None:
+            scheduler = megatron_lm_prepare_scheduler(self, optimizer, scheduler)
+
+        if model is not None:
+            model = MegatronEngine(self, model, optimizer, scheduler)
+        if optimizer is not None:
+            optimizer = MegatronLMOptimizerWrapper(optimizer)
+        if scheduler is not None:
+            scheduler = MegatronLMSchedulerWrapper(scheduler, optimizer)
+
+        for i in range(len(result)):
+            if isinstance(result[i], torch.nn.Module):
+                result[i] = model
+            elif isinstance(result[i], torch.optim.Optimizer):
+                result[i] = optimizer
+            elif isinstance(result[i], MegatronLMDummyScheduler):
+                result[i] = scheduler
+        if model is not None:
+            self._models.append(model)
+        if optimizer is not None:
+            self._optimizers.append(optimizer)
+        if scheduler is not None:
+            self._schedulers.append(scheduler)
+        if len(self._models) > 1:
+            raise AssertionError(
+                "You can't use same `Accelerator()` instance with multiple models when using Megatron-LM"
+            )
+        return tuple(result)
+
     def prepare_data_loader(self, data_loader: torch.utils.data.DataLoader, device_placement=None):
         """
         Prepares a PyTorch DataLoader for training in any distributed setup. It is recommended to use
@@ -1001,6 +1163,8 @@ def backward(self, loss, **kwargs):
             loss = loss / self.gradient_accumulation_steps
         if self.distributed_type == DistributedType.DEEPSPEED:
             self.deepspeed_engine_wrapped.backward(loss, **kwargs)
+        elif self.distributed_type == DistributedType.MEGATRON_LM:
+            return
         elif self.scaler is not None:
             self.scaler.scale(loss).backward(**kwargs)
         else:
@@ -1312,6 +1476,10 @@ def save_state(self, output_dir: str):
                 ckpt_id = f"{MODEL_NAME}" if i == 0 else f"{MODEL_NAME}_{i}"
                 model.save_checkpoint(output_dir, ckpt_id)
                 logger.info(f"DeepSpeed Model and Optimizer saved to output dir {os.path.join(output_dir, ckpt_id)}")
+            elif self.distributed_type == DistributedType.MEGATRON_LM:
+                logger.info("Saving Megatron-LM Model, Optimizer and Scheduler")
+                model.save_checkpoint(output_dir)
+                logger.info(f"Megatron-LM Model , Optimizer and Scheduler saved to output dir {output_dir}")
             else:
                 weights.append(self.get_state_dict(model, unwrap=False))
 
@@ -1322,7 +1490,7 @@ def save_state(self, output_dir: str):
                 logger.info("Saving FSDP Optimizer")
                 self.state.fsdp_plugin.save_optimizer(self, opt, self._models[i], output_dir, i)
                 logger.info(f"FSDP Optimizer saved to output dir {output_dir}")
-        elif self.distributed_type != DistributedType.DEEPSPEED:
+        elif self.distributed_type not in [DistributedType.DEEPSPEED, DistributedType.MEGATRON_LM]:
             optimizers = self._optimizers
 
         # Save the lr schedulers taking care of DeepSpeed nuances
@@ -1332,7 +1500,7 @@ def save_state(self, output_dir: str):
                 if isinstance(scheduler, DeepSpeedSchedulerWrapper):
                     continue
                 schedulers.append(scheduler)
-        else:
+        elif self.distributed_type not in [DistributedType.MEGATRON_LM]:
             schedulers = self._schedulers
 
         save_location = save_accelerator_state(
@@ -1374,6 +1542,10 @@ def load_state(self, input_dir: str):
                 ckpt_id = f"{MODEL_NAME}" if i == 0 else f"{MODEL_NAME}_{i}"
                 model.load_checkpoint(input_dir, ckpt_id)
                 logger.info(f"DeepSpeed Model and Optimizer loaded from input dir {os.path.join(input_dir, ckpt_id)}")
+            elif self.distributed_type == DistributedType.MEGATRON_LM:
+                logger.info("Loading Megatron-LM Model, Optimizer and Scheduler")
+                model.load_checkpoint(input_dir)
+                logger.info(f"Megatron-LM Model , Optimizer and Scheduler loaded from input dir {input_dir}")
             else:
                 models.append(model)
 
@@ -1384,7 +1556,7 @@ def load_state(self, input_dir: str):
                 logger.info("Loading FSDP Optimizer")
                 self.state.fsdp_plugin.load_optimizer(self, opt, self._models[i], input_dir, i)
                 logger.info(f"FSDP Optimizer loaded from input dir {input_dir}")
-        elif self.distributed_type != DistributedType.DEEPSPEED:
+        elif self.distributed_type not in [DistributedType.DEEPSPEED, DistributedType.MEGATRON_LM]:
             optimizers = self._optimizers
 
         # Load the lr schedulers taking care of DeepSpeed nuances
@@ -1394,7 +1566,7 @@ def load_state(self, input_dir: str):
                 if isinstance(scheduler, DeepSpeedSchedulerWrapper):
                     continue
                 schedulers.append(scheduler)
-        else:
+        elif self.distributed_type not in [DistributedType.MEGATRON_LM]:
             schedulers = self._schedulers
 
         load_accelerator_state(input_dir, models, optimizers, schedulers, self.state.process_index, self.scaler)
diff --git a/src/accelerate/checkpointing.py b/src/accelerate/checkpointing.py
index d5e816aab..00dfce73e 100644
--- a/src/accelerate/checkpointing.py
+++ b/src/accelerate/checkpointing.py
@@ -155,15 +155,18 @@ def load_accelerator_state(input_dir, models, optimizers, schedulers, process_in
         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")
+    try:
+        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")
+    except:
+        logger.info("Could not load random states")
 
 
 def save_custom_state(obj, path, index: int = 0):
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index e86a0abff..0fa3ceab2 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -275,6 +275,67 @@ def get_cluster_input():
                 default="FULL_STATE_DICT",
             )
 
+    megatron_lm_config = {}
+    if distributed_type in [DistributedType.MULTI_GPU]:
+        use_megatron_lm = _ask_field(
+            "Do you want to use Megatron-LM ? [yes/NO]: ",
+            _convert_yes_no_to_bool,
+            default=False,
+            error_message="Please enter yes or no.",
+        )
+        if use_megatron_lm:
+            distributed_type = DistributedType.MEGATRON_LM
+        if distributed_type == DistributedType.MEGATRON_LM:
+            prefix = "megatron_lm_"
+            megatron_lm_config[prefix + "tp_degree"] = _ask_field(
+                "What is the Tensor Parallelism degree/size? [1]:",
+                lambda x: int(x),
+                default=1,
+                error_message="Please enter an integer.",
+            )
+            if megatron_lm_config[prefix + "tp_degree"] > 1:
+                megatron_lm_config[prefix + "sequence_parallelism"] = _ask_field(
+                    "Do you want to enable Sequence Parallelism? [YES/no]: ",
+                    _convert_yes_no_to_bool,
+                    default=True,
+                    error_message="Please enter yes or no.",
+                )
+
+            megatron_lm_config[prefix + "pp_degree"] = _ask_field(
+                "What is the Pipeline Parallelism degree/size? [1]:",
+                lambda x: int(x),
+                default=1,
+                error_message="Please enter an integer.",
+            )
+            if megatron_lm_config[prefix + "pp_degree"] > 1:
+                megatron_lm_config[prefix + "num_micro_batches"] = _ask_field(
+                    "What is the number of micro-batches? [1]:",
+                    lambda x: int(x),
+                    default=1,
+                    error_message="Please enter an integer.",
+                )
+
+            megatron_lm_config[prefix + "recompute_activations"] = _ask_field(
+                "Do you want to enable selective activation recomputation? [YES/no]: ",
+                _convert_yes_no_to_bool,
+                default=True,
+                error_message="Please enter yes or no.",
+            )
+
+            megatron_lm_config[prefix + "use_distributed_optimizer"] = _ask_field(
+                "Do you want to use distributed optimizer "
+                "which shards optimizer state and gradients across data pralellel ranks? [YES/no]: ",
+                _convert_yes_no_to_bool,
+                default=True,
+                error_message="Please enter yes or no.",
+            )
+
+            megatron_lm_config[prefix + "gradient_clipping"] = _ask_field(
+                "What is the gradient clipping value based on global L2 Norm (0 to disable)? [1.0]: ",
+                lambda x: float(x),
+                default=1.0,
+            )
+
     if distributed_type == DistributedType.TPU:
         main_training_function = _ask_field(
             "What is the name of the function in your script that should be launched in all parallel scripts? [main]: ",
@@ -295,7 +356,7 @@ def get_cluster_input():
             default=1,
             error_message="Please enter an integer.",
         )
-    elif distributed_type in [DistributedType.FSDP, DistributedType.DEEPSPEED]:
+    elif distributed_type in [DistributedType.FSDP, DistributedType.DEEPSPEED, DistributedType.MEGATRON_LM]:
         num_processes = _ask_field(
             "How many GPU(s) should be used for distributed training? [1]:",
             lambda x: int(x),
@@ -343,6 +404,7 @@ def get_cluster_input():
         main_training_function=main_training_function,
         deepspeed_config=deepspeed_config,
         fsdp_config=fsdp_config,
+        megatron_lm_config=megatron_lm_config,
         use_cpu=use_cpu,
         rdzv_backend=rdzv_backend,
         same_network=same_network,
diff --git a/src/accelerate/commands/config/config_args.py b/src/accelerate/commands/config/config_args.py
index 43faf45ca..159f951dd 100644
--- a/src/accelerate/commands/config/config_args.py
+++ b/src/accelerate/commands/config/config_args.py
@@ -146,6 +146,8 @@ class ClusterConfig(BaseConfig):
     deepspeed_config: dict = None
     # args for fsdp
     fsdp_config: dict = None
+    # args for megatron_lm
+    megatron_lm_config: dict = None
     # args for TPU
     downcast_bf16: bool = False
 
@@ -154,6 +156,8 @@ def __post_init__(self):
             self.deepspeed_config = {}
         if self.fsdp_config is None:
             self.fsdp_config = {}
+        if self.megatron_lm_config is None:
+            self.megatron_lm_config = {}
         return super().__post_init__()
 
 
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index f348bb0a6..c297223cb 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -83,6 +83,8 @@ def launch_command_parser(subparsers=None):
         action="store_true",
         help="Whether or not this should use MPS-enabled GPU device on MacOS machines.",
     )
+
+    # deepspeed args
     parser.add_argument(
         "--use_deepspeed",
         default=False,
@@ -163,6 +165,8 @@ def launch_command_parser(subparsers=None):
         type=str,
         help="DeepSpeed multi-node launcher to use.",
     )
+
+    # fsdp args
     parser.add_argument(
         "--use_fsdp",
         default=False,
@@ -236,6 +240,62 @@ def launch_command_parser(subparsers=None):
         type=str,
         help="This argument is deprecated. Use `fsdp_transformer_layer_cls_to_wrap` instead.",
     )
+
+    # megatron_lm args
+    parser.add_argument(
+        "--use_megatron_lm",
+        default=False,
+        action="store_true",
+        help="Whether to use Megatron-LM.",
+    )
+    parser.add_argument(
+        "--megatron_lm_tp_degree",
+        type=int,
+        default=1,
+        help="Megatron-LM's Tensor Parallelism (TP) degree. (useful only when `use_megatron_lm` flag is passed).",
+    )
+    parser.add_argument(
+        "--megatron_lm_pp_degree",
+        type=int,
+        default=1,
+        help="Megatron-LM's Pipeline Parallelism (PP) degree. (useful only when `use_megatron_lm` flag is passed).",
+    )
+    parser.add_argument(
+        "--megatron_lm_num_micro_batches",
+        type=int,
+        default=None,
+        help="Megatron-LM's number of micro batches when PP degree > 1. (useful only when `use_megatron_lm` flag is passed).",
+    )
+    parser.add_argument(
+        "--megatron_lm_sequence_parallelism",
+        default=None,
+        type=str,
+        help="Decides Whether (true|false) to enable Sequence Parallelism when TP degree > 1. "
+        "(useful only when `use_megatron_lm` flag is passed).",
+    )
+    parser.add_argument(
+        "--megatron_lm_recompute_activations",
+        default=None,
+        type=str,
+        help="Decides Whether (true|false) to enable Selective Activation Recomputation. "
+        "(useful only when `use_megatron_lm` flag is passed).",
+    )
+    parser.add_argument(
+        "--megatron_lm_use_distributed_optimizer",
+        default=None,
+        type=str,
+        help="Decides Whether (true|false) to use distributed optimizer "
+        "which shards optimizer state and gradients across Data Pralellel (DP) ranks. "
+        "(useful only when `use_megatron_lm` flag is passed).",
+    )
+    parser.add_argument(
+        "--megatron_lm_gradient_clipping",
+        default=1.0,
+        type=float,
+        help="Megatron-LM's gradient clipping value based on global L2 Norm (0 to disable). "
+        "(useful only when `use_megatron_lm` flag is passed).",
+    )
+
     parser.add_argument(
         "--tpu", default=False, action="store_true", help="Whether or not this should launch a TPU training."
     )
@@ -487,6 +547,22 @@ def multi_gpu_launcher(args):
             current_env["FSDP_BACKWARD_PREFETCH"] = str(args.fsdp_backward_prefetch_policy)
         if args.fsdp_state_dict_type is not None:
             current_env["FSDP_STATE_DICT_TYPE"] = str(args.fsdp_state_dict_type)
+
+    if args.use_megatron_lm:
+        prefix = "MEGATRON_LM_"
+        current_env["USE_MEGATRON_LM"] = "true"
+        current_env[prefix + "TP_DEGREE"] = str(args.megatron_lm_tp_degree)
+        current_env[prefix + "PP_DEGREE"] = str(args.megatron_lm_pp_degree)
+        current_env[prefix + "GRADIENT_CLIPPING"] = str(args.megatron_lm_gradient_clipping)
+        if args.megatron_lm_num_micro_batches is not None:
+            current_env[prefix + "NUM_MICRO_BATCHES"] = str(args.megatron_lm_num_micro_batches)
+        if args.megatron_lm_sequence_parallelism is not None:
+            current_env[prefix + "SEQUENCE_PARALLELISM"] = str(args.megatron_lm_sequence_parallelism)
+        if args.megatron_lm_recompute_activations is not None:
+            current_env[prefix + "RECOMPUTE_ACTIVATIONS"] = str(args.megatron_lm_recompute_activations)
+        if args.megatron_lm_use_distributed_optimizer is not None:
+            current_env[prefix + "USE_DISTRIBUTED_OPTIMIZER"] = str(args.megatron_lm_use_distributed_optimizer)
+
     current_env["OMP_NUM_THREADS"] = str(args.num_cpu_threads_per_process)
     if is_torch_version("<", "1.9.0"):
         raise NotImplementedError("Multi-node training requires pytorch>=1.9.0")
@@ -825,12 +901,14 @@ def launch_command(args):
             and not args.use_deepspeed
             and not args.use_fsdp
             and not args.use_mps_device
+            and not args.use_megatron_lm
         ):
             args.use_deepspeed = defaults.distributed_type == DistributedType.DEEPSPEED
             args.multi_gpu = defaults.distributed_type == DistributedType.MULTI_GPU
             args.tpu = defaults.distributed_type == DistributedType.TPU
             args.use_fsdp = defaults.distributed_type == DistributedType.FSDP
             args.use_mps_device = defaults.distributed_type == DistributedType.MPS
+            args.use_megatron_lm = defaults.distributed_type == DistributedType.MEGATRON_LM
         if not args.use_mps_device:
             if args.gpu_ids is None:
                 if defaults.gpu_ids is not None:
@@ -851,6 +929,8 @@ def launch_command(args):
                         if "fsdp" not in arg_to_set:
                             arg_to_set = "fsdp_" + arg_to_set
                         setattr(args, arg_to_set, defaults.fsdp_config[k])
+                    for k in defaults.megatron_lm_config:
+                        setattr(args, k, defaults.megatron_lm_config[k])
                     continue
 
                 # Those args are handled separately
@@ -901,6 +981,8 @@ def launch_command(args):
         deepspeed_launcher(args)
     elif args.use_fsdp and not args.cpu:
         multi_gpu_launcher(args)
+    elif args.use_megatron_lm and not args.cpu:
+        multi_gpu_launcher(args)
     elif args.multi_gpu and not args.cpu:
         multi_gpu_launcher(args)
     elif args.tpu and not args.cpu:
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index fac1b50e2..c4a015156 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -625,7 +625,7 @@ def prepare_data_loader(
     sampler_is_batch_sampler = False
     generator = getattr(dataloader, "generator", None)
     # No change if no multiprocess
-    if num_processes != 1 and not dispatch_batches:
+    if (num_processes != 1 or state.distributed_type == DistributedType.MEGATRON_LM) and not dispatch_batches:
         if isinstance(new_dataset, IterableDataset):
             if getattr(dataloader.dataset, "generator", None) is not None:
                 generator = dataloader.dataset.generator
diff --git a/src/accelerate/logging.py b/src/accelerate/logging.py
index 2c67e24eb..dca093215 100644
--- a/src/accelerate/logging.py
+++ b/src/accelerate/logging.py
@@ -15,6 +15,7 @@
 import logging
 
 from .state import AcceleratorState
+from .utils import DistributedType
 
 
 class MultiProcessAdapter(logging.LoggerAdapter):
@@ -28,7 +29,12 @@ class MultiProcessAdapter(logging.LoggerAdapter):
     @staticmethod
     def _should_log(main_process_only):
         "Check if log should be performed"
-        return not main_process_only or (main_process_only and AcceleratorState().local_process_index == 0)
+        state = AcceleratorState()
+        if state.distributed_type != DistributedType.MEGATRON_LM:
+            process_index_flag = state.local_process_index == 0
+        else:
+            process_index_flag = state.process_index == state.num_processes - 1
+        return not main_process_only or (main_process_only and process_index_flag)
 
     def log(self, level, msg, *args, **kwargs):
         """
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index a700cfad6..82ea288f0 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -70,6 +70,7 @@ def __init__(
         cpu: bool = False,
         deepspeed_plugin=None,
         fsdp_plugin=None,
+        megatron_lm_plugin=None,
         _from_accelerator: bool = False,
         **kwargs,
     ):
@@ -162,6 +163,10 @@ def __init__(
                     if self.mixed_precision != "no":
                         fsdp_plugin.set_mixed_precision(self.mixed_precision)
                     self.fsdp_plugin = fsdp_plugin
+                if os.environ.get("USE_MEGATRON_LM", "false") == "true":
+                    self.distributed_type = DistributedType.MEGATRON_LM
+                    megatron_lm_plugin.set_mixed_precision(self.mixed_precision)
+                    self.megatron_lm_plugin = megatron_lm_plugin
             elif get_int_from_env(["PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE", "WORLD_SIZE"], 1) > 1:
                 self.distributed_type = DistributedType.MULTI_CPU
                 if is_ccl_available() and get_int_from_env(["CCL_WORKER_COUNT"], 0) > 0:
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
index 4bc71ba83..b472ec556 100644
--- a/src/accelerate/utils/__init__.py
+++ b/src/accelerate/utils/__init__.py
@@ -13,6 +13,7 @@
     InitProcessGroupKwargs,
     KwargsHandler,
     LoggerType,
+    MegatronLMPlugin,
     PrecisionType,
     RNGType,
     SageMakerDistributedType,
@@ -28,6 +29,7 @@
     is_comet_ml_available,
     is_datasets_available,
     is_deepspeed_available,
+    is_megatron_lm_available,
     is_rich_available,
     is_sagemaker_available,
     is_tensorboard_available,
@@ -94,6 +96,23 @@
     )
 
 from .launch import PrepareForLaunch, _filter_args, get_launch_prefix
+from .megatron_lm import (
+    AbstractTrainStep,
+    BertTrainStep,
+    GPTTrainStep,
+    MegatronEngine,
+    MegatronLMDummyDataLoader,
+    MegatronLMDummyScheduler,
+    MegatronLMOptimizerWrapper,
+    MegatronLMSchedulerWrapper,
+    T5TrainStep,
+    avg_losses_across_data_parallel_group,
+)
+from .megatron_lm import initialize as megatron_lm_initialize
+from .megatron_lm import prepare_data_loader as megatron_lm_prepare_data_loader
+from .megatron_lm import prepare_model as megatron_lm_prepare_model
+from .megatron_lm import prepare_optimizer as megatron_lm_prepare_optimizer
+from .megatron_lm import prepare_scheduler as megatron_lm_prepare_scheduler
 from .memory import find_executable_batch_size
 from .other import (
     extract_model_from_parallel,
diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py
index 9f6e30bec..8227fa5bc 100644
--- a/src/accelerate/utils/dataclasses.py
+++ b/src/accelerate/utils/dataclasses.py
@@ -16,6 +16,7 @@
 General namespace and dataclass related classes
 """
 
+import argparse
 import copy
 import enum
 import functools
@@ -24,7 +25,8 @@
 import warnings
 from dataclasses import dataclass, field
 from datetime import timedelta
-from typing import Any, Callable, Iterable, Optional
+from distutils.util import strtobool
+from typing import Any, Callable, Dict, Iterable, List, Optional
 
 import torch
 
@@ -127,6 +129,7 @@ class DistributedType(str, enum.Enum):
     FSDP = "FSDP"
     TPU = "TPU"
     MPS = "MPS"
+    MEGATRON_LM = "MEGATRON_LM"
 
 
 class SageMakerDistributedType(str, enum.Enum):
@@ -644,3 +647,351 @@ def load_optimizer(self, accelerator, optimizer, model, input_dir, optimizer_ind
         # called from all ranks, though only rank0 has a valid param for full_osd
         sharded_osd = FSDP.scatter_full_optim_state_dict(full_osd, model)
         optimizer.load_state_dict(sharded_osd)
+
+
+@dataclass
+class MegatronLMPlugin:
+    """
+    Plugin for Megatron-LM to enable tensor, pipeline, sequence and data parallelism. Also to enable selective
+    activation recomputation and optimized fused kernels.
+    """
+
+    tp_degree: int = field(default=None, metadata={"help": "tensor parallelism degree."})
+    pp_degree: int = field(default=None, metadata={"help": "pipeline parallelism degree."})
+    num_micro_batches: int = field(default=None, metadata={"help": "number of micro-batches."})
+    gradient_clipping: float = field(
+        default=None, metadata={"help": "gradient clipping value based on global L2 Norm (0 to disable)"}
+    )
+    sequence_parallelism: bool = field(
+        default=None,
+        metadata={"help": "enable sequence parallelism"},
+    )
+    recompute_activation: bool = field(
+        default=None,
+        metadata={"help": "enable selective activation recomputation"},
+    )
+    use_distributed_optimizer: bool = field(
+        default=None,
+        metadata={"help": "enable distributed optimizer"},
+    )
+    pipeline_model_parallel_split_rank: int = field(
+        default=None, metadata={"help": "Rank where encoder and decoder should be split."}
+    )
+    num_layers_per_virtual_pipeline_stage: int = field(
+        default=None, metadata={"help": "Number of layers per virtual pipeline stage."}
+    )
+    is_train_batch_min: str = field(
+        default=True,
+        metadata={"help": "If both train & eval dataloaders are specified, this will decide the micro_batch_size"},
+    )
+    train_iters: int = field(
+        default=None,
+        metadata={
+            "help": "Total number of iterations to train over all training runs. "
+            "Note that either train-iters or train-samples should be provided when using `MegatronLMDummyScheduler`"
+        },
+    )
+    train_samples: int = field(
+        default=None,
+        metadata={
+            "help": "Total number of samples to train over all training runs. "
+            "Note that either train-iters or train-samples should be provided when using `MegatronLMDummyScheduler`"
+        },
+    )
+    weight_decay_incr_style: str = field(
+        default="constant",
+        metadata={"help": 'Weight decay increment function. choices=["constant", "linear", "cosine"]. '},
+    )
+    start_weight_decay: float = field(
+        default=None,
+        metadata={"help": "Initial weight decay coefficient for L2 regularization."},
+    )
+    end_weight_decay: float = field(
+        default=None,
+        metadata={"help": "End of run weight decay coefficient for L2 regularization."},
+    )
+    lr_decay_style: str = field(
+        default="linear",
+        metadata={"help": "Learning rate decay function. choices=['constant', 'linear', 'cosine']."},
+    )
+    lr_decay_iters: int = field(
+        default=None,
+        metadata={"help": "Number of iterations for learning rate decay. If None defaults to `train_iters`."},
+    )
+    lr_decay_samples: int = field(
+        default=None,
+        metadata={"help": "Number of samples for learning rate decay. If None defaults to `train_samples`."},
+    )
+    lr_warmup_iters: int = field(
+        default=None,
+        metadata={"help": "number of iterations to linearly warmup learning rate over."},
+    )
+    lr_warmup_samples: int = field(
+        default=None,
+        metadata={"help": "number of samples to linearly warmup learning rate over."},
+    )
+    lr_warmup_fraction: float = field(
+        default=None,
+        metadata={"help": "fraction of lr-warmup-(iters/samples) to linearly warmup learning rate over."},
+    )
+    min_lr: float = field(
+        default=0,
+        metadata={"help": "Minumum value for learning rate. The scheduler clip values below this threshold."},
+    )
+    consumed_samples: List[int] = field(
+        default=None,
+        metadata={
+            "help": "Number of samples consumed in the same order as the dataloaders to `accelerator.prepare` call."
+        },
+    )
+    no_wd_decay_cond: Optional[Callable] = field(default=None, metadata={"help": "Condition to disable weight decay."})
+    scale_lr_cond: Optional[Callable] = field(default=None, metadata={"help": "Condition to scale learning rate."})
+    lr_mult: float = field(default=1.0, metadata={"help": "Learning rate multiplier."})
+    megatron_dataset_flag: bool = field(
+        default=False,
+        metadata={"help": "Whether the format of dataset follows Megatron-LM Indexed/Cached/MemoryMapped format."},
+    )
+    seq_length: int = field(
+        default=None,
+        metadata={"help": "Maximum sequence length to process."},
+    )
+    encoder_seq_length: int = field(
+        default=None,
+        metadata={"help": "Maximum sequence length to process for the encoder."},
+    )
+    decoder_seq_length: int = field(
+        default=None,
+        metadata={"help": "Maximum sequence length to process for the decoder."},
+    )
+    tensorboard_dir: str = field(
+        default=None,
+        metadata={"help": "Path to save tensorboard logs."},
+    )
+    set_all_logging_options: bool = field(
+        default=False,
+        metadata={"help": "Whether to set all logging options."},
+    )
+    custom_train_step_class: Optional[Any] = field(
+        default=None,
+        metadata={"help": "Custom train step class."},
+    )
+    custom_train_step_kwargs: Optional[Dict[str, Any]] = field(
+        default=None,
+        metadata={"help": "Custom train step kwargs."},
+    )
+    eval_iters: int = field(
+        default=100, metadata={"help": "Number of iterations to run for evaluation validation/test for."}
+    )
+    eval_interval: int = field(
+        default=1000, metadata={"help": "Interval between running evaluation on validation set."}
+    )
+
+    def __post_init__(self):
+        prefix = "MEGATRON_LM_"
+        if self.tp_degree is None:
+            self.tp_degree = int(os.environ.get(prefix + "TP_DEGREE", 1))
+        if self.pp_degree is None:
+            self.pp_degree = int(os.environ.get(prefix + "PP_DEGREE", 1))
+        if self.num_micro_batches is None:
+            self.num_micro_batches = int(os.environ.get(prefix + "NUM_MICRO_BATCHES", 1))
+        if self.gradient_clipping is None:
+            self.gradient_clipping = float(os.environ.get(prefix + "GRADIENT_CLIPPING", 1.0))
+        if self.recompute_activation is None:
+            self.recompute_activation = strtobool(os.environ.get(prefix + "RECOMPUTE_ACTIVATION", "False")) == 1
+        if self.use_distributed_optimizer is None:
+            self.use_distributed_optimizer = (
+                strtobool(os.environ.get(prefix + "USE_DISTRIBUTED_OPTIMIZER", "False")) == 1
+            )
+        if self.sequence_parallelism is None:
+            self.sequence_parallelism = strtobool(os.environ.get(prefix + "SEQUENCE_PARALLELISM", "False")) == 1
+
+        if self.pp_degree > 1 or self.use_distributed_optimizer:
+            self.DDP_impl = "local"
+        else:
+            self.DDP_impl = "torch"
+
+        if self.consumed_samples is not None:
+            if len(self.consumed_samples) == 1:
+                self.consumed_samples.extend([0, 0])
+            elif len(self.consumed_samples) == 2:
+                self.consumed_samples.append(0)
+
+        self.megatron_lm_default_args = {
+            "tensor_model_parallel_size": self.tp_degree,
+            "pipeline_model_parallel_size": self.pp_degree,
+            "pipeline_model_parallel_split_rank": self.pipeline_model_parallel_split_rank,
+            "num_layers_per_virtual_pipeline_stage": self.num_layers_per_virtual_pipeline_stage,
+            "DDP_impl": self.DDP_impl,
+            "use_distributed_optimizer": self.use_distributed_optimizer,
+            "sequence_parallel": self.sequence_parallelism,
+            "clip_grad": self.gradient_clipping,
+            "num_micro_batches": self.num_micro_batches,
+            "consumed_samples": self.consumed_samples,
+            "no_wd_decay_cond": self.no_wd_decay_cond,
+            "scale_lr_cond": self.scale_lr_cond,
+            "lr_mult": self.lr_mult,
+            "megatron_dataset_flag": self.megatron_dataset_flag,
+            "eval_iters": self.eval_iters,
+            "eval_interval": self.eval_interval,
+        }
+        if self.recompute_activation:
+            self.megatron_lm_default_args["recompute_granularity"] = "selective"
+        if self.tensorboard_dir is not None:
+            self.megatron_lm_default_args["tensorboard_dir"] = self.tensorboard_dir
+            if self.set_all_logging_options:
+                self.set_tensorboard_logging_options()
+
+    def set_network_size_args(self, model, batch_data=None):
+        # Check if the model is either BERT, GPT or T5 else raise error
+        # set 'num_layers', 'hidden_size', 'num_attention_heads', 'max_position_embeddings'
+        if "megatron-bert" in model.config.model_type.lower():
+            model_type_name = "bert"
+            num_layers = model.config.num_hidden_layers
+            hidden_size = model.config.hidden_size
+            num_attention_heads = model.config.num_attention_heads
+            max_position_embeddings = model.config.max_position_embeddings
+            num_labels = model.config.num_labels
+            orig_vocab_size = model.config.vocab_size
+            if "maskedlm" in model.__class__.__name__.lower():
+                pretraining_flag = True
+            if self.seq_length is not None:
+                if self.encoder_seq_length is not None:
+                    warnings.warn("Both `seq_length` and `encoder_seq_length` are set. Using `encoder_seq_length`.")
+                self.seq_length = self.encoder_seq_length
+            elif self.encoder_seq_length is not None:
+                self.seq_length = self.encoder_seq_length
+            elif batch_data is not None:
+                self.seq_length = batch_data["input_ids"].shape[1]
+            else:
+                self.seq_length = max_position_embeddings
+            self.megatron_lm_default_args["seq_length"] = self.seq_length
+        elif "gpt2" in model.config.model_type.lower():
+            model_type_name = "gpt"
+            num_layers = model.config.n_layer
+            hidden_size = model.config.n_embd
+            num_attention_heads = model.config.n_head
+            max_position_embeddings = model.config.n_positions
+            orig_vocab_size = model.config.vocab_size
+            pretraining_flag = True
+            if self.seq_length is not None:
+                if self.decoder_seq_length is not None:
+                    warnings.warn("Both `seq_length` and `decoder_seq_length` are set. Using `decoder_seq_length`.")
+                self.seq_length = self.decoder_seq_length
+            elif self.decoder_seq_length is not None:
+                self.seq_length = self.decoder_seq_length
+            elif batch_data is not None:
+                self.seq_length = batch_data["input_ids"].shape[1]
+            else:
+                self.seq_length = max_position_embeddings
+            self.megatron_lm_default_args["seq_length"] = self.seq_length
+        elif "t5" in model.config.model_type.lower():
+            model_type_name = "t5"
+            num_layers = model.config.num_layers
+            hidden_size = model.config.d_model
+            num_attention_heads = model.config.num_heads
+            max_position_embeddings = model.config.n_positions if hasattr(model.config, "n_positions") else 1024
+            orig_vocab_size = model.config.vocab_size
+            pretraining_flag = True
+            if self.encoder_seq_length is None:
+                if batch_data is not None:
+                    self.encoder_seq_length = batch_data["input_ids"].shape[1]
+                else:
+                    self.encoder_seq_length = max_position_embeddings
+            if self.decoder_seq_length is None:
+                if batch_data is not None:
+                    self.decoder_seq_length = batch_data["labels"].shape[1]
+                else:
+                    self.decoder_seq_length = max_position_embeddings
+
+            self.megatron_lm_default_args["encoder_seq_length"] = self.encoder_seq_length
+            self.megatron_lm_default_args["decoder_seq_length"] = self.decoder_seq_length
+        else:
+            raise ValueError(
+                "🤗 Accelerate Megatron-LM integration supports only BERT, GPT and T5 model. "
+                "Please check the model you are using is one of those."
+            )
+
+        self.megatron_lm_default_args["model_type_name"] = model_type_name
+        self.megatron_lm_default_args["num_layers"] = num_layers
+        self.megatron_lm_default_args["hidden_size"] = hidden_size
+        self.megatron_lm_default_args["num_attention_heads"] = num_attention_heads
+        self.megatron_lm_default_args["max_position_embeddings"] = max_position_embeddings
+        self.megatron_lm_default_args["pretraining_flag"] = pretraining_flag
+        self.megatron_lm_default_args["orig_vocab_size"] = orig_vocab_size
+        self.megatron_lm_default_args["model_return_dict"] = model.config.return_dict
+        if model_type_name == "bert":
+            self.megatron_lm_default_args["num_labels"] = num_labels
+
+    def set_mixed_precision(self, mixed_precision):
+        if mixed_precision == "fp16":
+            self.megatron_lm_default_args["fp16"] = True
+        elif mixed_precision == "bf16":
+            self.megatron_lm_default_args["bf16"] = True
+            self.DDP_impl = "local"
+            self.megatron_lm_default_args["DDP_impl"] = self.DDP_impl
+
+    def set_training_args(self, micro_batch_size, dp_degree):
+        self.data_parallel_size = dp_degree
+        self.micro_batch_size = micro_batch_size
+        self.global_batch_size = dp_degree * micro_batch_size * self.num_micro_batches
+        self.megatron_lm_default_args["data_parallel_size"] = self.data_parallel_size
+        self.megatron_lm_default_args["micro_batch_size"] = self.micro_batch_size
+        self.megatron_lm_default_args["global_batch_size"] = self.global_batch_size
+
+    def set_optimizer_type(self, optimizer):
+        optimizer_name = optimizer.__class__.__name__.lower()
+        if "adam" in optimizer_name:
+            self.megatron_lm_default_args["optimizer"] = "adam"
+            self.megatron_lm_default_args["adam_beta1"] = optimizer.defaults["betas"][0]
+            self.megatron_lm_default_args["adam_beta2"] = optimizer.defaults["betas"][1]
+            self.megatron_lm_default_args["adam_eps"] = optimizer.defaults["eps"]
+        elif "sgd" in optimizer_name:
+            self.megatron_lm_default_args["optimizer"] = "sgd"
+            self.megatron_lm_default_args["sgd_momentum"] = optimizer.defaults["momentum"]
+        else:
+            raise ValueError(f"Optimizer {optimizer_name} is not supported by Megatron-LM")
+
+        self.megatron_lm_default_args["lr"] = optimizer.defaults["lr"]
+        self.megatron_lm_default_args["weight_decay"] = optimizer.defaults["weight_decay"]
+
+    def set_scheduler_args(self, scheduler):
+        if self.train_iters is None:
+            self.train_iters = scheduler.total_num_steps // self.megatron_lm_default_args["data_parallel_size"]
+            if self.train_samples is not None:
+                self.train_samples = None
+                warnings.warn(
+                    "Ignoring `train_samples` as `train_iters` based on scheduler is being used for training."
+                )
+        if self.lr_warmup_iters is None:
+            self.lr_warmup_iters = scheduler.warmup_num_steps // self.megatron_lm_default_args["data_parallel_size"]
+            if self.lr_warmup_samples is not None:
+                warnings.warn(
+                    "Ignoring `lr_warmup_samples` as `lr_warmup_iters` based on scheduler is being used for training."
+                )
+            self.lr_warmup_samples = 0
+
+        self.megatron_lm_default_args["train_iters"] = self.train_iters
+        self.megatron_lm_default_args["lr_warmup_iters"] = self.lr_warmup_iters
+        self.megatron_lm_default_args["train_samples"] = self.train_samples
+        self.megatron_lm_default_args["lr_warmup_samples"] = self.lr_warmup_samples
+        self.megatron_lm_default_args["lr_decay_iters"] = self.lr_decay_iters
+        self.megatron_lm_default_args["lr_decay_samples"] = self.lr_decay_samples
+        self.megatron_lm_default_args["lr_warmup_fraction"] = self.lr_warmup_fraction
+        self.megatron_lm_default_args["lr_decay_style"] = self.lr_decay_style
+        self.megatron_lm_default_args["weight_decay_incr_style"] = self.weight_decay_incr_style
+        self.megatron_lm_default_args["start_weight_decay"] = self.start_weight_decay
+        self.megatron_lm_default_args["end_weight_decay"] = self.end_weight_decay
+        self.megatron_lm_default_args["min_lr"] = self.min_lr
+
+    def set_tensorboard_logging_options(self):
+        from megatron.arguments import _add_logging_args
+
+        parser = argparse.ArgumentParser()
+        parser = _add_logging_args(parser)
+        logging_args = parser.parse_known_args()
+        self.dataset_args = vars(logging_args[0])
+        for key, value in self.dataset_args.items():
+            if key.startswith("log_"):
+                self.megatron_lm_default_args[key] = True
+            elif key.startswith("no_log_"):
+                self.megatron_lm_default_args[key.replace("no_", "")] = True
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index eb5fda468..58009d03b 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -87,6 +87,10 @@ def is_bf16_available(ignore_tpu=False):
     return False
 
 
+def is_megatron_lm_available():
+    return importlib.util.find_spec("megatron") is not None
+
+
 def is_transformers_available():
     return importlib.util.find_spec("transformers") is not None
 
diff --git a/src/accelerate/utils/megatron_lm.py b/src/accelerate/utils/megatron_lm.py
new file mode 100644
index 000000000..cb988e360
--- /dev/null
+++ b/src/accelerate/utils/megatron_lm.py
@@ -0,0 +1,1199 @@
+# 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 math
+from abc import ABC
+from functools import partial
+
+import torch
+import torch.nn.functional as F
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from ..optimizer import AcceleratedOptimizer
+from ..scheduler import AcceleratedScheduler
+from .imports import is_megatron_lm_available, is_transformers_available
+from .operations import send_to_device
+
+
+if is_transformers_available():
+    from transformers.modeling_outputs import (
+        CausalLMOutputWithCrossAttentions,
+        Seq2SeqLMOutput,
+        SequenceClassifierOutput,
+    )
+
+
+if is_megatron_lm_available():
+    from megatron import (
+        get_args,
+        get_num_microbatches,
+        get_tensorboard_writer,
+        get_timers,
+        mpu,
+        print_rank_0,
+        print_rank_last,
+    )
+    from megatron.arguments import _add_data_args, _add_validation_args, parse_args, validate_args
+    from megatron.checkpointing import load_args_from_checkpoint, load_checkpoint, save_checkpoint
+    from megatron.data.data_samplers import MegatronPretrainingRandomSampler, MegatronPretrainingSampler
+    from megatron.global_vars import set_global_variables
+    from megatron.initialize import (
+        _compile_dependencies,
+        _init_autoresume,
+        _set_random_seed,
+        set_jit_fusion_options,
+        write_args_to_tensorboard,
+    )
+    from megatron.model import BertModel, GPTModel, ModelType, T5Model
+    from megatron.model.classification import Classification
+    from megatron.optimizer import get_megatron_optimizer
+    from megatron.schedules import get_forward_backward_func
+    from megatron.tokenizer.tokenizer import _vocab_size_with_padding
+    from megatron.training import get_model, get_optimizer_param_scheduler, training_log
+    from megatron.utils import (
+        average_losses_across_data_parallel_group,
+        calc_params_l2_norm,
+        get_ltor_masks_and_position_ids,
+    )
+
+
+# model utilities
+def model_provider_func(pre_process=True, post_process=True, add_encoder=True, add_decoder=True):
+    """Build the model."""
+    args = get_args()
+    mode = "pre-training" if args.pretraining_flag else "fine-tuning"
+    if args.rank == 0:
+        print(f"Building {args.model_type_name} model in the {mode} mode.")
+        print(
+            "The Megatron LM model weights are initialized at random in `accelerator.prepare`. "
+            "Please use `accelerator.load_checkpoint` to load a pre-trained checkpoint matching the distributed setup."
+        )
+    if args.model_type_name == "bert":
+        if args.pretraining_flag:
+            num_tokentypes = 2 if args.bert_binary_head else 0
+            model = BertModel(
+                num_tokentypes=num_tokentypes,
+                add_binary_head=args.bert_binary_head,
+                parallel_output=True,
+                pre_process=pre_process,
+                post_process=post_process,
+            )
+        else:
+
+            model = Classification(
+                num_classes=args.num_labels, num_tokentypes=2, pre_process=pre_process, post_process=post_process
+            )
+    elif args.model_type_name == "gpt":
+        model = GPTModel(num_tokentypes=0, parallel_output=True, pre_process=pre_process, post_process=post_process)
+    elif args.model_type_name == "t5":
+        model = T5Model(
+            num_tokentypes=0,
+            parallel_output=True,
+            pre_process=pre_process,
+            post_process=post_process,
+            add_encoder=add_encoder,
+            add_decoder=add_decoder,
+        )
+    else:
+        raise ValueError(f"Unsupported model type: {args.model_type_name}")
+    return model
+
+
+def prepare_model(accelerator):
+    accelerator.print("Preparing model")
+    args = get_args()
+    if args.model_type_name == "bert" or args.model_type_name == "gpt":
+        model_type = ModelType.encoder_or_decoder
+    elif args.model_type_name == "t5":
+        model_type = ModelType.encoder_and_decoder
+        if args.pipeline_model_parallel_split_rank is None and args.pipeline_model_parallel_size > 1:
+            args.pipeline_model_parallel_split_rank = args.pipeline_model_parallel_size // 2
+    model = get_model(model_provider_func, model_type)
+    return model
+
+
+# dataloader utilities
+class MegatronLMDummyDataLoader:
+    """
+    Dummy dataloader presents model parameters or param groups, this is primarily used to follow conventional training
+
+    Args:
+        **dataset_kwargs: Megatron data arguments.
+    """
+
+    def __init__(self, **dataset_kwargs):
+        parser = argparse.ArgumentParser()
+        parser = _add_data_args(parser)
+        parser = _add_validation_args(parser)
+        data_args = parser.parse_known_args()
+        self.dataset_args = vars(data_args[0])
+        self.dataset_args.update(dataset_kwargs)
+        self.dataset_args["megatron_dataset_flag"] = True
+
+    def set_megatron_data_args(self):
+        args = get_args()
+        for key, value in self.dataset_args.items():
+            setattr(args, key, value)
+
+    def get_train_valid_test_datasets_provider(self):
+        def train_valid_test_datasets_provider(train_val_test_num_samples):
+            """Build train, valid, and test datasets."""
+            args = get_args()
+            dataset_args = {
+                "data_prefix": args.data_path,
+                "data_impl": args.data_impl,
+                "splits_string": args.split,
+                "train_valid_test_num_samples": train_val_test_num_samples,
+                "skip_warmup": (not args.mmap_warmup),
+                "seed": args.seed,
+            }
+            if args.model_type_name == "bert":
+                dataset_args.update(
+                    {
+                        "max_seq_length": args.seq_length,
+                        "masked_lm_prob": args.mask_prob,
+                        "short_seq_prob": args.short_seq_prob,
+                        "binary_head": args.bert_binary_head,
+                    }
+                )
+            elif args.model_type_name == "gpt":
+                dataset_args.update(
+                    {
+                        "seq_length": args.seq_length,
+                    }
+                )
+            elif args.model_type_name == "t5":
+                dataset_args.update(
+                    {
+                        "max_seq_length": args.encoder_seq_length,
+                        "max_seq_length_dec": args.decoder_seq_length,
+                        "masked_lm_prob": args.mask_prob,
+                        "short_seq_prob": args.short_seq_prob,
+                        "dataset_type": "t5",
+                    }
+                )
+            else:
+                raise ValueError(f"Unsupported model type: {args.model_type_name}")
+            if args.model_type_name == "gpt":
+                from megatron.data.gpt_dataset import build_train_valid_test_datasets
+            else:
+                from megatron.data.dataset_utils import build_train_valid_test_datasets
+            train_ds, valid_ds, test_ds = build_train_valid_test_datasets(**dataset_args)
+            return train_ds, valid_ds, test_ds
+
+        return train_valid_test_datasets_provider
+
+    def build_pretraining_data_loader(self, dataset, consumed_samples):
+        if dataset is None:
+            return None
+        args = get_args()
+        micro_batch_size = args.micro_batch_size * args.num_micro_batches
+
+        # Megatron sampler
+        if args.dataloader_type == "single":
+            batch_sampler = MegatronPretrainingSampler(
+                total_samples=len(dataset),
+                consumed_samples=consumed_samples,
+                micro_batch_size=micro_batch_size,
+                data_parallel_rank=mpu.get_data_parallel_rank(),
+                data_parallel_size=mpu.get_data_parallel_world_size(),
+            )
+        elif args.dataloader_type == "cyclic":
+            batch_sampler = MegatronPretrainingRandomSampler(
+                dataset,
+                total_samples=len(dataset),
+                consumed_samples=consumed_samples,
+                micro_batch_size=micro_batch_size,
+                data_parallel_rank=mpu.get_data_parallel_rank(),
+                data_parallel_size=mpu.get_data_parallel_world_size(),
+                data_sharding=args.data_sharding,
+            )
+        else:
+            raise Exception("{} dataloader type is not supported.".format(args.dataloader_type))
+
+        # Torch dataloader.
+        return torch.utils.data.DataLoader(
+            dataset, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True
+        )
+
+    def build_train_valid_test_data_iterators(self):
+        def cyclic_iter(iter):
+            while True:
+                for x in iter:
+                    yield x
+
+        args = get_args()
+
+        (train_dataloader, valid_dataloader, test_dataloader) = (None, None, None)
+
+        print_rank_0("> building train, validation, and test datasets ...")
+
+        # Backward compatibility, assume fixed batch size.
+        if args.iteration > 0 and args.consumed_train_samples == 0:
+            assert args.train_samples is None, "only backward compatiblity support for iteration-based training"
+            args.consumed_train_samples = args.iteration * args.global_batch_size
+        if args.iteration > 0 and args.consumed_valid_samples == 0:
+            if args.train_samples is None:
+                args.consumed_valid_samples = (
+                    (args.iteration // args.eval_interval) * args.eval_iters * args.global_batch_size
+                )
+
+        # Data loader only on rank 0 of each model parallel group.
+        if mpu.get_tensor_model_parallel_rank() == 0:
+
+            # Number of train/valid/test samples.
+            if args.train_samples:
+                train_samples = args.train_samples
+            else:
+                train_samples = args.train_iters * args.global_batch_size
+            eval_iters = (args.train_iters // args.eval_interval + 1) * args.eval_iters
+            test_iters = args.eval_iters
+            train_val_test_num_samples = [
+                train_samples,
+                eval_iters * args.global_batch_size,
+                test_iters * args.global_batch_size,
+            ]
+            print_rank_0(" > datasets target sizes (minimum size):")
+            print_rank_0("    train:      {}".format(train_val_test_num_samples[0]))
+            print_rank_0("    validation: {}".format(train_val_test_num_samples[1]))
+            print_rank_0("    test:       {}".format(train_val_test_num_samples[2]))
+
+            # Build the datasets.
+            train_valid_test_datasets_provider = self.get_train_valid_test_datasets_provider()
+            train_ds, valid_ds, test_ds = train_valid_test_datasets_provider(train_val_test_num_samples)
+
+            # Build dataloders.
+            train_dataloader = self.build_pretraining_data_loader(train_ds, args.consumed_train_samples)
+            valid_dataloader = self.build_pretraining_data_loader(valid_ds, args.consumed_valid_samples)
+            test_dataloader = self.build_pretraining_data_loader(test_ds, 0)
+
+            # Flags to know if we need to do training/validation/testing.
+            do_train = train_dataloader is not None and args.train_iters > 0
+            do_valid = valid_dataloader is not None and args.eval_iters > 0
+            do_test = test_dataloader is not None and args.eval_iters > 0
+            # Need to broadcast num_tokens and num_type_tokens.
+            flags = torch.cuda.LongTensor([int(do_train), int(do_valid), int(do_test)])
+        else:
+            flags = torch.cuda.LongTensor([0, 0, 0])
+
+        # Broadcast num tokens.
+        torch.distributed.broadcast(
+            flags, mpu.get_tensor_model_parallel_src_rank(), group=mpu.get_tensor_model_parallel_group()
+        )
+        args.do_train = flags[0].item()
+        args.do_valid = flags[1].item()
+        args.do_test = flags[2].item()
+
+        # Build iterators.
+        dl_type = args.dataloader_type
+        assert dl_type in ["single", "cyclic"]
+
+        if train_dataloader is not None:
+            train_data_iterator = (
+                iter(train_dataloader) if dl_type == "single" else iter(cyclic_iter(train_dataloader))
+            )
+        else:
+            train_data_iterator = None
+
+        if valid_dataloader is not None:
+            valid_data_iterator = (
+                iter(valid_dataloader) if dl_type == "single" else iter(cyclic_iter(valid_dataloader))
+            )
+        else:
+            valid_data_iterator = None
+
+        if test_dataloader is not None:
+            test_data_iterator = iter(test_dataloader) if dl_type == "single" else iter(cyclic_iter(test_dataloader))
+        else:
+            test_data_iterator = None
+
+        return train_data_iterator, valid_data_iterator, test_data_iterator
+
+
+def prepare_data_loader(accelerator, dataloader):
+    accelerator.print("Preparing dataloader")
+    args = get_args()
+    if not args.megatron_dataset_flag:
+        from ..data_loader import _PYTORCH_DATALOADER_KWARGS, prepare_data_loader
+
+        args = get_args()
+        micro_batch_size = args.micro_batch_size * args.num_micro_batches
+        kwargs = {k: getattr(dataloader, k, _PYTORCH_DATALOADER_KWARGS[k]) for k in _PYTORCH_DATALOADER_KWARGS}
+        if kwargs["batch_size"] is None:
+            if isinstance(kwargs["sampler"], torch.utils.data.BatchSampler):
+                kwargs["sampler"].batch_size = micro_batch_size
+            else:
+                del kwargs["sampler"]
+                del kwargs["shuffle"]
+                del kwargs["batch_size"]
+                kwargs["batch_sampler"].batch_size = micro_batch_size
+        else:
+            del kwargs["batch_sampler"]
+            kwargs["batch_size"] = micro_batch_size
+
+        dataloader = torch.utils.data.DataLoader(dataloader.dataset, **kwargs)
+        return prepare_data_loader(
+            dataloader,
+            accelerator.device,
+            num_processes=mpu.get_data_parallel_world_size(),
+            process_index=mpu.get_data_parallel_rank(),
+            split_batches=accelerator.split_batches,
+            put_on_device=True,
+            rng_types=accelerator.rng_types.copy(),
+            dispatch_batches=accelerator.dispatch_batches,
+        )
+    else:
+        if args.consumed_samples is not None:
+            (
+                args.consumed_train_samples,
+                args.consumed_valid_samples,
+                args.consumed_test_samples,
+            ) = args.consumed_samples
+        else:
+            args.consumed_train_samples, args.consumed_valid_samples, args.consumed_test_samples = 0, 0, 0
+        (
+            train_data_iterator,
+            valid_data_iterator,
+            test_data_iterator,
+        ) = dataloader.build_train_valid_test_data_iterators()
+        return train_data_iterator, valid_data_iterator, test_data_iterator
+
+
+# optimizer utilities
+class MegatronLMOptimizerWrapper(AcceleratedOptimizer):
+    def __init__(self, optimizer):
+        super().__init__(optimizer, device_placement=False, scaler=None)
+
+    def zero_grad(self, set_to_none=None):
+        pass  # `model(**batch)` is doing that automatically. Therefore, it's implementation is not needed
+
+    def step(self):
+        pass  # `model(**batch)` is doing that automatically. Therefore, it's implementation is not needed
+
+    @property
+    def step_was_skipped(self):
+        """Whether or not the optimizer step was done, or skipped because of gradient overflow."""
+        return self.optimizer.skipped_iter
+
+
+def prepare_optimizer(accelerator, model):
+    accelerator.print("Preparing optimizer")
+    args = get_args()
+    optimizer = get_megatron_optimizer(model, args.no_wd_decay_cond, args.scale_lr_cond, args.lr_mult)
+    return optimizer
+
+
+# scheduler utilities
+class MegatronLMDummyScheduler:
+    """
+    Dummy scheduler presents model parameters or param groups, this is primarily used to follow conventional training
+    loop when scheduler config is specified in the deepspeed config file.
+
+    Args:
+        optimizer (`torch.optim.optimizer.Optimizer`):
+            The optimizer to wrap.
+        total_num_steps (int):
+            Total number of steps.
+        warmup_num_steps (int):
+            Number of steps for warmup.
+        **kwargs:
+            Other arguments.
+    """
+
+    def __init__(self, optimizer, total_num_steps=None, warmup_num_steps=0, **kwargs):
+        self.optimizer = optimizer
+        self.total_num_steps = total_num_steps
+        self.warmup_num_steps = warmup_num_steps
+        self.kwargs = kwargs
+
+
+class MegatronLMSchedulerWrapper(AcceleratedScheduler):
+    def __init__(self, scheduler, optimizers):
+        super().__init__(scheduler, optimizers)
+
+    def step(self, *args, **kwargs):
+        return  # `model(**batch)` is doing that automatically. Therefore, it's implementation is not needed
+
+
+def prepare_scheduler(accelerator, optimizer, scheduler):
+    accelerator.print("Preparing scheduler")
+    scheduler = get_optimizer_param_scheduler(optimizer)
+    return scheduler
+
+
+class AbstractTrainStep(ABC):
+    """Abstract class for batching, forward pass and loss handler."""
+
+    def __init__(self, name):
+        super().__init__()
+        self.name = name
+
+    def get_batch_func(self):
+        pass
+
+    def get_forward_step_func(self):
+        pass
+
+    def get_loss_func(self):
+        pass
+
+
+class BertTrainStep(AbstractTrainStep):
+    """
+    Bert train step class.
+
+    Args:
+        args (`argparse.Namespace`): Megatron-LM arguments.
+    """
+
+    def __init__(self, args):
+        super().__init__("BertTrainStep")
+        self.get_batch = self.get_batch_func(args.megatron_dataset_flag)
+        self.loss_func = self.get_loss_func(args.pretraining_flag, args.num_labels)
+        self.forward_step = self.get_forward_step_func(args.pretraining_flag, args.bert_binary_head)
+        if not args.model_return_dict:
+            self.model_output_class = None
+        else:
+            self.model_output_class = SequenceClassifierOutput
+
+    def get_batch_func(self, megatron_dataset_flag):
+        def get_batch_megatron(data_iterator):
+            """Build the batch."""
+
+            # Items and their type.
+            keys = ["text", "types", "labels", "is_random", "loss_mask", "padding_mask"]
+            datatype = torch.int64
+
+            # Broadcast data.
+            if data_iterator is not None:
+                data = next(data_iterator)
+            else:
+                data = None
+            data_b = mpu.broadcast_data(keys, data, datatype)
+
+            # Unpack.
+            tokens = data_b["text"].long()
+            types = data_b["types"].long()
+            sentence_order = data_b["is_random"].long()
+            loss_mask = data_b["loss_mask"].float()
+            lm_labels = data_b["labels"].long()
+            padding_mask = data_b["padding_mask"].long()
+
+            return tokens, types, sentence_order, loss_mask, lm_labels, padding_mask
+
+        def get_batch_transformer(data_iterator):
+            """Build the batch."""
+            data = next(data_iterator)
+            data = send_to_device(data, torch.cuda.current_device())
+
+            # Unpack.
+            tokens = data["input_ids"].long()
+            padding_mask = data["attention_mask"].long()
+            if "token_type_ids" in data:
+                types = data["token_type_ids"].long()
+            else:
+                types = None
+            if "labels" in data:
+                lm_labels = data["labels"].long()
+                loss_mask = (data["labels"] != -100).to(torch.float)
+            else:
+                lm_labels = None
+                loss_mask = None
+            if "next_sentence_label" in data:
+                sentence_order = data["next_sentence_label"].long()
+            else:
+                sentence_order = None
+
+            return tokens, types, sentence_order, loss_mask, lm_labels, padding_mask
+
+        if megatron_dataset_flag:
+            return get_batch_megatron
+        else:
+            return get_batch_transformer
+
+    def get_loss_func(self, pretraining_flag, num_labels):
+        def loss_func_pretrain(loss_mask, sentence_order, output_tensor):
+            lm_loss_, sop_logits = output_tensor
+
+            lm_loss_ = lm_loss_.float()
+            loss_mask = loss_mask.float()
+            lm_loss = torch.sum(lm_loss_.view(-1) * loss_mask.reshape(-1)) / loss_mask.sum()
+
+            if sop_logits is not None:
+                sop_loss = F.cross_entropy(sop_logits.view(-1, 2).float(), sentence_order.view(-1), ignore_index=-1)
+                sop_loss = sop_loss.float()
+                loss = lm_loss + sop_loss
+                averaged_losses = average_losses_across_data_parallel_group([lm_loss, sop_loss])
+                return loss, {"lm loss": averaged_losses[0], "sop loss": averaged_losses[1]}
+
+            else:
+                loss = lm_loss
+                averaged_losses = average_losses_across_data_parallel_group([lm_loss])
+                return loss, {"lm loss": averaged_losses[0]}
+
+        def loss_func_finetune(labels, logits):
+            if num_labels == 1:
+                #  We are doing regression
+                loss_fct = MSELoss()
+                loss = loss_fct(logits.view(-1), labels.view(-1))
+            elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(logits.view(-1, num_labels), labels.view(-1))
+            else:
+                loss_fct = BCEWithLogitsLoss()
+                loss = loss_fct(logits, labels)
+            averaged_losses = average_losses_across_data_parallel_group([loss])
+            return loss, {"loss": averaged_losses[0]}
+
+        if pretraining_flag:
+            return loss_func_pretrain
+        else:
+            return loss_func_finetune
+
+    def get_forward_step_func(self, pretraining_flag, bert_binary_head):
+        def forward_step(data_iterator, model):
+            """Forward step."""
+            tokens, types, sentence_order, loss_mask, labels, padding_mask = self.get_batch(data_iterator)
+            if not bert_binary_head:
+                types = None
+            # Forward pass through the model.
+            if pretraining_flag:
+                output_tensor = model(tokens, padding_mask, tokentype_ids=types, lm_labels=labels)
+                return output_tensor, partial(self.loss_func, loss_mask, sentence_order)
+            else:
+                logits = model(tokens, padding_mask, tokentype_ids=types)
+                return logits, partial(self.loss_func, labels)
+
+        return forward_step
+
+
+class GPTTrainStep(AbstractTrainStep):
+    """
+    GPT train step class.
+
+    Args:
+        args (`argparse.Namespace`): Megatron-LM arguments.
+    """
+
+    def __init__(self, args):
+        super().__init__("GPTTrainStep")
+        self.get_batch = self.get_batch_func(args.megatron_dataset_flag)
+        self.loss_func = self.get_loss_func()
+        self.forward_step = self.get_forward_step_func()
+        self.eod_token = args.padded_vocab_size - 1
+        if not args.model_return_dict:
+            self.model_output_class = None
+        else:
+            self.model_output_class = CausalLMOutputWithCrossAttentions
+
+    def get_batch_func(self, megatron_dataset_flag):
+        def get_batch_megatron(data_iterator):
+            """Generate a batch"""
+            args = get_args()
+
+            # Items and their type.
+            keys = ["text"]
+            datatype = torch.int64
+
+            # Broadcast data.
+            if data_iterator is not None:
+                data = next(data_iterator)
+            else:
+                data = None
+            data_b = mpu.broadcast_data(keys, data, datatype)
+
+            # Unpack.
+            tokens_ = data_b["text"].long()
+            labels = tokens_[:, 1:].contiguous()
+            tokens = tokens_[:, :-1].contiguous()
+
+            # Get the masks and postition ids.
+            attention_mask, loss_mask, position_ids = get_ltor_masks_and_position_ids(
+                tokens, self.eod_token, args.reset_position_ids, args.reset_attention_mask, args.eod_mask_loss
+            )
+
+            return tokens, labels, loss_mask, attention_mask, position_ids
+
+        def get_batch_transformer(data_iterator):
+            data = next(data_iterator)
+            data = {"input_ids": data["input_ids"]}
+            data = send_to_device(data, torch.cuda.current_device())
+
+            tokens_ = data["input_ids"].long()
+            padding = torch.zeros((tokens_.shape[0], 1), dtype=tokens_.dtype, device=tokens_.device) + self.eod_token
+            tokens_ = torch.concat([tokens_, padding], dim=1)
+            labels = tokens_[:, 1:].contiguous()
+            tokens = tokens_[:, :-1].contiguous()
+            # Get the masks and postition ids.
+            attention_mask, loss_mask, position_ids = get_ltor_masks_and_position_ids(
+                tokens, self.eod_token, False, False, True
+            )
+            return tokens, labels, loss_mask, attention_mask, position_ids
+
+        if megatron_dataset_flag:
+            return get_batch_megatron
+        else:
+            return get_batch_transformer
+
+    def get_loss_func(self):
+        def loss_func(loss_mask, output_tensor):
+            losses = output_tensor.float()
+            loss_mask = loss_mask.view(-1).float()
+            loss = torch.sum(losses.view(-1) * loss_mask) / loss_mask.sum()
+
+            # Reduce loss for logging.
+            averaged_loss = average_losses_across_data_parallel_group([loss])
+
+            return loss, {"lm loss": averaged_loss[0]}
+
+        return loss_func
+
+    def get_forward_step_func(self):
+        def forward_step(data_iterator, model):
+            """Forward step."""
+            # Get the batch.
+            tokens, labels, loss_mask, attention_mask, position_ids = self.get_batch(data_iterator)
+            output_tensor = model(tokens, position_ids, attention_mask, labels=labels)
+
+            return output_tensor, partial(self.loss_func, loss_mask)
+
+        return forward_step
+
+
+class T5TrainStep(AbstractTrainStep):
+    """
+    T5 train step class.
+
+    Args:
+        args (`argparse.Namespace`): Megatron-LM arguments.
+    """
+
+    def __init__(self, args):
+        super().__init__("T5TrainStep")
+        self.get_batch = self.get_batch_func(args.megatron_dataset_flag)
+        self.loss_func = self.get_loss_func()
+        self.forward_step = self.get_forward_step_func()
+        if not args.model_return_dict:
+            self.model_output_class = None
+        else:
+            self.model_output_class = Seq2SeqLMOutput
+
+    @staticmethod
+    def attn_mask_postprocess(attention_mask):
+        # We create a 3D attention mask from a 2D tensor mask.
+        # [b, 1, s]
+        attention_mask_b1s = attention_mask.unsqueeze(1)
+        # [b, s, 1]
+        attention_mask_bs1 = attention_mask.unsqueeze(2)
+        # [b, s, s]
+        attention_mask_bss = attention_mask_b1s * attention_mask_bs1
+        # Convert attention mask to binary:
+        extended_attention_mask = attention_mask_bss < 0.5
+        return extended_attention_mask
+
+    @staticmethod
+    def get_decoder_mask(seq_length, device):
+        attention_mask = torch.tril(torch.ones((1, seq_length, seq_length), device=device))
+        attention_mask = attention_mask < 0.5
+        return attention_mask
+
+    @staticmethod
+    def get_enc_dec_mask(attention_mask, dec_seq_length, device):
+        batch_size, _ = attention_mask.shape
+        # We create a 3D attention mask from a 2D tensor mask.
+        # [b, 1, s]
+        attention_mask_b1s = attention_mask.unsqueeze(1)
+        # [b, s, 1]
+        attention_mask_bs1 = torch.ones((batch_size, dec_seq_length, 1), device=device)
+        attention_mask_bss = attention_mask_bs1 * attention_mask_b1s
+        extended_attention_mask = attention_mask_bss < 0.5
+        return extended_attention_mask
+
+    def get_batch_func(self, megatron_dataset_flag):
+        def get_batch_megatron(data_iterator):
+            """Build the batch."""
+
+            keys = ["text_enc", "text_dec", "labels", "loss_mask", "enc_mask", "dec_mask", "enc_dec_mask"]
+            datatype = torch.int64
+
+            # Broadcast data.
+            if data_iterator is not None:
+                data = next(data_iterator)
+            else:
+                data = None
+            data_b = mpu.broadcast_data(keys, data, datatype)
+
+            # Unpack.
+            tokens_enc = data_b["text_enc"].long()
+            tokens_dec = data_b["text_dec"].long()
+            labels = data_b["labels"].long()
+            loss_mask = data_b["loss_mask"].float()
+
+            enc_mask = data_b["enc_mask"] < 0.5
+            dec_mask = data_b["dec_mask"] < 0.5
+            enc_dec_mask = data_b["enc_dec_mask"] < 0.5
+
+            return tokens_enc, tokens_dec, loss_mask, labels, enc_mask, dec_mask, enc_dec_mask
+
+        def get_batch_transformer(data_iterator):
+            """Build the batch."""
+            data = next(data_iterator)
+            data = send_to_device(data, torch.cuda.current_device())
+
+            tokens_enc = data["input_ids"].long()
+            labels = data["labels"].long()
+            loss_mask = (labels != -100).to(torch.float)
+            if "decoder_input_ids" in data:
+                tokens_dec = data["decoder_input_ids"].long()
+            else:
+                tokens_dec = labels.new_zeros(labels.shape, device=labels.device, dtype=torch.long)
+                tokens_dec[..., 1:] = labels[..., :-1].clone()
+                tokens_dec[..., 0] = 0
+                tokens_dec.masked_fill_(tokens_dec == -100, 0)
+            enc_mask = T5TrainStep.attn_mask_postprocess(data["attention_mask"].long())
+            dec_mask = T5TrainStep.get_decoder_mask(tokens_dec.shape[1], tokens_dec.device)
+            enc_dec_mask = T5TrainStep.get_enc_dec_mask(
+                data["attention_mask"].long(), tokens_dec.shape[1], tokens_dec.device
+            )
+
+            return tokens_enc, tokens_dec, loss_mask, labels, enc_mask, dec_mask, enc_dec_mask
+
+        if megatron_dataset_flag:
+            return get_batch_megatron
+        else:
+            return get_batch_transformer
+
+    def get_loss_func(self):
+        def loss_func(loss_mask, output_tensor):
+            lm_loss_ = output_tensor.float()
+            lm_loss = torch.sum(lm_loss_.view(-1) * loss_mask.reshape(-1)) / loss_mask.sum()
+
+            loss = lm_loss
+            averaged_losses = average_losses_across_data_parallel_group([lm_loss])
+
+            return loss, {"lm loss": averaged_losses[0]}
+
+        return loss_func
+
+    def get_forward_step_func(self):
+        def forward_step(data_iterator, model):
+            """Forward step."""
+            # Get the batch.
+            tokens_enc, tokens_dec, loss_mask, lm_labels, enc_mask, dec_mask, enc_dec_mask = self.get_batch(
+                data_iterator
+            )
+            # Forward model lm_labels
+            output_tensor = model(
+                tokens_enc, tokens_dec, enc_mask, dec_mask, enc_dec_mask, tokentype_ids=None, lm_labels=lm_labels
+            )
+
+            return output_tensor, partial(self.loss_func, loss_mask)
+
+        return forward_step
+
+
+# intialize megatron setup
+def initialize(accelerator, extra_args_provider=None, args_defaults={}):
+    accelerator.print("Initializing Megatron-LM")
+    assert torch.cuda.is_available(), "Megatron requires CUDA."
+
+    # Parse arguments
+    args = parse_args(extra_args_provider, ignore_unknown_args=True)
+
+    # Set defaults
+    for key, value in args_defaults.items():
+        if getattr(args, key, None) is not None:
+            if args.rank == 0:
+                print(
+                    "WARNING: overriding default arguments for {key}:{v} \
+                        with {key}:{v2}".format(
+                        key=key, v=getattr(args, key), v2=value
+                    ),
+                    flush=True,
+                )
+        setattr(args, key, value)
+
+    if args.use_checkpoint_args or args_defaults.get("use_checkpoint_args", False):
+        assert args.load is not None, "--use-checkpoints-args requires --load argument"
+        load_args_from_checkpoint(args)
+
+    validate_args(args)
+
+    # set global args, build tokenizer, and set adlr-autoresume,
+    # tensorboard-writer, and timers.
+    set_global_variables(args)
+
+    # torch.distributed initialization
+    def finish_mpu_init():
+        args = get_args()
+        # Pytorch distributed.
+        device_count = torch.cuda.device_count()
+        args.rank = torch.distributed.get_rank()
+        args.world_size = torch.distributed.get_world_size()
+        if device_count > 0:
+            device = args.rank % device_count
+            if args.local_rank is not None:
+                assert args.local_rank == device, "expected local-rank to be the same as rank % device-count."
+            else:
+                args.local_rank = device
+
+            # Set the tensor model-parallel, pipeline model-parallel, and
+            # data-parallel communicators.
+            if mpu.model_parallel_is_initialized():
+                print("model parallel is already initialized")
+            else:
+                mpu.initialize_model_parallel(
+                    args.tensor_model_parallel_size,
+                    args.pipeline_model_parallel_size,
+                    args.virtual_pipeline_model_parallel_size,
+                    args.pipeline_model_parallel_split_rank,
+                )
+
+        # Random seeds for reproducibility.
+        if args.rank == 0:
+            print("> setting random seeds to {} ...".format(args.seed))
+        _set_random_seed(args.seed, args.data_parallel_random_init)
+
+    args = get_args()
+
+    # Megatron's MPU is the master. Complete initialization right away.
+    finish_mpu_init()
+
+    # Autoresume.
+    _init_autoresume()
+
+    # Compile dependencies.
+    _compile_dependencies()
+
+    # Set pytorch JIT layer fusion options and warmup JIT functions.
+    set_jit_fusion_options()
+    args = get_args()
+    args.padded_vocab_size = _vocab_size_with_padding(args.orig_vocab_size, args)
+    if args.model_type_name == "bert" and args.pretraining_flag and args.num_labels == 2:
+        args.bert_binary_head = True
+    else:
+        args.bert_binary_head = False
+    args.iteration = 0
+
+
+class MegatronEngine(torch.nn.Module):
+    """
+    Megatron-LM model wrapper
+
+    Args:
+        accelerator (:class:`~accelerate.Accelerator`): The accelerator object to use.
+        model: Megatron-LM model
+        optimizer: Megatron-LM optimizer
+        lr_scheduler: Megatron-LM lr scheduler
+    """
+
+    def __init__(self, accelerator, model, optimizer, scheduler):
+        super(MegatronEngine, self).__init__()
+        self.module = model
+        self.base_model = model[0]
+        self.optimizer = optimizer
+        self.scheduler = scheduler
+        args = get_args()
+        if accelerator.state.megatron_lm_plugin.custom_train_step_class is not None:
+            self.train_step_handler = accelerator.state.megatron_lm_plugin.custom_train_step_class(
+                args, **accelerator.state.megatron_lm_plugin.custom_train_step_kwargs
+            )
+        elif args.model_type_name == "bert":
+            self.train_step_handler = BertTrainStep(args)
+        elif args.model_type_name == "gpt":
+            self.train_step_handler = GPTTrainStep(args)
+        elif args.model_type_name == "t5":
+            self.train_step_handler = T5TrainStep(args)
+        else:
+            raise ValueError(f"Unsupported model type: {args.model_type_name}")
+        self.optimizer.skipped_iter = False
+
+        # Tracking loss.
+        self.total_loss_dict = {}
+        self.eval_total_loss_dict = {}
+        self.iteration = 0
+        self.report_memory_flag = True
+        if args.tensorboard_dir is not None:
+            write_args_to_tensorboard()
+
+    def train(self):
+        for model_module in self.module:
+            model_module.train()
+        self.log_eval_results()
+
+    def eval(self):
+        for model_module in self.module:
+            model_module.eval()
+
+    def train_step(self, **batch_data):
+        """
+        Training step for Megatron-LM
+
+        Args:
+            batch_data (:obj:`dict`): The batch data to train on.
+        """
+
+        args = get_args()
+        timers = get_timers()
+
+        if len(batch_data) > 0:
+            data_chunks = []
+            if args.num_micro_batches > 1:
+                for i in range(0, args.num_micro_batches):
+                    data_chunks.append(
+                        {
+                            k: v[i * args.micro_batch_size : (i + 1) * args.micro_batch_size]
+                            for k, v in batch_data.items()
+                        }
+                    )
+            else:
+                data_chunks = [batch_data]
+
+        if len(self.module) > 1:
+            batch_data_iterator = (
+                [iter(data_chunks) for _ in range(len(self.module))]
+                if len(batch_data) > 0
+                else [None] * len(self.module)
+            )
+        else:
+            batch_data_iterator = iter(data_chunks) if len(batch_data) > 0 else None
+
+        # Set grad to zero.
+        if args.DDP_impl == "local" and args.use_contiguous_buffers_in_local_ddp:
+            for partition in self.module:
+                partition.zero_grad_buffer()
+        self.optimizer.zero_grad()
+
+        # Forward pass.
+        forward_backward_func = get_forward_backward_func()
+        losses_reduced = forward_backward_func(
+            self.train_step_handler.forward_step,
+            batch_data_iterator,
+            self.module,
+            self.optimizer,
+            None,
+            forward_only=False,
+        )
+
+        # Empty unused memory.
+        if args.empty_unused_memory_level >= 1:
+            torch.cuda.empty_cache()
+
+        # Reduce gradients.
+        timers("backward-reduce-model-grads").start()
+        self.optimizer.reduce_model_grads(args, timers)
+        timers("backward-reduce-model-grads").stop()
+
+        # Update parameters.
+        timers("optimizer").start()
+        update_successful, grad_norm, num_zeros_in_grad = self.optimizer.step(args, timers)
+        timers("optimizer").stop()
+
+        # Gather params.
+        if update_successful:
+            timers("backward-gather-model-params").start()
+            self.optimizer.gather_model_params(args, timers)
+            timers("backward-gather-model-params").stop()
+
+        # Update learning rate.
+        if update_successful:
+            if self.scheduler is not None:
+                increment = get_num_microbatches() * args.micro_batch_size * args.data_parallel_size
+                self.scheduler.step(increment=increment)
+            skipped_iter = 0
+        else:
+            skipped_iter = 1
+
+        self.optimizer.skipped_iter = not update_successful
+
+        # Empty unused memory.
+        if args.empty_unused_memory_level >= 2:
+            torch.cuda.empty_cache()
+
+        args.consumed_train_samples += (
+            mpu.get_data_parallel_world_size() * args.micro_batch_size * get_num_microbatches()
+        )
+
+        if mpu.is_pipeline_last_stage(ignore_virtual=True):
+            # Average loss across microbatches.
+            loss_reduced = {}
+            for key in losses_reduced[0]:
+                losses_reduced_for_key = [x[key] for x in losses_reduced]
+                loss_reduced[key] = sum(losses_reduced_for_key) / len(losses_reduced_for_key)
+            return loss_reduced, skipped_iter, grad_norm, num_zeros_in_grad
+        return {}, skipped_iter, grad_norm, num_zeros_in_grad
+
+    def eval_step(self, **batch_data):
+        """
+        Evaluation step for Megatron-LM
+
+        Args:
+            batch_data (:obj:`dict`): The batch data to evaluate on.
+        """
+
+        args = get_args()
+        data_chunks = []
+        if args.num_micro_batches > 1:
+            for i in range(0, args.num_micro_batches):
+                data_chunks.append(
+                    {k: v[i * args.micro_batch_size : (i + 1) * args.micro_batch_size] for k, v in batch_data.items()}
+                )
+        else:
+            data_chunks = [batch_data]
+
+        if len(self.module) > 1:
+            batch_data_iterator = [iter(data_chunks) for _ in range(len(self.module))]
+        else:
+            batch_data_iterator = iter(data_chunks)
+        forward_backward_func = get_forward_backward_func()
+        loss_dicts = forward_backward_func(
+            self.train_step_handler.forward_step,
+            batch_data_iterator,
+            self.module,
+            optimizer=None,
+            timers=None,
+            forward_only=True,
+        )
+        # Empty unused memory
+        if args.empty_unused_memory_level >= 1:
+            torch.cuda.empty_cache()
+
+        args.consumed_valid_samples += (
+            mpu.get_data_parallel_world_size() * args.micro_batch_size * get_num_microbatches()
+        )
+
+        if mpu.is_pipeline_last_stage(ignore_virtual=True):
+            # Average loss across microbatches.
+            loss_reduced = {}
+            for key in loss_dicts[0]:
+                losses_reduced_for_key = [x[key] for x in loss_dicts]
+                loss_reduced[key] = sum(losses_reduced_for_key) / len(losses_reduced_for_key)
+            return loss_reduced
+        else:
+            return {}
+
+    def forward(self, **batch_data):
+        # During training, we use train_step()
+        # model(**batch_data) performs following operations by delegating it to `self.train_step`:
+        # 1. Prepare **batch_data for Tendor, Pipeline and Model Parallelism
+        # 2. Set grad to zero.
+        # 3. forward pass and backward pass using Pipeline Parallelism
+        # 4. Empty unused memory.
+        # 5. Reduce gradients.
+        # 6. Update parameters.
+        # 7. Gather params when using Distributed Optimizer (Data Parallelism).
+        # 8. Update learning rate if scheduler is specified.
+        # 9. Empty unused memory.
+        # 10. Average loss across microbatches and across DP ranks.
+        #
+        # During evaluation, we use eval_step()
+        args = get_args()
+        if self.module[0].training:
+            loss_dict, skipped_iter, grad_norm, num_zeros_in_grad = self.train_step(**batch_data)
+            self.iteration += 1
+            if args.tensorboard_dir is not None:
+                # Logging.
+                loss_scale = self.optimizer.get_loss_scale().item()
+                params_norm = None
+                if args.log_params_norm:
+                    params_norm = calc_params_l2_norm(self.model)
+                self.report_memory_flag = training_log(
+                    loss_dict,
+                    self.total_loss_dict,
+                    self.optimizer.param_groups[0]["lr"],
+                    self.iteration,
+                    loss_scale,
+                    self.report_memory_flag,
+                    skipped_iter,
+                    grad_norm,
+                    params_norm,
+                    num_zeros_in_grad,
+                )
+        else:
+            loss_dict = self.eval_step(**batch_data)
+            if args.tensorboard_dir is not None:
+                for key in loss_dict:
+                    self.eval_total_loss_dict[key] = (
+                        self.eval_total_loss_dict.get(key, torch.cuda.FloatTensor([0.0])) + loss_dict[key]
+                    )
+                    self.eval_total_loss_dict[key + "_num_iters"] = self.eval_total_loss_dict.get(
+                        key + "_num_iters", torch.cuda.FloatTensor([0.0])
+                    ) + torch.cuda.FloatTensor([1.0])
+
+        loss = torch.tensor(0.0, device=args.local_rank)
+        for key in loss_dict:
+            loss += loss_dict[key]
+        # loss = reduce(loss)
+        if self.train_step_handler.model_output_class is not None:
+            return self.train_step_handler.model_output_class(loss=loss)
+        return loss
+
+    def log_eval_results(self):
+        args = get_args()
+        if args.tensorboard_dir is None or self.iteration == 0:
+            return
+        args = get_args()
+        writer = get_tensorboard_writer()
+        string = f"validation loss at iteration {self.iteration} | "
+        for key in self.eval_total_loss_dict:
+            if key.endswith("_num_iters"):
+                continue
+            value = self.eval_total_loss_dict[key] / self.eval_total_loss_dict[key + "_num_iters"]
+            string += f"{key} value: {value} | "
+            ppl = math.exp(min(20, value.item()))
+            if args.pretraining_flag:
+                string += f"{key} PPL: {ppl} | "
+            if writer:
+                writer.add_scalar(f"{key} validation", value.item(), self.iteration)
+                if args.pretraining_flag:
+                    writer.add_scalar(f"{key} validation ppl", ppl, self.iteration)
+
+        length = len(string) + 1
+        print_rank_last("-" * length)
+        print_rank_last(string)
+        print_rank_last("-" * length)
+        self.eval_total_loss_dict = {}
+
+    def save_checkpoint(self, output_dir):
+        self.log_eval_results()
+        args = get_args()
+        args.save = output_dir
+        torch.distributed.barrier()
+        save_checkpoint(self.iteration, self.module, self.optimizer, self.scheduler)
+        torch.distributed.barrier()
+
+    def load_checkpoint(self, input_dir):
+        args = get_args()
+        args.load = input_dir
+        args.consumed_train_samples = 0
+        args.consumed_valid_samples = 0
+        torch.distributed.barrier()
+        iteration = load_checkpoint(self.module, self.optimizer, self.scheduler)
+        torch.distributed.barrier()
+        self.iteration = iteration
+        if args.fp16 and self.iteration == 0:
+            self.optimizer.reload_model_params()
+
+
+# other utilities
+def avg_losses_across_data_parallel_group(losses):
+    """
+    Average losses across data parallel group.
+
+    Args:
+        losses (List[Tensor]): List of losses to average across data parallel group.
+    """
+
+    return average_losses_across_data_parallel_group(losses)
diff --git a/src/accelerate/utils/operations.py b/src/accelerate/utils/operations.py
index 3ba46a03f..aef1fd589 100644
--- a/src/accelerate/utils/operations.py
+++ b/src/accelerate/utils/operations.py
@@ -224,6 +224,7 @@ def gather(tensor):
         DistributedType.DEEPSPEED,
         DistributedType.MULTI_GPU,
         DistributedType.FSDP,
+        DistributedType.MEGATRON_LM,
     ]:
         return _gpu_gather(tensor)
     elif AcceleratorState().distributed_type == DistributedType.MULTI_CPU:
@@ -261,6 +262,7 @@ def gather_object(object: Any):
         DistributedType.DEEPSPEED,
         DistributedType.MULTI_GPU,
         DistributedType.FSDP,
+        DistributedType.MEGATRON_LM,
     ]:
         return _gpu_gather_object(object)
     elif AcceleratorState().distributed_type == DistributedType.MULTI_CPU:
@@ -304,6 +306,7 @@ def broadcast(tensor, from_process: int = 0):
         DistributedType.DEEPSPEED,
         DistributedType.MULTI_GPU,
         DistributedType.FSDP,
+        DistributedType.MEGATRON_LM,
     ]:
         return _gpu_broadcast(tensor, src=from_process)
     elif AcceleratorState().distributed_type == DistributedType.MULTI_CPU:
@@ -332,6 +335,7 @@ def broadcast_object_list(object_list, from_process: int = 0):
         DistributedType.DEEPSPEED,
         DistributedType.MULTI_GPU,
         DistributedType.FSDP,
+        DistributedType.MEGATRON_LM,
     ]:
         torch.distributed.broadcast_object_list(object_list, src=from_process)
     elif AcceleratorState().distributed_type == DistributedType.MULTI_CPU:
@@ -452,6 +456,7 @@ def _reduce_across_processes(tensor, reduction="mean"):
             DistributedType.DEEPSPEED,
             DistributedType.MULTI_GPU,
             DistributedType.FSDP,
+            DistributedType.MEGATRON_LM,
         ]:
             torch.distributed.all_reduce(cloned_tensor, ReduceOp.SUM)
             return cloned_tensor
diff --git a/src/accelerate/utils/random.py b/src/accelerate/utils/random.py
index e95ed03bf..8dc149a02 100644
--- a/src/accelerate/utils/random.py
+++ b/src/accelerate/utils/random.py
@@ -64,7 +64,12 @@ def synchronize_rng_state(rng_type: Optional[RNGType] = None, generator: Optiona
     state = AcceleratorState()
     if state.distributed_type == DistributedType.TPU:
         rng_state = xm.mesh_reduce("random_seed", rng_state, lambda x: x[0])
-    elif state.distributed_type in [DistributedType.DEEPSPEED, DistributedType.MULTI_GPU, DistributedType.FSDP]:
+    elif state.distributed_type in [
+        DistributedType.DEEPSPEED,
+        DistributedType.MULTI_GPU,
+        DistributedType.FSDP,
+        DistributedType.MEGATRON_LM,
+    ]:
         rng_state = rng_state.to(state.device)
         torch.distributed.broadcast(rng_state, 0)
         rng_state = rng_state.cpu()
diff --git a/tests/test_examples.py b/tests/test_examples.py
index 9197070c8..05688f2b5 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -39,6 +39,7 @@
     "automatic_gradient_accumulation.py",
     "fsdp_with_peak_mem_tracking.py",
     "deepspeed_with_config_support.py",
+    "megatron_lm_gpt_pretraining.py",
 ]
 
 

```suggestion
        self.megatron_lm_default_args = {
```
Also you use `megatronlm_plugin` with no `_` between `megatron` and `lm`., just flagging this so you stay consistent (I prefer `megatron_lm` personally.
@@ -643,3 +645,142 @@ def load_optimizer(self, accelerator, optimizer, model, input_dir, optimizer_ind
         # called from all ranks, though only rank0 has a valid param for full_osd
         sharded_osd = FSDP.scatter_full_optim_state_dict(full_osd, model)
         optimizer.load_state_dict(sharded_osd)
+
+
+@dataclass
+class MegatronLMPlugin:
+    """
+    Plugin for Megatron-LM to enable tensor, pipeline, sequence and data parallelism. Also to enable selective
+    activation recomputation and optimized fused kernels.
+    """
+
+    tp_degree: int = field(default=None, metadata={"help": "tensor parallelism degree."})
+    pp_degree: int = field(default=None, metadata={"help": "pipeline parallelism degree."})
+    num_micro_batches: int = field(default=None, metadata={"help": "number of micro-batches."})
+    gradient_clipping: float = field(
+        default=None, metadata={"help": "gradient clipping value based on global L2 Norm (0 to disable)"}
+    )
+    sequence_parallelism: bool = field(
+        default=None,
+        metadata={"help": "enable sequence parallelism"},
+    )
+    recompute_activation: bool = field(
+        default=None,
+        metadata={"help": "enable selective activation recomputation"},
+    )
+    use_distributed_optimizer: bool = field(
+        default=None,
+        metadata={"help": "enable distributed optimizer"},
+    )
+    pipeline_model_parallel_split_rank: int = field(
+        default=None, metadata={"help": "Rank where encoder and decoder should be split."}
+    )
+    num_layers_per_virtual_pipeline_stage: int = field(
+        default=None, metadata={"help": "Number of layers per virtual pipeline stage."}
+    )
+    is_train_batch_min: str = field(
+        default=True,
+        metadata={"help": "If both train & eval dataloaders are specified, this will decide the micro_batch_size"},
+    )
+
+    def __post_init__(self):
+        prefix = "MEGATRON_"
+        if self.tp_degree is None:
+            self.tp_degree = int(os.environ.get(prefix + "TP_DEGREE", 1))
+        if self.pp_degree is None:
+            self.pp_degree = int(os.environ.get(prefix + "PP_DEGREE", 1))
+        if self.num_micro_batches is None:
+            self.num_micro_batches = int(os.environ.get(prefix + "NUM_MICRO_BATCHES", 1))
+        if self.gradient_clipping is None:
+            self.gradient_clipping = float(os.environ.get(prefix + "GRADIENT_CLIPPING", 1.0))
+        if self.recompute_activation is None:
+            self.recompute_activation = strtobool(os.environ.get(prefix + "RECOMPUTE_ACTIVATION", "False")) == 1
+        if self.use_distributed_optimizer is None:
+            self.use_distributed_optimizer = (
+                strtobool(os.environ.get(prefix + "USE_DISTRIBUTED_OPTIMIZER", "False")) == 1
+            )
+        if self.sequence_parallelism is None:
+            self.sequence_parallelism = strtobool(os.environ.get(prefix + "SEQUENCE_PARALLELISM", "False")) == 1
+
+        if self.pp_degree > 1:
+            self.DDP_impl = "local"
+        else:
+            self.DDP_impl = "torch"
+
+        self.megtron_lm_default_args = {
Done.
@@ -643,3 +645,142 @@ def load_optimizer(self, accelerator, optimizer, model, input_dir, optimizer_ind
         # called from all ranks, though only rank0 has a valid param for full_osd
         sharded_osd = FSDP.scatter_full_optim_state_dict(full_osd, model)
         optimizer.load_state_dict(sharded_osd)
+
+
+@dataclass
+class MegatronLMPlugin:
+    """
+    Plugin for Megatron-LM to enable tensor, pipeline, sequence and data parallelism. Also to enable selective
+    activation recomputation and optimized fused kernels.
+    """
+
+    tp_degree: int = field(default=None, metadata={"help": "tensor parallelism degree."})
+    pp_degree: int = field(default=None, metadata={"help": "pipeline parallelism degree."})
+    num_micro_batches: int = field(default=None, metadata={"help": "number of micro-batches."})
+    gradient_clipping: float = field(
+        default=None, metadata={"help": "gradient clipping value based on global L2 Norm (0 to disable)"}
+    )
+    sequence_parallelism: bool = field(
+        default=None,
+        metadata={"help": "enable sequence parallelism"},
+    )
+    recompute_activation: bool = field(
+        default=None,
+        metadata={"help": "enable selective activation recomputation"},
+    )
+    use_distributed_optimizer: bool = field(
+        default=None,
+        metadata={"help": "enable distributed optimizer"},
+    )
+    pipeline_model_parallel_split_rank: int = field(
+        default=None, metadata={"help": "Rank where encoder and decoder should be split."}
+    )
+    num_layers_per_virtual_pipeline_stage: int = field(
+        default=None, metadata={"help": "Number of layers per virtual pipeline stage."}
+    )
+    is_train_batch_min: str = field(
+        default=True,
+        metadata={"help": "If both train & eval dataloaders are specified, this will decide the micro_batch_size"},
+    )
+
+    def __post_init__(self):
+        prefix = "MEGATRON_"
+        if self.tp_degree is None:
+            self.tp_degree = int(os.environ.get(prefix + "TP_DEGREE", 1))
+        if self.pp_degree is None:
+            self.pp_degree = int(os.environ.get(prefix + "PP_DEGREE", 1))
+        if self.num_micro_batches is None:
+            self.num_micro_batches = int(os.environ.get(prefix + "NUM_MICRO_BATCHES", 1))
+        if self.gradient_clipping is None:
+            self.gradient_clipping = float(os.environ.get(prefix + "GRADIENT_CLIPPING", 1.0))
+        if self.recompute_activation is None:
+            self.recompute_activation = strtobool(os.environ.get(prefix + "RECOMPUTE_ACTIVATION", "False")) == 1
+        if self.use_distributed_optimizer is None:
+            self.use_distributed_optimizer = (
+                strtobool(os.environ.get(prefix + "USE_DISTRIBUTED_OPTIMIZER", "False")) == 1
+            )
+        if self.sequence_parallelism is None:
+            self.sequence_parallelism = strtobool(os.environ.get(prefix + "SEQUENCE_PARALLELISM", "False")) == 1
+
+        if self.pp_degree > 1:
+            self.DDP_impl = "local"
+        else:
+            self.DDP_impl = "torch"
+
+        self.megtron_lm_default_args = {
I don't believe all of those should be changed just for MegatronLM.
@@ -1032,7 +1174,7 @@ def wait_for_everyone(self):
         """
         wait_for_everyone()
 
-    @on_main_process
+    @on_last_process
Here the message should be more in the lines of "The Megatron-LM only supports those"
@@ -643,3 +645,291 @@ def load_optimizer(self, accelerator, optimizer, model, input_dir, optimizer_ind
         # called from all ranks, though only rank0 has a valid param for full_osd
         sharded_osd = FSDP.scatter_full_optim_state_dict(full_osd, model)
         optimizer.load_state_dict(sharded_osd)
+
+
+@dataclass
+class MegatronLMPlugin:
+    """
+    Plugin for Megatron-LM to enable tensor, pipeline, sequence and data parallelism. Also to enable selective
+    activation recomputation and optimized fused kernels.
+    """
+
+    tp_degree: int = field(default=None, metadata={"help": "tensor parallelism degree."})
+    pp_degree: int = field(default=None, metadata={"help": "pipeline parallelism degree."})
+    num_micro_batches: int = field(default=None, metadata={"help": "number of micro-batches."})
+    gradient_clipping: float = field(
+        default=None, metadata={"help": "gradient clipping value based on global L2 Norm (0 to disable)"}
+    )
+    sequence_parallelism: bool = field(
+        default=None,
+        metadata={"help": "enable sequence parallelism"},
+    )
+    recompute_activation: bool = field(
+        default=None,
+        metadata={"help": "enable selective activation recomputation"},
+    )
+    use_distributed_optimizer: bool = field(
+        default=None,
+        metadata={"help": "enable distributed optimizer"},
+    )
+    pipeline_model_parallel_split_rank: int = field(
+        default=None, metadata={"help": "Rank where encoder and decoder should be split."}
+    )
+    num_layers_per_virtual_pipeline_stage: int = field(
+        default=None, metadata={"help": "Number of layers per virtual pipeline stage."}
+    )
+    is_train_batch_min: str = field(
+        default=True,
+        metadata={"help": "If both train & eval dataloaders are specified, this will decide the micro_batch_size"},
+    )
+    train_iters: int = field(
+        default=None,
+        metadata={
+            "help": "Total number of iterations to train over all training runs. "
+            "Note that either train-iters or train-samples should be provided when using `MegatronLMDummyScheduler`"
+        },
+    )
+    train_samples: int = field(
+        default=None,
+        metadata={
+            "help": "Total number of samples to train over all training runs. "
+            "Note that either train-iters or train-samples should be provided when using `MegatronLMDummyScheduler`"
+        },
+    )
+    weight_decay_incr_style: str = field(
+        default="constant",
+        metadata={
+            "help": 'Weight decay increment function. choices=["constant", "linear", "cosine"]. '
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    start_weight_decay: float = field(
+        default=None,
+        metadata={
+            "help": "Initial weight decay coefficient for L2 regularization. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    end_weight_decay: float = field(
+        default=None,
+        metadata={
+            "help": "End of run weight decay coefficient for L2 regularization. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    lr_decay_style: str = field(
+        default="linear",
+        metadata={
+            "help": "Learning rate decay function. choices=['constant', 'linear', 'cosine']. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    lr_decay_iters: int = field(
+        default=None,
+        metadata={
+            "help": "Number of iterations for learning rate decay. If None defaults to `train_iters`. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    lr_decay_samples: int = field(
+        default=None,
+        metadata={
+            "help": "Number of samples for learning rate decay. If None defaults to `train_samples`. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    lr_warmup_iters: int = field(
+        default=None,
+        metadata={
+            "help": "number of iterations to linearly warmup learning rate over. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    lr_warmup_samples: int = field(
+        default=None,
+        metadata={
+            "help": "number of samples to linearly warmup learning rate over. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    lr_warmup_fraction: float = field(
+        default=None,
+        metadata={
+            "help": "fraction of lr-warmup-(iters/samples) to linearly warmup learning rate over. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    min_lr: float = field(
+        default=None,
+        metadata={
+            "help": "Minumum value for learning rate. The scheduler clip values below this threshold. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    consumed_samples: List[int] = field(
+        default=None,
+        metadata={
+            "help": "Number of samples consumed in the same order as the dataloaders to `accelerator.prepare` call."
+        },
+    )
+    no_wd_decay_cond: Optional[Callable] = field(default=None, metadata={"help": "Condition to disable weight decay."})
+    scale_lr_cond: Optional[Callable] = field(default=None, metadata={"help": "Condition to scale learning rate."})
+    lr_mult: float = field(default=1.0, metadata={"help": "Learning rate multiplier."})
+    megatron_dataset_flag: bool = field(
+        default=False,
+        metadata={"help": "Whether the format of dataset follows Megatron-LM Indexed/Cached/MemoryMapped format."},
+    )
+
+    def __post_init__(self):
+        prefix = "MEGATRON_"
+        if self.tp_degree is None:
+            self.tp_degree = int(os.environ.get(prefix + "TP_DEGREE", 1))
+        if self.pp_degree is None:
+            self.pp_degree = int(os.environ.get(prefix + "PP_DEGREE", 1))
+        if self.num_micro_batches is None:
+            self.num_micro_batches = int(os.environ.get(prefix + "NUM_MICRO_BATCHES", 1))
+        if self.gradient_clipping is None:
+            self.gradient_clipping = float(os.environ.get(prefix + "GRADIENT_CLIPPING", 1.0))
+        if self.recompute_activation is None:
+            self.recompute_activation = strtobool(os.environ.get(prefix + "RECOMPUTE_ACTIVATION", "False")) == 1
+        if self.use_distributed_optimizer is None:
+            self.use_distributed_optimizer = (
+                strtobool(os.environ.get(prefix + "USE_DISTRIBUTED_OPTIMIZER", "False")) == 1
+            )
+        if self.sequence_parallelism is None:
+            self.sequence_parallelism = strtobool(os.environ.get(prefix + "SEQUENCE_PARALLELISM", "False")) == 1
+
+        if self.pp_degree > 1:
+            self.DDP_impl = "local"
+        else:
+            self.DDP_impl = "torch"
+
+        if self.consumed_samples is not None:
+            if len(self.consumed_samples) == 1:
+                self.consumed_samples.extend([0, 0])
+            elif len(self.consumed_samples) == 2:
+                self.consumed_samples.append(0)
+
+        self.megatron_lm_default_args = {
+            "tensor_model_parallel_size": self.tp_degree,
+            "pipeline_model_parallel_size": self.pp_degree,
+            "pipeline_model_parallel_split_rank": self.pipeline_model_parallel_split_rank,
+            "num_layers_per_virtual_pipeline_stage": self.num_layers_per_virtual_pipeline_stage,
+            "DDP_impl": self.DDP_impl,
+            "use_distributed_optimizer": self.use_distributed_optimizer,
+            "recompute_activations": self.recompute_activation,
+            "sequence_parallel": self.sequence_parallelism,
+            "clip_grad": self.gradient_clipping,
+            "num_micro_batches": self.num_micro_batches,
+            "consumed_samples": self.consumed_samples,
+            "no_wd_decay_cond": self.no_wd_decay_cond,
+            "scale_lr_cond": self.scale_lr_cond,
+            "lr_mult": self.lr_mult,
+            "megatron_dataset_flag": self.megatron_dataset_flag,
+        }
+
+    def set_network_size_args(self, model):
+        # Check if the model is either BERT, GPT or T5 else raise error
+        # set 'num_layers', 'hidden_size', 'num_attention_heads', 'max_position_embeddings'
+        if "bert" in model.__class__.__name__.lower():
+            model_type_name = "bert"
+            num_layers = model.config.num_hidden_layers
+            hidden_size = model.config.hidden_size
+            num_attention_heads = model.config.num_attention_heads
+            max_position_embeddings = model.config.max_position_embeddings
+            num_labels = model.config.num_labels
+            orig_vocab_size = model.config.vocab_size
+            if "maskedlm" in model.__class__.__name__.lower():
+                pretraining_flag = True
+
+        elif "gpt" in model.__class__.__name__.lower():
+            model_type_name = "gpt"
+            num_layers = model.config.n_layer
+            hidden_size = model.config.n_embd
+            num_attention_heads = model.config.n_head
+            max_position_embeddings = model.config.n_positions
+            orig_vocab_size = model.config.vocab_size
+            pretraining_flag = True
+        elif "t5" in model.__class__.__name__.lower():
+            model_type_name = "t5"
+            num_layers = model.config.num_layers
+            hidden_size = model.config.d_model
+            num_attention_heads = model.config.num_heads
+            max_position_embeddings = model.config.n_positions
+            orig_vocab_size = model.config.vocab_size
+            pretraining_flag = True
+        else:
+            raise ValueError("Model is not BERT, GPT or T5. Please check the model you are using.")
You should check the `model_type` attribute instead of the class name IMO.
@@ -643,3 +645,291 @@ def load_optimizer(self, accelerator, optimizer, model, input_dir, optimizer_ind
         # called from all ranks, though only rank0 has a valid param for full_osd
         sharded_osd = FSDP.scatter_full_optim_state_dict(full_osd, model)
         optimizer.load_state_dict(sharded_osd)
+
+
+@dataclass
+class MegatronLMPlugin:
+    """
+    Plugin for Megatron-LM to enable tensor, pipeline, sequence and data parallelism. Also to enable selective
+    activation recomputation and optimized fused kernels.
+    """
+
+    tp_degree: int = field(default=None, metadata={"help": "tensor parallelism degree."})
+    pp_degree: int = field(default=None, metadata={"help": "pipeline parallelism degree."})
+    num_micro_batches: int = field(default=None, metadata={"help": "number of micro-batches."})
+    gradient_clipping: float = field(
+        default=None, metadata={"help": "gradient clipping value based on global L2 Norm (0 to disable)"}
+    )
+    sequence_parallelism: bool = field(
+        default=None,
+        metadata={"help": "enable sequence parallelism"},
+    )
+    recompute_activation: bool = field(
+        default=None,
+        metadata={"help": "enable selective activation recomputation"},
+    )
+    use_distributed_optimizer: bool = field(
+        default=None,
+        metadata={"help": "enable distributed optimizer"},
+    )
+    pipeline_model_parallel_split_rank: int = field(
+        default=None, metadata={"help": "Rank where encoder and decoder should be split."}
+    )
+    num_layers_per_virtual_pipeline_stage: int = field(
+        default=None, metadata={"help": "Number of layers per virtual pipeline stage."}
+    )
+    is_train_batch_min: str = field(
+        default=True,
+        metadata={"help": "If both train & eval dataloaders are specified, this will decide the micro_batch_size"},
+    )
+    train_iters: int = field(
+        default=None,
+        metadata={
+            "help": "Total number of iterations to train over all training runs. "
+            "Note that either train-iters or train-samples should be provided when using `MegatronLMDummyScheduler`"
+        },
+    )
+    train_samples: int = field(
+        default=None,
+        metadata={
+            "help": "Total number of samples to train over all training runs. "
+            "Note that either train-iters or train-samples should be provided when using `MegatronLMDummyScheduler`"
+        },
+    )
+    weight_decay_incr_style: str = field(
+        default="constant",
+        metadata={
+            "help": 'Weight decay increment function. choices=["constant", "linear", "cosine"]. '
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    start_weight_decay: float = field(
+        default=None,
+        metadata={
+            "help": "Initial weight decay coefficient for L2 regularization. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    end_weight_decay: float = field(
+        default=None,
+        metadata={
+            "help": "End of run weight decay coefficient for L2 regularization. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    lr_decay_style: str = field(
+        default="linear",
+        metadata={
+            "help": "Learning rate decay function. choices=['constant', 'linear', 'cosine']. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    lr_decay_iters: int = field(
+        default=None,
+        metadata={
+            "help": "Number of iterations for learning rate decay. If None defaults to `train_iters`. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    lr_decay_samples: int = field(
+        default=None,
+        metadata={
+            "help": "Number of samples for learning rate decay. If None defaults to `train_samples`. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    lr_warmup_iters: int = field(
+        default=None,
+        metadata={
+            "help": "number of iterations to linearly warmup learning rate over. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    lr_warmup_samples: int = field(
+        default=None,
+        metadata={
+            "help": "number of samples to linearly warmup learning rate over. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    lr_warmup_fraction: float = field(
+        default=None,
+        metadata={
+            "help": "fraction of lr-warmup-(iters/samples) to linearly warmup learning rate over. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    min_lr: float = field(
+        default=None,
+        metadata={
+            "help": "Minumum value for learning rate. The scheduler clip values below this threshold. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    consumed_samples: List[int] = field(
+        default=None,
+        metadata={
+            "help": "Number of samples consumed in the same order as the dataloaders to `accelerator.prepare` call."
+        },
+    )
+    no_wd_decay_cond: Optional[Callable] = field(default=None, metadata={"help": "Condition to disable weight decay."})
+    scale_lr_cond: Optional[Callable] = field(default=None, metadata={"help": "Condition to scale learning rate."})
+    lr_mult: float = field(default=1.0, metadata={"help": "Learning rate multiplier."})
+    megatron_dataset_flag: bool = field(
+        default=False,
+        metadata={"help": "Whether the format of dataset follows Megatron-LM Indexed/Cached/MemoryMapped format."},
+    )
+
+    def __post_init__(self):
+        prefix = "MEGATRON_"
+        if self.tp_degree is None:
+            self.tp_degree = int(os.environ.get(prefix + "TP_DEGREE", 1))
+        if self.pp_degree is None:
+            self.pp_degree = int(os.environ.get(prefix + "PP_DEGREE", 1))
+        if self.num_micro_batches is None:
+            self.num_micro_batches = int(os.environ.get(prefix + "NUM_MICRO_BATCHES", 1))
+        if self.gradient_clipping is None:
+            self.gradient_clipping = float(os.environ.get(prefix + "GRADIENT_CLIPPING", 1.0))
+        if self.recompute_activation is None:
+            self.recompute_activation = strtobool(os.environ.get(prefix + "RECOMPUTE_ACTIVATION", "False")) == 1
+        if self.use_distributed_optimizer is None:
+            self.use_distributed_optimizer = (
+                strtobool(os.environ.get(prefix + "USE_DISTRIBUTED_OPTIMIZER", "False")) == 1
+            )
+        if self.sequence_parallelism is None:
+            self.sequence_parallelism = strtobool(os.environ.get(prefix + "SEQUENCE_PARALLELISM", "False")) == 1
+
+        if self.pp_degree > 1:
+            self.DDP_impl = "local"
+        else:
+            self.DDP_impl = "torch"
+
+        if self.consumed_samples is not None:
+            if len(self.consumed_samples) == 1:
+                self.consumed_samples.extend([0, 0])
+            elif len(self.consumed_samples) == 2:
+                self.consumed_samples.append(0)
+
+        self.megatron_lm_default_args = {
+            "tensor_model_parallel_size": self.tp_degree,
+            "pipeline_model_parallel_size": self.pp_degree,
+            "pipeline_model_parallel_split_rank": self.pipeline_model_parallel_split_rank,
+            "num_layers_per_virtual_pipeline_stage": self.num_layers_per_virtual_pipeline_stage,
+            "DDP_impl": self.DDP_impl,
+            "use_distributed_optimizer": self.use_distributed_optimizer,
+            "recompute_activations": self.recompute_activation,
+            "sequence_parallel": self.sequence_parallelism,
+            "clip_grad": self.gradient_clipping,
+            "num_micro_batches": self.num_micro_batches,
+            "consumed_samples": self.consumed_samples,
+            "no_wd_decay_cond": self.no_wd_decay_cond,
+            "scale_lr_cond": self.scale_lr_cond,
+            "lr_mult": self.lr_mult,
+            "megatron_dataset_flag": self.megatron_dataset_flag,
+        }
+
+    def set_network_size_args(self, model):
+        # Check if the model is either BERT, GPT or T5 else raise error
+        # set 'num_layers', 'hidden_size', 'num_attention_heads', 'max_position_embeddings'
+        if "bert" in model.__class__.__name__.lower():
To be cleaned up?
@@ -0,0 +1,821 @@
+# 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
+from abc import ABC
+from functools import partial
+
+import torch
+import torch.nn.functional as F
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from ..optimizer import AcceleratedOptimizer
+from ..scheduler import AcceleratedScheduler
+from .imports import is_megatron_lm_available, is_transformers_available
+from .operations import recursively_apply
+
+
+if is_transformers_available():
+    from transformers.modeling_outputs import (
+        CausalLMOutputWithCrossAttentions,
+        Seq2SeqLMOutput,
+        SequenceClassifierOutput,
+    )
+    from transformers.models.bert.modeling_bert import BertForPreTrainingOutput
+
+
+if is_megatron_lm_available():
+    from megatron import get_args, get_num_microbatches, get_timers, mpu
+    from megatron.arguments import _add_data_args
+    from megatron.data.data_samplers import build_pretraining_data_loader
+    from megatron.data.dataset_utils import build_train_valid_test_datasets
+    from megatron.initialize import (
+        get_tensor_model_parallel_group,
+        get_tensor_model_parallel_src_rank,
+        initialize_megatron,
+        set_jit_fusion_options,
+    )
+    from megatron.model import BertModel, GPTModel, ModelType, T5Model
+    from megatron.model.classification import Classification
+    from megatron.model.module import MegatronModule
+    from megatron.optimizer import get_megatron_optimizer
+    from megatron.schedules import get_forward_backward_func
+    from megatron.tokenizer.tokenizer import _vocab_size_with_padding
+    from megatron.training import build_train_valid_test_data_iterators, get_model, get_optimizer_param_scheduler
+    from megatron.utils import average_losses_across_data_parallel_group, get_ltor_masks_and_position_ids
+
+
+def model_provider_func(accelerator, pre_process=True, post_process=True, add_encoder=True, add_decoder=True):
+    """Build the model."""
+    args = get_args()
+    mode = "pre-training" if args.pretraining_flag else "fine-tuning"
+    accelerator.print(f"Building {args.model_type_name} model in the {mode} mode.")
+    if args.model_type_name == "bert":
+        if args.pretraining_flag:
+            num_tokentypes = 2 if args.bert_binary_head else 0
+            model = BertModel(
+                num_tokentypes=num_tokentypes,
+                add_binary_head=args.bert_binary_head,
+                parallel_output=True,
+                pre_process=pre_process,
+                post_process=post_process,
+            )
+        else:
+
+            model = Classification(
+                num_classes=args.num_labels, num_tokentypes=2, pre_process=pre_process, post_process=post_process
+            )
+    elif args.model_type_name == "gpt2":
+        model = GPTModel(num_tokentypes=0, parallel_output=True, pre_process=pre_process, post_process=post_process)
+    elif args.model_type_name == "t5":
+        model = T5Model(
+            num_tokentypes=0,
+            parallel_output=True,
+            pre_process=pre_process,
+            post_process=post_process,
+            add_encoder=add_encoder,
+            add_decoder=add_decoder,
+        )
+    else:
+        raise ValueError(f"Unknown model type: {args.model_type_name}")
+    return model
+
+
+def prepare_data_loader(accelerator, dataloader, consumed_samples_index=-1, consumed_samples=0):
+    accelerator.print("Preparing dataloader")
+    args = get_args()
+    if not args.megatron_dataset_flag:
+        if args.consumed_samples is not None:
+            consumed_samples = args.consumed_samples[consumed_samples_index]
+        return build_pretraining_data_loader(dataloader.dataset, consumed_samples)
+    else:
+        if args.consumed_samples is not None:
+            (
+                args.consumed_train_samples,
+                args.consumed_valid_samples,
+                args.consumed_test_samples,
+            ) = args.consumed_samples
+        else:
+            args.consumed_train_samples, args.consumed_valid_samples, args.consumed_test_samples = 0, 0, 0
+        train_data_iterator, valid_data_iterator, test_data_iterator = build_train_valid_test_data_iterators(
+            MegatronLMDummyDataLoader.get_train_valid_test_datasets_provider
+        )
+        return train_data_iterator, valid_data_iterator, test_data_iterator
+
+
+def prepare_model(accelerator):
+    accelerator.print("Preparing model")
+    args = get_args()
+    if args.model_type_name == "bert" or args.model_type_name == "gpt":
+        model_type = ModelType.encoder_or_decoder
+    elif args.model_type_name == "t5":
+        model_type = ModelType.encoder_and_decoder
+    model = get_model(model_provider_func, model_type)
+    return model
+
+
+def prepare_optimizer(accelerator, model):
+    accelerator.print("Preparing optimizer")
+    args = get_args()
+    optimizer = get_megatron_optimizer(model, args.no_wd_decay_cond, args.scale_lr_cond, args.lr_mult)
+    return optimizer
+
+
+def prepare_scheduler(accelerator, optimizer, scheduler, is_dummy_scheduler):
+    accelerator.print("Preparing scheduler")
+    if is_dummy_scheduler:
+        scheduler = get_optimizer_param_scheduler(optimizer)
+    else:
+        scheduler.optimizer = optimizer
+        if isinstance(scheduler, torch.optim.lr_scheduler.LambdaLR):
+            scheduler = scheduler.__class__(optimizer, scheduler.lr_lambdas[0])
+    return scheduler
+
+
+def initialize(accelerator, extra_args_provider=None, args_defaults={}):
+    accelerator.print("Initializing Megatron-LM")
+    # Initalize and get arguments
+    initialize_megatron(extra_args_provider=extra_args_provider, args_defaults=args_defaults, ignore_unknown_args=True)
+    # Set pytorch JIT layer fusion options and warmup JIT functions.
+    set_jit_fusion_options()
+    args = get_args()
+    args.padded_vocab_size = _vocab_size_with_padding(args.orig_vocab_size, args)
+    if args.model_type_name == "bert" and args.pretraining_flag and args.num_labels == 2:
+        args.bert_binary_head = True
+    # if args.virtual_pipeline_model_parallel_size is not None:
+    #     raise Exception("Virtual pipeline model parallelism is not supported in Accelerate")
The aliases should be done in the `.utils` module: if we have another integration like this, we won't be able to have both names.
@@ -71,6 +73,21 @@
         DummyScheduler,
     )
 
+if is_megatron_lm_available():
+    from .utils import (
+        MegatronEngine,
+        MegatronLMDummyDataLoader,
+        MegatronLMDummyScheduler,
+        MegatronLMOptimizerWrapper,
+        MegatronLMSchedulerWrapper,
+    )
+    from .utils import initialize as megatron_lm_initialize
+    from .utils import prepare_data_loader as megatron_lm_prepare_data_loader
+    from .utils import prepare_model as megatron_lm_prepare_model
+    from .utils import prepare_optimizer as megatron_lm_prepare_optimizer
+    from .utils import prepare_scheduler as megatron_lm_prepare_scheduler
AS said above, those should probably have the `megatron_lm` prefix here.
@@ -94,6 +96,23 @@
     )
 
 from .launch import PrepareForLaunch, _filter_args, get_launch_prefix
+from .megatron_lm import (
+    AbstractTrainStep,
+    BertTrainStep,
+    GPTTrainStep,
+    MegatronEngine,
+    MegatronLMDummyDataLoader,
+    MegatronLMDummyScheduler,
+    MegatronLMOptimizerWrapper,
+    MegatronLMSchedulerWrapper,
+    T5TrainStep,
+    avg_losses_across_data_parallel_group,
+    initialize,
+    prepare_data_loader,
+    prepare_model,
+    prepare_optimizer,
+    prepare_scheduler,
Nit: in a followup PR it would be nice to have all of those in one constant.
@@ -224,6 +224,7 @@ def gather(tensor):
         DistributedType.DEEPSPEED,
         DistributedType.MULTI_GPU,
         DistributedType.FSDP,
+        DistributedType.MEGATRON_LM,
Done.
@@ -71,6 +73,21 @@
         DummyScheduler,
     )
 
+if is_megatron_lm_available():
+    from .utils import (
+        MegatronEngine,
+        MegatronLMDummyDataLoader,
+        MegatronLMDummyScheduler,
+        MegatronLMOptimizerWrapper,
+        MegatronLMSchedulerWrapper,
+    )
+    from .utils import initialize as megatron_lm_initialize
+    from .utils import prepare_data_loader as megatron_lm_prepare_data_loader
+    from .utils import prepare_model as megatron_lm_prepare_model
+    from .utils import prepare_optimizer as megatron_lm_prepare_optimizer
+    from .utils import prepare_scheduler as megatron_lm_prepare_scheduler
Done.
@@ -94,6 +96,23 @@
     )
 
 from .launch import PrepareForLaunch, _filter_args, get_launch_prefix
+from .megatron_lm import (
+    AbstractTrainStep,
+    BertTrainStep,
+    GPTTrainStep,
+    MegatronEngine,
+    MegatronLMDummyDataLoader,
+    MegatronLMDummyScheduler,
+    MegatronLMOptimizerWrapper,
+    MegatronLMSchedulerWrapper,
+    T5TrainStep,
+    avg_losses_across_data_parallel_group,
+    initialize,
+    prepare_data_loader,
+    prepare_model,
+    prepare_optimizer,
+    prepare_scheduler,
Can be for another PR, but would be nice to explain the difference between those kinds of parallelism and what is parallelized exactly.
@@ -0,0 +1,507 @@
+<!--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.
+-->
+
+
+# Megatron-LM
+
+[Megatron-LM](https://github.com/NVIDIA/Megatron-LM) enables training large transformer language models at scale.
+It provides efficient tensor, pipeline and sequence based model parallelism for pre-training transformer based
+Language Models such as [GPT](https://arxiv.org/abs/2005.14165) (Decoder Only), [BERT](https://arxiv.org/pdf/1810.04805.pdf) (Encoder Only) and [T5](https://arxiv.org/abs/1910.10683) (Encoder-Decoder).
+For detailed information and how things work behind the scene please refer the github [repo](https://github.com/NVIDIA/Megatron-LM).
+
+## What is integrated?
+
+Accelerate integrates following feature of Megatron-LM to enable large scale pre-training/finetuning
+of BERT (Encoder), GPT (Decoder) or T5 models (Encoder and Decoder):
+
+a. **Tensor Parallelism (TP)**: Reduces memory footprint without much additional communication on intra-node ranks
+
+b. **Pipeline Parallelism (PP)**: Reduces memory footprint and enables large scale training via inter-node parallelization. 
+Reduces the bubble of naive PP via PipeDream-Flush schedule/1F1B schedule and Interleaved 1F1B schedule. 
+Layers are distributed uniformly across PP stages.
+
+c. **Sequence Parallelism (SP)**: Reduces memory footprint without any additional communication. Only applicable when using TP.
+
+d. **Data Parallelism (DP)** via Distributed Optimizer: Reduces memory footprint by sharding optimizer states and gradients across DP ranks.
PR #758 to resolve this
@@ -224,6 +224,7 @@ def gather(tensor):
         DistributedType.DEEPSPEED,
         DistributedType.MULTI_GPU,
         DistributedType.FSDP,
+        DistributedType.MEGATRON_LM,
PR #759 to address this.
@@ -0,0 +1,507 @@
+<!--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.
+-->
+
+
+# Megatron-LM
+
+[Megatron-LM](https://github.com/NVIDIA/Megatron-LM) enables training large transformer language models at scale.
+It provides efficient tensor, pipeline and sequence based model parallelism for pre-training transformer based
+Language Models such as [GPT](https://arxiv.org/abs/2005.14165) (Decoder Only), [BERT](https://arxiv.org/pdf/1810.04805.pdf) (Encoder Only) and [T5](https://arxiv.org/abs/1910.10683) (Encoder-Decoder).
+For detailed information and how things work behind the scene please refer the github [repo](https://github.com/NVIDIA/Megatron-LM).
+
+## What is integrated?
+
+Accelerate integrates following feature of Megatron-LM to enable large scale pre-training/finetuning
+of BERT (Encoder), GPT (Decoder) or T5 models (Encoder and Decoder):
+
+a. **Tensor Parallelism (TP)**: Reduces memory footprint without much additional communication on intra-node ranks
+
+b. **Pipeline Parallelism (PP)**: Reduces memory footprint and enables large scale training via inter-node parallelization. 
+Reduces the bubble of naive PP via PipeDream-Flush schedule/1F1B schedule and Interleaved 1F1B schedule. 
+Layers are distributed uniformly across PP stages.
+
+c. **Sequence Parallelism (SP)**: Reduces memory footprint without any additional communication. Only applicable when using TP.
+
+d. **Data Parallelism (DP)** via Distributed Optimizer: Reduces memory footprint by sharding optimizer states and gradients across DP ranks.
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/667
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 4991ddbe4..353ccf894 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -140,7 +140,7 @@ class Accelerator:
             A list of `KwargHandler` to customize how the objects related to distributed training or mixed precision
             are created. See [kwargs](kwargs) for more information.
 
-    **Attributes:**
+    **Available attributes:**
 
         - **device** (`torch.device`) -- The device to use.
         - **distributed_type** ([`~utils.DistributedType`]) -- The distributed training configuration.
@@ -463,6 +463,28 @@ def no_sync(self, model):
         Args:
             model (`torch.nn.Module`):
                 PyTorch Module that was prepared with `Accelerator.prepare`
+
+        Example:
+
+        ```python
+        >>> from accelerate import Accelerator
+
+        >>> accelerator = Accelerator()
+        >>> dataloader, model, optimizer = accelerator.prepare(dataloader, model, optimizer)
+        >>> input_a = next(iter(dataloader))
+        >>> input_b = next(iter(dataloader))
+
+        >>> with accelerator.no_sync():
+        ...     outputs = model(input_a)
+        ...     loss = loss_func(outputs)
+        ...     accelerator.backward(loss)
+        ...     # No synchronization across processes, only accumulate gradients
+        >>> outputs = model(input_b)
+        >>> accelerator.backward(loss)
+        >>> # Synchronization across all processes
+        >>> optimizer.step()
+        >>> optimizer.zero_grad()
+        ```
         """
         context = contextlib.nullcontext
         if self.use_distributed:
@@ -492,6 +514,24 @@ def accumulate(self, model):
         Args:
             model (`torch.nn.Module`):
                 PyTorch Module that was prepared with `Accelerator.prepare`
+
+        Example:
+
+        ```python
+        >>> from accelerate import Accelerator
+
+        >>> accelerator = Accelerator(gradient_accumulation_steps=2)
+        >>> dataloader, model, optimizer, scheduler = accelerator.prepare(dataloader, model, optimizer, scheduler)
+
+        >>> with accelerator.accumulate():
+        ...     for input, output in dataloader:
+        ...         outputs = model(input)
+        ...         loss = loss_func(outputs)
+        ...         loss.backward()
+        ...         optimizer.step()
+        ...         scheduler.step()
+        ...         optimizer.zero_grad()
+        ```
         """
         self._do_sync()
         if self.sync_gradients:
@@ -873,7 +913,10 @@ def prepare_scheduler(self, scheduler):
 
     def backward(self, loss, **kwargs):
         """
-        Use `accelerator.backward(loss)` in lieu of `loss.backward()`.
+        Scales the gradients in accordance to `Accelerator.gradient_accumulation_steps` and calls the correct
+        `backward()` based on the configuration.
+
+        Should be used in lieu of `loss.backward()`.
         """
         loss /= self.gradient_accumulation_steps
         if self.distributed_type == DistributedType.DEEPSPEED:
@@ -906,6 +949,24 @@ def unscale_gradients(self, optimizer=None):
     def clip_grad_norm_(self, parameters, max_norm, norm_type=2):
         """
         Should be used in place of `torch.nn.utils.clip_grad_norm_`.
+
+        Example:
+
+        ```python
+        >>> from accelerate import Accelerator
+
+        >>> accelerator = Accelerator(gradient_accumulation_steps=2)
+        >>> dataloader, model, optimizer, scheduler = accelerator.prepare(dataloader, model, optimizer, scheduler)
+
+        >>> for (input, target) in dataloader:
+        ...     optimizer.zero_grad()
+        ...     output = model(input)
+        ...     loss = loss_func(output, target)
+        ...     accelerator.backward(loss)
+        ...     if accelerator.sync_gradients:
+        ...         accelerator.clip_grad_norm_(model.parameters(), max_grad_norm)
+        ...     optimizer.step()
+        ```
         """
         if self.distributed_type == DistributedType.FSDP:
             self.unscale_gradients()
@@ -923,6 +984,24 @@ def clip_grad_norm_(self, parameters, max_norm, norm_type=2):
     def clip_grad_value_(self, parameters, clip_value):
         """
         Should be used in place of `torch.nn.utils.clip_grad_value_`.
+
+        Example:
+
+        ```python
+        >>> from accelerate import Accelerator
+
+        >>> accelerator = Accelerator(gradient_accumulation_steps=2)
+        >>> dataloader, model, optimizer, scheduler = accelerator.prepare(dataloader, model, optimizer, scheduler)
+
+        >>> for (input, target) in dataloader:
+        ...     optimizer.zero_grad()
+        ...     output = model(input)
+        ...     loss = loss_func(output, target)
+        ...     accelerator.backward(loss)
+        ...     if accelerator.sync_gradients:
+        ...         accelerator.clip_grad_value_(model.parameters(), clip_value)
+        ...     optimizer.step()
+        ```
         """
         if self.distributed_type in [DistributedType.DEEPSPEED, DistributedType.FSDP]:
             raise Exception("DeepSpeed and FSDP  do not support `clip_grad_value_`. Use `clip_grad_norm_` instead.")
@@ -1124,6 +1203,13 @@ def save_state(self, output_dir: str):
         """
         Saves the current states of the model, optimizer, scaler, RNG generators, and registered objects.
 
+        <Tip>
+
+        Should only be used when wanting to save a checkpoint during training and restoring the state in the same
+        environment.
+
+        </Tip>
+
         Args:
             output_dir (`str` or `os.PathLike`):
                 The name of the folder to save all relevant weights and states.
@@ -1179,6 +1265,12 @@ def load_state(self, input_dir: str):
         """
         Loads the current states of the model, optimizer, scaler, RNG generators, and registered objects.
 
+        <Tip>
+
+        Should only be used in conjunction with [`Accelerator.save_state`].
+
+        </Tip>
+
         Args:
             input_dir (`str` or `os.PathLike`):
                 The name of the folder all relevant weights and states were saved in.
diff --git a/src/accelerate/logging.py b/src/accelerate/logging.py
index 2128cdca7..2c67e24eb 100644
--- a/src/accelerate/logging.py
+++ b/src/accelerate/logging.py
@@ -49,15 +49,20 @@ def get_logger(name: str):
 
     If a log should be called on all processes, pass `main_process_only=False`
 
-    E.g.
-    ```python
-    logger.info("My log", main_process_only=False)
-    logger.debug("My log", main_process_only=False)
-    ```
-
     Args:
         name (`str`):
             The name for the logger, such as `__file__`
+
+    Example:
+
+    ```python
+    >>> from accelerate.logging import get_logger
+
+    >>> logger = get_logger(__name__)
+
+    >>> logger.info("My log", main_process_only=False)
+    >>> logger.debug("My log", main_process_only=True)
+    ```
     """
     logger = logging.getLogger(name)
     return MultiProcessAdapter(logger, {})
diff --git a/src/accelerate/scheduler.py b/src/accelerate/scheduler.py
index 75695091e..4ae91a657 100644
--- a/src/accelerate/scheduler.py
+++ b/src/accelerate/scheduler.py
@@ -28,7 +28,7 @@ class AcceleratedScheduler:
     to avoid making a scheduler step too fast when gradients went overflow and there was no training step (in mixed
     precision training)
 
-    When performing gradient accumulation scheduler lengths should not be changed accordingly, accelerate will always
+    When performing gradient accumulation scheduler lengths should not be changed accordingly, Accelerate will always
     step the scheduler to account for it.
 
     Args:
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index 051cf8bc6..a700cfad6 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -50,14 +50,14 @@ class AcceleratorState:
     """
     Singleton class that has information about the current training environment.
 
-    **Attributes:**
+    **Available attributes:**
 
         - **device** (`torch.device`) -- The device to use.
         - **distributed_type** ([`~accelerate.state.DistributedType`]) -- The type of distributed environment currently
           in use.
         - **local_process_index** (`int`) -- The index of the current process on the current server.
-        - **mixed_precision** (`str`) -- Whether or not the current script will use mixed precision. If you are using
-          mixed precision, define if you want to use FP16 or BF16 (bfloat16) as the floating point.
+        - **mixed_precision** (`str`) -- Whether or not the current script will use mixed precision, and if so the type
+          of mixed precision being performed.
         - **num_processes** (`int`) -- The number of processes currently launched in parallel.
         - **process_index** (`int`) -- The index of the current process.
     """
@@ -278,10 +278,11 @@ class GradientState:
     """
     Singleton class that has information related to gradient synchronization for gradient accumulation
 
-    **Attributes:**
+    **Available attributes:**
 
         - **end_of_dataloader** (`bool`) -- Whether we have reached the end the current dataloader
         - **remainder** (`int`) -- The number of extra samples that were added from padding the dataloader
+        - **sync_gradients** (`bool`) -- Whether the gradients should be synced across all devices
     """
 
     _shared_state = {}
@@ -310,5 +311,5 @@ def _set_end_of_dataloader(self, end_of_dataloader):
         self.end_of_dataloader = end_of_dataloader
 
     def _set_remainder(self, remainder):
-        "Private function that sets the number of remaining samples at the end of the dataloader"
+        "Private function that sets the number of remaining samples at the end of the dataloader. Users should not have to call this."
         self.remainder = remainder

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/666
diff --git a/README.md b/README.md
index f15f442a7..deddf2986 100644
--- a/README.md
+++ b/README.md
@@ -243,3 +243,16 @@ pip install accelerate
 - FP16 with native AMP (apex on the roadmap)
 - DeepSpeed support (Experimental)
 - PyTorch Fully Sharded Data Parallel (FSDP) support (Experimental)
+
+## Citing 🤗 Accelerate
+
+If you use 🤗 Accelerate in your publication, please cite it by using the following BibTeX entry.
+
+```bibtex
+@Misc{accelerate,
+  title =        {Accelerate: Training and inference at scale made simple, efficient and adaptable.},
+  author =       {Sylvain Gugger, Lysandre Debut, Thomas Wolf, Philipp Schmid, Zachary Mueller, Sourab Mangrulkar},
+  howpublished = {\url{https://github.com/huggingface/accelerate}},
+  year =         {2022}
+}
+```

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/660
diff --git a/src/accelerate/tracking.py b/src/accelerate/tracking.py
index 839a7a2a9..028e431e1 100644
--- a/src/accelerate/tracking.py
+++ b/src/accelerate/tracking.py
@@ -16,9 +16,12 @@
 # Provide a project dir name, then each type of logger gets stored in project/{`logging_dir`}
 
 import os
+import time
 from abc import ABCMeta, abstractmethod, abstractproperty
 from typing import List, Optional, Union
 
+import yaml
+
 from .logging import get_logger
 from .utils import LoggerType, is_comet_ml_available, is_tensorboard_available, is_wandb_available
 
@@ -142,7 +145,8 @@ def tracker(self):
 
     def store_init_configuration(self, values: dict):
         """
-        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment.
+        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment. Stores the
+        hyperparameters in a yaml file for future use.
 
         Args:
             values (Dictionary `str` to `bool`, `str`, `float` or `int`):
@@ -151,7 +155,16 @@ def store_init_configuration(self, values: dict):
         """
         self.writer.add_hparams(values, metric_dict={})
         self.writer.flush()
-        logger.info("Stored initial configuration hyperparameters to TensorBoard")
+        project_run_name = time.time()
+        dir_name = os.path.join(self.logging_dir, str(project_run_name))
+        os.makedirs(dir_name, exist_ok=True)
+        with open(os.path.join(dir_name, "hparams.yml"), "w") as outfile:
+            try:
+                yaml.dump(values, outfile)
+            except yaml.representer.RepresenterError:
+                logger.error("Serialization to store hyperparameters failed")
+                raise
+        logger.info("Stored initial configuration hyperparameters to TensorBoard and hparams yaml file")
 
     def log(self, values: dict, step: Optional[int] = None, **kwargs):
         """

Perhaps it may be good to make use of the `logging_dir` to send the yaml to? Or should these be stored separately normally.

(This part):
```python
writer =  tensorboard.SummaryWriter(self.logging_dir, **kwargs)
```

@@ -151,7 +153,9 @@ def store_init_configuration(self, values: dict):
         """
         self.writer.add_hparams(values, metric_dict={})
         self.writer.flush()
-        logger.info("Stored initial configuration hyperparameters to TensorBoard")
+        with open("hparams.yaml", "w") as outfile:
So, should I save the yaml file in the `logging_dir`?
@@ -151,7 +153,9 @@ def store_init_configuration(self, values: dict):
         """
         self.writer.add_hparams(values, metric_dict={})
         self.writer.flush()
-        logger.info("Stored initial configuration hyperparameters to TensorBoard")
+        with open("hparams.yaml", "w") as outfile:
Yep, as this is where users would expect it to wind up landing. E.g

logging_dir/hparams.yaml
@@ -151,7 +153,9 @@ def store_init_configuration(self, values: dict):
         """
         self.writer.add_hparams(values, metric_dict={})
         self.writer.flush()
-        logger.info("Stored initial configuration hyperparameters to TensorBoard")
+        with open("hparams.yaml", "w") as outfile:
Yeah makes sense. I have made the changes now.
@@ -151,7 +153,9 @@ def store_init_configuration(self, values: dict):
         """
         self.writer.add_hparams(values, metric_dict={})
         self.writer.flush()
-        logger.info("Stored initial configuration hyperparameters to TensorBoard")
+        with open("hparams.yaml", "w") as outfile:
Would we want to use pathlib `Path(self.logging_dir) / 'hparams.yaml'` or `os.path.join(self.logging_dir, 'hparams.yaml')` instead here for sanity purposes?
@@ -151,7 +153,9 @@ def store_init_configuration(self, values: dict):
         """
         self.writer.add_hparams(values, metric_dict={})
         self.writer.flush()
-        logger.info("Stored initial configuration hyperparameters to TensorBoard")
+        with open("self.logging_dir/hparams.yaml", "w") as outfile:
Do we need to check if all keys/vals are serializable here, or are we going to let the error yaml throws be the one we present to users in this case? (not sure if this is done elsewhere in the codebase)
@@ -151,7 +153,9 @@ def store_init_configuration(self, values: dict):
         """
         self.writer.add_hparams(values, metric_dict={})
         self.writer.flush()
-        logger.info("Stored initial configuration hyperparameters to TensorBoard")
+        with open("self.logging_dir/hparams.yaml", "w") as outfile:
+            yaml.dump(values, outfile)
We do that somewhat here: https://github.com/huggingface/accelerate/blob/main/src/accelerate/tracking.py#L170-L176

Maybe a try/catch that alerts the user that non-serializable bits were tried to be stored? Not sure what the try/catch error would be for that, otherwise I'm open to it just throwing the yaml error if we cannot :) 
@@ -151,7 +153,9 @@ def store_init_configuration(self, values: dict):
         """
         self.writer.add_hparams(values, metric_dict={})
         self.writer.flush()
-        logger.info("Stored initial configuration hyperparameters to TensorBoard")
+        with open("self.logging_dir/hparams.yaml", "w") as outfile:
+            yaml.dump(values, outfile)
could catch a `yaml.representer.RepresenterError` and return the message maybe?
@@ -151,7 +153,9 @@ def store_init_configuration(self, values: dict):
         """
         self.writer.add_hparams(values, metric_dict={})
         self.writer.flush()
-        logger.info("Stored initial configuration hyperparameters to TensorBoard")
+        with open("self.logging_dir/hparams.yaml", "w") as outfile:
+            yaml.dump(values, outfile)
Along with more descriptive explanation of what happened under the hood accelerate...otherwise its the same as letting yaml handle it. 

I'm fine with either solution, just wanted to bring it up!
@@ -151,7 +153,9 @@ def store_init_configuration(self, values: dict):
         """
         self.writer.add_hparams(values, metric_dict={})
         self.writer.flush()
-        logger.info("Stored initial configuration hyperparameters to TensorBoard")
+        with open("self.logging_dir/hparams.yaml", "w") as outfile:
+            yaml.dump(values, outfile)
We should be careful here that `hparams.yaml` is saved separately for separate experiments. Should make sure that when you run 2 separate runs, you get 2 separate entries in the hparams dashboard of tensorboard. 

If we save directly to `logging_dir/hparams.yaml` we would be overwriting every time, eh? 
@@ -151,7 +153,9 @@ def store_init_configuration(self, values: dict):
         """
         self.writer.add_hparams(values, metric_dict={})
         self.writer.flush()
-        logger.info("Stored initial configuration hyperparameters to TensorBoard")
+        with open("hparams.yaml", "w") as outfile:
I have changed the hardcoded path here.
@@ -151,7 +153,9 @@ def store_init_configuration(self, values: dict):
         """
         self.writer.add_hparams(values, metric_dict={})
         self.writer.flush()
-        logger.info("Stored initial configuration hyperparameters to TensorBoard")
+        with open("self.logging_dir/hparams.yaml", "w") as outfile:
I have added the try/catch error. 
@@ -151,7 +153,9 @@ def store_init_configuration(self, values: dict):
         """
         self.writer.add_hparams(values, metric_dict={})
         self.writer.flush()
-        logger.info("Stored initial configuration hyperparameters to TensorBoard")
+        with open("self.logging_dir/hparams.yaml", "w") as outfile:
+            yaml.dump(values, outfile)
I am not sure where to give a descriptive explanation. Should I add it as a comment or should I define it in the PR?
@@ -151,7 +153,9 @@ def store_init_configuration(self, values: dict):
         """
         self.writer.add_hparams(values, metric_dict={})
         self.writer.flush()
-        logger.info("Stored initial configuration hyperparameters to TensorBoard")
+        with open("self.logging_dir/hparams.yaml", "w") as outfile:
+            yaml.dump(values, outfile)
The default tensorboard summary writer uses time.time() to create a folder for every run. So, I added the same for hparams.yml file. Hope this is okay.
@@ -151,7 +153,9 @@ def store_init_configuration(self, values: dict):
         """
         self.writer.add_hparams(values, metric_dict={})
         self.writer.flush()
-        logger.info("Stored initial configuration hyperparameters to TensorBoard")
+        with open("self.logging_dir/hparams.yaml", "w") as outfile:
We should change this to actually raise the error, otherwise this will silently fail.

```suggestion
            except yaml.representer.RepresenterError:
                logger.error("Serialization to store hyperparameters failed")
                raise
```
I think just doing `logger.error` here is fine enough
@@ -151,7 +154,13 @@ def store_init_configuration(self, values: dict):
         """
         self.writer.add_hparams(values, metric_dict={})
         self.writer.flush()
-        logger.info("Stored initial configuration hyperparameters to TensorBoard")
+        project_run_name = time.time()
+        with open(os.path.join(self.logging_dir, project_run_name,  "hparams.yml"), "w") as outfile:
+            try:
+                yaml.dump(values, outfile)
+            except yaml.representer.RepresenterError:
+                logger.info("Serialization to store hyperparmeters failed")
Doing so in the `logger.error` (see my suggestions) is good enough :) 
@@ -151,7 +153,9 @@ def store_init_configuration(self, values: dict):
         """
         self.writer.add_hparams(values, metric_dict={})
         self.writer.flush()
-        logger.info("Stored initial configuration hyperparameters to TensorBoard")
+        with open("self.logging_dir/hparams.yaml", "w") as outfile:
+            yaml.dump(values, outfile)
Yes, I have made the changes.
@@ -151,7 +154,13 @@ def store_init_configuration(self, values: dict):
         """
         self.writer.add_hparams(values, metric_dict={})
         self.writer.flush()
-        logger.info("Stored initial configuration hyperparameters to TensorBoard")
+        project_run_name = time.time()
+        with open(os.path.join(self.logging_dir, project_run_name,  "hparams.yml"), "w") as outfile:
+            try:
+                yaml.dump(values, outfile)
+            except yaml.representer.RepresenterError:
+                logger.info("Serialization to store hyperparmeters failed")
_The documentation is not available anymore as the PR was closed or merged._
cc @nateraw if you have any comments 🤗 
@Shreyz-max cast `run_name` to a string on the yaml part and it should solve the error. E.g. `str(run_name)`

Great work on this, were almost there!
Yeah, sorry about the mistake. I did not notice this at all.
I suppose I need to create the folder.
Great job! CC @nateraw for one last look, but lg2m!
https://api.github.com/repos/huggingface/accelerate/pulls/657
diff --git a/docs/source/concept_guides/gradient_synchronization.mdx b/docs/source/concept_guides/gradient_synchronization.mdx
index 02b5adf0d..ea4de3d72 100644
--- a/docs/source/concept_guides/gradient_synchronization.mdx
+++ b/docs/source/concept_guides/gradient_synchronization.mdx
@@ -114,4 +114,6 @@ for batch in dataloader:
         outputs = model(inputs)
         loss = loss_function(outputs, targets)
         accelerator.backward(loss)
-```
\ No newline at end of file
+```
+
+As a result, you should either use *`accelerator.accumulate` or `accelerator.no_sync`* when it comes to API choice. 
\ No newline at end of file
diff --git a/docs/source/usage_guides/gradient_accumulation.mdx b/docs/source/usage_guides/gradient_accumulation.mdx
index bd06efed9..798e15847 100644
--- a/docs/source/usage_guides/gradient_accumulation.mdx
+++ b/docs/source/usage_guides/gradient_accumulation.mdx
@@ -126,3 +126,5 @@ for batch in training_dataloader:
         scheduler.step()
         optimizer.zero_grad()
 ```
+
+To learn more about what magic this wraps around, read the [Gradient Synchronization concept guide](/concept_guides/gradient_synchronization)
\ No newline at end of file

```suggestion
As a result, you should either use *`accelerator.accumulate` or `accelerator.no_sync`* when it comes to API choice. 
```
@@ -114,4 +114,6 @@ for batch in dataloader:
         outputs = model(inputs)
         loss = loss_function(outputs, targets)
         accelerator.backward(loss)
-```
\ No newline at end of file
+```
+
+As a result, you should either use *one or the other* when it comes to API choice. 
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_656). All of your documentation changes will be reflected on that endpoint.
https://api.github.com/repos/huggingface/accelerate/pulls/656
diff --git a/src/accelerate/tracking.py b/src/accelerate/tracking.py
index 028e431e1..5f189c326 100644
--- a/src/accelerate/tracking.py
+++ b/src/accelerate/tracking.py
@@ -134,8 +134,8 @@ def __init__(self, run_name: str, logging_dir: Optional[Union[str, os.PathLike]]
         self.run_name = run_name
         self.logging_dir = os.path.join(logging_dir, run_name)
         self.writer = tensorboard.SummaryWriter(self.logging_dir, **kwargs)
-        logger.info(f"Initialized TensorBoard project {self.run_name} logging to {self.logging_dir}")
-        logger.info(
+        logger.debug(f"Initialized TensorBoard project {self.run_name} logging to {self.logging_dir}")
+        logger.debug(
             "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
         )
 
@@ -164,7 +164,7 @@ def store_init_configuration(self, values: dict):
             except yaml.representer.RepresenterError:
                 logger.error("Serialization to store hyperparameters failed")
                 raise
-        logger.info("Stored initial configuration hyperparameters to TensorBoard and hparams yaml file")
+        logger.debug("Stored initial configuration hyperparameters to TensorBoard and hparams yaml file")
 
     def log(self, values: dict, step: Optional[int] = None, **kwargs):
         """
@@ -188,14 +188,14 @@ def log(self, values: dict, step: Optional[int] = None, **kwargs):
             elif isinstance(v, dict):
                 self.writer.add_scalars(k, v, global_step=step, **kwargs)
         self.writer.flush()
-        logger.info("Successfully logged to TensorBoard")
+        logger.debug("Successfully logged to TensorBoard")
 
     def finish(self):
         """
         Closes `TensorBoard` writer
         """
         self.writer.close()
-        logger.info("TensorBoard writer closed")
+        logger.debug("TensorBoard writer closed")
 
 
 class WandBTracker(GeneralTracker):
@@ -215,8 +215,8 @@ class WandBTracker(GeneralTracker):
     def __init__(self, run_name: str, **kwargs):
         self.run_name = run_name
         self.run = wandb.init(project=self.run_name, **kwargs)
-        logger.info(f"Initialized WandB project {self.run_name}")
-        logger.info(
+        logger.debug(f"Initialized WandB project {self.run_name}")
+        logger.debug(
             "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
         )
 
@@ -234,7 +234,7 @@ def store_init_configuration(self, values: dict):
                 `str`, `float`, `int`, or `None`.
         """
         wandb.config.update(values)
-        logger.info("Stored initial configuration hyperparameters to WandB")
+        logger.debug("Stored initial configuration hyperparameters to WandB")
 
     def log(self, values: dict, step: Optional[int] = None, **kwargs):
         """
@@ -250,14 +250,14 @@ def log(self, values: dict, step: Optional[int] = None, **kwargs):
                 Additional key word arguments passed along to the `wandb.log` method.
         """
         self.run.log(values, step=step, **kwargs)
-        logger.info("Successfully logged to WandB")
+        logger.debug("Successfully logged to WandB")
 
     def finish(self):
         """
         Closes `wandb` writer
         """
         self.run.finish()
-        logger.info("WandB run closed")
+        logger.debug("WandB run closed")
 
 
 class CometMLTracker(GeneralTracker):
@@ -279,8 +279,8 @@ class CometMLTracker(GeneralTracker):
     def __init__(self, run_name: str, **kwargs):
         self.run_name = run_name
         self.writer = Experiment(project_name=run_name, **kwargs)
-        logger.info(f"Initialized CometML project {self.run_name}")
-        logger.info(
+        logger.debug(f"Initialized CometML project {self.run_name}")
+        logger.debug(
             "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
         )
 
@@ -298,7 +298,7 @@ def store_init_configuration(self, values: dict):
                 `str`, `float`, `int`, or `None`.
         """
         self.writer.log_parameters(values)
-        logger.info("Stored initial configuration hyperparameters to CometML")
+        logger.debug("Stored initial configuration hyperparameters to CometML")
 
     def log(self, values: dict, step: Optional[int] = None, **kwargs):
         """
@@ -323,14 +323,14 @@ def log(self, values: dict, step: Optional[int] = None, **kwargs):
                 self.writer.log_other(k, v, **kwargs)
             elif isinstance(v, dict):
                 self.writer.log_metrics(v, step=step, **kwargs)
-        logger.info("Successfully logged to CometML")
+        logger.debug("Successfully logged to CometML")
 
     def finish(self):
         """
         Closes `comet-ml` writer
         """
         self.writer.end()
-        logger.info("CometML run closed")
+        logger.debug("CometML run closed")
 
 
 LOGGER_TYPE_TO_CLASS = {"tensorboard": TensorBoardTracker, "wandb": WandBTracker, "comet_ml": CometMLTracker}
@@ -384,6 +384,6 @@ def filter_trackers(
                                     )
                             loggers.append(log_type)
                         else:
-                            logger.info(f"Tried adding logger {log_type}, but package is unavailable in the system.")
+                            logger.debug(f"Tried adding logger {log_type}, but package is unavailable in the system.")
 
     return loggers

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/655
diff --git a/src/accelerate/tracking.py b/src/accelerate/tracking.py
index 839a7a2a9..08b988250 100644
--- a/src/accelerate/tracking.py
+++ b/src/accelerate/tracking.py
@@ -20,7 +20,7 @@
 from typing import List, Optional, Union
 
 from .logging import get_logger
-from .utils import LoggerType, is_comet_ml_available, is_tensorboard_available, is_wandb_available
+from .utils import LoggerType, is_aim_available, is_comet_ml_available, is_tensorboard_available, is_wandb_available
 
 
 _available_trackers = []
@@ -40,6 +40,11 @@
 
     _available_trackers.append(LoggerType.COMETML)
 
+if is_aim_available():
+    from aim import Run
+
+    _available_trackers.append(LoggerType.AIM)
+
 
 logger = get_logger(__name__)
 
@@ -320,7 +325,72 @@ def finish(self):
         logger.info("CometML run closed")
 
 
-LOGGER_TYPE_TO_CLASS = {"tensorboard": TensorBoardTracker, "wandb": WandBTracker, "comet_ml": CometMLTracker}
+class AimTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `aim`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run.
+        kwargs:
+            Additional key word arguments passed along to the `Run.__init__` method.
+    """
+
+    name = "aim"
+    requires_logging_directory = True
+
+    def __init__(self, run_name: str, logging_dir: Optional[Union[str, os.PathLike]] = ".", **kwargs):
+        self.run_name = run_name
+        self.writer = Run(repo=logging_dir, **kwargs)
+        self.writer.name = self.run_name
+        logger.debug(f"Initialized Aim project {self.run_name}")
+        logger.debug(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    @property
+    def tracker(self):
+        return self.writer
+
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment.
+
+        Args:
+            values (`dict`):
+                Values to be stored as initial hyperparameters as key-value pairs.
+        """
+        self.writer["hparams"] = values
+
+    def log(self, values: dict, step: Optional[int], **kwargs):
+        """
+        Logs `values` to the current run.
+
+        Args:
+            values (`dict`):
+                Values to be logged as key-value pairs.
+            step (`int`, *optional*):
+                The run step. If included, the log will be affiliated with this step.
+            kwargs:
+                Additional key word arguments passed along to the `Run.track` method.
+        """
+        # Note: replace this with the dictionary support when merged
+        for key, value in values.items():
+            self.writer.track(value, name=key, step=step, **kwargs)
+
+    def finish(self):
+        """
+        Closes `aim` writer
+        """
+        self.writer.close()
+
+
+LOGGER_TYPE_TO_CLASS = {
+    "aim": AimTracker,
+    "comet_ml": CometMLTracker,
+    "tensorboard": TensorBoardTracker,
+    "wandb": WandBTracker,
+}
 
 
 def filter_trackers(
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
index 2b8f566f0..ae5af1c3f 100644
--- a/src/accelerate/utils/__init__.py
+++ b/src/accelerate/utils/__init__.py
@@ -20,6 +20,7 @@
 )
 from .imports import (
     get_ccl_version,
+    is_aim_available,
     is_apex_available,
     is_bf16_available,
     is_boto3_available,
diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py
index 5f425f200..9f6e30bec 100644
--- a/src/accelerate/utils/dataclasses.py
+++ b/src/accelerate/utils/dataclasses.py
@@ -196,6 +196,7 @@ class LoggerType(BaseEnum):
     """
 
     ALL = "all"
+    AIM = "aim"
     TENSORBOARD = "tensorboard"
     WANDB = "wandb"
     COMETML = "comet_ml"
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index 074d02e4a..6239e3deb 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -93,6 +93,10 @@ def is_datasets_available():
     return importlib.util.find_spec("datasets") is not None
 
 
+def is_aim_available():
+    return importlib.util.find_spec("aim") is not None
+
+
 def is_tensorboard_available():
     return importlib.util.find_spec("tensorboard") is not None or importlib.util.find_spec("tensorboardX") is not None
 

@muellerzr the `finalize` method is more for the SDK internal use (it handles storage indexing), which will become a private method in the upcoming releases.
The resources cleanup is done in `close` method, which should be used to properly finalize the run:
```py
self.writer.close()
```
@@ -320,7 +325,74 @@ def finish(self):
         logger.info("CometML run closed")
 
 
-LOGGER_TYPE_TO_CLASS = {"tensorboard": TensorBoardTracker, "wandb": WandBTracker, "comet_ml": CometMLTracker}
+class AimTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `aim`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run.
+        kwargs:
+            Additional key word arguments passed along to the `Run.__init__` method.
+    """
+
+    name = "aim"
+    requires_logging_directory = True
+
+    def __init__(self, run_name: str, logging_dir: Optional[Union[str, os.PathLike]] = ".", **kwargs):
+        self.run_name = run_name
+        self.writer = Run(repo=logging_dir, experiment=run_name, **kwargs)
+        logger.info(f"Initialized Aim project {self.run_name}")
+        logger.info(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    @property
+    def tracker(self):
+        return self.writer
+
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment.
+
+        Args:
+            values (`dict`):
+                Values to be stored as initial hyperparameters as key-value pairs.
+        """
+        self.writer["hparams"].update(values)
+
+    def log(self, values: dict, step: Optional[int], **kwargs):
+        """
+        Logs `values` to the current run.
+
+        Args:
+            values (`dict`):
+                Values to be logged as key-value pairs.
+            step (`int`, *optional*):
+                The run step. If included, the log will be affiliated with this step.
+            kwargs (`dict`, *optional*):
+                Additional key word arguments passed along to the `Run.track` method. Valid keys include `context` and
+                `epoch`.
+        """
+        # Note: replace this with the dictionary support when merged
+        context = kwargs.pop("context", {})
+        epoch = kwargs.pop("epoch", None)
+        for key, value in values.items():
+            self.writer.track(value, key, step=step, epoch=epoch, context=context)
+
+    def finish(self):
+        """
+        Closes `aim` writer
+        """
+        self.writer.finalize()
What about passing the name argument as a keyword argument? I guess it will be more readable and will minimize the risk of forward incompatibilities. Thoughts? 🤔 
@@ -320,7 +325,74 @@ def finish(self):
         logger.info("CometML run closed")
 
 
-LOGGER_TYPE_TO_CLASS = {"tensorboard": TensorBoardTracker, "wandb": WandBTracker, "comet_ml": CometMLTracker}
+class AimTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `aim`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run.
+        kwargs:
+            Additional key word arguments passed along to the `Run.__init__` method.
+    """
+
+    name = "aim"
+    requires_logging_directory = True
+
+    def __init__(self, run_name: str, logging_dir: Optional[Union[str, os.PathLike]] = ".", **kwargs):
+        self.run_name = run_name
+        self.writer = Run(repo=logging_dir, experiment=run_name, **kwargs)
+        logger.info(f"Initialized Aim project {self.run_name}")
+        logger.info(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    @property
+    def tracker(self):
+        return self.writer
+
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment.
+
+        Args:
+            values (`dict`):
+                Values to be stored as initial hyperparameters as key-value pairs.
+        """
+        self.writer["hparams"].update(values)
+
+    def log(self, values: dict, step: Optional[int], **kwargs):
+        """
+        Logs `values` to the current run.
+
+        Args:
+            values (`dict`):
+                Values to be logged as key-value pairs.
+            step (`int`, *optional*):
+                The run step. If included, the log will be affiliated with this step.
+            kwargs (`dict`, *optional*):
+                Additional key word arguments passed along to the `Run.track` method. Valid keys include `context` and
+                `epoch`.
+        """
+        # Note: replace this with the dictionary support when merged
+        context = kwargs.pop("context", {})
+        epoch = kwargs.pop("epoch", None)
+        for key, value in values.items():
+            self.writer.track(value, key, step=step, epoch=epoch, context=context)
The experiment is mainly used to group related runs together (e.g. often used in sweeps, hparams search). The experiment is not a mandatory argument, it defaults to `"default"`.

I would recommend the `init` method to receive the repo path and the run name and pass the rest of arguments as kwargs to the Aim Run:
```py
def __init__(self, run_name: str, logging_dir: Optional[Union[str, os.PathLike]] = ".", **kwargs):
    ...

    # Init the run - specify the logs dir and pass kwargs (including experiment) 
    self.writer = Run(repo=logging_dir, **kwargs)

    # Assign the run name
    self.writer.name = run_name
```


@@ -320,7 +325,74 @@ def finish(self):
         logger.info("CometML run closed")
 
 
-LOGGER_TYPE_TO_CLASS = {"tensorboard": TensorBoardTracker, "wandb": WandBTracker, "comet_ml": CometMLTracker}
+class AimTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `aim`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run.
+        kwargs:
+            Additional key word arguments passed along to the `Run.__init__` method.
+    """
+
+    name = "aim"
+    requires_logging_directory = True
+
+    def __init__(self, run_name: str, logging_dir: Optional[Union[str, os.PathLike]] = ".", **kwargs):
If you're worried about that, we can absolutely do so. Easy to achieve!
@@ -320,7 +325,74 @@ def finish(self):
         logger.info("CometML run closed")
 
 
-LOGGER_TYPE_TO_CLASS = {"tensorboard": TensorBoardTracker, "wandb": WandBTracker, "comet_ml": CometMLTracker}
+class AimTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `aim`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run.
+        kwargs:
+            Additional key word arguments passed along to the `Run.__init__` method.
+    """
+
+    name = "aim"
+    requires_logging_directory = True
+
+    def __init__(self, run_name: str, logging_dir: Optional[Union[str, os.PathLike]] = ".", **kwargs):
+        self.run_name = run_name
+        self.writer = Run(repo=logging_dir, experiment=run_name, **kwargs)
+        logger.info(f"Initialized Aim project {self.run_name}")
+        logger.info(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    @property
+    def tracker(self):
+        return self.writer
+
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment.
+
+        Args:
+            values (`dict`):
+                Values to be stored as initial hyperparameters as key-value pairs.
+        """
+        self.writer["hparams"].update(values)
+
+    def log(self, values: dict, step: Optional[int], **kwargs):
+        """
+        Logs `values` to the current run.
+
+        Args:
+            values (`dict`):
+                Values to be logged as key-value pairs.
+            step (`int`, *optional*):
+                The run step. If included, the log will be affiliated with this step.
+            kwargs (`dict`, *optional*):
+                Additional key word arguments passed along to the `Run.track` method. Valid keys include `context` and
+                `epoch`.
+        """
+        # Note: replace this with the dictionary support when merged
+        context = kwargs.pop("context", {})
+        epoch = kwargs.pop("epoch", None)
+        for key, value in values.items():
+            self.writer.track(value, key, step=step, epoch=epoch, context=context)
Interesting. I've been doing it the opposite way in my experiments. So e.g. if I have 3 trials of X hyperparameter, I'd just name it `run_hparam_x` *not* `run_hyparam_x_1`?
@@ -320,7 +325,74 @@ def finish(self):
         logger.info("CometML run closed")
 
 
-LOGGER_TYPE_TO_CLASS = {"tensorboard": TensorBoardTracker, "wandb": WandBTracker, "comet_ml": CometMLTracker}
+class AimTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `aim`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run.
+        kwargs:
+            Additional key word arguments passed along to the `Run.__init__` method.
+    """
+
+    name = "aim"
+    requires_logging_directory = True
+
+    def __init__(self, run_name: str, logging_dir: Optional[Union[str, os.PathLike]] = ".", **kwargs):
@gorarakelyan I'd recommend updating this example, as this is where I got it from :) Will use close

https://colab.research.google.com/drive/14rIAjpEyklf5fSMiRbyZs6iYG7IVibcI?usp=sharing#scrollTo=qOu7icths6n5
@@ -320,7 +325,74 @@ def finish(self):
         logger.info("CometML run closed")
 
 
-LOGGER_TYPE_TO_CLASS = {"tensorboard": TensorBoardTracker, "wandb": WandBTracker, "comet_ml": CometMLTracker}
+class AimTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `aim`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run.
+        kwargs:
+            Additional key word arguments passed along to the `Run.__init__` method.
+    """
+
+    name = "aim"
+    requires_logging_directory = True
+
+    def __init__(self, run_name: str, logging_dir: Optional[Union[str, os.PathLike]] = ".", **kwargs):
+        self.run_name = run_name
+        self.writer = Run(repo=logging_dir, experiment=run_name, **kwargs)
+        logger.info(f"Initialized Aim project {self.run_name}")
+        logger.info(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    @property
+    def tracker(self):
+        return self.writer
+
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment.
+
+        Args:
+            values (`dict`):
+                Values to be stored as initial hyperparameters as key-value pairs.
+        """
+        self.writer["hparams"].update(values)
+
+    def log(self, values: dict, step: Optional[int], **kwargs):
+        """
+        Logs `values` to the current run.
+
+        Args:
+            values (`dict`):
+                Values to be logged as key-value pairs.
+            step (`int`, *optional*):
+                The run step. If included, the log will be affiliated with this step.
+            kwargs (`dict`, *optional*):
+                Additional key word arguments passed along to the `Run.track` method. Valid keys include `context` and
+                `epoch`.
+        """
+        # Note: replace this with the dictionary support when merged
+        context = kwargs.pop("context", {})
+        epoch = kwargs.pop("epoch", None)
+        for key, value in values.items():
+            self.writer.track(value, key, step=step, epoch=epoch, context=context)
+
+    def finish(self):
+        """
+        Closes `aim` writer
+        """
+        self.writer.finalize()
@muellerzr Yup, makes sense. What about assigning trial hashes (or ids) to run names and the sweep name to experiment? Like this:

| Run name  | Experiment |
| --- | --- |
| trial_1  | run_hparam_x  |
| trial_2  | run_hparam_x  |
| trial_3  | run_hparam_x  |
| --- | --- |
| trial_1  | run_hparam_y  |
| trial_2  | run_hparam_y  |
| trial_3  | run_hparam_y  |
@@ -320,7 +325,74 @@ def finish(self):
         logger.info("CometML run closed")
 
 
-LOGGER_TYPE_TO_CLASS = {"tensorboard": TensorBoardTracker, "wandb": WandBTracker, "comet_ml": CometMLTracker}
+class AimTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `aim`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run.
+        kwargs:
+            Additional key word arguments passed along to the `Run.__init__` method.
+    """
+
+    name = "aim"
+    requires_logging_directory = True
+
+    def __init__(self, run_name: str, logging_dir: Optional[Union[str, os.PathLike]] = ".", **kwargs):
Sounds good! I'll switch it to this as that makes perfect sense. 
@@ -320,7 +325,74 @@ def finish(self):
         logger.info("CometML run closed")
 
 
-LOGGER_TYPE_TO_CLASS = {"tensorboard": TensorBoardTracker, "wandb": WandBTracker, "comet_ml": CometMLTracker}
+class AimTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `aim`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run.
+        kwargs:
+            Additional key word arguments passed along to the `Run.__init__` method.
+    """
+
+    name = "aim"
+    requires_logging_directory = True
+
+    def __init__(self, run_name: str, logging_dir: Optional[Union[str, os.PathLike]] = ".", **kwargs):
@muellerzr The last minor thing - the `run_hash` is a unique auto-generated id of a run. It would be better to store the name as a `run.name`. Like this:
```py
self.writer = Run(repo=logging_dir, **kwargs)
self.writer.name = run_name
```
@@ -320,7 +325,71 @@ def finish(self):
         logger.info("CometML run closed")
 
 
-LOGGER_TYPE_TO_CLASS = {"tensorboard": TensorBoardTracker, "wandb": WandBTracker, "comet_ml": CometMLTracker}
+class AimTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `aim`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run.
+        kwargs:
+            Additional key word arguments passed along to the `Run.__init__` method.
+    """
+
+    name = "aim"
+    requires_logging_directory = True
+
+    def __init__(self, run_name: str, logging_dir: Optional[Union[str, os.PathLike]] = ".", **kwargs):
+        self.run_name = run_name
+        self.writer = Run(run_hash=run_name, repo=logging_dir, **kwargs)
Fixed, thanks!
@@ -320,7 +325,71 @@ def finish(self):
         logger.info("CometML run closed")
 
 
-LOGGER_TYPE_TO_CLASS = {"tensorboard": TensorBoardTracker, "wandb": WandBTracker, "comet_ml": CometMLTracker}
+class AimTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `aim`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run.
+        kwargs:
+            Additional key word arguments passed along to the `Run.__init__` method.
+    """
+
+    name = "aim"
+    requires_logging_directory = True
+
+    def __init__(self, run_name: str, logging_dir: Optional[Union[str, os.PathLike]] = ".", **kwargs):
+        self.run_name = run_name
+        self.writer = Run(run_hash=run_name, repo=logging_dir, **kwargs)
Found one issue during testing. As Aim Run doesn't have an `update` method, let's just use a simple assignment here:
```py
self.writer["hparams"] = values
```
@@ -320,7 +325,72 @@ def finish(self):
         logger.info("CometML run closed")
 
 
-LOGGER_TYPE_TO_CLASS = {"tensorboard": TensorBoardTracker, "wandb": WandBTracker, "comet_ml": CometMLTracker}
+class AimTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `aim`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run.
+        kwargs:
+            Additional key word arguments passed along to the `Run.__init__` method.
+    """
+
+    name = "aim"
+    requires_logging_directory = True
+
+    def __init__(self, run_name: str, logging_dir: Optional[Union[str, os.PathLike]] = ".", **kwargs):
+        self.run_name = run_name
+        self.writer = Run(repo=logging_dir, **kwargs)
+        self.writer.name = self.run_name
+        logger.info(f"Initialized Aim project {self.run_name}")
+        logger.info(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    @property
+    def tracker(self):
+        return self.writer
+
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment.
+
+        Args:
+            values (`dict`):
+                Values to be stored as initial hyperparameters as key-value pairs.
+        """
+        self.writer["hparams"].update(values)
Makes sense, will update with this. Thanks for checking!
@@ -320,7 +325,72 @@ def finish(self):
         logger.info("CometML run closed")
 
 
-LOGGER_TYPE_TO_CLASS = {"tensorboard": TensorBoardTracker, "wandb": WandBTracker, "comet_ml": CometMLTracker}
+class AimTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `aim`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run.
+        kwargs:
+            Additional key word arguments passed along to the `Run.__init__` method.
+    """
+
+    name = "aim"
+    requires_logging_directory = True
+
+    def __init__(self, run_name: str, logging_dir: Optional[Union[str, os.PathLike]] = ".", **kwargs):
+        self.run_name = run_name
+        self.writer = Run(repo=logging_dir, **kwargs)
+        self.writer.name = self.run_name
+        logger.info(f"Initialized Aim project {self.run_name}")
+        logger.info(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    @property
+    def tracker(self):
+        return self.writer
+
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment.
+
+        Args:
+            values (`dict`):
+                Values to be stored as initial hyperparameters as key-value pairs.
+        """
+        self.writer["hparams"].update(values)
Now that #655 was merged, this should be adapted.
@@ -320,7 +325,72 @@ def finish(self):
         logger.info("CometML run closed")
 
 
-LOGGER_TYPE_TO_CLASS = {"tensorboard": TensorBoardTracker, "wandb": WandBTracker, "comet_ml": CometMLTracker}
+class AimTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `aim`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run.
+        kwargs:
+            Additional key word arguments passed along to the `Run.__init__` method.
+    """
+
+    name = "aim"
+    requires_logging_directory = True
+
+    def __init__(self, run_name: str, logging_dir: Optional[Union[str, os.PathLike]] = ".", **kwargs):
+        self.run_name = run_name
+        self.writer = Run(repo=logging_dir, **kwargs)
+        self.writer.name = self.run_name
+        logger.info(f"Initialized Aim project {self.run_name}")
+        logger.info(
_The documentation is not available anymore as the PR was closed or merged._
@muellerzr the tracker looks great. Is there any example training script, which I can run to test out the integration? Thanks a million!
@gorarakelyan you can follow the directions here for installing the git version of accelerate + configuring it for your system:
https://huggingface.co/docs/accelerate/basic_tutorials/install

And then feel free to run the `nlp_example.py` script found in the examples directory here: https://github.com/huggingface/accelerate/blob/main/examples/nlp_example.py

Just make sure to also install `transformers` and `datasets` first :) 
@muellerzr thanks, will test it out soon! 🤗 
https://api.github.com/repos/huggingface/accelerate/pulls/649
diff --git a/README.md b/README.md
index c67d5dd50..f15f442a7 100644
--- a/README.md
+++ b/README.md
@@ -196,7 +196,7 @@ from accelerate import notebook_launcher
 notebook_launcher(training_function)
 ```
 
-An example can be found in [this notebook](https://github.com/huggingface/notebooks/blob/master/examples/accelerate/simple_nlp_example.ipynb). [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/notebooks/blob/master/examples/accelerate/simple_nlp_example.ipynb)
+An example can be found in [this notebook](https://github.com/huggingface/notebooks/blob/main/examples/accelerate_examples/simple_nlp_example.ipynb). [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/notebooks/blob/main/examples/accelerate_examples/simple_nlp_example.ipynb)
 
 ## Why should I use 🤗 Accelerate?
 

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/645
diff --git a/examples/by_feature/deepspeed_with_config_support.py b/examples/by_feature/deepspeed_with_config_support.py
index fd8a8fa82..89234a911 100755
--- a/examples/by_feature/deepspeed_with_config_support.py
+++ b/examples/by_feature/deepspeed_with_config_support.py
@@ -588,14 +588,12 @@ def group_texts(examples):
         checkpointing_steps = None
 
     # We need to initialize the trackers we use, and also store our configuration.
-    # We initialize the trackers only on main process because `accelerator.log`
-    # only logs on main process and we don't want empty logs/runs on other processes.
+    # The trackers initializes automatically on the main process.
     if args.with_tracking:
-        if accelerator.is_main_process:
-            experiment_config = vars(args)
-            # TensorBoard cannot log Enums, need the raw value
-            experiment_config["lr_scheduler_type"] = experiment_config["lr_scheduler_type"].value
-            accelerator.init_trackers("clm_no_trainer", experiment_config)
+        experiment_config = vars(args)
+        # TensorBoard cannot log Enums, need the raw value
+        experiment_config["lr_scheduler_type"] = experiment_config["lr_scheduler_type"].value
+        accelerator.init_trackers("clm_no_trainer", experiment_config)
 
     # Train!
     total_batch_size = args.per_device_train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
diff --git a/examples/by_feature/fsdp_with_peak_mem_tracking.py b/examples/by_feature/fsdp_with_peak_mem_tracking.py
index 274ccdfce..4ae8e917e 100644
--- a/examples/by_feature/fsdp_with_peak_mem_tracking.py
+++ b/examples/by_feature/fsdp_with_peak_mem_tracking.py
@@ -113,7 +113,7 @@ def training_function(config, args):
     batch_size = int(config["batch_size"])
 
     # We need to initialize the trackers we use, and also store our configuration
-    if args.with_tracking and accelerator.is_main_process:
+    if args.with_tracking:
         experiment_config = vars(args)
         accelerator.init_trackers("fsdp_glue_no_trainer", experiment_config)
 
diff --git a/examples/by_feature/tracking.py b/examples/by_feature/tracking.py
index 76ad7a64b..e4467697c 100644
--- a/examples/by_feature/tracking.py
+++ b/examples/by_feature/tracking.py
@@ -166,8 +166,8 @@ def training_function(config, args):
     )
 
     # New Code #
-    # We need to initalize the trackers we use. Overall configurations can also be stored
-    if args.with_tracking and accelerator.is_main_process:
+    # We need to initialize the trackers we use. Overall configurations can also be stored
+    if args.with_tracking:
         run = os.path.split(__file__)[-1].split(".")[0]
         accelerator.init_trackers(run, config)
 
diff --git a/examples/complete_cv_example.py b/examples/complete_cv_example.py
index 1b86e66ee..880981594 100644
--- a/examples/complete_cv_example.py
+++ b/examples/complete_cv_example.py
@@ -103,7 +103,7 @@ def training_function(config, args):
         checkpointing_steps = None
 
     # We need to initialize the trackers we use, and also store our configuration
-    if args.with_tracking and accelerator.is_main_process:
+    if args.with_tracking:
         run = os.path.split(__file__)[-1].split(".")[0]
         accelerator.init_trackers(run, config)
 
diff --git a/examples/complete_nlp_example.py b/examples/complete_nlp_example.py
index 93e6764f6..559a5c963 100644
--- a/examples/complete_nlp_example.py
+++ b/examples/complete_nlp_example.py
@@ -75,7 +75,7 @@ def training_function(config, args):
     batch_size = int(config["batch_size"])
 
     # We need to initialize the trackers we use, and also store our configuration
-    if args.with_tracking and accelerator.is_main_process:
+    if args.with_tracking:
         run = os.path.split(__file__)[-1].split(".")[0]
         accelerator.init_trackers(run, config)
 

_The documentation is not available anymore as the PR was closed or merged._
Thanks! Do you want to try to update the complete examples in this too? (The last spot would be in the docs, but I can update that :) )
Sure, I will update all the example scripts in the `examples` folder in this PR :)
@Gladiator07 great! It should just be the complete_nlp_example and complete_cv_example
@muellerzr but many of the example scripts have the `is_main_process` check (for instance [here](https://github.com/huggingface/accelerate/blob/a3d94916a87d07471865261d015128870f513c38/examples/by_feature/deepspeed_with_config_support.py#L594)). Should I update them as well or just `complete_nlp_example` and `complete_cv_example` ?
Ah, great point! Yes in that case those would be great to have done as well! :)
Hi @muellerzr, I have removed the `is_main_process` check from all example scripts. Please check it once.
https://api.github.com/repos/huggingface/accelerate/pulls/643
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 8261524fd..4991ddbe4 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -1032,6 +1032,7 @@ def wait_for_everyone(self):
         """
         wait_for_everyone()
 
+    @on_main_process
     def init_trackers(self, project_name: str, config: Optional[dict] = None, init_kwargs: Optional[dict] = {}):
         """
         Initializes a run for all trackers stored in `self.log_with`, potentially with starting configurations

_The documentation is not available anymore as the PR was closed or merged._
@Gladiator07 could you adjust the following code in the examples as well? Since this removes the need for the check:

https://github.com/huggingface/accelerate/blob/main/examples/by_feature/tracking.py#L168-L172
@muellerzr adjust as in? Should I remove the `accelerator.is_main_process` check?
@Gladiator07 correct. We're putting this in place so that is no longer needed, so the examples should show this :)
Cool, doing this in a separate PR.
https://api.github.com/repos/huggingface/accelerate/pulls/642
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 806f0094d..3fec513d2 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -11,6 +11,10 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
+        pytorch-version: [
+          latest,
+          minimum
+        ]
         test-kind: [
           test_prod,
           test_core,
@@ -43,6 +47,7 @@ jobs:
         if [[ ${{ matrix.test-kind }} = test_prod ]]; then pip install -e .[test_prod]; fi
         if [[ ${{ matrix.test-kind }} != test_prod ]]; then pip install -e .[testing,test_trackers]; fi
         if [[ ${{ matrix.test-kind }} = test_rest ]]; then pip uninstall comet_ml -y; fi
+        if [[ ${{ matrix.pytorch-version }} = minimum ]]; then pip install torch==1.6.0; fi
     
     - name: Run Tests
       run: |
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 353ccf894..27db19b9b 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -279,9 +279,7 @@ def __init__(
         self.native_amp = False
         err = "{mode} mixed precision requires {requirement}"
         if self.state.mixed_precision == "fp16":
-            self.native_amp = is_torch_version(">=", "1.6")
-            if not self.native_amp:
-                raise ValueError(err.format(mode="fp16", requirement="PyTorch >= 1.6"))
+            self.native_amp = True
             if not torch.cuda.is_available() and not parse_flag_from_env("USE_MPS_DEVICE"):
                 raise ValueError(err.format(mode="fp16", requirement="a GPU"))
             kwargs = self.scaler_handler.to_kwargs() if self.scaler_handler is not None else {}
@@ -314,7 +312,7 @@ def __init__(
         # RNG Types
         self.rng_types = rng_types
         if self.rng_types is None:
-            self.rng_types = ["torch"] if is_torch_version("<=", "1.5.1") else ["generator"]
+            self.rng_types = ["generator"]
 
     @property
     def use_distributed(self):
diff --git a/src/accelerate/big_modeling.py b/src/accelerate/big_modeling.py
index 311111ad4..6643fa579 100644
--- a/src/accelerate/big_modeling.py
+++ b/src/accelerate/big_modeling.py
@@ -29,6 +29,7 @@
     load_checkpoint_in_model,
     offload_state_dict,
 )
+from .utils.versions import is_torch_version
 
 
 @contextmanager
@@ -59,6 +60,8 @@ def init_empty_weights(include_buffers: bool = False):
 
     </Tip>
     """
+    if not is_torch_version(">=", "1.9.0"):
+        raise NotImplementedError("Initializing empty weights to a meta device requires torch >= 1.9.0")
     old_register_parameter = nn.Module.register_parameter
     if include_buffers:
         old_register_buffer = nn.Module.register_buffer
@@ -114,6 +117,8 @@ def cpu_offload(
             called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
             `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
     """
+    if not is_torch_version(">=", "1.9.0"):
+        raise NotImplementedError("CPU offloading requires torch >= 1.9.0")
     if execution_device is None:
         execution_device = next(iter(model.parameters())).device
     if state_dict is None:
@@ -157,6 +162,8 @@ def disk_offload(
             called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
             `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
     """
+    if not is_torch_version(">=", "1.9.0"):
+        raise NotImplementedError("Disk offloading requires torch >= 1.9.0")
     if not os.path.isdir(offload_dir) or not os.path.isfile(os.path.join(offload_dir, "index.json")):
         offload_state_dict(offload_dir, model.state_dict())
     if execution_device is None:
@@ -208,6 +215,8 @@ def dispatch_model(
             called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
             `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
     """
+    if not is_torch_version(">=", "1.9.0"):
+        raise NotImplementedError("Model dispatching requires torch >= 1.9.0")
     # Error early if the device map is incomplete.
     check_device_map(model, device_map)
 
@@ -304,6 +313,8 @@ def load_checkpoint_and_dispatch(
             called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
             `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
     """
+    if not is_torch_version(">=", "1.9.0"):
+        raise NotImplementedError("Loading and dispatching requires torch >= 1.9.0")
     if isinstance(device_map, str) and device_map not in ["auto", "balanced", "balanced_low_0", "sequential"]:
         raise ValueError(
             "If passing a string for `device_map`, please choose 'auto', 'balanced', 'balanced_low_0' or "
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index b8152b302..9236ee310 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -75,11 +75,11 @@ def total_dataset_length(self):
     "timeout": 0,
     "worker_init_fn": None,
     "multiprocessing_context": None,
+    "generator": None,
 }
 
 # kwargs added after by version
 _PYTORCH_DATALOADER_ADDITIONAL_KWARGS = {
-    "1.6.0": {"generator": None},
     "1.7.0": {"prefetch_factor": 2, "persistent_workers": False},
 }
 
@@ -412,7 +412,7 @@ def __init__(self, dataset, split_batches: bool = False, _drop_last: bool = Fals
         self.split_batches = split_batches
         if is_torch_version("<", "1.8.0"):
             raise ImportError(
-                "Using `DataLoaderDispatcher` requires PyTorch 1.8.0 minimum. You have {torch.__version__}."
+                f"Using `DataLoaderDispatcher` requires PyTorch 1.8.0 minimum. You have {torch.__version__}."
             )
         if shuffle:
             torch.utils.data.graph_settings.apply_shuffle_settings(dataset, shuffle=shuffle)
diff --git a/src/accelerate/launchers.py b/src/accelerate/launchers.py
index 8ddc9af34..dcdead956 100644
--- a/src/accelerate/launchers.py
+++ b/src/accelerate/launchers.py
@@ -20,7 +20,7 @@
 import torch
 
 from .state import AcceleratorState
-from .utils import PrecisionType, PrepareForLaunch, is_torch_version, patch_environment
+from .utils import PrecisionType, PrepareForLaunch, patch_environment
 
 
 def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, mixed_precision="no", use_port="29500"):
@@ -90,12 +90,6 @@ def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, mix
 
         if num_processes > 1:
             # Multi-GPU launch
-            if is_torch_version("<", "1.5.0"):
-                raise ImportError(
-                    "Using `notebook_launcher` for distributed training on GPUs require torch >= 1.5.0, got "
-                    f"{torch.__version__}."
-                )
-
             from torch.multiprocessing import start_processes
 
             if len(AcceleratorState._shared_state) > 0:
@@ -154,12 +148,6 @@ def debug_launcher(function, args=(), num_processes=2):
         num_processes (`int`, *optional*, defaults to 2):
             The number of processes to use for training.
     """
-    if is_torch_version("<", "1.5.0"):
-        raise ImportError(
-            "Using `debug_launcher` for distributed training on GPUs require torch >= 1.5.0, got "
-            f"{torch.__version__}."
-        )
-
     from torch.multiprocessing import start_processes
 
     with tempfile.NamedTemporaryFile() as tmp_file:
diff --git a/src/accelerate/test_utils/__init__.py b/src/accelerate/test_utils/__init__.py
index 43d5ed7ef..faf8a7da8 100644
--- a/src/accelerate/test_utils/__init__.py
+++ b/src/accelerate/test_utils/__init__.py
@@ -10,6 +10,7 @@
     require_huggingface_suite,
     require_multi_gpu,
     require_single_gpu,
+    require_torch_min_version,
     require_tpu,
     skip,
     slow,
diff --git a/src/accelerate/test_utils/scripts/test_script.py b/src/accelerate/test_utils/scripts/test_script.py
index 07d794b42..6897d9084 100644
--- a/src/accelerate/test_utils/scripts/test_script.py
+++ b/src/accelerate/test_utils/scripts/test_script.py
@@ -46,10 +46,9 @@ def rng_sync_check():
     if state.distributed_type == DistributedType.MULTI_GPU:
         synchronize_rng_states(["cuda"])
         assert are_the_same_tensors(torch.cuda.get_rng_state()), "RNG states improperly synchronized on GPU."
-    if is_torch_version(">=", "1.6.0"):
-        generator = torch.Generator()
-        synchronize_rng_states(["generator"], generator=generator)
-        assert are_the_same_tensors(generator.get_state()), "RNG states improperly synchronized in generator."
+    generator = torch.Generator()
+    synchronize_rng_states(["generator"], generator=generator)
+    assert are_the_same_tensors(generator.get_state()), "RNG states improperly synchronized in generator."
 
     if state.local_process_index == 0:
         print("All rng are properly synched.")
@@ -339,7 +338,7 @@ def main():
     if state.local_process_index == 0:
         print("\n**DataLoader integration test**")
     dl_preparation_check()
-    if state.distributed_type != DistributedType.TPU:
+    if state.distributed_type != DistributedType.TPU and is_torch_version(">=", "1.8.0"):
         central_dl_preparation_check()
 
     # Trainings are not exactly the same in DeepSpeed and CPU mode
diff --git a/src/accelerate/test_utils/testing.py b/src/accelerate/test_utils/testing.py
index 381f92c58..94e13412a 100644
--- a/src/accelerate/test_utils/testing.py
+++ b/src/accelerate/test_utils/testing.py
@@ -20,6 +20,7 @@
 import tempfile
 import unittest
 from distutils.util import strtobool
+from functools import partial
 from pathlib import Path
 from typing import List, Union
 from unittest import mock
@@ -132,6 +133,16 @@ def require_fsdp(test_case):
     return unittest.skipUnless(is_torch_version(">=", "1.12.0"), "test requires torch version >= 1.12.0")(test_case)
 
 
+def require_torch_min_version(test_case=None, version=None):
+    """
+    Decorator marking that a test requires a particular torch version to be tested. These tests are skipped when an
+    installed torch version is less than the required one.
+    """
+    if test_case is None:
+        return partial(require_torch_min_version, version=version)
+    return unittest.skipUnless(is_torch_version(">=", version), f"test requires torch version >= {version}")(test_case)
+
+
 def require_tensorboard(test_case):
     """
     Decorator marking a test that requires tensorboard installed. These tests are skipped when tensorboard isn't
diff --git a/tests/test_big_modeling.py b/tests/test_big_modeling.py
index 4f3054b73..4f738f4d8 100644
--- a/tests/test_big_modeling.py
+++ b/tests/test_big_modeling.py
@@ -27,7 +27,7 @@
     load_checkpoint_and_dispatch,
 )
 from accelerate.hooks import remove_hook_from_submodules
-from accelerate.test_utils import require_cuda, require_multi_gpu, slow
+from accelerate.test_utils import require_cuda, require_multi_gpu, require_torch_min_version, slow
 from accelerate.utils import offload_state_dict
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
@@ -79,6 +79,7 @@ def forward(self, x):
         return self.linear4(self.linear3(self.batchnorm(self.linear2(self.linear1(x)))))
 
 
+@require_torch_min_version(version="1.9.0")
 class BigModelingTester(unittest.TestCase):
     def test_init_empty_weights(self):
         # base use
diff --git a/tests/test_hooks.py b/tests/test_hooks.py
index 2e0519668..9d48db9e1 100644
--- a/tests/test_hooks.py
+++ b/tests/test_hooks.py
@@ -27,7 +27,7 @@
     remove_hook_from_module,
     remove_hook_from_submodules,
 )
-from accelerate.test_utils import require_multi_gpu
+from accelerate.test_utils import require_multi_gpu, require_torch_min_version
 
 
 class ModelForTest(nn.Module):
@@ -51,6 +51,7 @@ def post_forward(self, module, output):
         return output + 1
 
 
+@require_torch_min_version(version="1.9.0")
 class HooksModelTester(unittest.TestCase):
     def test_add_and_remove_hooks(self):
         test_model = ModelForTest()
diff --git a/tests/test_metrics.py b/tests/test_metrics.py
index b26ef00b2..9dd098456 100644
--- a/tests/test_metrics.py
+++ b/tests/test_metrics.py
@@ -26,11 +26,13 @@
     require_huggingface_suite,
     require_multi_gpu,
     require_single_gpu,
+    require_torch_min_version,
 )
 from accelerate.utils import get_launch_prefix, patch_environment
 
 
 @require_huggingface_suite
+@require_torch_min_version(version="1.8.0")
 class MetricTester(unittest.TestCase):
     def setUp(self):
         mod_file = inspect.getfile(accelerate.test_utils)
diff --git a/tests/test_modeling_utils.py b/tests/test_modeling_utils.py
index f5c36ee3d..1c6f6088d 100644
--- a/tests/test_modeling_utils.py
+++ b/tests/test_modeling_utils.py
@@ -21,6 +21,7 @@
 import torch.nn as nn
 
 from accelerate.test_utils import require_cuda, require_multi_gpu
+from accelerate.test_utils.testing import require_torch_min_version
 from accelerate.utils.modeling import (
     check_device_map,
     clean_device_map,
@@ -45,6 +46,7 @@ def forward(self, x):
         return self.linear2(self.batchnorm(self.linear1(x)))
 
 
+@require_torch_min_version(version="1.9.0")
 class ModelingUtilsTester(unittest.TestCase):
     def check_set_module_tensor_for_device(self, model, device1, device2):
         self.assertEqual(model.linear1.weight.device, torch.device(device1))

It should be possible to import those regardless of the version (and then the actual use of it would trigger the error).
@@ -16,7 +15,18 @@
     InitProcessGroupKwargs,
     find_executable_batch_size,
     infer_auto_device_map,
+    is_torch_version,
     load_checkpoint_in_model,
     rich,
     synchronize_rng_states,
 )
+
+
+if is_torch_version(">=", "1.9.0"):
+    from .big_modeling import (
+        cpu_offload,
+        disk_offload,
+        dispatch_model,
+        init_empty_weights,
+        load_checkpoint_and_dispatch,
+    )
Let's leave the type hint but as a forward reference?
@@ -48,7 +49,7 @@ def save_accelerator_state(
     optimizers: list,
     schedulers: list,
     process_index: int,
-    scaler: GradScaler = None,
+    scaler=None,
```suggestion
def require_torch_min_version(test_case=None, version=None):
```
Naming isn't completely right IMO.
@@ -132,6 +133,16 @@ def require_fsdp(test_case):
     return unittest.skipUnless(is_torch_version(">=", "1.12.0"), "test requires torch version >= 1.12.0")(test_case)
 
 
+def require_torch_version(test_case=None, version=None):
Not needed since we're raising the min torch
@@ -48,7 +49,7 @@ def save_accelerator_state(
     optimizers: list,
     schedulers: list,
     process_index: int,
-    scaler: GradScaler = None,
+    scaler=None,
Fixed
@@ -16,7 +15,18 @@
     InitProcessGroupKwargs,
     find_executable_batch_size,
     infer_auto_device_map,
+    is_torch_version,
     load_checkpoint_in_model,
     rich,
     synchronize_rng_states,
 )
+
+
+if is_torch_version(">=", "1.9.0"):
+    from .big_modeling import (
+        cpu_offload,
+        disk_offload,
+        dispatch_model,
+        init_empty_weights,
+        load_checkpoint_and_dispatch,
+    )
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/638
diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py
index cbbaeae16..5f425f200 100644
--- a/src/accelerate/utils/dataclasses.py
+++ b/src/accelerate/utils/dataclasses.py
@@ -60,6 +60,8 @@ class DistributedDataParallelKwargs(KwargsHandler):
 
     `gradient_as_bucket_view` is only available in PyTorch 1.7.0 and later versions.
 
+    `static_graph` is only available in PyTorch 1.11.0 and later versions.
+
     </Tip>"""
 
     dim: int = 0
@@ -68,6 +70,7 @@ class DistributedDataParallelKwargs(KwargsHandler):
     find_unused_parameters: bool = False
     check_reduction: bool = False
     gradient_as_bucket_view: bool = False
+    static_graph: bool = False
 
 
 @dataclass

_The documentation is not available anymore as the PR was closed or merged._
Thanks! Could we add a warning to this please similar to the other one that exists for these kwargs in the docstring mentioning its only available on pytorch>=1.11?
Sure, done
@rom1504 looks like the last thing we need is to run `make style; make quality` :) 
@muellerzr fixed
Great! Thanks!!
https://api.github.com/repos/huggingface/accelerate/pulls/637
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 11aadc8e2..fcc80bb7f 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -26,7 +26,6 @@
 from typing import Dict, List
 
 import torch
-import torch.distributed.run as distrib_run
 
 import psutil
 from accelerate.commands.config import default_config_file, load_config_from_file
@@ -50,6 +49,10 @@
 from rich.logging import RichHandler
 
 
+if is_torch_version(">=", "1.9.0"):
+    import torch.distributed.run as distrib_run
+
+
 FORMAT = "%(message)s"
 logging.basicConfig(format=FORMAT, datefmt="[%X]", handlers=[RichHandler()])
 

_The documentation is not available anymore as the PR was closed or merged._
Thanks for fixing!
https://api.github.com/repos/huggingface/accelerate/pulls/634
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index 1ba5839e2..65c0e56b2 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -37,6 +37,7 @@ def get_cluster_input():
     num_machines = 1
     main_process_ip = None
     main_process_port = None
+    rdzv_backend = "static"
     if distributed_type in [DistributedType.MULTI_GPU, DistributedType.MULTI_CPU]:
         num_machines = _ask_field(
             "How many different machines will you use (use more than 1 for multi-node training)? [1]: ",
@@ -56,6 +57,9 @@ def get_cluster_input():
                 "What is the port you will use to communicate with the main process? ",
                 lambda x: int(x),
             )
+            rdzv_backend = _ask_field(
+                "What rendezvous backend will you use? ('static', 'c10d', ...)", default="static"
+            )
 
     if distributed_type == DistributedType.NO:
         use_cpu = _ask_field(
@@ -323,4 +327,5 @@ def get_cluster_input():
         deepspeed_config=deepspeed_config,
         fsdp_config=fsdp_config,
         use_cpu=use_cpu,
+        rdzv_backend=rdzv_backend,
     )
diff --git a/src/accelerate/commands/config/config_args.py b/src/accelerate/commands/config/config_args.py
index 11ca0ca90..e78aad181 100644
--- a/src/accelerate/commands/config/config_args.py
+++ b/src/accelerate/commands/config/config_args.py
@@ -137,6 +137,7 @@ class ClusterConfig(BaseConfig):
     num_machines: int = 1
     main_process_ip: Optional[str] = None
     main_process_port: Optional[int] = None
+    rdzv_backend: Optional[str] = "static"
     main_training_function: str = "main"
 
     # args for deepspeed_plugin
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index ea27afb75..11aadc8e2 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -26,13 +26,13 @@
 from typing import Dict, List
 
 import torch
+import torch.distributed.run as distrib_run
 
 import psutil
 from accelerate.commands.config import default_config_file, load_config_from_file
 from accelerate.commands.config.config_args import SageMakerConfig
 from accelerate.state import get_int_from_env
 from accelerate.utils import (
-    TORCH_LAUNCH_PARAMS,
     ComputeEnvironment,
     DistributedType,
     PrecisionType,
@@ -50,9 +50,6 @@
 from rich.logging import RichHandler
 
 
-if is_torch_version(">=", "1.9.0"):
-    import torch.distributed.run as distrib_run
-
 FORMAT = "%(message)s"
 logging.basicConfig(format=FORMAT, datefmt="[%X]", handlers=[RichHandler()])
 
@@ -379,8 +376,7 @@ def multi_gpu_launcher(args):
         setattr(args, "nproc_per_node", str(num_processes // num_machines))
         setattr(args, "nnodes", str(num_machines))
         setattr(args, "node_rank", str(args.machine_rank))
-        setattr(args, "master_addr", str(args.main_process_ip))
-        setattr(args, "master_port", str(args.main_process_port))
+        setattr(args, "rdzv_endpoint", f"{args.main_process_ip}:{args.main_process_port}")
     else:
         setattr(args, "nproc_per_node", str(num_processes))
         if args.main_process_port is not None:
@@ -451,33 +447,19 @@ def multi_gpu_launcher(args):
         if args.fsdp_state_dict_type is not None:
             current_env["FSDP_STATE_DICT_TYPE"] = str(args.fsdp_state_dict_type)
     current_env["OMP_NUM_THREADS"] = str(args.num_cpu_threads_per_process)
-    if is_torch_version(">=", "1.9.0"):
-        debug = getattr(args, "debug", False)
-        args = _filter_args(args)
-        with patch_environment(**current_env):
-            console = get_console()
-
-            try:
-                distrib_run.run(args)
-            except:
-                if debug:
-                    console.print("\n[bold red]Using --debug, `torch.distributed` Stack Trace:[/bold red]")
-                    console.print_exception(suppress=[__file__], show_locals=False)
-    else:
-        # We still have to use subprocess, the user won't get a clean traceback as a result
-        cmd = get_launch_prefix()
-        for k, v in vars(args).items():
-            if k in TORCH_LAUNCH_PARAMS and v:
-                param = [f"--{k}"]
-                if type(v) != bool:
-                    param.append(v)
-                cmd.extend(param)
-        cmd.append(args.training_script)
-        cmd.extend(args.training_script_args)
-        process = subprocess.Popen(cmd, env=current_env)
-        process.wait()
-        if process.returncode != 0:
-            raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)
+    if is_torch_version("<", "1.9.0"):
+        raise NotImplementedError("Multi-node training requires pytorch>=1.9.0")
+
+    debug = getattr(args, "debug", False)
+    args = _filter_args(args)
+    with patch_environment(**current_env):
+        try:
+            distrib_run.run(args)
+        except:
+            if debug:
+                console = get_console()
+                console.print("\n[bold red]Using --debug, `torch.distributed` Stack Trace:[/bold red]")
+                console.print_exception(suppress=[__file__], show_locals=False)
 
 
 def deepspeed_launcher(args):

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/631
diff --git a/examples/complete_cv_example.py b/examples/complete_cv_example.py
index 99d109bd1..1b86e66ee 100644
--- a/examples/complete_cv_example.py
+++ b/examples/complete_cv_example.py
@@ -105,8 +105,6 @@ def training_function(config, args):
     # We need to initialize the trackers we use, and also store our configuration
     if args.with_tracking and accelerator.is_main_process:
         run = os.path.split(__file__)[-1].split(".")[0]
-        if args.logging_dir:
-            run = os.path.join(args.logging_dir, run)
         accelerator.init_trackers(run, config)
 
     # Grab all the image filenames
diff --git a/examples/complete_nlp_example.py b/examples/complete_nlp_example.py
index 094f6cfb5..93e6764f6 100644
--- a/examples/complete_nlp_example.py
+++ b/examples/complete_nlp_example.py
@@ -77,8 +77,6 @@ def training_function(config, args):
     # We need to initialize the trackers we use, and also store our configuration
     if args.with_tracking and accelerator.is_main_process:
         run = os.path.split(__file__)[-1].split(".")[0]
-        if args.logging_dir:
-            run = os.path.join(args.logging_dir, run)
         accelerator.init_trackers(run, config)
 
     tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index 4e0900eb4..051cf8bc6 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -251,7 +251,7 @@ def __repr__(self):
         if self.distributed_type == DistributedType.DEEPSPEED:
             repr += f"ds_config: {self.deepspeed_plugin.deepspeed_config}\n"
         else:
-            f"Mixed precision type: {mixed_precision}\n"
+            repr += f"Mixed precision type: {mixed_precision}\n"
         return repr
 
     # For backward compatibility

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/630
diff --git a/docs/source/usage_guides/mps.mdx b/docs/source/usage_guides/mps.mdx
index 43a00f104..3648bdce8 100644
--- a/docs/source/usage_guides/mps.mdx
+++ b/docs/source/usage_guides/mps.mdx
@@ -72,8 +72,9 @@ accelerate launch /examples/cv_example.py --data_dir images
 
 ## A few caveats to be aware of
 
-1. For `nlp_example.py` the metrics are too bad when compared to CPU-only training. 
-This means certain operations in BERT model are going wrong using mps device and this needs to be fixed by PyTorch.
+1. We strongly recommend to install PyTorch >= 1.13 (nightly version at the time of writing) on your MacOS machine. 
+It has major fixes related to model correctness and performance improvements for transformer based models.
+Please refer to https://github.com/pytorch/pytorch/issues/82707 for more details.
 2. Distributed setups `gloo` and `nccl` are not working with `mps` device. 
 This means that currently only single GPU of `mps` device type can be used.
 

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/629
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 7cb4529d9..541bda48e 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -367,7 +367,7 @@ def multi_gpu_launcher(args):
     if num_machines > 1:
         setattr(args, "nproc_per_node", str(num_processes // num_machines))
         setattr(args, "nnodes", str(num_machines))
-        setattr(args, "machine_rank", str(args.machine_rank))
+        setattr(args, "node_rank", str(args.machine_rank))
         setattr(args, "master_addr", str(args.main_process_ip))
         setattr(args, "master_port", str(args.main_process_port))
     else:
@@ -441,16 +441,16 @@ def multi_gpu_launcher(args):
             current_env["FSDP_STATE_DICT_TYPE"] = str(args.fsdp_state_dict_type)
     current_env["OMP_NUM_THREADS"] = str(args.num_cpu_threads_per_process)
     if is_torch_version(">=", "1.9.0"):
-        distrib_args = _filter_args(args)
+        args = _filter_args(args)
         with patch_environment(**current_env):
-            distrib_run.run(distrib_args)
+            distrib_run.run(args)
     else:
         # We still have to use subprocess, the user won't get a clean traceback as a result
         cmd = get_launch_prefix()
         for k, v in vars(args).items():
             if k in TORCH_LAUNCH_PARAMS and v:
                 param = [f"--{k}"]
-                if not v:
+                if type(v) != bool:
                     param.append(v)
                 cmd.extend(param)
         cmd.append(args.training_script)
diff --git a/src/accelerate/utils/launch.py b/src/accelerate/utils/launch.py
index c6605ee26..1a452ca63 100644
--- a/src/accelerate/utils/launch.py
+++ b/src/accelerate/utils/launch.py
@@ -46,16 +46,12 @@ def _filter_args(args):
     Filters out all `accelerate` specific args
     """
     distrib_args = distrib_run.get_args_parser()
-    known_args, _ = distrib_args.parse_known_args()
-    for arg in list(vars(args).keys()):
-        if arg not in vars(known_args).keys():
-            delattr(args, arg)
-    distrib_args = distrib_run.parse_args(vars(args))
+    new_args, _ = distrib_args.parse_known_args()
+
     for key, value in vars(args).items():
-        setattr(distrib_args, key, value)
-    if is_torch_version("<", "1.9.0"):
-        setattr(distrib_args, "use_env", True)
-    return distrib_args
+        if key in vars(new_args).keys():
+            setattr(new_args, key, value)
+    return new_args
 
 
 class PrepareForLaunch:

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/627
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index c7f235843..8196413c5 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -37,7 +37,6 @@
     PrecisionType,
     PrepareForLaunch,
     _filter_args,
-    get_launch_prefix,
     is_deepspeed_available,
     is_rich_available,
     is_sagemaker_available,
@@ -477,55 +476,58 @@ def multi_gpu_launcher(args):
 def deepspeed_launcher(args):
     if not is_deepspeed_available():
         raise ImportError("DeepSpeed is not installed => run `pip3 install deepspeed` or build it from source.")
-    cmd = ["deepspeed", "--no_local_rank"]
-    if args.num_machines > 1:
-        if args.deepspeed_multinode_launcher == DEEPSPEED_MULTINODE_LAUNCHERS[1]:
-            cmd = get_launch_prefix()
+    num_processes = getattr(args, "num_processes")
+    num_machines = getattr(args, "num_machines")
+    main_process_ip = getattr(args, "main_process_ip")
+    main_process_port = getattr(args, "main_process_port")
+    if num_machines > 1 and args.deepspeed_multinode_launcher != DEEPSPEED_MULTINODE_LAUNCHERS[1]:
+        cmd = ["deepspeed", "--no_local_rank"]
+        cmd.extend(["--hostfile", str(args.deepspeed_hostfile), "--launcher", str(args.deepspeed_multinode_launcher)])
+        if args.deepspeed_exclusion_filter is not None:
             cmd.extend(
                 [
-                    "--nproc_per_node",
-                    str(args.num_processes // args.num_machines),
-                    "--nnodes",
-                    str(args.num_machines),
-                    "--node_rank",
-                    str(args.machine_rank),
-                    "--master_addr",
-                    args.main_process_ip,
-                    "--master_port",
-                    str(args.main_process_port),
+                    "--exclude",
+                    str(args.deepspeed_exclusion_filter),
                 ]
             )
-        else:
+        elif args.deepspeed_inclusion_filter is not None:
             cmd.extend(
-                ["--hostfile", str(args.deepspeed_hostfile), "--launcher", str(args.deepspeed_multinode_launcher)]
+                [
+                    "--include",
+                    str(args.deepspeed_inclusion_filter),
+                ]
             )
-            if args.deepspeed_exclusion_filter is not None:
-                cmd.extend(
-                    [
-                        "--exclude",
-                        str(args.deepspeed_exclusion_filter),
-                    ]
-                )
-            elif args.deepspeed_inclusion_filter is not None:
-                cmd.extend(
-                    [
-                        "--include",
-                        str(args.deepspeed_inclusion_filter),
-                    ]
-                )
-            else:
-                cmd.extend(["--num_gpus", str(args.num_processes // args.num_machines)])
+        else:
+            cmd.extend(["--num_gpus", str(args.num_processes // args.num_machines)])
+
+        if args.module and args.no_python:
+            raise ValueError("--module and --no_python cannot be used together")
+        elif args.module:
+            cmd.append("--module")
+        elif args.no_python:
+            cmd.append("--no_python")
+        cmd.append(args.training_script)
+        cmd.extend(args.training_script_args)
+    elif num_machines > 1 and args.deepspeed_multinode_launcher == DEEPSPEED_MULTINODE_LAUNCHERS[1]:
+        setattr(args, "nproc_per_node", str(num_processes // num_machines))
+        setattr(args, "nnodes", str(num_machines))
+        setattr(args, "node_rank", int(args.machine_rank))
+        if getattr(args, "same_network"):
+            setattr(args, "master_addr", str(main_process_ip))
+            setattr(args, "master_port", str(main_process_port))
+        else:
+            setattr(args, "rdzv_endpoint", f"{main_process_ip}:{main_process_port}")
     else:
-        cmd.extend(["--num_gpus", str(args.num_processes)])
+        setattr(args, "nproc_per_node", str(num_processes))
+        if main_process_port is not None:
+            setattr(args, "master_port", str(main_process_port))
 
     if args.module and args.no_python:
         raise ValueError("--module and --no_python cannot be used together")
     elif args.module:
-        cmd.append("--module")
+        setattr(args, "module", True)
     elif args.no_python:
-        cmd.append("--no_python")
-    cmd.append(args.training_script)
-    cmd.extend(args.training_script_args)
+        setattr(args, "no_python", True)
 
     current_env = os.environ.copy()
     try:
@@ -558,10 +560,24 @@ def deepspeed_launcher(args):
                     continue
                 f.write(f"{key}={value}\n")
 
-    process = subprocess.Popen(cmd, env=current_env)
-    process.wait()
-    if process.returncode != 0:
-        raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)
+        process = subprocess.Popen(cmd, env=current_env)
+        process.wait()
+        if process.returncode != 0:
+            raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)
+    else:
+        if is_torch_version("<", "1.9.0"):
+            raise NotImplementedError("Multi-node training requires pytorch>=1.9.0")
+
+        debug = getattr(args, "debug", False)
+        args = _filter_args(args)
+        with patch_environment(**current_env):
+            try:
+                distrib_run.run(args)
+            except:
+                if debug:
+                    console = get_console()
+                    console.print("\n[bold red]Using --debug, `torch.distributed` Stack Trace:[/bold red]")
+                    console.print_exception(suppress=[__file__], show_locals=False)
 
 
 def tpu_launcher(args):

_The documentation is not available anymore as the PR was closed or merged._
> There is a bit too much in this PR to wrap my head around. Can we split it between multiGPU launcher fixes, DeepSpeed launcher fixes and other fixes? Thanks!

1. MultiGPU launcher fixes and simplification was put in another PR by Zach #627 
2. Minor other fixes are in #630 
3. This will make deepspeed launcher updates to remove a call to subprocess
https://api.github.com/repos/huggingface/accelerate/pulls/626
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 7cb4529d9..117932f2b 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -334,6 +334,8 @@ def simple_launcher(args):
     current_env = os.environ.copy()
     current_env["USE_CPU"] = str(args.cpu or args.use_cpu)
     current_env["USE_MPS_DEVICE"] = str(args.use_mps_device)
+    if args.use_mps_device:
+        current_env["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
     if args.num_machines > 1:
         current_env["MASTER_ADDR"] = args.main_process_ip
         current_env["MASTER_PORT"] = str(args.main_process_port)
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index bd339048f..4e0900eb4 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import os
+import warnings
 from distutils.util import strtobool
 
 import torch
@@ -221,6 +222,14 @@ def __init__(
                                 "and/or you do not have an MPS-enabled device on this machine."
                             )
                     else:
+                        from .utils import is_torch_version
+
+                        if not is_torch_version(">", "1.12.0"):
+                            warnings.warn(
+                                "We strongly recommend to install PyTorch >= 1.13 (nightly version at the time of writing) on your MacOS machine. "
+                                "It has major fixes related to model correctness and performance improvements for transformer based models. "
+                                "Please refer to https://github.com/pytorch/pytorch/issues/82707 for more details."
+                            )
                         self.device = torch.device("mps")
                 elif cpu or not torch.cuda.is_available():
                     self.device = torch.device("cpu")

That warning may not age well once PyTorch 1.13 is released. Maybe:
```suggestion
                                "We strongly recommend to install PyTorch >= 1.13 (nightly version at the time of writing) on your MacOS machine. "
```
@@ -221,6 +222,14 @@ def __init__(
                                 "and/or you do not have an MPS-enabled device on this machine."
                             )
                     else:
+                        from .utils import is_torch_version
+
+                        if not is_torch_version(">", "1.12.0"):
+                            warnings.warn(
+                                "We strongly recommend to install the latest PyTorch nightly version (1.13.0.dev...) on your MacOS machine. "
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/625
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 8e372adea..7cb4529d9 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -32,19 +32,26 @@
 from accelerate.commands.config.config_args import SageMakerConfig
 from accelerate.state import get_int_from_env
 from accelerate.utils import (
+    TORCH_LAUNCH_PARAMS,
     ComputeEnvironment,
     DistributedType,
     PrecisionType,
     PrepareForLaunch,
+    _filter_args,
     get_launch_prefix,
     is_deepspeed_available,
     is_sagemaker_available,
+    is_torch_version,
     patch_environment,
 )
 from accelerate.utils.constants import DEEPSPEED_MULTINODE_LAUNCHERS
 from accelerate.utils.dataclasses import SageMakerDistributedType
 
 
+if is_torch_version(">=", "1.9.0"):
+    import torch.distributed.run as distrib_run
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -355,43 +362,32 @@ def simple_launcher(args):
 
 
 def multi_gpu_launcher(args):
-    cmd = get_launch_prefix()
-    if args.num_machines > 1:
-        cmd.extend(
-            [
-                "--nproc_per_node",
-                str(args.num_processes // args.num_machines),
-                "--nnodes",
-                str(args.num_machines),
-                "--node_rank",
-                str(args.machine_rank),
-                "--master_addr",
-                args.main_process_ip,
-                "--master_port",
-                str(args.main_process_port),
-            ]
-        )
+    num_processes = getattr(args, "num_processes")
+    num_machines = getattr(args, "num_machines")
+    if num_machines > 1:
+        setattr(args, "nproc_per_node", str(num_processes // num_machines))
+        setattr(args, "nnodes", str(num_machines))
+        setattr(args, "machine_rank", str(args.machine_rank))
+        setattr(args, "master_addr", str(args.main_process_ip))
+        setattr(args, "master_port", str(args.main_process_port))
     else:
-        cmd.extend(["--nproc_per_node", str(args.num_processes)])
+        setattr(args, "nproc_per_node", str(num_processes))
         if args.main_process_port is not None:
-            cmd.extend(["--master_port", str(args.main_process_port)])
+            setattr(args, "master_port", str(args.main_process_port))
 
     if args.module and args.no_python:
         raise ValueError("--module and --no_python cannot be used together")
     elif args.module:
-        cmd.append("--module")
+        setattr(args, "module", True)
     elif args.no_python:
-        cmd.append("--no_python")
-    cmd.append(args.training_script)
-    cmd.extend(args.training_script_args)
+        setattr(args, "no_python", True)
 
     current_env = os.environ.copy()
+    mixed_precision = args.mixed_precision.lower()
     try:
-        mixed_precision = PrecisionType(args.mixed_precision.lower())
+        mixed_precision = PrecisionType(mixed_precision)
     except ValueError:
-        raise ValueError(
-            f"Unknown mixed_precision mode: {args.mixed_precision.lower()}. Choose between {PrecisionType.list()}."
-        )
+        raise ValueError(f"Unknown mixed_precision mode: {mixed_precision}. Choose between {PrecisionType.list()}.")
 
     if args.fp16:
         warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', DeprecationWarning)
@@ -444,10 +440,25 @@ def multi_gpu_launcher(args):
         if args.fsdp_state_dict_type is not None:
             current_env["FSDP_STATE_DICT_TYPE"] = str(args.fsdp_state_dict_type)
     current_env["OMP_NUM_THREADS"] = str(args.num_cpu_threads_per_process)
-    process = subprocess.Popen(cmd, env=current_env)
-    process.wait()
-    if process.returncode != 0:
-        raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)
+    if is_torch_version(">=", "1.9.0"):
+        distrib_args = _filter_args(args)
+        with patch_environment(**current_env):
+            distrib_run.run(distrib_args)
+    else:
+        # We still have to use subprocess, the user won't get a clean traceback as a result
+        cmd = get_launch_prefix()
+        for k, v in vars(args).items():
+            if k in TORCH_LAUNCH_PARAMS and v:
+                param = [f"--{k}"]
+                if not v:
+                    param.append(v)
+                cmd.extend(param)
+        cmd.append(args.training_script)
+        cmd.extend(args.training_script_args)
+        process = subprocess.Popen(cmd, env=current_env)
+        process.wait()
+        if process.returncode != 0:
+            raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)
 
 
 def deepspeed_launcher(args):
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
index db21afc3a..2b8f566f0 100644
--- a/src/accelerate/utils/__init__.py
+++ b/src/accelerate/utils/__init__.py
@@ -2,7 +2,7 @@
 # There's no way to ignore "F401 '...' imported but unused" warnings in this
 # module, but to preserve other warnings. So, don't check this module at all
 
-from .constants import MODEL_NAME, OPTIMIZER_NAME, RNG_STATE_NAME, SCALER_NAME, SCHEDULER_NAME
+from .constants import MODEL_NAME, OPTIMIZER_NAME, RNG_STATE_NAME, SCALER_NAME, SCHEDULER_NAME, TORCH_LAUNCH_PARAMS
 from .dataclasses import (
     ComputeEnvironment,
     DeepSpeedPlugin,
@@ -91,7 +91,7 @@
         HfDeepSpeedConfig,
     )
 
-from .launch import PrepareForLaunch, get_launch_prefix
+from .launch import PrepareForLaunch, _filter_args, get_launch_prefix
 from .memory import find_executable_batch_size
 from .other import (
     extract_model_from_parallel,
diff --git a/src/accelerate/utils/constants.py b/src/accelerate/utils/constants.py
index 4e7c71853..934923bef 100644
--- a/src/accelerate/utils/constants.py
+++ b/src/accelerate/utils/constants.py
@@ -31,3 +31,30 @@
 DEEPSPEED_MULTINODE_LAUNCHERS = ["pdsh", "standard", "openmpi", "mvapich"]
 
 STR_OPERATION_TO_FUNC = {">": op.gt, ">=": op.ge, "==": op.eq, "!=": op.ne, "<=": op.le, "<": op.lt}
+
+# These are the args for `torch.distributed.launch` for pytorch < 1.9
+TORCH_LAUNCH_PARAMS = [
+    "nnodes",
+    "nproc_per_node",
+    "rdzv_backend",
+    "rdzv_endpoint",
+    "rdzv_id",
+    "rdzv_conf",
+    "standalone",
+    "max_restarts",
+    "monitor_interval",
+    "start_method",
+    "role",
+    "module",
+    "m",
+    "no_python",
+    "run_path",
+    "log_dir",
+    "r",
+    "redirects",
+    "t",
+    "tee",
+    "node_rank",
+    "master_addr",
+    "master_port",
+]
diff --git a/src/accelerate/utils/launch.py b/src/accelerate/utils/launch.py
index c74185c5d..c6605ee26 100644
--- a/src/accelerate/utils/launch.py
+++ b/src/accelerate/utils/launch.py
@@ -21,6 +21,12 @@
 from .dataclasses import DistributedType
 
 
+if is_torch_version(">=", "1.9.0"):
+    import torch.distributed.run as distrib_run
+else:
+    import torch.distributed.launch as distrib_run
+
+
 def get_launch_prefix():
     """
     Grabs the correct launcher for starting a distributed command, such as either `torchrun`, `python -m
@@ -35,6 +41,23 @@ def get_launch_prefix():
     return cmd
 
 
+def _filter_args(args):
+    """
+    Filters out all `accelerate` specific args
+    """
+    distrib_args = distrib_run.get_args_parser()
+    known_args, _ = distrib_args.parse_known_args()
+    for arg in list(vars(args).keys()):
+        if arg not in vars(known_args).keys():
+            delattr(args, arg)
+    distrib_args = distrib_run.parse_args(vars(args))
+    for key, value in vars(args).items():
+        setattr(distrib_args, key, value)
+    if is_torch_version("<", "1.9.0"):
+        setattr(distrib_args, "use_env", True)
+    return distrib_args
+
+
 class PrepareForLaunch:
     """
     Prepare a function that will launched in a distributed setup.

Clean!
@@ -35,6 +36,21 @@ def get_launch_prefix():
     return cmd
 
 
+def _filter_args(args):
Please use a regular list here, it's going to be easier to maintain (in terms of diff)
@@ -31,3 +31,8 @@
 DEEPSPEED_MULTINODE_LAUNCHERS = ["pdsh", "standard", "openmpi", "mvapich"]
 
 STR_OPERATION_TO_FUNC = {">": op.gt, ">=": op.ge, "==": op.eq, "!=": op.ne, "<=": op.le, "<": op.lt}
+
+# These are the args for `torch.distributed.launch` for pytorch < 1.9
+TORCH_LAUNCH_PARAMS = """nnodes,nproc_per_node,rdzv_backend,rdzv_endpoint,rdzv_id,rdzv_conf,standalone,max_restarts,monitor_interval,start_method,role,module,m,no_python,run_path,log_dir,r,redirects,t,tee,node_rank,master_addr,master_port""".split(
+    ","
+)
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/623
diff --git a/examples/README.md b/examples/README.md
index 150afb75c..c4f050b3f 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -136,7 +136,7 @@ To run it in each of these various modes, use the following commands:
         ```
 - single GPU:
     ```bash
-    python ./nlp_example.py  # from a server with a GPU
+    python ./cv_example.py  # from a server with a GPU
     ```
 - with fp16 (mixed-precision)
     * from any server by passing `fp16=True` to the `Accelerator`.

Thanks for the fix!
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/622
diff --git a/docs/source/usage_guides/gradient_accumulation.mdx b/docs/source/usage_guides/gradient_accumulation.mdx
index f9f5dc62a..bd06efed9 100644
--- a/docs/source/usage_guides/gradient_accumulation.mdx
+++ b/docs/source/usage_guides/gradient_accumulation.mdx
@@ -31,7 +31,6 @@ model.to(device)
 gradient_accumulation_steps = 2
 
 for index, batch in enumerate(training_dataloader):
-    optimizer.zero_grad()
     inputs, targets = batch
     inputs = inputs.to(device)
     targets = targets.to(device)
@@ -42,6 +41,7 @@ for index, batch in enumerate(training_dataloader):
     if (index + 1) % gradient_accumulation_steps == 0:
         optimizer.step()
         scheduler.step()
+        optimizer.zero_grad()
 ```
 
 ## Converting it to 🤗 Accelerate
@@ -57,7 +57,6 @@ First the code shown earlier will be converted to utilize 🤗 Accelerate withou
 + )
 
   for index, batch in enumerate(training_dataloader):
-      optimizer.zero_grad()
       inputs, targets = batch
 -     inputs = inputs.to(device)
 -     targets = targets.to(device)
@@ -68,6 +67,7 @@ First the code shown earlier will be converted to utilize 🤗 Accelerate withou
       if (index+1) % gradient_accumulation_steps == 0:
           optimizer.step()
           scheduler.step()
+          optimizer.zero_grad()
 ```
 
 <Tip warning={true}>
@@ -94,7 +94,6 @@ You just wrap it around the entire training part of our code:
 - for index, batch in enumerate(training_dataloader):
 + for batch in training_dataloader:
 +     with accelerator.accumulate(model):
-          optimizer.zero_grad()
           inputs, targets = batch
           outputs = model(inputs)
 ```
@@ -107,6 +106,7 @@ You can remove all the special checks for the step number and the loss adjustmen
 - if (index+1) % gradient_accumulation_steps == 0:
   optimizer.step()
   scheduler.step()
+  optimizer.zero_grad()
 ```
 
 As you can see the [`Accelerator`] is able to keep track of the batch number you are on and it will automatically know whether to step through the prepared optimizer and how to adjust the loss. 
@@ -118,11 +118,11 @@ Below is the finished implementation for performing gradient accumulation with 
 ```python
 for batch in training_dataloader:
     with accelerator.accumulate(model):
-        optimizer.zero_grad()
         inputs, targets = batch
         outputs = model(inputs)
         loss = loss_function(outputs, targets)
         accelerator.backward(loss)
         optimizer.step()
         scheduler.step()
+        optimizer.zero_grad()
 ```

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/620
diff --git a/examples/README.md b/examples/README.md
index 271ffd3c5..150afb75c 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -184,6 +184,10 @@ To run it in each of these various modes, use the following commands:
     * In PyTorch:
         Add an `xmp.spawn` line in your script as you usually do.
 
+### Simple vision example (GANs)
+
+- [huggan project](https://github.com/huggingface/community-events/tree/main/huggan)
+    
 ## Finer Examples
 
 While the first two scripts are extremely barebones when it comes to what you can do with accelerate, more advanced features are documented in two other locations.

```suggestion
- [huggan project](https://github.com/huggingface/community-events/tree/main/huggan)
```
@@ -184,6 +184,10 @@ To run it in each of these various modes, use the following commands:
     * In PyTorch:
         Add an `xmp.spawn` line in your script as you usually do.
 
+### Simple vision example (GANs)
+
+- Please see: https://github.com/huggingface/community-events/tree/main/huggan
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/619
diff --git a/setup.py b/setup.py
index 8247aa5fe..4bc977029 100644
--- a/setup.py
+++ b/setup.py
@@ -50,7 +50,7 @@
         ]
     },
     python_requires=">=3.7.0",
-    install_requires=["numpy>=1.17", "packaging>=20.0", "psutil", "pyyaml", "torch>=1.4.0"],
+    install_requires=["numpy>=1.17", "packaging>=20.0", "psutil", "pyyaml", "torch>=1.4.0", "rich"],
     extras_require=extras,
     classifiers=[
         "Development Status :: 5 - Production/Stable",
diff --git a/src/accelerate/__init__.py b/src/accelerate/__init__.py
index 145b0aa27..03a7198ac 100644
--- a/src/accelerate/__init__.py
+++ b/src/accelerate/__init__.py
@@ -17,5 +17,6 @@
     find_executable_batch_size,
     infer_auto_device_map,
     load_checkpoint_in_model,
+    rich,
     synchronize_rng_states,
 )
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index fcb26228a..ea27afb75 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -46,11 +46,15 @@
 )
 from accelerate.utils.constants import DEEPSPEED_MULTINODE_LAUNCHERS
 from accelerate.utils.dataclasses import SageMakerDistributedType
+from rich import get_console
+from rich.logging import RichHandler
 
 
 if is_torch_version(">=", "1.9.0"):
     import torch.distributed.run as distrib_run
 
+FORMAT = "%(message)s"
+logging.basicConfig(format=FORMAT, datefmt="[%X]", handlers=[RichHandler()])
 
 logger = logging.getLogger(__name__)
 
@@ -301,7 +305,12 @@ def launch_command_parser(subparsers=None):
         "--aws_secret_access_key",
         type=str,
         default=None,
-        help="The AWS_SECRET_ACCESS_KEY used to launch the Amazon SageMaker training job",
+        help="The AWS_SECRET_ACCESS_KEY used to launch the Amazon SageMaker training job.",
+    )
+    parser.add_argument(
+        "--debug",
+        action="store_true",
+        help="Whether to print out the torch.distributed stack trace when something fails.",
     )
     parser.add_argument(
         "training_script",
@@ -443,9 +452,17 @@ def multi_gpu_launcher(args):
             current_env["FSDP_STATE_DICT_TYPE"] = str(args.fsdp_state_dict_type)
     current_env["OMP_NUM_THREADS"] = str(args.num_cpu_threads_per_process)
     if is_torch_version(">=", "1.9.0"):
+        debug = getattr(args, "debug", False)
         args = _filter_args(args)
         with patch_environment(**current_env):
-            distrib_run.run(args)
+            console = get_console()
+
+            try:
+                distrib_run.run(args)
+            except:
+                if debug:
+                    console.print("\n[bold red]Using --debug, `torch.distributed` Stack Trace:[/bold red]")
+                    console.print_exception(suppress=[__file__], show_locals=False)
     else:
         # We still have to use subprocess, the user won't get a clean traceback as a result
         cmd = get_launch_prefix()
diff --git a/src/accelerate/utils/launch.py b/src/accelerate/utils/launch.py
index 1a452ca63..31d71ad21 100644
--- a/src/accelerate/utils/launch.py
+++ b/src/accelerate/utils/launch.py
@@ -23,8 +23,6 @@
 
 if is_torch_version(">=", "1.9.0"):
     import torch.distributed.run as distrib_run
-else:
-    import torch.distributed.launch as distrib_run
 
 
 def get_launch_prefix():
diff --git a/src/accelerate/utils/rich.py b/src/accelerate/utils/rich.py
new file mode 100644
index 000000000..7a54c8b5a
--- /dev/null
+++ b/src/accelerate/utils/rich.py
@@ -0,0 +1,18 @@
+# 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.
+
+from rich.traceback import install
+
+
+install(show_locals=False)

I think this is a leftover of your tests?
@@ -97,6 +98,8 @@ def collate_fn(examples):
 def training_function(config, args):
     # Initialize accelerator
     accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
+    if accelerator.process_index == 1:
+        raise ValueError()
You put `rich` in the core install here but use an `is_rich_available` later on.
@@ -50,7 +50,7 @@
         ]
     },
     python_requires=">=3.7.0",
-    install_requires=["numpy>=1.17", "packaging>=20.0", "psutil", "pyyaml", "torch>=1.4.0"],
+    install_requires=["numpy>=1.17", "packaging>=20.0", "psutil", "pyyaml", "torch>=1.4.0", "rich"],
Should be its own list in a constant probably, so it's easier to read/add/see diff
@@ -444,10 +441,29 @@ def multi_gpu_launcher(args):
         if args.fsdp_state_dict_type is not None:
             current_env["FSDP_STATE_DICT_TYPE"] = str(args.fsdp_state_dict_type)
     current_env["OMP_NUM_THREADS"] = str(args.num_cpu_threads_per_process)
-    process = subprocess.Popen(cmd, env=current_env)
-    process.wait()
-    if process.returncode != 0:
-        raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)
+    # use torchrun directly here and try to call it
+
+    with patch_environment(**current_env):
+        import torch.distributed.run as distrib_run
+        debug = getattr(args, "debug", False)
+        keys = """nnodes,nproc_per_node,rdzv_backend,rdzv_endpoint,rdzv_id,rdzv_conf,standalone,max_restarts,monitor_interval,start_method,role,module,m,no_python,run_path,log_dir,r,redirects,t,tee,node_rank,master_addr,master_port,training_script,training_script_args""".split(",")
Thanks, fixed this now! 🤗 
@@ -97,6 +98,8 @@ def collate_fn(examples):
 def training_function(config, args):
     # Initialize accelerator
     accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
+    if accelerator.process_index == 1:
+        raise ValueError()
As discussed offline, opted to make rich a true dep in Accelerate
@@ -50,7 +50,7 @@
         ]
     },
     python_requires=">=3.7.0",
-    install_requires=["numpy>=1.17", "packaging>=20.0", "psutil", "pyyaml", "torch>=1.4.0"],
+    install_requires=["numpy>=1.17", "packaging>=20.0", "psutil", "pyyaml", "torch>=1.4.0", "rich"],
_The documentation is not available anymore as the PR was closed or merged._
Nice utility! LGTM 🤗, I think Sylvain's suggestion would be a really good way to use this feature. Also, I think this won't result in rich traceback if the issue was in creating an `Accelerate` object itself, right? if so, should that be mentioned as part of the documentation?
<s>In Multi GPU launcher, if one isn't using `--debug` then they won't see any error stack trace at all (which might be hard to read but important for root cause and resolution). This makes `--debug` sort of mandatory instead of optional. Can the user have normal error traceback if `--debug` isn't given? Please let me know if I am missing something or way off. </s>

Went over it again with a sample code, 
```python
from rich.traceback import 
install install(show_locals=True)
``` 
makes rich the default traceback handler. `--debug` meaning is no more what Sylvain was referring to, hence my confusion. 
https://api.github.com/repos/huggingface/accelerate/pulls/613
diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml
index c0c9c62f2..cb0a7ee42 100644
--- a/.github/workflows/nightly.yml
+++ b/.github/workflows/nightly.yml
@@ -57,10 +57,16 @@ jobs:
           git fetch && git checkout ${{ github.sha }}
           pip install -e . --no-deps
 
-      - name: Run test on GPUs
+      - name: Run core and big modeling tests on GPUs
         run: |
           source activate accelerate
-          make test
+          make test_big_modeling
+          make test_core
+
+      - name: Run Integration tests on GPUs
+        run: |
+          source activate accelerate
+          make test_integrations
 
       - name: Run examples on GPUs
         run: |
diff --git a/Makefile b/Makefile
index 100a1484a..c1c79d6e8 100644
--- a/Makefile
+++ b/Makefile
@@ -44,6 +44,8 @@ test_examples:
 	python -m pytest -s -v ./tests/test_examples.py
 
 # Broken down example tests for the CI runners
+test_integrations:
+	python -m pytest -s -v ./tests/deepspeed ./tests/fsdp
 test_example_differences:
 	python -m pytest -s -v ./tests/test_examples.py::ExampleDifferenceTests
 
diff --git a/tests/deepspeed/test_deepspeed.py b/tests/deepspeed/test_deepspeed.py
index 30eb2a4d8..1be794421 100644
--- a/tests/deepspeed/test_deepspeed.py
+++ b/tests/deepspeed/test_deepspeed.py
@@ -35,7 +35,6 @@
     require_cuda,
     require_deepspeed,
     require_multi_gpu,
-    skip,
     slow,
 )
 from accelerate.test_utils.training import RegressionDataset
@@ -697,7 +696,6 @@ def test_performance(self):
             with patch_environment(omp_num_threads=1):
                 execute_subprocess_async(cmd_stage, env=os.environ.copy())
 
-    @skip
     def test_checkpointing(self):
         self.test_file_path = os.path.join(self.test_scripts_folder, "test_checkpointing.py")
         cmd = [

_The documentation is not available anymore as the PR was closed or merged._
Thank you 😄. This is neat and finally all the tests would run as expected 🤗. 
https://api.github.com/repos/huggingface/accelerate/pulls/612
diff --git a/.github/workflows/check_dependencies.yml b/.github/workflows/build_and_run_tests.yml
similarity index 82%
rename from .github/workflows/check_dependencies.yml
rename to .github/workflows/build_and_run_tests.yml
index f8ac1f492..a5378353d 100644
--- a/.github/workflows/check_dependencies.yml
+++ b/.github/workflows/build_and_run_tests.yml
@@ -10,7 +10,7 @@ env:
   GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
 
 jobs:
-  check-for-setup:
+  check-for-source:
     runs-on: ubuntu-latest
     name: Check if setup was changed
     outputs:
@@ -28,7 +28,7 @@ jobs:
         id: was_changed
         run: |
           for file in ${{ steps.changed-files.outputs.all_changed_files }}; do
-            if [ `basename "${file}"` = "setup.py" ]; then
+            if [ `basename "${file}"` == "setup.py" ]; then
               echo ::set-output name=changed::"1"
             fi
           done
@@ -36,10 +36,10 @@ jobs:
   build-docker-containers:
     needs: check-for-setup
     if: (github.event_name == 'push') && (needs.check-for-setup.outputs.changed == '1')
-    uses: ./.github/workflows/build-docker-images.yml
+    uses: ./.github/workflows/build_docker_images.yml
     secrets: inherit
 
-  run-tests:
+  run-merge-tests:
     needs: build-docker-containers
     if: always()
-    uses: ./.github/workflows/on-merge.yml
\ No newline at end of file
+    uses: ./.github/workflows/run_merge_tests.yml
\ No newline at end of file
diff --git a/.github/workflows/build-docker-images.yml b/.github/workflows/build_docker_images.yml
similarity index 100%
rename from .github/workflows/build-docker-images.yml
rename to .github/workflows/build_docker_images.yml
diff --git a/.github/workflows/on-merge.yml b/.github/workflows/run_merge_tests.yml
similarity index 100%
rename from .github/workflows/on-merge.yml
rename to .github/workflows/run_merge_tests.yml
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index d39d36951..5d8ff0c1b 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -1,6 +1,13 @@
 name: Run Tests
 
-on: [pull_request]
+on:
+  pull_request:
+    paths:
+      - "src/**"
+      - "tests/**"
+      - ".github/**"
+      - "setup.py"
+    types: [opened, synchronize, reopened]
 
 env:
   HF_HOME: ~/hf_cache

I feel this name should be `Checks if setup.py were changed`, as I can only see it checks that file.

And on merge event to `main`, it seems the tests are always run.
@@ -10,9 +10,9 @@ env:
   GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
 
 jobs:
-  check-for-setup:
+  check-for-source:
     runs-on: ubuntu-latest
-    name: Check if setup was changed
+    name: Checks if source code or tests were changed
_The documentation is not available anymore as the PR was closed or merged._
So the purpose of `check-for-source` in `build_and_run_tests.yml` and `test.yml` are different?

- on PR request: check source file, and run the tests only if any changed found.
- on merge (to main): check `setup.py` to build the image if necessary. Then always run the tests?
I am probably wrong above, as I see you have `needs: build-docker-containers`. 
But I am confused a bit: currently (without this PR), what happens (on merge) if `setup.py` is not changed, but some `.py` files change?

I will review in more detail tomorrow.

@ydshieh if some `.py` was changed but not the setup then the tests are ran but Docker images won't be rebuilt
Thanks, @muellerzr . So in this PR, what happens (on merge) if no `.py` file changes at all (`setup.py` not change + no other `.py` change)? If I understand, the tests still run. Is this the expected behavior?
My only remaining question is as above: on merge, should we skip the tests if no `.py` file changes at all. Otherwise, all LGTM.
@ydshieh (sorry this got buried in my priorities):

Re: on merge, yes it is because we never know if some other dependency has a regression issue or what. So we're still made aware of when this happens post merge and can immediately know it's unrelated to something directly in that PR without it silently not failing until the nightlies
https://api.github.com/repos/huggingface/accelerate/pulls/611
diff --git a/docs/source/index.mdx b/docs/source/index.mdx
index 3f0f58802..1664fd7a2 100644
--- a/docs/source/index.mdx
+++ b/docs/source/index.mdx
@@ -55,7 +55,7 @@ accelerate launch {my_script.py}
       ><div class="w-full text-center bg-gradient-to-br from-blue-400 to-blue-500 rounded-lg py-1.5 font-semibold mb-5 text-white text-lg leading-relaxed">Tutorials</div>
       <p class="text-gray-700">Learn the basics and become familiar with using 🤗 Accelerate. Start here if you are using 🤗 Accelerate for the first time!</p>
     </a>
-    <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="/docs/accelerate/utility_guides/gradient_accumulation"
+    <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="/docs/accelerate/usage_guides/gradient_accumulation"
       ><div class="w-full text-center bg-gradient-to-br from-indigo-400 to-indigo-500 rounded-lg py-1.5 font-semibold mb-5 text-white text-lg leading-relaxed">How-to guides</div>
       <p class="text-gray-700">Practical guides to help you achieve a specific goal. Take a look at these guides to learn how to use 🤗 Accelerate to solve real-world problems.</p>
     </a>
@@ -68,4 +68,4 @@ accelerate launch {my_script.py}
       <p class="text-gray-700">Technical descriptions of how 🤗 Accelerate classes and methods work.</p>
     </a>
   </div>
-</div>
\ No newline at end of file
+</div>

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/610
diff --git a/src/accelerate/utils/operations.py b/src/accelerate/utils/operations.py
index 42868a0a5..3ba46a03f 100644
--- a/src/accelerate/utils/operations.py
+++ b/src/accelerate/utils/operations.py
@@ -103,7 +103,7 @@ def recursively_apply(func, data, *args, test_type=is_torch_tensor, error_on_oth
     return data
 
 
-def send_to_device(tensor, device):
+def send_to_device(tensor, device, non_blocking=False):
     """
     Recursively sends the elements in a nested list/tuple/dictionary of tensors to a given device.
 
@@ -117,13 +117,16 @@ def send_to_device(tensor, device):
         The same data structure as `tensor` with all tensors sent to the proper device.
     """
 
-    def _send_to_device(t, device):
-        return t.to(device)
+    def _send_to_device(t, device, non_blocking):
+        try:
+            return t.to(device, non_blocking=non_blocking)
+        except TypeError:  # .to() doesn't accept non_blocking as kwarg
+            return t.to(device)
 
     def _has_to_method(t):
         return hasattr(t, "to")
 
-    return recursively_apply(_send_to_device, tensor, device, test_type=_has_to_method)
+    return recursively_apply(_send_to_device, tensor, device, non_blocking, test_type=_has_to_method)
 
 
 def get_data_structure(data):

So far all objects in Pytorch which accept `.to()` method accept the `non_blocking` kwarg or at least kwargs in general, but are we sure this is going to be always the case?

- [torch.jit.ScriptModule.to](https://pytorch.org/docs/stable/generated/torch.jit.ScriptModule.html?highlight=#torch.jit.ScriptModule.to) (Python method, in ScriptModule)
- [torch.nn.Module.to](https://pytorch.org/docs/stable/generated/torch.nn.Module.html?highlight=#torch.nn.Module.to) (Python method, in Module)
- [torch.nn.utils.rnn.PackedSequence.to](https://pytorch.org/docs/stable/generated/torch.nn.utils.rnn.PackedSequence.html?highlight=#torch.nn.utils.rnn.PackedSequence.to) (Python method, in PackedSequence)
- [torch.Tensor.to](https://pytorch.org/docs/stable/generated/torch.Tensor.to.html?highlight=#torch.Tensor.to) (Python method, in torch.Tensor.to)

cc @sgugger 
@@ -117,13 +117,13 @@ def send_to_device(tensor, device):
         The same data structure as `tensor` with all tensors sent to the proper device.
     """
 
-    def _send_to_device(t, device):
-        return t.to(device)
+    def _send_to_device(t, device, non_blocking):
+        return t.to(device, non_blocking=non_blocking)
 
     def _has_to_method(t):
         return hasattr(t, "to")
This is for our own objects FYI, `BatchEncoding` in Transformers, which does not accept the `non_blocking` argument :-/
Maybe we can inspect the signature and only pass it when we see it?
@@ -117,13 +117,13 @@ def send_to_device(tensor, device):
         The same data structure as `tensor` with all tensors sent to the proper device.
     """
 
-    def _send_to_device(t, device):
-        return t.to(device)
+    def _send_to_device(t, device, non_blocking):
+        return t.to(device, non_blocking=non_blocking)
 
     def _has_to_method(t):
         return hasattr(t, "to")
I tried using `t.to.__code__` and `inspect.signature(t.to)` and both give errors. Do you have a better way to get kwargs of a `builtin_function_or_method` @sgugger ?
@@ -117,13 +117,13 @@ def send_to_device(tensor, device):
         The same data structure as `tensor` with all tensors sent to the proper device.
     """
 
-    def _send_to_device(t, device):
-        return t.to(device)
+    def _send_to_device(t, device, non_blocking):
+        return t.to(device, non_blocking=non_blocking)
 
     def _has_to_method(t):
         return hasattr(t, "to")
Mmm, `inspect.signature(t.to)` should give you access to the keyword arguments. You can find them in the `.parameters` attribute, which will be a dict-like object.
@@ -117,13 +117,13 @@ def send_to_device(tensor, device):
         The same data structure as `tensor` with all tensors sent to the proper device.
     """
 
-    def _send_to_device(t, device):
-        return t.to(device)
+    def _send_to_device(t, device, non_blocking):
+        return t.to(device, non_blocking=non_blocking)
 
     def _has_to_method(t):
         return hasattr(t, "to")
I get this error when trying `inspect.signature(t.to)` where `t` is a tensor
```
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/home/nouamane/miniconda3/envs/hf/lib/python3.9/inspect.py", line 3111, in signature
    return Signature.from_callable(obj, follow_wrapped=follow_wrapped)
  File "/home/nouamane/miniconda3/envs/hf/lib/python3.9/inspect.py", line 2860, in from_callable
    return _signature_from_callable(obj, sigcls=cls,
  File "/home/nouamane/miniconda3/envs/hf/lib/python3.9/inspect.py", line 2327, in _signature_from_callable
    return _signature_from_builtin(sigcls, obj,
  File "/home/nouamane/miniconda3/envs/hf/lib/python3.9/inspect.py", line 2145, in _signature_from_builtin
    raise ValueError("no signature found for builtin {!r}".format(func))
ValueError: no signature found for builtin <built-in method to of Tensor object at 0x7f7e98c57720>
```
@@ -117,13 +117,13 @@ def send_to_device(tensor, device):
         The same data structure as `tensor` with all tensors sent to the proper device.
     """
 
-    def _send_to_device(t, device):
-        return t.to(device)
+    def _send_to_device(t, device, non_blocking):
+        return t.to(device, non_blocking=non_blocking)
 
     def _has_to_method(t):
         return hasattr(t, "to")
I think this is related to the warning in [inspect's docs](https://docs.python.org/3/library/inspect.html#inspect.signature)

> Note: Some callables may not be introspectable in certain implementations of Python. For example, in CPython, some built-in functions defined in C provide no metadata about their arguments.
@@ -117,13 +117,13 @@ def send_to_device(tensor, device):
         The same data structure as `tensor` with all tensors sent to the proper device.
     """
 
-    def _send_to_device(t, device):
-        return t.to(device)
+    def _send_to_device(t, device, non_blocking):
+        return t.to(device, non_blocking=non_blocking)
 
     def _has_to_method(t):
         return hasattr(t, "to")
Maybe we can isolate tensors in the function then? Something like:
```
if isinstance(t, torch.Tensor):
    return t.to(device, non_blocking=non_blocking)
else:
    # inspect signature here in a try/except block
```
@@ -117,13 +117,13 @@ def send_to_device(tensor, device):
         The same data structure as `tensor` with all tensors sent to the proper device.
     """
 
-    def _send_to_device(t, device):
-        return t.to(device)
+    def _send_to_device(t, device, non_blocking):
+        return t.to(device, non_blocking=non_blocking)
 
     def _has_to_method(t):
         return hasattr(t, "to")
Wouldn't this be enough, or should I inspect signature?
  ```python
  def _send_to_device(t, device, non_blocking):
        if isinstance(t, torch.Tensor):
            return t.to(device, non_blocking=non_blocking)
        else:
            try:
                return t.to(device, non_blocking=non_blocking)
            except TypeError:
                return t.to(device)
```
@@ -117,13 +117,13 @@ def send_to_device(tensor, device):
         The same data structure as `tensor` with all tensors sent to the proper device.
     """
 
-    def _send_to_device(t, device):
-        return t.to(device)
+    def _send_to_device(t, device, non_blocking):
+        return t.to(device, non_blocking=non_blocking)
 
     def _has_to_method(t):
         return hasattr(t, "to")
Oh in this case, you don't even need the first test. Are you sure about the type of `TypeError`?
@@ -117,13 +117,13 @@ def send_to_device(tensor, device):
         The same data structure as `tensor` with all tensors sent to the proper device.
     """
 
-    def _send_to_device(t, device):
-        return t.to(device)
+    def _send_to_device(t, device, non_blocking):
+        return t.to(device, non_blocking=non_blocking)
 
     def _has_to_method(t):
         return hasattr(t, "to")
Searching on google for `python unexpected argument`,  it's usually `TypeError: ... got an unexpected keyword`. I'm not sure if it's always the case though
@@ -117,13 +117,13 @@ def send_to_device(tensor, device):
         The same data structure as `tensor` with all tensors sent to the proper device.
     """
 
-    def _send_to_device(t, device):
-        return t.to(device)
+    def _send_to_device(t, device, non_blocking):
+        return t.to(device, non_blocking=non_blocking)
 
     def _has_to_method(t):
         return hasattr(t, "to")
Seems to work for C defined methods as well:
```python
>> t.__abs__(a=5)
Traceback (most recent call last):
  File "<string>", line 1, in <module>
TypeError: Tensor.abs() takes no keyword arguments
```


```python
>> t.to(a=5)
Traceback (most recent call last):
  File "<string>", line 1, in <module>
TypeError: to() received an invalid combination of arguments - got (a=int, ), but expected one of:
 * (torch.device device, torch.dtype dtype, bool non_blocking, bool copy, *, torch.memory_format memory_format)
 * (torch.dtype dtype, bool non_blocking, bool copy, *, torch.memory_format memory_format)
 * (Tensor tensor, bool non_blocking, bool copy, *, torch.memory_format memory_format)
```
@@ -117,13 +117,13 @@ def send_to_device(tensor, device):
         The same data structure as `tensor` with all tensors sent to the proper device.
     """
 
-    def _send_to_device(t, device):
-        return t.to(device)
+    def _send_to_device(t, device, non_blocking):
+        return t.to(device, non_blocking=non_blocking)
 
     def _has_to_method(t):
         return hasattr(t, "to")
Nice!
@@ -117,13 +117,13 @@ def send_to_device(tensor, device):
         The same data structure as `tensor` with all tensors sent to the proper device.
     """
 
-    def _send_to_device(t, device):
-        return t.to(device)
+    def _send_to_device(t, device, non_blocking):
+        return t.to(device, non_blocking=non_blocking)
 
     def _has_to_method(t):
         return hasattr(t, "to")
@sgugger the docs say that `tensor` is a torch.Tensor here, which isn't necessary the case iiuc. We can probably update the docs as well
@@ -103,7 +103,7 @@ def recursively_apply(func, data, *args, test_type=is_torch_tensor, error_on_oth
     return data
 
 
-def send_to_device(tensor, device):
+def send_to_device(tensor, device, non_blocking=False):
     """
     Recursively sends the elements in a nested list/tuple/dictionary of tensors to a given device.
True
@@ -103,7 +103,7 @@ def recursively_apply(func, data, *args, test_type=is_torch_tensor, error_on_oth
     return data
 
 
-def send_to_device(tensor, device):
+def send_to_device(tensor, device, non_blocking=False):
     """
     Recursively sends the elements in a nested list/tuple/dictionary of tensors to a given device.
We could leave that for another PR 
@@ -103,7 +103,7 @@ def recursively_apply(func, data, *args, test_type=is_torch_tensor, error_on_oth
     return data
 
 
-def send_to_device(tensor, device):
+def send_to_device(tensor, device, non_blocking=False):
     """
     Recursively sends the elements in a nested list/tuple/dictionary of tensors to a given device.
_The documentation is not available anymore as the PR was closed or merged._
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/accelerate/blob/main/CONTRIBUTING.md) are likely to be ignored.
Hi @sgugger, was there a reason that this was closed and not merged? I came across this when I was about to request the same thing. 
@Chris-hughes10 it was automatically closed by the stale bot due to the absence of activity. @NouamaneTazi is still investigating whether this could have some negative impact sometimes.
### Should we make .to(non_blocking=True) the default in **accelerate**?
Let’s use this script to find out

```python
import torch
if __name__ == '__main__':
    seed = 0
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    stream = torch.cuda.current_stream()

    x = torch.rand(32, 256, 220, 220).cuda()

    t = (x.min() - x.max()).to(torch.device("cpu"), non_blocking=False) # try with False then True
    # t = (x.min() - x.max()).to("cuda:1", non_blocking=True) # try GPU0 to GPU1 copy

    print(stream.query()) # False - Checks if all the work submitted has been completed.
    print(t)
    stream.synchronize() # wait for stream to finish the work
    
    print(stream.query()) # True - work done
    print(t)
```

<details>
<summary>Copy to CPU with non_blocking=False   (default)</summary>

In this case `.to()`  adds a `cudaStreamSynchronize` op which makes the CPU use the correct value of the tensor when printing
![image](https://user-images.githubusercontent.com/29777165/194038267-91033b35-fc18-4a0d-9af1-cfb1e382721e.png)

</details>

<details>
<summary>Copy to CPU with non_blocking=True  </summary>


In this case the CPU submits the kernels for `.to()`  to the GPU then moves on to perform the print operation which uses an incorrect value for the tensor `tensor(0.)`  **(The dangerous part)** 
![image](https://user-images.githubusercontent.com/29777165/194038545-88833ffc-5141-4fd4-a503-c140202f0bed.png)

</details>

<details>
<summary>Copy to another GPU with non_blocking=True </summary>

It seems that the `non_blocking` here doesn’t do much (we get basically the same thing using `non_blocking=True` ). In both cases we have GPU 1 waiting for GPU 0 to finish working on the tensor, and THEN copy it to GPU 1. And finally the CPU prints the tensor that’s now located on GPU 1
In this case `.to()`  creates a `cudaStreamWaitEvent` event (figure 2) which makes GPU 1 waits for GPU 0. I made [an issue](https://discuss.pytorch.org/t/to-cuda-1-non-blocking-true-creates-cudastreamwaitevent/162296) on Pytorch’s forums to investigate why is this the case

![image](https://user-images.githubusercontent.com/29777165/194038566-56a7c178-d659-4c3f-a759-d94da8c0f152.png)

![image](https://user-images.githubusercontent.com/29777165/194038587-5a089b5c-af42-4bab-86b1-516e2ebe3300.png)

</details>

**tldr;** `non_blocking` could be a game changer in using your GPUs efficiently. 
- **Good use scenario:** Copying your data from CPU to GPU in a non_blocking way then running your model which exists on GPU (this would make the CPU launch the copy kernel, then moves on to queuing other kernels in your model on the GPU. As opposed to waiting for the copy to end, and only then launching kernels from the model). [Example from Pytorch's repo](https://github.com/pytorch/examples/blob/main/imagenet/main.py#L321-L326).
- **Bad use scenario:** Copying your data from CPU to GPU in a non_blocking way then start some operations **on CPU** that would use the non-ready tensors. (could be in `if` statements, or simple arithmetics...)


=> It’s good to support that argument in accelerate but it’s better to keep the default as it is, just like it’s the case in Pytorch 
@NouamaneTazi These are really interesting insights! If you are looking at inspecting signatures, do you think it would be too much complexity to set non-blocking automatically based on whether it is CPU -> GPU transfer or GPU -> CPU transfer by inspecting the tensor's current device?

If you like the idea, perhaps this would be a different feature though, toggled by an accelerator flag?
@Chris-hughes10  CPU -> GPU and GPU -> CPU both lead to the same issues as mentioned above. Only GPU -> GPU is the safe operation but as I said above, it seems that it requires the two GPUs synchronization whether we set `non_blocking=True` or not
> @Chris-hughes10 CPU -> GPU and GPU -> CPU both lead to the same issues as mentioned above. Only GPU -> GPU is the safe operation but as I said above, it seems that it requires the two GPUs synchronization whether we set `non_blocking=True` or not

Apologies, I misread the post above. Please ignore my previous suggestion!
https://api.github.com/repos/huggingface/accelerate/pulls/607
diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
index 188800527..467a757b4 100644
--- a/src/accelerate/utils/modeling.py
+++ b/src/accelerate/utils/modeling.py
@@ -677,7 +677,7 @@ def load_checkpoint_in_model(
         else:
             for param_name, param in checkpoint.items():
                 module_name = param_name
-                if dtype is not None:
+                if dtype is not None and not str(param.dtype).startswith(("torch.uint", "torch.int", "torch.bool")):
                     param = param.to(dtype)
                 while len(module_name) > 0 and module_name not in device_map:
                     module_name = ".".join(module_name.split(".")[:-1])

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/606
diff --git a/src/accelerate/tracking.py b/src/accelerate/tracking.py
index 673975fc7..839a7a2a9 100644
--- a/src/accelerate/tracking.py
+++ b/src/accelerate/tracking.py
@@ -209,7 +209,7 @@ def __init__(self, run_name: str, **kwargs):
 
     @property
     def tracker(self):
-        return self.run.run
+        return self.run
 
     def store_init_configuration(self, values: dict):
         """

_The documentation is not available anymore as the PR was closed or merged._
Thanks!
https://api.github.com/repos/huggingface/accelerate/pulls/605
diff --git a/tests/deepspeed/test_deepspeed.py b/tests/deepspeed/test_deepspeed.py
index a8c1aedac..30eb2a4d8 100644
--- a/tests/deepspeed/test_deepspeed.py
+++ b/tests/deepspeed/test_deepspeed.py
@@ -35,6 +35,7 @@
     require_cuda,
     require_deepspeed,
     require_multi_gpu,
+    skip,
     slow,
 )
 from accelerate.test_utils.training import RegressionDataset
@@ -642,7 +643,9 @@ def setUp(self):
             "deepspeed_stage_1_fp16": 1600,
             "deepspeed_stage_2_fp16": 2500,
             "deepspeed_stage_3_zero_init_fp16": 2800,
-            "deepspeed_stage_3_cpu_offload_fp16": 1900,
+            # Disabling below test as it overwhelms the RAM memory usage
+            # on CI self-hosted runner leading to tests getting killed.
+            # "deepspeed_stage_3_cpu_offload_fp16": 1900,
         }
         self.n_train = 160
         self.n_val = 160
@@ -694,6 +697,7 @@ def test_performance(self):
             with patch_environment(omp_num_threads=1):
                 execute_subprocess_async(cmd_stage, env=os.environ.copy())
 
+    @skip
     def test_checkpointing(self):
         self.test_file_path = os.path.join(self.test_scripts_folder, "test_checkpointing.py")
         cmd = [
diff --git a/tests/fsdp/test_fsdp.py b/tests/fsdp/test_fsdp.py
index 628e377c3..249d2b692 100644
--- a/tests/fsdp/test_fsdp.py
+++ b/tests/fsdp/test_fsdp.py
@@ -191,7 +191,9 @@ def setUp(self):
             "multi_gpu_fp16": 3200,
             "fsdp_shard_grad_op_transformer_based_wrap_fp16": 2000,
             "fsdp_full_shard_transformer_based_wrap_fp16": 1900,
-            "fsdp_full_shard_cpu_offload_transformer_based_wrap_fp32": 1500,  # fp16 was leading to indefinite hang
+            # Disabling below test as it overwhelms the RAM memory usage
+            # on CI self-hosted runner leading to tests getting killed.
+            # "fsdp_full_shard_cpu_offload_transformer_based_wrap_fp32": 1500,  # fp16 was leading to indefinite hang
         }
         self.n_train = 160
         self.n_val = 160

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/604
diff --git a/docs/source/concept_guides/training_tpu.mdx b/docs/source/concept_guides/training_tpu.mdx
index 32736ce31..7fe54b14a 100644
--- a/docs/source/concept_guides/training_tpu.mdx
+++ b/docs/source/concept_guides/training_tpu.mdx
@@ -89,7 +89,7 @@ like:
 ProcessExitedException: process 0 terminated with signal SIGSEGV
 ```
 
-This error is *extremely* cryptic but the basic explaination is you ran out of system RAM. You can avoid this entirely by reconfiguring the training function to 
+This error is *extremely* cryptic but the basic explanation is you ran out of system RAM. You can avoid this entirely by reconfiguring the training function to 
 accept a single `model` argument, and declare it in an outside cell:
 
 ```python
@@ -137,7 +137,7 @@ accelerator = Accelerator(mixed_precision="bf16")
 By default this will cast `torch.float` and `torch.double` to `bfloat16` on TPUs. 
 The specific configuration being set is an environmental variable of `XLA_USE_BF16` is set to `1`.
 
-There is a futher configuration you can perform which is setting the `XLA_DOWNCAST_BF16` environmental variable. If set to `1`, then 
+There is a further configuration you can perform which is setting the `XLA_DOWNCAST_BF16` environmental variable. If set to `1`, then 
 `torch.float` is `bfloat16` and `torch.double` is `float32`.
 
 This is performed in the `Accelerator` object when passing `downcast_bf16=True`:
diff --git a/docs/source/quicktour.mdx b/docs/source/quicktour.mdx
index 812b20226..6d92e7958 100644
--- a/docs/source/quicktour.mdx
+++ b/docs/source/quicktour.mdx
@@ -391,7 +391,7 @@ and [`~Accelerator.clip_grad_value_`] respectively.
 
 If you are running your training in Mixed Precision with 🤗 Accelerate, you will get the best result with your loss being
 computed inside your model (like in Transformer models for instance). Every computation outside of the model will be
-executed in full precision (which is generally what you want for loss computation, expecially if it involves a
+executed in full precision (which is generally what you want for loss computation, especially if it involves a
 softmax). However you might want to put your loss computation inside the *accelerator.autocast* context manager:
 
 ```
diff --git a/docs/source/usage_guides/big_modeling.mdx b/docs/source/usage_guides/big_modeling.mdx
index 243b354d0..9b3252809 100644
--- a/docs/source/usage_guides/big_modeling.mdx
+++ b/docs/source/usage_guides/big_modeling.mdx
@@ -213,7 +213,7 @@ You can let 🤗 Accelerate handle the device map computation by setting `device
 
 </Tip>
 
-All the options will produce the same result when you don't have enough GPU memory to accomodate the whole model (which is to fit everything that can on the GPU, then offload weights on the CPU or even on the disk if there is not enough RAM). 
+All the options will produce the same result when you don't have enough GPU memory to accommodate the whole model (which is to fit everything that can on the GPU, then offload weights on the CPU or even on the disk if there is not enough RAM). 
 
 When you have more GPU memory available than the model size, here the difference between each option:
 - `"auto"` and `"balanced"` evenly split the model on all available GPUs, making it possible for you to use a batch size greater than 1.
diff --git a/docs/source/usage_guides/checkpoint.mdx b/docs/source/usage_guides/checkpoint.mdx
index c818c9847..7d6bbbf99 100644
--- a/docs/source/usage_guides/checkpoint.mdx
+++ b/docs/source/usage_guides/checkpoint.mdx
@@ -13,7 +13,7 @@ specific language governing permissions and limitations under the License.
 # Checkpointing
 
 When training a PyTorch model with 🤗 Accelerate, you may often want to save and continue a state of training. Doing so requires
-saving and loading the model, optimizer, RNG generators, and the GradScaler. Inside 🤗 Accelerate are two convience functions to achieve this quickly:
+saving and loading the model, optimizer, RNG generators, and the GradScaler. Inside 🤗 Accelerate are two convenience functions to achieve this quickly:
 - Use [`~Accelerator.save_state`] for saving everything mentioned above to a folder location
 - Use [`~Accelerator.load_state`] for loading everything stored from an earlier `save_state`
 
diff --git a/docs/source/usage_guides/deepspeed.mdx b/docs/source/usage_guides/deepspeed.mdx
index 8a4238732..29561c77b 100644
--- a/docs/source/usage_guides/deepspeed.mdx
+++ b/docs/source/usage_guides/deepspeed.mdx
@@ -68,7 +68,7 @@ Inference:
 
 ## How it works?
 
-**Pre-Requisites**: Install DeepSpeed version >=0.6.5. Please refer to the [DeepSpeed Insallation details](https://github.com/microsoft/DeepSpeed#installation)
+**Pre-Requisites**: Install DeepSpeed version >=0.6.5. Please refer to the [DeepSpeed Installation details](https://github.com/microsoft/DeepSpeed#installation)
 for more information.
 
 We will first look at easy to use integration via `accelerate config`. 
@@ -383,13 +383,13 @@ We will look at the changes needed in the code when using these.
    ```
    b. Custom Optim + Custom Scheduler: The case when both `optimizer` and `scheduler` keys are absent in the DeepSpeed config file.
    In this situation, no code changes are needed from the user and this is the case when using integration via DeepSpeed Plugin.
-   In the above example we can see that the code reamins unchanged if the `optimizer` and `scheduler` keys are absent in the DeepSpeed config file.
+   In the above example we can see that the code remains unchanged if the `optimizer` and `scheduler` keys are absent in the DeepSpeed config file.
 
    c. Custom Optim + DS Scheduler: The case when only `scheduler` key is present in the DeepSpeed config file. 
    In this situation, user has to use `accelerate.utils.DummyScheduler` to replace the PyTorch/Custom scheduler in their code. 
 
    d. DS Optim + Custom Scheduler: The case when only `optimizer` key is present in the DeepSpeed config file. 
-   This will result in an error because one can only use DS Scheduler when using DS Optim.
+   This will result in an error because you can only use DS Scheduler when using DS Optim.
 
 2. Notice the `auto` values in the above example DeepSpeed config files. These are automatically handled by `prepare` method 
 based on model, dataloaders, dummy optimizer and dummy schedulers provided to `prepare` method. 
@@ -435,7 +435,7 @@ ZeRO Stage-3 has 2 options:
        logging.warning(f"Failure {status_msg}")
    ``` 
    This will create ZeRO model and optimizer partitions along with `zero_to_fp32.py` script in checkpoint directory.
-   One can use this script to do offline consolidation.  
+   You can use this script to do offline consolidation.  
    It requires no configuration files or GPUs. Here is an example of its usage:  
    ```bash
    $ cd /path/to/checkpoint_dir
@@ -444,14 +444,14 @@ ZeRO Stage-3 has 2 options:
    Detected checkpoint of type zero stage 3, world_size: 2
    Saving fp32 state dict to pytorch_model.bin (total_numel=60506624)
    ```
-   To get 32bit model for saving/inference, one can do the following:
+   To get 32bit model for saving/inference, you can perform:
    ```python
    from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
 
    unwrapped_model = accelerator.unwrap_model(model)
    fp32_model = load_state_dict_from_zero_checkpoint(unwrapped_model, checkpoint_dir)
    ```
-   If only interested in state_dict, one can do the following:
+   If you are only interested in the `state_dict`, you can do the following:
    ```python
    from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
 
@@ -462,7 +462,7 @@ ZeRO Stage-3 has 2 options:
 ## ZeRO Inference
 DeepSpeed ZeRO Inference supports ZeRO stage 3 with ZeRO-Infinity. 
 It uses the same ZeRO protocol as training, but it doesn't use an optimizer and a lr scheduler and only stage 3 is relevant.
-With accelerate integration, one has to just prepare model and dataloader as shown below:
+With accelerate integration, you just need to prepare the model and dataloader as shown below:
 
 ```python
 model, eval_dataloader = accelerator.prepare(model, eval_dataloader)
diff --git a/docs/source/usage_guides/tracking.mdx b/docs/source/usage_guides/tracking.mdx
index 7c61b21e9..cc5c17418 100644
--- a/docs/source/usage_guides/tracking.mdx
+++ b/docs/source/usage_guides/tracking.mdx
@@ -104,7 +104,7 @@ Every tracker must implement three functions and have three properties:
     - This should be implemented as a `@property` function 
     - Should return the internal tracking mechanism the library uses, such as the `run` object for `wandb`.
 
-A brief example can be seen below with an integration with Weights and Biases, containing only the relevent information:
+A brief example can be seen below with an integration with Weights and Biases, containing only the relevant information:
 ```python
 from accelerate.tracking import GeneralTracker
 from typing import Optional

```suggestion
To introduce special behavior in your script for TPUs you can check the `distributed_type` of your
```
@@ -261,7 +261,7 @@ lot of time. In practice, that means you must take special care to have all your
 shape (so no dynamic padding for instance if you are in an NLP problem) and should not use layers with for loops that
 have different lengths depending on the inputs (such as an LSTM) or the training will be excruciatingly slow.
 
-To introduce special behavior in your script for TPUs you can check the `distributed_type` of your
+To introduce special behaviour in your script for TPUs you can check the `distributed_type` of your
```suggestion
Internally, the library works by first analyzing the environment in which the script is launched to determine which
```
@@ -447,7 +447,7 @@ will be added in a next version.
 
 ## Internal mechanism
 
-Internally, the library works by first analyzing the environment in which the script is launched to determine which
+Internally, the library works by first analysing the environment in which the script is launched to determine which
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/603
diff --git a/src/accelerate/test_utils/scripts/test_checkpointing.py b/src/accelerate/test_utils/scripts/external_deps/test_checkpointing.py
similarity index 100%
rename from src/accelerate/test_utils/scripts/test_checkpointing.py
rename to src/accelerate/test_utils/scripts/external_deps/test_checkpointing.py
diff --git a/src/accelerate/test_utils/scripts/test_peak_memory_usage.py b/src/accelerate/test_utils/scripts/external_deps/test_peak_memory_usage.py
similarity index 100%
rename from src/accelerate/test_utils/scripts/test_peak_memory_usage.py
rename to src/accelerate/test_utils/scripts/external_deps/test_peak_memory_usage.py
diff --git a/src/accelerate/test_utils/scripts/test_performance.py b/src/accelerate/test_utils/scripts/external_deps/test_performance.py
similarity index 100%
rename from src/accelerate/test_utils/scripts/test_performance.py
rename to src/accelerate/test_utils/scripts/external_deps/test_performance.py
diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py
index 350d95320..c394b50cc 100644
--- a/src/accelerate/utils/dataclasses.py
+++ b/src/accelerate/utils/dataclasses.py
@@ -383,7 +383,7 @@ def deepspeed_config_process(self, prefix="", mismatches=None, config=None, must
     def set_mixed_precision(self, mixed_precision):
         ds_config = self.deepspeed_config
         if mixed_precision == "fp16" and "fp16" not in ds_config and "bf16" not in ds_config:
-            ds_config.update({"fp16": {"enabled": True}})
+            ds_config.update({"fp16": {"enabled": True, "auto_cast": True}})
         elif mixed_precision == "bf16" and "fp16" not in ds_config and "bf16" not in ds_config:
             ds_config.update({"bf16": {"enabled": True}})
 
diff --git a/tests/deepspeed/test_deepspeed.py b/tests/deepspeed/test_deepspeed.py
index 6b37eb93e..a8c1aedac 100644
--- a/tests/deepspeed/test_deepspeed.py
+++ b/tests/deepspeed/test_deepspeed.py
@@ -594,7 +594,9 @@ def test_autofill_dsconfig(self):
 
     def test_basic_run(self):
         mod_file = inspect.getfile(accelerate.test_utils)
-        test_file_path = os.path.sep.join(mod_file.split(os.path.sep)[:-1] + ["scripts", "test_performance.py"])
+        test_file_path = os.path.sep.join(
+            mod_file.split(os.path.sep)[:-1] + ["scripts", "external_deps", "test_performance.py"]
+        )
         with tempfile.TemporaryDirectory() as dirpath:
             cmd = [
                 "accelerate",
@@ -634,7 +636,7 @@ def setUp(self):
 
         self.stages = [1, 2, 3]
         self.zero3_offload_config = False
-        self.performance_lower_bound = 0.83
+        self.performance_lower_bound = 0.82
         self.peak_memory_usage_upper_bound = {
             "multi_gpu_fp16": 3200,
             "deepspeed_stage_1_fp16": 1600,
@@ -646,7 +648,7 @@ def setUp(self):
         self.n_val = 160
 
         mod_file = inspect.getfile(accelerate.test_utils)
-        self.test_scripts_folder = os.path.sep.join(mod_file.split(os.path.sep)[:-1] + ["scripts"])
+        self.test_scripts_folder = os.path.sep.join(mod_file.split(os.path.sep)[:-1] + ["scripts", "external_deps"])
 
     def test_performance(self):
         self.test_file_path = os.path.join(self.test_scripts_folder, "test_performance.py")
diff --git a/tests/fsdp/test_fsdp.py b/tests/fsdp/test_fsdp.py
index 8ad088c04..628e377c3 100644
--- a/tests/fsdp/test_fsdp.py
+++ b/tests/fsdp/test_fsdp.py
@@ -182,7 +182,7 @@ def test_cpu_offload(self):
 class FSDPIntegrationTest(TempDirTestCase):
     def setUp(self):
         super().setUp()
-        self.performance_lower_bound = 0.83
+        self.performance_lower_bound = 0.82
         self.performance_configs = [
             "fsdp_shard_grad_op_transformer_based_wrap",
             "fsdp_full_shard_transformer_based_wrap",
@@ -197,7 +197,7 @@ def setUp(self):
         self.n_val = 160
 
         mod_file = inspect.getfile(accelerate.test_utils)
-        self.test_scripts_folder = os.path.sep.join(mod_file.split(os.path.sep)[:-1] + ["scripts"])
+        self.test_scripts_folder = os.path.sep.join(mod_file.split(os.path.sep)[:-1] + ["scripts", "external_deps"])
 
     def test_performance(self):
         self.test_file_path = os.path.join(self.test_scripts_folder, "test_performance.py")

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/602
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 599e710f0..3930745dc 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -339,17 +339,6 @@ def simple_launcher(args):
         mixed_precision = "fp16"
 
     current_env["MIXED_PRECISION"] = str(mixed_precision)
-    if args.num_cpu_threads_per_process is None:
-        local_size = get_int_from_env(
-            ["MPI_LOCALNRANKS", "OMPI_COMM_WORLD_LOCAL_SIZE", "MV2_COMM_WORLD_LOCAL_SIZE"], 1
-        )
-        args.num_cpu_threads_per_process = int(psutil.cpu_count(logical=False) / local_size)
-        if args.num_cpu_threads_per_process == 0:
-            args.num_cpu_threads_per_process = 1
-        logger.info(
-            f"num_cpu_threads_per_process unset, we set it at {args.num_cpu_threads_per_process} to improve oob performance."
-        )
-
     current_env["OMP_NUM_THREADS"] = str(args.num_cpu_threads_per_process)
 
     process = subprocess.Popen(cmd, env=current_env)
@@ -447,9 +436,6 @@ def multi_gpu_launcher(args):
             current_env["FSDP_BACKWARD_PREFETCH"] = str(args.fsdp_backward_prefetch_policy)
         if args.fsdp_state_dict_type is not None:
             current_env["FSDP_STATE_DICT_TYPE"] = str(args.fsdp_state_dict_type)
-    if args.num_cpu_threads_per_process is None:
-        args.num_cpu_threads_per_process = 1
-        logger.info(f"num_cpu_threads_per_process unset, we set it at {args.num_cpu_threads_per_process}.")
     current_env["OMP_NUM_THREADS"] = str(args.num_cpu_threads_per_process)
     process = subprocess.Popen(cmd, env=current_env)
     process.wait()
@@ -803,6 +789,17 @@ def launch_command(args):
                 if "--num_processes" in warn:
                     warned[i] = warn.replace("`1`", f"`{args.num_processes}`")
 
+    if args.num_cpu_threads_per_process is None:
+        local_size = get_int_from_env(
+            ["MPI_LOCALNRANKS", "OMPI_COMM_WORLD_LOCAL_SIZE", "MV2_COMM_WORLD_LOCAL_SIZE"], 1
+        )
+        args.num_cpu_threads_per_process = int(psutil.cpu_count(logical=False) / local_size)
+        if args.num_cpu_threads_per_process == 0:
+            args.num_cpu_threads_per_process = 1
+        warned.append(
+            f"\t`--num_cpu_threads_per_process` was set to `{args.num_cpu_threads_per_process}` to improve out-of-box performance"
+        )
+
     if any(warned):
         message = "The following values were not passed to `accelerate launch` and had defaults used instead:\n"
         message += "\n".join(warned)

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/598
diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml
index 0da843ec2..84c4e03f6 100644
--- a/docs/source/_toctree.yml
+++ b/docs/source/_toctree.yml
@@ -33,6 +33,8 @@
     title: How to avoid CUDA Out-of-Memory
   - local: usage_guides/sagemaker
     title: Using Accelerate on SageMaker
+  - local: usage_guides/mps
+    title: How to use Apple Silicon M1 GPUs
   title: How-To Guides
 - sections:
   - local: concept_guides/gradient_synchronization
diff --git a/docs/source/usage_guides/mps.mdx b/docs/source/usage_guides/mps.mdx
new file mode 100644
index 000000000..43a00f104
--- /dev/null
+++ b/docs/source/usage_guides/mps.mdx
@@ -0,0 +1,81 @@
+<!--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.
+-->
+
+# Accelerated PyTorch Training on Mac
+
+With PyTorch v1.12 release, developers and researchers can take advantage of Apple silicon GPUs for significantly faster model training. 
+This unlocks the ability to perform machine learning workflows like prototyping and fine-tuning locally, right on Mac.
+Apple's Metal Performance Shaders (MPS) as a backend for PyTorch enables this and can be used via the new `"mps"` device. 
+This will map computational graphs and primitives on the MPS Graph framework and tuned kernels provided by MPS.
+For more information please refer official documents [Introducing Accelerated PyTorch Training on Mac](https://pytorch.org/blog/introducing-accelerated-pytorch-training-on-mac/)
+and [MPS BACKEND](https://pytorch.org/docs/stable/notes/mps.html).
+
+### Benefits of Training and Inference using Apple M1 Chips
+
+1. Enables users to train larger networks or batch sizes locally
+2. Reduces data retrieval latency and provides the GPU with direct access to the full memory store due to unified memory architecture. 
+Therefore, improving end-to-end performance.
+3. Reduces costs associated with cloud-based development or the need for additional local GPUs.
+
+**Pre-requisites**: To install torch with mps support, 
+please follow this nice medium article [GPU-Acceleration Comes to PyTorch on M1 Macs](https://medium.com/towards-data-science/gpu-acceleration-comes-to-pytorch-on-m1-macs-195c399efcc1).
+
+
+## How it works out of the box
+
+On your machine(s) just run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked, specifically choose `MPS` for the query:
+
+```
+ Which type of machine are you using?. 
+ ```
+
+This will generate a config file that will be used automatically to properly set 
+the default options when doing `accelerate launch`, such as the one shown below:
+
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config: {}
+distributed_type: MPS
+downcast_bf16: 'no'
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: 'no'
+num_machines: 1
+num_processes: 1
+use_cpu: false
+```
+
+After this configuration has been made, here is how you run the CV example 
+(from the root of the repo) with MPS enabled:
+
+```bash
+accelerate launch /examples/cv_example.py --data_dir images
+```
+
+## A few caveats to be aware of
+
+1. For `nlp_example.py` the metrics are too bad when compared to CPU-only training. 
+This means certain operations in BERT model are going wrong using mps device and this needs to be fixed by PyTorch.
+2. Distributed setups `gloo` and `nccl` are not working with `mps` device. 
+This means that currently only single GPU of `mps` device type can be used.
+
+Finally, please, remember that, 🤗 `Accelerate` only integrates MPS backend, therefore if you
+have any problems or questions with regards to MPS backend usage, please, file an issue with [PyTorch GitHub](https://github.com/pytorch/pytorch/issues).
\ No newline at end of file
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index bbd8db435..9cd4b7da3 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -29,7 +29,7 @@
 from .logging import get_logger
 from .optimizer import AcceleratedOptimizer
 from .scheduler import AcceleratedScheduler
-from .state import AcceleratorState, GradientState
+from .state import AcceleratorState, GradientState, parse_flag_from_env
 from .tracking import LOGGER_TYPE_TO_CLASS, GeneralTracker, filter_trackers
 from .utils import (
     MODEL_NAME,
@@ -282,7 +282,7 @@ def __init__(
             self.native_amp = is_torch_version(">=", "1.6")
             if not self.native_amp:
                 raise ValueError(err.format(mode="fp16", requirement="PyTorch >= 1.6"))
-            if not torch.cuda.is_available():
+            if not torch.cuda.is_available() and not parse_flag_from_env("USE_MPS_DEVICE"):
                 raise ValueError(err.format(mode="fp16", requirement="a GPU"))
             kwargs = self.scaler_handler.to_kwargs() if self.scaler_handler is not None else {}
             if self.distributed_type == DistributedType.FSDP:
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index ba677e3eb..1ba5839e2 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -28,9 +28,9 @@
 
 def get_cluster_input():
     distributed_type = _ask_field(
-        "Which type of machine are you using? ([0] No distributed training, [1] multi-CPU, [2] multi-GPU, [3] TPU): ",
+        "Which type of machine are you using? ([0] No distributed training, [1] multi-CPU, [2] multi-GPU, [3] TPU [4] MPS): ",
         _convert_distributed_mode,
-        error_message="Please enter 0, 1, 2 or 3.",
+        error_message="Please enter 0, 1, 2, 3 or 4.",
     )
 
     machine_rank = 0
diff --git a/src/accelerate/commands/config/config_utils.py b/src/accelerate/commands/config/config_utils.py
index fe72be267..9dd1f4c99 100644
--- a/src/accelerate/commands/config/config_utils.py
+++ b/src/accelerate/commands/config/config_utils.py
@@ -37,7 +37,7 @@ def _convert_compute_environment(value):
 
 def _convert_distributed_mode(value):
     value = int(value)
-    return DistributedType(["NO", "MULTI_CPU", "MULTI_GPU", "TPU"][value])
+    return DistributedType(["NO", "MULTI_CPU", "MULTI_GPU", "TPU", "MPS"][value])
 
 
 def _convert_sagemaker_distributed_mode(value):
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 599e710f0..863f903c1 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -63,6 +63,12 @@ def launch_command_parser(subparsers=None):
         action="store_true",
         help="Whether or not this should launch a distributed GPU training.",
     )
+    parser.add_argument(
+        "--use_mps_device",
+        default=False,
+        action="store_true",
+        help="Whether or not this should use MPS-enabled GPU device on MacOS machines.",
+    )
     parser.add_argument(
         "--use_deepspeed",
         default=False,
@@ -320,6 +326,7 @@ def simple_launcher(args):
 
     current_env = os.environ.copy()
     current_env["USE_CPU"] = str(args.cpu or args.use_cpu)
+    current_env["USE_MPS_DEVICE"] = str(args.use_mps_device)
     if args.num_machines > 1:
         current_env["MASTER_ADDR"] = args.main_process_ip
         current_env["MASTER_PORT"] = str(args.main_process_port)
@@ -752,11 +759,18 @@ def launch_command(args):
     # Get the default from the config file.
     if args.config_file is not None or os.path.isfile(default_config_file) and not args.cpu:
         defaults = load_config_from_file(args.config_file)
-        if not args.multi_gpu and not args.tpu and not args.use_deepspeed and not args.use_fsdp:
+        if (
+            not args.multi_gpu
+            and not args.tpu
+            and not args.use_deepspeed
+            and not args.use_fsdp
+            and not args.use_mps_device
+        ):
             args.use_deepspeed = defaults.distributed_type == DistributedType.DEEPSPEED
             args.multi_gpu = defaults.distributed_type == DistributedType.MULTI_GPU
             args.tpu = defaults.distributed_type == DistributedType.TPU
             args.use_fsdp = defaults.distributed_type == DistributedType.FSDP
+            args.use_mps_device = defaults.distributed_type == DistributedType.MPS
         if defaults.compute_environment == ComputeEnvironment.LOCAL_MACHINE:
             # Update args with the defaults
             for name, attr in defaults.__dict__.items():
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index a4f6bac2a..bd339048f 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -206,7 +206,26 @@ def __init__(
                 self.distributed_type = DistributedType.NO
                 self.num_processes = 1
                 self.process_index = self.local_process_index = 0
-                self.device = torch.device("cuda" if torch.cuda.is_available() and not cpu else "cpu")
+                if parse_flag_from_env("USE_MPS_DEVICE") and not cpu:
+                    if not torch.backends.mps.is_available():
+                        if not torch.backends.mps.is_built():
+                            raise AssertionError(
+                                "MPS not available because the current PyTorch install was not "
+                                "built with MPS enabled. Please install torch version >=1.12.0 on "
+                                "your Apple silicon Mac running macOS 12.3 or later with a native "
+                                "version (arm64) of Python"
+                            )
+                        else:
+                            raise AssertionError(
+                                "MPS not available because the current MacOS version is not 12.3+ "
+                                "and/or you do not have an MPS-enabled device on this machine."
+                            )
+                    else:
+                        self.device = torch.device("mps")
+                elif cpu or not torch.cuda.is_available():
+                    self.device = torch.device("cpu")
+                else:
+                    self.device = torch.device("cuda")
                 self.mixed_precision = mixed_precision
             self.initialized = True
 
diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py
index 350d95320..50c80adb1 100644
--- a/src/accelerate/utils/dataclasses.py
+++ b/src/accelerate/utils/dataclasses.py
@@ -123,6 +123,7 @@ class DistributedType(str, enum.Enum):
     DEEPSPEED = "DEEPSPEED"
     FSDP = "FSDP"
     TPU = "TPU"
+    MPS = "MPS"
 
 
 class SageMakerDistributedType(str, enum.Enum):

Can we do something similar to this [here](https://github.com/huggingface/accelerate/blob/main/src/accelerate/test_utils/testing.py#L94-L99) so that cuda tests can also be ran? As so far tests are only running on the CPU here 😄 
@@ -272,7 +272,7 @@ def __init__(
             self.native_amp = is_torch_version(">=", "1.6")
             if not self.native_amp:
                 raise ValueError(err.format(mode="fp16", requirement="PyTorch >= 1.6"))
-            if not torch.cuda.is_available():
+            if not torch.cuda.is_available() and os.environ.get("USE_MPS_DEVICE", "false") == "false":
There might be other truthy values we're not catching here, can se use something like `parse_flag_from_env`?
@@ -272,7 +272,7 @@ def __init__(
             self.native_amp = is_torch_version(">=", "1.6")
             if not self.native_amp:
                 raise ValueError(err.format(mode="fp16", requirement="PyTorch >= 1.6"))
-            if not torch.cuda.is_available():
+            if not torch.cuda.is_available() and os.environ.get("USE_MPS_DEVICE", "false") == "false":
```suggestion
    current_env["USE_MPS_DEVICE"] = str(args.use_mps_device)
```
We can't rely on users always setting lowercased values, so we should handle any casing.
@@ -320,6 +326,7 @@ def simple_launcher(args):
 
     current_env = os.environ.copy()
     current_env["USE_CPU"] = str(args.cpu or args.use_cpu)
+    current_env["USE_MPS_DEVICE"] = str(args.use_mps_device).lower()
Give instructions here maybe?
@@ -208,7 +208,24 @@ def __init__(
                 self.distributed_type = DistributedType.NO
                 self.num_processes = 1
                 self.process_index = self.local_process_index = 0
-                self.device = torch.device("cuda" if torch.cuda.is_available() and not cpu else "cpu")
+                if os.environ.get("USE_MPS_DEVICE", "false") == "true" and not cpu:
+                    if not torch.backends.mps.is_available():
+                        if not torch.backends.mps.is_built():
+                            raise AssertionError(
+                                "MPS not available because the current PyTorch install was not "
+                                "built with MPS enabled."
Hello, `use_mps_device` is a boolean argument using argparse action, so I think this should be fine.
@@ -320,6 +326,7 @@ def simple_launcher(args):
 
     current_env = os.environ.copy()
     current_env["USE_CPU"] = str(args.cpu or args.use_cpu)
+    current_env["USE_MPS_DEVICE"] = str(args.use_mps_device).lower()
Done.
@@ -208,7 +208,24 @@ def __init__(
                 self.distributed_type = DistributedType.NO
                 self.num_processes = 1
                 self.process_index = self.local_process_index = 0
-                self.device = torch.device("cuda" if torch.cuda.is_available() and not cpu else "cpu")
+                if os.environ.get("USE_MPS_DEVICE", "false") == "true" and not cpu:
+                    if not torch.backends.mps.is_available():
+                        if not torch.backends.mps.is_built():
+                            raise AssertionError(
+                                "MPS not available because the current PyTorch install was not "
+                                "built with MPS enabled."
Done.
@@ -320,6 +326,7 @@ def simple_launcher(args):
 
     current_env = os.environ.copy()
     current_env["USE_CPU"] = str(args.cpu or args.use_cpu)
+    current_env["USE_MPS_DEVICE"] = str(args.use_mps_device).lower()
Done.
@@ -272,7 +272,7 @@ def __init__(
             self.native_amp = is_torch_version(">=", "1.6")
             if not self.native_amp:
                 raise ValueError(err.format(mode="fp16", requirement="PyTorch >= 1.6"))
-            if not torch.cuda.is_available():
+            if not torch.cuda.is_available() and os.environ.get("USE_MPS_DEVICE", "false") == "false":
Hello, I didn't understand. Let's take this offline to better understand.
@@ -272,7 +272,7 @@ def __init__(
             self.native_amp = is_torch_version(">=", "1.6")
             if not self.native_amp:
                 raise ValueError(err.format(mode="fp16", requirement="PyTorch >= 1.6"))
-            if not torch.cuda.is_available():
+            if not torch.cuda.is_available() and os.environ.get("USE_MPS_DEVICE", "false") == "false":
Specify what specific question related to MPS needs to be answered here :) 
@@ -0,0 +1,93 @@
+<!--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.
+-->
+
+# Accelerated PyTorch Training on Mac
+With PyTorch v1.12 release, developers and researchers can take advantage of Apple silicon GPUs for significantly faster model training. 
+This unlocks the ability to perform machine learning workflows like prototyping and fine-tuning locally, right on Mac.
+Apple's Metal Performance Shaders (MPS) as a backend for PyTorch enables this. 
+New device `mps` maps computational graphs and primitives on the MPS Graph framework and tuned kernels provided by MPS.
+For more information please refer Official documents [Introducing Accelerated PyTorch Training on Mac](https://pytorch.org/blog/introducing-accelerated-pytorch-training-on-mac/)
+and [MPS BACKEND](https://pytorch.org/docs/stable/notes/mps.html).
+
+### Benefits of Training and Inference using Apple M1 Chips
+1. Enables users to train larger networks or batch sizes locally
+2. Reduces data retrieval latency and provides the GPU with direct access to the full memory store due to unified memory architecture. 
+Therefore, improving end-to-end performance.
+3. Reduces costs associated with cloud-based development or the need for additional local GPUs.
+
+### Pre-requisites: Installing torch with mps support
+Following this great medium article [GPU-Acceleration Comes to PyTorch on M1 Macs](https://medium.com/towards-data-science/gpu-acceleration-comes-to-pytorch-on-m1-macs-195c399efcc1),
+below are the steps to make sure torch is installed with `mps` support
+```
+# installing torch with m1 support on mac
+# install latest python from https://www.python.org/downloads/release/python-3106/
+# check the platform
+import platform
+platform.platform()
+'macOS-12.5-arm64-arm-64bit' 
+# This is compatible as the macOS version is above 12.3 and it is the ARM64 native version of Python
+# install torch 1.12 via the below command
+# pip3 install torch torchvision torchaudio
+# test the `mps` device support
+import torch
+torch.has_mps
+True
+a = torch.Tensor([10,11])
+a.to("mps")
+/Users/mac/ml/lib/python3.10/site-packages/torch/_tensor_str.py:103: UserWarning: The operator 'aten::bitwise_and.Tensor_out' is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at  /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/mps/MPSFallback.mm:11.)
+  nonzero_finite_vals = torch.masked_select(tensor_view, torch.isfinite(tensor_view) & tensor_view.ne(0))
+tensor([10.0000, 11.0000], device='mps:0')
+```
+
+## How it works out of the box
+On your machine(s) just run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked. This will generate a config file that will be used automatically to properly set the
+default options when doing
Resolving based on Sylvains note about the tests
@@ -272,7 +272,7 @@ def __init__(
             self.native_amp = is_torch_version(">=", "1.6")
             if not self.native_amp:
                 raise ValueError(err.format(mode="fp16", requirement="PyTorch >= 1.6"))
-            if not torch.cuda.is_available():
+            if not torch.cuda.is_available() and os.environ.get("USE_MPS_DEVICE", "false") == "false":
Add a line under each header, I believe this is why the URLs aren't showing correctly and why everything's formatting seems to be broken: https://moon-ci-docs.huggingface.co/docs/accelerate/pr_596/en/usage_guides/mps#accelerated-pytorch-training-on-mac
@@ -0,0 +1,93 @@
+<!--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.
+-->
+
+# Accelerated PyTorch Training on Mac
+With PyTorch v1.12 release, developers and researchers can take advantage of Apple silicon GPUs for significantly faster model training. 
```suggestion
Apple's Metal Performance Shaders (MPS) as a backend for PyTorch enables this and can be used via the new `"mps"` device. This will map computational graphs and primitives on the MPS Graph framework and tuned kernels provided by MPS.
```
@@ -0,0 +1,93 @@
+<!--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.
+-->
+
+# Accelerated PyTorch Training on Mac
+With PyTorch v1.12 release, developers and researchers can take advantage of Apple silicon GPUs for significantly faster model training. 
+This unlocks the ability to perform machine learning workflows like prototyping and fine-tuning locally, right on Mac.
+Apple's Metal Performance Shaders (MPS) as a backend for PyTorch enables this. 
+New device `mps` maps computational graphs and primitives on the MPS Graph framework and tuned kernels provided by MPS.
Done.
@@ -0,0 +1,93 @@
+<!--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.
+-->
+
+# Accelerated PyTorch Training on Mac
+With PyTorch v1.12 release, developers and researchers can take advantage of Apple silicon GPUs for significantly faster model training. 
Done

@@ -0,0 +1,93 @@
+<!--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.
+-->
+
+# Accelerated PyTorch Training on Mac
+With PyTorch v1.12 release, developers and researchers can take advantage of Apple silicon GPUs for significantly faster model training. 
+This unlocks the ability to perform machine learning workflows like prototyping and fine-tuning locally, right on Mac.
+Apple's Metal Performance Shaders (MPS) as a backend for PyTorch enables this. 
+New device `mps` maps computational graphs and primitives on the MPS Graph framework and tuned kernels provided by MPS.
+For more information please refer Official documents [Introducing Accelerated PyTorch Training on Mac](https://pytorch.org/blog/introducing-accelerated-pytorch-training-on-mac/)
+and [MPS BACKEND](https://pytorch.org/docs/stable/notes/mps.html).
+
+### Benefits of Training and Inference using Apple M1 Chips
+1. Enables users to train larger networks or batch sizes locally
+2. Reduces data retrieval latency and provides the GPU with direct access to the full memory store due to unified memory architecture. 
+Therefore, improving end-to-end performance.
+3. Reduces costs associated with cloud-based development or the need for additional local GPUs.
+
+### Pre-requisites: Installing torch with mps support
+Following this great medium article [GPU-Acceleration Comes to PyTorch on M1 Macs](https://medium.com/towards-data-science/gpu-acceleration-comes-to-pytorch-on-m1-macs-195c399efcc1),
+below are the steps to make sure torch is installed with `mps` support
+```
+# installing torch with m1 support on mac
+# install latest python from https://www.python.org/downloads/release/python-3106/
+# check the platform
+import platform
+platform.platform()
+'macOS-12.5-arm64-arm-64bit' 
+# This is compatible as the macOS version is above 12.3 and it is the ARM64 native version of Python
+# install torch 1.12 via the below command
+# pip3 install torch torchvision torchaudio
+# test the `mps` device support
+import torch
+torch.has_mps
+True
+a = torch.Tensor([10,11])
+a.to("mps")
+/Users/mac/ml/lib/python3.10/site-packages/torch/_tensor_str.py:103: UserWarning: The operator 'aten::bitwise_and.Tensor_out' is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at  /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/mps/MPSFallback.mm:11.)
+  nonzero_finite_vals = torch.masked_select(tensor_view, torch.isfinite(tensor_view) & tensor_view.ne(0))
+tensor([10.0000, 11.0000], device='mps:0')
+```
+
+## How it works out of the box
+On your machine(s) just run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked. This will generate a config file that will be used automatically to properly set the
+default options when doing
We can break this up a bit to make it more digestible/readable + syntax highlighting

```suggestion
```python
# installing torch with m1 support on mac
# install latest python from https://www.python.org/downloads/release/python-3106/
# check the platform
import platform
platform.platform()
```
```python out
'macOS-12.5-arm64-arm-64bit' 
```
```python
# This is compatible as the macOS version is above 12.3 and it is the ARM64 native version of Python
# install torch 1.12 via the below command
# pip3 install torch torchvision torchaudio
# test the `mps` device support
import torch
torch.has_mps
```
```python out
True
```
```python
a = torch.Tensor([10,11])
a.to("mps")
```
```python out
/Users/mac/ml/lib/python3.10/site-packages/torch/_tensor_str.py:103: UserWarning: The operator 'aten::bitwise_and.Tensor_out' is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at  /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/mps/MPSFallback.mm:11.)
  nonzero_finite_vals = torch.masked_select(tensor_view, torch.isfinite(tensor_view) & tensor_view.ne(0))
tensor([10.0000, 11.0000], device='mps:0')
```
```
@@ -0,0 +1,100 @@
+<!--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.
+-->
+
+# Accelerated PyTorch Training on Mac
+
+With PyTorch v1.12 release, developers and researchers can take advantage of Apple silicon GPUs for significantly faster model training. 
+This unlocks the ability to perform machine learning workflows like prototyping and fine-tuning locally, right on Mac.
+Apple's Metal Performance Shaders (MPS) as a backend for PyTorch enables this and can be used via the new `"mps"` device. 
+This will map computational graphs and primitives on the MPS Graph framework and tuned kernels provided by MPS.
+For more information please refer official documents [Introducing Accelerated PyTorch Training on Mac](https://pytorch.org/blog/introducing-accelerated-pytorch-training-on-mac/)
+and [MPS BACKEND](https://pytorch.org/docs/stable/notes/mps.html).
+
+### Benefits of Training and Inference using Apple M1 Chips
+
+1. Enables users to train larger networks or batch sizes locally
+2. Reduces data retrieval latency and provides the GPU with direct access to the full memory store due to unified memory architecture. 
+Therefore, improving end-to-end performance.
+3. Reduces costs associated with cloud-based development or the need for additional local GPUs.
+
+### Pre-requisites: Installing torch with mps support
+
+Following this great medium article [GPU-Acceleration Comes to PyTorch on M1 Macs](https://medium.com/towards-data-science/gpu-acceleration-comes-to-pytorch-on-m1-macs-195c399efcc1),
+below are the steps to make sure torch is installed with `mps` support.
+
+```
+# installing torch with m1 support on mac
+# install latest python from https://www.python.org/downloads/release/python-3106/
+# check the platform
+import platform
+platform.platform()
+'macOS-12.5-arm64-arm-64bit' 
+# This is compatible as the macOS version is above 12.3 and it is the ARM64 native version of Python
+# install torch 1.12 via the below command
+# pip3 install torch torchvision torchaudio
+# test the `mps` device support
+import torch
+torch.has_mps
+True
+a = torch.Tensor([10,11])
+a.to("mps")
+/Users/mac/ml/lib/python3.10/site-packages/torch/_tensor_str.py:103: UserWarning: The operator 'aten::bitwise_and.Tensor_out' is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at  /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/mps/MPSFallback.mm:11.)
+  nonzero_finite_vals = torch.masked_select(tensor_view, torch.isfinite(tensor_view) & tensor_view.ne(0))
+tensor([10.0000, 11.0000], device='mps:0')
+```
I can't make a suggestion for this one, but I recommend something like this:

and answer the questions asked, specifically choose `MPS` for the query:
```
Which type of machine are you using?
```
This will generate a config file that will be used automatically to properly set the
default options when doing
@@ -0,0 +1,100 @@
+<!--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.
+-->
+
+# Accelerated PyTorch Training on Mac
+
+With PyTorch v1.12 release, developers and researchers can take advantage of Apple silicon GPUs for significantly faster model training. 
+This unlocks the ability to perform machine learning workflows like prototyping and fine-tuning locally, right on Mac.
+Apple's Metal Performance Shaders (MPS) as a backend for PyTorch enables this and can be used via the new `"mps"` device. 
+This will map computational graphs and primitives on the MPS Graph framework and tuned kernels provided by MPS.
+For more information please refer official documents [Introducing Accelerated PyTorch Training on Mac](https://pytorch.org/blog/introducing-accelerated-pytorch-training-on-mac/)
+and [MPS BACKEND](https://pytorch.org/docs/stable/notes/mps.html).
+
+### Benefits of Training and Inference using Apple M1 Chips
+
+1. Enables users to train larger networks or batch sizes locally
+2. Reduces data retrieval latency and provides the GPU with direct access to the full memory store due to unified memory architecture. 
+Therefore, improving end-to-end performance.
+3. Reduces costs associated with cloud-based development or the need for additional local GPUs.
+
+### Pre-requisites: Installing torch with mps support
+
+Following this great medium article [GPU-Acceleration Comes to PyTorch on M1 Macs](https://medium.com/towards-data-science/gpu-acceleration-comes-to-pytorch-on-m1-macs-195c399efcc1),
+below are the steps to make sure torch is installed with `mps` support.
+
+```
+# installing torch with m1 support on mac
+# install latest python from https://www.python.org/downloads/release/python-3106/
+# check the platform
+import platform
+platform.platform()
+'macOS-12.5-arm64-arm-64bit' 
+# This is compatible as the macOS version is above 12.3 and it is the ARM64 native version of Python
+# install torch 1.12 via the below command
+# pip3 install torch torchvision torchaudio
+# test the `mps` device support
+import torch
+torch.has_mps
+True
+a = torch.Tensor([10,11])
+a.to("mps")
+/Users/mac/ml/lib/python3.10/site-packages/torch/_tensor_str.py:103: UserWarning: The operator 'aten::bitwise_and.Tensor_out' is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at  /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/mps/MPSFallback.mm:11.)
+  nonzero_finite_vals = torch.masked_select(tensor_view, torch.isfinite(tensor_view) & tensor_view.ne(0))
+tensor([10.0000, 11.0000], device='mps:0')
+```
+
+## How it works out of the box
+
+On your machine(s) just run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked, specifically choose `MPS` for the query `Which type of machine are you using?` . 
+This will generate a config file that will be used automatically to properly set the
+default options when doing
Move this section to after you've ran `accelerate config` and state that this is what the generated configuration looks like
@@ -0,0 +1,100 @@
+<!--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.
+-->
+
+# Accelerated PyTorch Training on Mac
+
+With PyTorch v1.12 release, developers and researchers can take advantage of Apple silicon GPUs for significantly faster model training. 
+This unlocks the ability to perform machine learning workflows like prototyping and fine-tuning locally, right on Mac.
+Apple's Metal Performance Shaders (MPS) as a backend for PyTorch enables this and can be used via the new `"mps"` device. 
+This will map computational graphs and primitives on the MPS Graph framework and tuned kernels provided by MPS.
+For more information please refer official documents [Introducing Accelerated PyTorch Training on Mac](https://pytorch.org/blog/introducing-accelerated-pytorch-training-on-mac/)
+and [MPS BACKEND](https://pytorch.org/docs/stable/notes/mps.html).
+
+### Benefits of Training and Inference using Apple M1 Chips
+
+1. Enables users to train larger networks or batch sizes locally
+2. Reduces data retrieval latency and provides the GPU with direct access to the full memory store due to unified memory architecture. 
+Therefore, improving end-to-end performance.
+3. Reduces costs associated with cloud-based development or the need for additional local GPUs.
+
+### Pre-requisites: Installing torch with mps support
+
+Following this great medium article [GPU-Acceleration Comes to PyTorch on M1 Macs](https://medium.com/towards-data-science/gpu-acceleration-comes-to-pytorch-on-m1-macs-195c399efcc1),
+below are the steps to make sure torch is installed with `mps` support.
+
+```
+# installing torch with m1 support on mac
+# install latest python from https://www.python.org/downloads/release/python-3106/
+# check the platform
+import platform
+platform.platform()
+'macOS-12.5-arm64-arm-64bit' 
+# This is compatible as the macOS version is above 12.3 and it is the ARM64 native version of Python
+# install torch 1.12 via the below command
+# pip3 install torch torchvision torchaudio
+# test the `mps` device support
+import torch
+torch.has_mps
+True
+a = torch.Tensor([10,11])
+a.to("mps")
+/Users/mac/ml/lib/python3.10/site-packages/torch/_tensor_str.py:103: UserWarning: The operator 'aten::bitwise_and.Tensor_out' is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at  /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/mps/MPSFallback.mm:11.)
+  nonzero_finite_vals = torch.masked_select(tensor_view, torch.isfinite(tensor_view) & tensor_view.ne(0))
+tensor([10.0000, 11.0000], device='mps:0')
+```
+
+## How it works out of the box
+
+On your machine(s) just run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked, specifically choose `MPS` for the query `Which type of machine are you using?` . 
+This will generate a config file that will be used automatically to properly set the
+default options when doing
+
+```bash
+accelerate launch my_script.py --args_to_my_script
+```
+
+For instance, here is how you would run the CV example (from the root of the repo) with MPS enabled:
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config: {}
+distributed_type: MPS
+downcast_bf16: 'no'
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: 'no'
+num_machines: 1
+num_processes: 1
+use_cpu: false
+```
I feel the above block is fine as it might be too distracting to have multiple blocks which can break the flow when the user is reading 
@@ -0,0 +1,100 @@
+<!--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.
+-->
+
+# Accelerated PyTorch Training on Mac
+
+With PyTorch v1.12 release, developers and researchers can take advantage of Apple silicon GPUs for significantly faster model training. 
+This unlocks the ability to perform machine learning workflows like prototyping and fine-tuning locally, right on Mac.
+Apple's Metal Performance Shaders (MPS) as a backend for PyTorch enables this and can be used via the new `"mps"` device. 
+This will map computational graphs and primitives on the MPS Graph framework and tuned kernels provided by MPS.
+For more information please refer official documents [Introducing Accelerated PyTorch Training on Mac](https://pytorch.org/blog/introducing-accelerated-pytorch-training-on-mac/)
+and [MPS BACKEND](https://pytorch.org/docs/stable/notes/mps.html).
+
+### Benefits of Training and Inference using Apple M1 Chips
+
+1. Enables users to train larger networks or batch sizes locally
+2. Reduces data retrieval latency and provides the GPU with direct access to the full memory store due to unified memory architecture. 
+Therefore, improving end-to-end performance.
+3. Reduces costs associated with cloud-based development or the need for additional local GPUs.
+
+### Pre-requisites: Installing torch with mps support
+
+Following this great medium article [GPU-Acceleration Comes to PyTorch on M1 Macs](https://medium.com/towards-data-science/gpu-acceleration-comes-to-pytorch-on-m1-macs-195c399efcc1),
+below are the steps to make sure torch is installed with `mps` support.
+
+```
+# installing torch with m1 support on mac
+# install latest python from https://www.python.org/downloads/release/python-3106/
+# check the platform
+import platform
+platform.platform()
+'macOS-12.5-arm64-arm-64bit' 
+# This is compatible as the macOS version is above 12.3 and it is the ARM64 native version of Python
+# install torch 1.12 via the below command
+# pip3 install torch torchvision torchaudio
+# test the `mps` device support
+import torch
+torch.has_mps
+True
+a = torch.Tensor([10,11])
+a.to("mps")
+/Users/mac/ml/lib/python3.10/site-packages/torch/_tensor_str.py:103: UserWarning: The operator 'aten::bitwise_and.Tensor_out' is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at  /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/mps/MPSFallback.mm:11.)
+  nonzero_finite_vals = torch.masked_select(tensor_view, torch.isfinite(tensor_view) & tensor_view.ne(0))
+tensor([10.0000, 11.0000], device='mps:0')
+```
But then it would be breaking the flow of commands. I have followed generic commands `accelerate config` and `accelerate launch script.py args` with specific example so that user can quickly get the gist. 
@@ -0,0 +1,100 @@
+<!--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.
+-->
+
+# Accelerated PyTorch Training on Mac
+
+With PyTorch v1.12 release, developers and researchers can take advantage of Apple silicon GPUs for significantly faster model training. 
+This unlocks the ability to perform machine learning workflows like prototyping and fine-tuning locally, right on Mac.
+Apple's Metal Performance Shaders (MPS) as a backend for PyTorch enables this and can be used via the new `"mps"` device. 
+This will map computational graphs and primitives on the MPS Graph framework and tuned kernels provided by MPS.
+For more information please refer official documents [Introducing Accelerated PyTorch Training on Mac](https://pytorch.org/blog/introducing-accelerated-pytorch-training-on-mac/)
+and [MPS BACKEND](https://pytorch.org/docs/stable/notes/mps.html).
+
+### Benefits of Training and Inference using Apple M1 Chips
+
+1. Enables users to train larger networks or batch sizes locally
+2. Reduces data retrieval latency and provides the GPU with direct access to the full memory store due to unified memory architecture. 
+Therefore, improving end-to-end performance.
+3. Reduces costs associated with cloud-based development or the need for additional local GPUs.
+
+### Pre-requisites: Installing torch with mps support
+
+Following this great medium article [GPU-Acceleration Comes to PyTorch on M1 Macs](https://medium.com/towards-data-science/gpu-acceleration-comes-to-pytorch-on-m1-macs-195c399efcc1),
+below are the steps to make sure torch is installed with `mps` support.
+
+```
+# installing torch with m1 support on mac
+# install latest python from https://www.python.org/downloads/release/python-3106/
+# check the platform
+import platform
+platform.platform()
+'macOS-12.5-arm64-arm-64bit' 
+# This is compatible as the macOS version is above 12.3 and it is the ARM64 native version of Python
+# install torch 1.12 via the below command
+# pip3 install torch torchvision torchaudio
+# test the `mps` device support
+import torch
+torch.has_mps
+True
+a = torch.Tensor([10,11])
+a.to("mps")
+/Users/mac/ml/lib/python3.10/site-packages/torch/_tensor_str.py:103: UserWarning: The operator 'aten::bitwise_and.Tensor_out' is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at  /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/mps/MPSFallback.mm:11.)
+  nonzero_finite_vals = torch.masked_select(tensor_view, torch.isfinite(tensor_view) & tensor_view.ne(0))
+tensor([10.0000, 11.0000], device='mps:0')
+```
+
+## How it works out of the box
+
+On your machine(s) just run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked, specifically choose `MPS` for the query `Which type of machine are you using?` . 
+This will generate a config file that will be used automatically to properly set the
+default options when doing
+
+```bash
+accelerate launch my_script.py --args_to_my_script
+```
+
+For instance, here is how you would run the CV example (from the root of the repo) with MPS enabled:
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config: {}
+distributed_type: MPS
+downcast_bf16: 'no'
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: 'no'
+num_machines: 1
+num_processes: 1
+use_cpu: false
+```
If you're using comments inside a code-block in the documentation, you are doing something wrong; you can put the comments as sentences **then** enter a code-block. Also you should use the doctest syntax to highlight the inputs and outputs, as it's not clear otherwise.
@@ -0,0 +1,100 @@
+<!--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.
+-->
+
+# Accelerated PyTorch Training on Mac
+
+With PyTorch v1.12 release, developers and researchers can take advantage of Apple silicon GPUs for significantly faster model training. 
+This unlocks the ability to perform machine learning workflows like prototyping and fine-tuning locally, right on Mac.
+Apple's Metal Performance Shaders (MPS) as a backend for PyTorch enables this and can be used via the new `"mps"` device. 
+This will map computational graphs and primitives on the MPS Graph framework and tuned kernels provided by MPS.
+For more information please refer official documents [Introducing Accelerated PyTorch Training on Mac](https://pytorch.org/blog/introducing-accelerated-pytorch-training-on-mac/)
+and [MPS BACKEND](https://pytorch.org/docs/stable/notes/mps.html).
+
+### Benefits of Training and Inference using Apple M1 Chips
+
+1. Enables users to train larger networks or batch sizes locally
+2. Reduces data retrieval latency and provides the GPU with direct access to the full memory store due to unified memory architecture. 
+Therefore, improving end-to-end performance.
+3. Reduces costs associated with cloud-based development or the need for additional local GPUs.
+
+### Pre-requisites: Installing torch with mps support
+
+Following this great medium article [GPU-Acceleration Comes to PyTorch on M1 Macs](https://medium.com/towards-data-science/gpu-acceleration-comes-to-pytorch-on-m1-macs-195c399efcc1),
+below are the steps to make sure torch is installed with `mps` support.
+
+```
+# installing torch with m1 support on mac
+# install latest python from https://www.python.org/downloads/release/python-3106/
+# check the platform
+import platform
+platform.platform()
+'macOS-12.5-arm64-arm-64bit' 
+# This is compatible as the macOS version is above 12.3 and it is the ARM64 native version of Python
+# install torch 1.12 via the below command
+# pip3 install torch torchvision torchaudio
+# test the `mps` device support
+import torch
+torch.has_mps
+True
+a = torch.Tensor([10,11])
+a.to("mps")
+/Users/mac/ml/lib/python3.10/site-packages/torch/_tensor_str.py:103: UserWarning: The operator 'aten::bitwise_and.Tensor_out' is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at  /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/mps/MPSFallback.mm:11.)
+  nonzero_finite_vals = torch.masked_select(tensor_view, torch.isfinite(tensor_view) & tensor_view.ne(0))
+tensor([10.0000, 11.0000], device='mps:0')
+```
I'd probably change it to be:

```
This will generate a config file that will be used automatically to properly set the default options when doing `accelerate launch`, such as the one shown below:
```
And just don't show the full `accelerate launch` snippet, since here we want to focus on the config not the launch.

The next part would then say:
```
After this configuration has been made, here is how you run the CV example (from the root of the repo) with MPS enabled:
```

This way each part is focused on what matters to that section
@@ -0,0 +1,100 @@
+<!--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.
+-->
+
+# Accelerated PyTorch Training on Mac
+
+With PyTorch v1.12 release, developers and researchers can take advantage of Apple silicon GPUs for significantly faster model training. 
+This unlocks the ability to perform machine learning workflows like prototyping and fine-tuning locally, right on Mac.
+Apple's Metal Performance Shaders (MPS) as a backend for PyTorch enables this and can be used via the new `"mps"` device. 
+This will map computational graphs and primitives on the MPS Graph framework and tuned kernels provided by MPS.
+For more information please refer official documents [Introducing Accelerated PyTorch Training on Mac](https://pytorch.org/blog/introducing-accelerated-pytorch-training-on-mac/)
+and [MPS BACKEND](https://pytorch.org/docs/stable/notes/mps.html).
+
+### Benefits of Training and Inference using Apple M1 Chips
+
+1. Enables users to train larger networks or batch sizes locally
+2. Reduces data retrieval latency and provides the GPU with direct access to the full memory store due to unified memory architecture. 
+Therefore, improving end-to-end performance.
+3. Reduces costs associated with cloud-based development or the need for additional local GPUs.
+
+### Pre-requisites: Installing torch with mps support
+
+Following this great medium article [GPU-Acceleration Comes to PyTorch on M1 Macs](https://medium.com/towards-data-science/gpu-acceleration-comes-to-pytorch-on-m1-macs-195c399efcc1),
+below are the steps to make sure torch is installed with `mps` support.
+
+```
+# installing torch with m1 support on mac
+# install latest python from https://www.python.org/downloads/release/python-3106/
+# check the platform
+import platform
+platform.platform()
+'macOS-12.5-arm64-arm-64bit' 
+# This is compatible as the macOS version is above 12.3 and it is the ARM64 native version of Python
+# install torch 1.12 via the below command
+# pip3 install torch torchvision torchaudio
+# test the `mps` device support
+import torch
+torch.has_mps
+True
+a = torch.Tensor([10,11])
+a.to("mps")
+/Users/mac/ml/lib/python3.10/site-packages/torch/_tensor_str.py:103: UserWarning: The operator 'aten::bitwise_and.Tensor_out' is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at  /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/mps/MPSFallback.mm:11.)
+  nonzero_finite_vals = torch.masked_select(tensor_view, torch.isfinite(tensor_view) & tensor_view.ne(0))
+tensor([10.0000, 11.0000], device='mps:0')
+```
+
+## How it works out of the box
+
+On your machine(s) just run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked, specifically choose `MPS` for the query `Which type of machine are you using?` . 
+This will generate a config file that will be used automatically to properly set the
+default options when doing
+
+```bash
+accelerate launch my_script.py --args_to_my_script
+```
+
+For instance, here is how you would run the CV example (from the root of the repo) with MPS enabled:
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config: {}
+distributed_type: MPS
+downcast_bf16: 'no'
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: 'no'
+num_machines: 1
+num_processes: 1
+use_cpu: false
+```
Oh, okay, then maybe I can just leave it at the link to the medium blogpost instead of summarizing it this way. 
@@ -0,0 +1,100 @@
+<!--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.
+-->
+
+# Accelerated PyTorch Training on Mac
+
+With PyTorch v1.12 release, developers and researchers can take advantage of Apple silicon GPUs for significantly faster model training. 
+This unlocks the ability to perform machine learning workflows like prototyping and fine-tuning locally, right on Mac.
+Apple's Metal Performance Shaders (MPS) as a backend for PyTorch enables this and can be used via the new `"mps"` device. 
+This will map computational graphs and primitives on the MPS Graph framework and tuned kernels provided by MPS.
+For more information please refer official documents [Introducing Accelerated PyTorch Training on Mac](https://pytorch.org/blog/introducing-accelerated-pytorch-training-on-mac/)
+and [MPS BACKEND](https://pytorch.org/docs/stable/notes/mps.html).
+
+### Benefits of Training and Inference using Apple M1 Chips
+
+1. Enables users to train larger networks or batch sizes locally
+2. Reduces data retrieval latency and provides the GPU with direct access to the full memory store due to unified memory architecture. 
+Therefore, improving end-to-end performance.
+3. Reduces costs associated with cloud-based development or the need for additional local GPUs.
+
+### Pre-requisites: Installing torch with mps support
+
+Following this great medium article [GPU-Acceleration Comes to PyTorch on M1 Macs](https://medium.com/towards-data-science/gpu-acceleration-comes-to-pytorch-on-m1-macs-195c399efcc1),
+below are the steps to make sure torch is installed with `mps` support.
+
+```
+# installing torch with m1 support on mac
+# install latest python from https://www.python.org/downloads/release/python-3106/
+# check the platform
+import platform
+platform.platform()
+'macOS-12.5-arm64-arm-64bit' 
+# This is compatible as the macOS version is above 12.3 and it is the ARM64 native version of Python
+# install torch 1.12 via the below command
+# pip3 install torch torchvision torchaudio
+# test the `mps` device support
+import torch
+torch.has_mps
+True
+a = torch.Tensor([10,11])
+a.to("mps")
+/Users/mac/ml/lib/python3.10/site-packages/torch/_tensor_str.py:103: UserWarning: The operator 'aten::bitwise_and.Tensor_out' is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at  /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/mps/MPSFallback.mm:11.)
+  nonzero_finite_vals = torch.masked_select(tensor_view, torch.isfinite(tensor_view) & tensor_view.ne(0))
+tensor([10.0000, 11.0000], device='mps:0')
+```
Done. Removed this entire block and just left it with a link to the medium blogpost explaining this in detail. 
@@ -0,0 +1,100 @@
+<!--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.
+-->
+
+# Accelerated PyTorch Training on Mac
+
+With PyTorch v1.12 release, developers and researchers can take advantage of Apple silicon GPUs for significantly faster model training. 
+This unlocks the ability to perform machine learning workflows like prototyping and fine-tuning locally, right on Mac.
+Apple's Metal Performance Shaders (MPS) as a backend for PyTorch enables this and can be used via the new `"mps"` device. 
+This will map computational graphs and primitives on the MPS Graph framework and tuned kernels provided by MPS.
+For more information please refer official documents [Introducing Accelerated PyTorch Training on Mac](https://pytorch.org/blog/introducing-accelerated-pytorch-training-on-mac/)
+and [MPS BACKEND](https://pytorch.org/docs/stable/notes/mps.html).
+
+### Benefits of Training and Inference using Apple M1 Chips
+
+1. Enables users to train larger networks or batch sizes locally
+2. Reduces data retrieval latency and provides the GPU with direct access to the full memory store due to unified memory architecture. 
+Therefore, improving end-to-end performance.
+3. Reduces costs associated with cloud-based development or the need for additional local GPUs.
+
+### Pre-requisites: Installing torch with mps support
+
+Following this great medium article [GPU-Acceleration Comes to PyTorch on M1 Macs](https://medium.com/towards-data-science/gpu-acceleration-comes-to-pytorch-on-m1-macs-195c399efcc1),
+below are the steps to make sure torch is installed with `mps` support.
+
+```
+# installing torch with m1 support on mac
+# install latest python from https://www.python.org/downloads/release/python-3106/
+# check the platform
+import platform
+platform.platform()
+'macOS-12.5-arm64-arm-64bit' 
+# This is compatible as the macOS version is above 12.3 and it is the ARM64 native version of Python
+# install torch 1.12 via the below command
+# pip3 install torch torchvision torchaudio
+# test the `mps` device support
+import torch
+torch.has_mps
+True
+a = torch.Tensor([10,11])
+a.to("mps")
+/Users/mac/ml/lib/python3.10/site-packages/torch/_tensor_str.py:103: UserWarning: The operator 'aten::bitwise_and.Tensor_out' is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at  /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/mps/MPSFallback.mm:11.)
+  nonzero_finite_vals = torch.masked_select(tensor_view, torch.isfinite(tensor_view) & tensor_view.ne(0))
+tensor([10.0000, 11.0000], device='mps:0')
+```
This should be rewritten or removed now that the PyTorch issue has been solved (in nightly at least), no?
@@ -0,0 +1,81 @@
+<!--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.
+-->
+
+# Accelerated PyTorch Training on Mac
+
+With PyTorch v1.12 release, developers and researchers can take advantage of Apple silicon GPUs for significantly faster model training. 
+This unlocks the ability to perform machine learning workflows like prototyping and fine-tuning locally, right on Mac.
+Apple's Metal Performance Shaders (MPS) as a backend for PyTorch enables this and can be used via the new `"mps"` device. 
+This will map computational graphs and primitives on the MPS Graph framework and tuned kernels provided by MPS.
+For more information please refer official documents [Introducing Accelerated PyTorch Training on Mac](https://pytorch.org/blog/introducing-accelerated-pytorch-training-on-mac/)
+and [MPS BACKEND](https://pytorch.org/docs/stable/notes/mps.html).
+
+### Benefits of Training and Inference using Apple M1 Chips
+
+1. Enables users to train larger networks or batch sizes locally
+2. Reduces data retrieval latency and provides the GPU with direct access to the full memory store due to unified memory architecture. 
+Therefore, improving end-to-end performance.
+3. Reduces costs associated with cloud-based development or the need for additional local GPUs.
+
+**Pre-requisites**: To install torch with mps support, 
+please follow this nice medium article [GPU-Acceleration Comes to PyTorch on M1 Macs](https://medium.com/towards-data-science/gpu-acceleration-comes-to-pytorch-on-m1-macs-195c399efcc1).
+
+
+## How it works out of the box
+
+On your machine(s) just run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked, specifically choose `MPS` for the query:
+
+```
+ Which type of machine are you using?. 
+ ```
+
+This will generate a config file that will be used automatically to properly set 
+the default options when doing `accelerate launch`, such as the one shown below:
+
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config: {}
+distributed_type: MPS
+downcast_bf16: 'no'
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: 'no'
+num_machines: 1
+num_processes: 1
+use_cpu: false
+```
+
+After this configuration has been made, here is how you run the CV example 
+(from the root of the repo) with MPS enabled:
+
+```bash
+accelerate launch /examples/cv_example.py --data_dir images
+```
+
+## A few caveats to be aware of
+
+1. For `nlp_example.py` the metrics are too bad when compared to CPU-only training. 
+This means certain operations in BERT model are going wrong using mps device and this needs to be fixed by PyTorch.
Done. Please refer this PR #629 
@@ -0,0 +1,81 @@
+<!--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.
+-->
+
+# Accelerated PyTorch Training on Mac
+
+With PyTorch v1.12 release, developers and researchers can take advantage of Apple silicon GPUs for significantly faster model training. 
+This unlocks the ability to perform machine learning workflows like prototyping and fine-tuning locally, right on Mac.
+Apple's Metal Performance Shaders (MPS) as a backend for PyTorch enables this and can be used via the new `"mps"` device. 
+This will map computational graphs and primitives on the MPS Graph framework and tuned kernels provided by MPS.
+For more information please refer official documents [Introducing Accelerated PyTorch Training on Mac](https://pytorch.org/blog/introducing-accelerated-pytorch-training-on-mac/)
+and [MPS BACKEND](https://pytorch.org/docs/stable/notes/mps.html).
+
+### Benefits of Training and Inference using Apple M1 Chips
+
+1. Enables users to train larger networks or batch sizes locally
+2. Reduces data retrieval latency and provides the GPU with direct access to the full memory store due to unified memory architecture. 
+Therefore, improving end-to-end performance.
+3. Reduces costs associated with cloud-based development or the need for additional local GPUs.
+
+**Pre-requisites**: To install torch with mps support, 
+please follow this nice medium article [GPU-Acceleration Comes to PyTorch on M1 Macs](https://medium.com/towards-data-science/gpu-acceleration-comes-to-pytorch-on-m1-macs-195c399efcc1).
+
+
+## How it works out of the box
+
+On your machine(s) just run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked, specifically choose `MPS` for the query:
+
+```
+ Which type of machine are you using?. 
+ ```
+
+This will generate a config file that will be used automatically to properly set 
+the default options when doing `accelerate launch`, such as the one shown below:
+
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config: {}
+distributed_type: MPS
+downcast_bf16: 'no'
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: 'no'
+num_machines: 1
+num_processes: 1
+use_cpu: false
+```
+
+After this configuration has been made, here is how you run the CV example 
+(from the root of the repo) with MPS enabled:
+
+```bash
+accelerate launch /examples/cv_example.py --data_dir images
+```
+
+## A few caveats to be aware of
+
+1. For `nlp_example.py` the metrics are too bad when compared to CPU-only training. 
+This means certain operations in BERT model are going wrong using mps device and this needs to be fixed by PyTorch.
_The documentation is not available anymore as the PR was closed or merged._
`notebook_launcher` doesn't seem to be modified in order to utilize this change that supports MPS device.
https://api.github.com/repos/huggingface/accelerate/pulls/596
diff --git a/setup.py b/setup.py
index a59d1ecaf..ec015f276 100644
--- a/setup.py
+++ b/setup.py
@@ -19,7 +19,7 @@
 extras["quality"] = ["black ~= 22.0", "isort >= 5.5.4", "flake8 >= 3.8.3", "hf-doc-builder >= 0.3.0"]
 extras["docs"] = []
 extras["test_prod"] = ["pytest", "pytest-xdist", "pytest-subtests", "parameterized"]
-extras["test_dev"] = ["datasets", "evaluate", "transformers", "scipy", "sklearn", "deepspeed", "tqdm"]
+extras["test_dev"] = ["datasets", "evaluate", "transformers", "scipy", "sklearn", "deepspeed<0.7.0", "tqdm"]
 extras["testing"] = extras["test_prod"] + extras["test_dev"]
 
 extras["test_trackers"] = ["wandb", "comet-ml", "tensorboard"]

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/595
diff --git a/docs/source/usage_guides/tracking.mdx b/docs/source/usage_guides/tracking.mdx
index 2fb68c0f2..7c61b21e9 100644
--- a/docs/source/usage_guides/tracking.mdx
+++ b/docs/source/usage_guides/tracking.mdx
@@ -13,7 +13,7 @@ specific language governing permissions and limitations under the License.
 # Tracking
 
 There are a large number of experiment tracking API's available, however getting them all to work with in a multi-processing environment can oftentimes be complex.
-🤗 Accelerate provides a general tracking API that can be used to log useful items during your script through [`~Accelerator.log`]
+🤗 Accelerate provides a general tracking API that can be used to log useful items during your script through [`Accelerator.log`]
 
 ## Integrated Trackers
 
@@ -33,19 +33,19 @@ accelerator = Accelerator(log_with="wandb")
 accelerator = Accelerator(log_with=["wandb", LoggerType.TENSORBOARD])
 ```
 
-At the start of your experiment [`~Accelerator.init_trackers`] should be used to setup your project, and potentially add any experiment hyperparameters to be logged:
+At the start of your experiment [`Accelerator.init_trackers`] should be used to setup your project, and potentially add any experiment hyperparameters to be logged:
 ```python
 hps = {"num_iterations": 5, "learning_rate": 1e-2}
 accelerator.init_trackers("my_project", config=hps)
 ```
 
-When you are ready to log any data, [`~Accelerator.log`] should be used.
+When you are ready to log any data, [`Accelerator.log`] should be used.
 A `step` can also be passed in to correlate the data with a particular step in the training loop.
 ```python
 accelerator.log({"train_loss": 1.12, "valid_loss": 0.8}, step=1)
 ```
 
-Once you've finished training, make sure to run [`~Accelerator.end_training`] so that all the trackers can run their finish functionalities if they have any.
+Once you've finished training, make sure to run [`Accelerator.end_training`] so that all the trackers can run their finish functionalities if they have any.
 ```python
 accelerator.end_training()
 ```
@@ -85,8 +85,8 @@ accelerator.end_training()
 
 ## Implementing Custom Trackers
 
-To implement a new tracker to be used in `Accelerator`, a new one can be made through implementing the [`~GeneralTracker`] class.
-Every tracker must implement three functions:
+To implement a new tracker to be used in `Accelerator`, a new one can be made through implementing the [`GeneralTracker`] class.
+Every tracker must implement three functions and have three properties:
   - `__init__`: 
     - Should store a `run_name` and initialize the tracker API of the integrated library. 
     - If a tracker stores their data locally (such as TensorBoard), a `logging_dir` parameter can be added.
@@ -95,6 +95,15 @@ Every tracker must implement three functions:
   - `log`: 
     - Should take in a `values` dictionary and a `step`, and should log them to the run
 
+  - `name` (`str`):
+    - A unique string name for the tracker, such as `"wandb"` for the wandb tracker. 
+    - This will be used for interacting with this tracker specifically
+  - `requires_logging_directory` (`bool`):
+    - Whether a `logging_dir` is needed for this particular tracker and if it uses one.
+  - `tracker`: 
+    - This should be implemented as a `@property` function 
+    - Should return the internal tracking mechanism the library uses, such as the `run` object for `wandb`.
+
 A brief example can be seen below with an integration with Weights and Biases, containing only the relevent information:
 ```python
 from accelerate.tracking import GeneralTracker
@@ -109,7 +118,11 @@ class MyCustomTracker(GeneralTracker):
 
     def __init__(self, run_name: str):
         self.run_name = run_name
-        wandb.init(self.run_name)
+        run = wandb.init(self.run_name)
+
+    @property
+    def tracker(self):
+        return self.run.run
 
     def store_init_configuration(self, values: dict):
         wandb.config(values)
@@ -118,7 +131,7 @@ class MyCustomTracker(GeneralTracker):
         wandb.log(values, step=step)
 ```
 
-When you are ready to build your `Accelerator` object, pass in an **instance** of your tracker to [`~Accelerator.log_with`] to have it automatically
+When you are ready to build your `Accelerator` object, pass in an **instance** of your tracker to [`Accelerator.log_with`] to have it automatically
 be used with the API:
 
 ```python
@@ -133,6 +146,30 @@ tracker = MyCustomTracker("some_run_name")
 accelerator = Accelerator(log_with=[tracker, "all"])
 ```
 
+## Accessing the internal tracker 
+
+If some custom interactions with a tracker might be wanted directly, you can quickly access one using the 
+[`Accelerator.get_tracker`] method. Just pass in the string corresponding to a tracker's `.name` attribute 
+and it will return that tracker on the main process.
+
+This example shows doing so with wandb:
+
+```python
+wandb_tracker = accelerator.get_tracker("wandb")
+```
+
+From there you can interact with `wandb`'s `run` object like normal:
+
+<Tip warning={true}>
+  Make sure to only interact with trackers on the main process!
+</Tip>
+
+
+```python
+if accelerator.is_main_process:
+    wandb_run.log_artifact(some_artifact_to_log)
+```
+
 ## When a wrapper cannot work
 
 If a library has an API that does not follow a strict `.log` with an overall dictionary such as Neptune.AI, logging can be done manually under an `if accelerator.is_main_process` statement:
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 093d4c00f..bbd8db435 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -1069,10 +1069,24 @@ def init_trackers(self, project_name: str, config: Optional[dict] = None, init_k
             for tracker in self.trackers:
                 tracker.store_init_configuration(config)
 
+    @on_main_process
+    def get_tracker(self, name: str):
+        """
+        Returns a `tracker` from `self.trackers` based on `name` on the main process only.
+
+        Args:
+            name (`str`):
+                The name of a tracker, corresponding to the `.name` property.
+        """
+        for tracker in self.trackers:
+            if tracker.name == name:
+                return tracker.tracker
+        raise ValueError(f"{name} is not an available tracker stored inside the `Accelerator`.")
+
     @on_main_process
     def log(self, values: dict, step: Optional[int] = None, log_kwargs: Optional[dict] = {}):
         """
-        Logs `values` to all stored trackers in `self.trackers`.
+        Logs `values` to all stored trackers in `self.trackers` on the main process only.
 
         Args:
             values (`dict`):
@@ -1092,7 +1106,7 @@ def log(self, values: dict, step: Optional[int] = None, log_kwargs: Optional[dic
     @on_main_process
     def end_training(self):
         """
-        Runs any special end training behaviors, such as stopping trackers
+        Runs any special end training behaviors, such as stopping trackers on the main process only.
         """
         for tracker in self.trackers:
             tracker.finish()
diff --git a/src/accelerate/tracking.py b/src/accelerate/tracking.py
index b89630a4e..673975fc7 100644
--- a/src/accelerate/tracking.py
+++ b/src/accelerate/tracking.py
@@ -103,6 +103,13 @@ def finish(self):
         """
         pass
 
+    @abstractproperty
+    def tracker(self):
+        """
+        Should return internal tracking mechanism used by a tracker class (such as the `run` for wandb)
+        """
+        pass
+
 
 class TensorBoardTracker(GeneralTracker):
     """
@@ -129,6 +136,10 @@ def __init__(self, run_name: str, logging_dir: Optional[Union[str, os.PathLike]]
             "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
         )
 
+    @property
+    def tracker(self):
+        return self.writer
+
     def store_init_configuration(self, values: dict):
         """
         Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment.
@@ -196,6 +207,10 @@ def __init__(self, run_name: str, **kwargs):
             "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
         )
 
+    @property
+    def tracker(self):
+        return self.run.run
+
     def store_init_configuration(self, values: dict):
         """
         Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment.
@@ -256,6 +271,10 @@ def __init__(self, run_name: str, **kwargs):
             "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
         )
 
+    @property
+    def tracker(self):
+        return self.writer
+
     def store_init_configuration(self, values: dict):
         """
         Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment.
diff --git a/tests/test_tracking.py b/tests/test_tracking.py
index c32502796..917ee329d 100644
--- a/tests/test_tracking.py
+++ b/tests/test_tracking.py
@@ -224,6 +224,10 @@ def __init__(self, dir: str):
         self.writer = csv.DictWriter(self.f, fieldnames=self._col_names)
         self.writer.writeheader()
 
+    @property
+    def tracker(self):
+        return self.writer
+
     def store_init_configuration(self, values: dict):
         logger.info("Call init")
         self.writer.writerow(values)

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/594
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index c2fd566b7..093d4c00f 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -412,7 +412,7 @@ def wrapper(self, *args, **kwargs):
 
     def on_local_process(local_process_idx):
         """
-        Run func on certain local process only
+        A decorator that will run the decorated function on a given local process index only.
         """
 
         def decorator(func):

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/591
diff --git a/examples/by_feature/checkpointing.py b/examples/by_feature/checkpointing.py
index 1b462a1f4..bffd843d8 100644
--- a/examples/by_feature/checkpointing.py
+++ b/examples/by_feature/checkpointing.py
@@ -236,23 +236,13 @@ def training_function(config, args):
                     accelerator.save_state(output_dir)
 
         model.eval()
-        samples_seen = 0
         for step, batch in enumerate(eval_dataloader):
             # We could avoid this line since we set the accelerator with `device_placement=True` (the default).
             batch.to(accelerator.device)
             with torch.no_grad():
                 outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1)
-            # It is slightly faster to call this once, than multiple times
-            predictions, references = accelerator.gather((predictions, batch["labels"]))
-            if accelerator.use_distributed:
-                if step == len(eval_dataloader) - 1:
-                    # Last batch needs to be truncated on distributed systems as it contains additional samples
-                    predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
-                    references = references[: len(eval_dataloader.dataset) - samples_seen]
-                else:
-                    # Otherwise we add the number of samples seen
-                    samples_seen += references.shape[0]
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
             metric.add_batch(
                 predictions=predictions,
                 references=references,
diff --git a/examples/by_feature/cross_validation.py b/examples/by_feature/cross_validation.py
index 6dd61bbf8..707c93c4c 100644
--- a/examples/by_feature/cross_validation.py
+++ b/examples/by_feature/cross_validation.py
@@ -203,23 +203,13 @@ def training_function(config, args):
                     optimizer.zero_grad()
 
             model.eval()
-            samples_seen = 0
             for step, batch in enumerate(eval_dataloader):
                 # We could avoid this line since we set the accelerator with `device_placement=True`.
                 batch.to(accelerator.device)
                 with torch.no_grad():
                     outputs = model(**batch)
                 predictions = outputs.logits.argmax(dim=-1)
-                # It is slightly faster to call this once, than multiple times
-                predictions, references = accelerator.gather((predictions, batch["labels"]))
-                if accelerator.use_distributed:
-                    if step == len(eval_dataloader) - 1:
-                        # Last batch needs to be truncated on distributed systems as it contains additional samples
-                        predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
-                        references = references[: len(eval_dataloader.dataset) - samples_seen]
-                    else:
-                        # Otherwise we add the number of samples seen
-                        samples_seen += references.shape[0]
+                predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
                 metric.add_batch(
                     predictions=predictions,
                     references=references,
diff --git a/examples/by_feature/fsdp_with_peak_mem_tracking.py b/examples/by_feature/fsdp_with_peak_mem_tracking.py
index a279e64eb..274ccdfce 100644
--- a/examples/by_feature/fsdp_with_peak_mem_tracking.py
+++ b/examples/by_feature/fsdp_with_peak_mem_tracking.py
@@ -274,23 +274,13 @@ def collate_fn(examples):
         # context manager to track the peak memory usage during the evaluation
         with TorchTracemalloc() as tracemalloc:
             model.eval()
-            samples_seen = 0
             for step, batch in enumerate(eval_dataloader):
                 # We could avoid this line since we set the accelerator with `device_placement=True`.
                 batch.to(accelerator.device)
                 with torch.no_grad():
                     outputs = model(**batch)
                 predictions = outputs.logits.argmax(dim=-1)
-                # It is slightly faster to call this once, than multiple times
-                predictions, references = accelerator.gather((predictions, batch["labels"]))
-                if accelerator.use_distributed:
-                    if step == len(eval_dataloader) - 1:
-                        # Last batch needs to be truncated on distributed systems as it contains additional samples
-                        predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
-                        references = references[: len(eval_dataloader.dataset) - samples_seen]
-                    else:
-                        # Otherwise we add the number of samples seen
-                        samples_seen += references.shape[0]
+                predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
                 metric.add_batch(
                     predictions=predictions,
                     references=references,
diff --git a/examples/by_feature/gradient_accumulation.py b/examples/by_feature/gradient_accumulation.py
index 69856c97b..170a885e3 100644
--- a/examples/by_feature/gradient_accumulation.py
+++ b/examples/by_feature/gradient_accumulation.py
@@ -170,23 +170,13 @@ def training_function(config, args):
                 optimizer.zero_grad()
 
         model.eval()
-        samples_seen = 0
         for step, batch in enumerate(eval_dataloader):
             # We could avoid this line since we set the accelerator with `device_placement=True`.
             batch.to(accelerator.device)
             with torch.no_grad():
                 outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1)
-            # It is slightly faster to call this once, than multiple times
-            predictions, references = accelerator.gather((predictions, batch["labels"]))
-            if accelerator.use_distributed:
-                if step == len(eval_dataloader) - 1:
-                    # Last batch needs to be truncated on distributed systems as it contains additional samples
-                    predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
-                    references = references[: len(eval_dataloader.dataset) - samples_seen]
-                else:
-                    # Otherwise we add the number of samples seen
-                    samples_seen += references.shape[0]
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
             metric.add_batch(
                 predictions=predictions,
                 references=references,
diff --git a/examples/by_feature/memory.py b/examples/by_feature/memory.py
index 91f3e41d1..9f0d55c69 100644
--- a/examples/by_feature/memory.py
+++ b/examples/by_feature/memory.py
@@ -182,23 +182,13 @@ def inner_training_loop(batch_size):
                     optimizer.zero_grad()
 
             model.eval()
-            samples_seen = 0
             for step, batch in enumerate(eval_dataloader):
                 # We could avoid this line since we set the accelerator with `device_placement=True`.
                 batch.to(accelerator.device)
                 with torch.no_grad():
                     outputs = model(**batch)
                 predictions = outputs.logits.argmax(dim=-1)
-                # It is slightly faster to call this once, than multiple times
-                predictions, references = accelerator.gather((predictions, batch["labels"]))
-                if accelerator.use_distributed:
-                    if step == len(eval_dataloader) - 1:
-                        # Last batch needs to be truncated on distributed systems as it contains additional samples
-                        predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
-                        references = references[: len(eval_dataloader.dataset) - samples_seen]
-                    else:
-                        # Otherwise we add the number of samples seen
-                        samples_seen += references.shape[0]
+                predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
                 metric.add_batch(
                     predictions=predictions,
                     references=references,
diff --git a/examples/by_feature/multi_process_metrics.py b/examples/by_feature/multi_process_metrics.py
index 522cc571b..cb9534c4a 100644
--- a/examples/by_feature/multi_process_metrics.py
+++ b/examples/by_feature/multi_process_metrics.py
@@ -192,6 +192,8 @@ def training_function(config, args):
                 else:
                     # Otherwise we add the number of samples seen
                     samples_seen += references.shape[0]
+            # All of this can be avoided if you use `Accelerator.gather_for_metrics` instead of `Accelerator.gather`:
+            # accelerator.gather_for_metrics((predictions, batch["labels"]))
             metric.add_batch(
                 predictions=predictions,
                 references=references,
diff --git a/examples/by_feature/tracking.py b/examples/by_feature/tracking.py
index d7248682d..76ad7a64b 100644
--- a/examples/by_feature/tracking.py
+++ b/examples/by_feature/tracking.py
@@ -194,23 +194,13 @@ def training_function(config, args):
                 optimizer.zero_grad()
 
         model.eval()
-        samples_seen = 0
         for step, batch in enumerate(eval_dataloader):
             # We could avoid this line since we set the accelerator with `device_placement=True` (the default).
             batch.to(accelerator.device)
             with torch.no_grad():
                 outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1)
-            # It is slightly faster to call this once, than multiple times
-            predictions, references = accelerator.gather((predictions, batch["labels"]))
-            if accelerator.use_distributed:
-                if step == len(eval_dataloader) - 1:
-                    # Last batch needs to be truncated on distributed systems as it contains additional samples
-                    predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
-                    references = references[: len(eval_dataloader.dataset) - samples_seen]
-                else:
-                    # Otherwise we add the number of samples seen
-                    samples_seen += references.shape[0]
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
             metric.add_batch(
                 predictions=predictions,
                 references=references,
diff --git a/examples/complete_cv_example.py b/examples/complete_cv_example.py
index f984406bf..99d109bd1 100644
--- a/examples/complete_cv_example.py
+++ b/examples/complete_cv_example.py
@@ -232,7 +232,6 @@ def training_function(config, args):
                     accelerator.save_state(output_dir)
         model.eval()
         accurate = 0
-        samples_seen = 0
         for step, batch in enumerate(eval_dataloader):
             # We could avoid this line since we set the accelerator with `device_placement=True`.
             batch = {k: v.to(accelerator.device) for k, v in batch.items()}
@@ -240,16 +239,7 @@ def training_function(config, args):
             with torch.no_grad():
                 outputs = model(inputs)
             predictions = outputs.argmax(dim=-1)
-            # It is slightly faster to call this once, than multiple times
-            predictions, references = accelerator.gather((predictions, batch["label"]))
-            if accelerator.use_distributed:
-                if step == len(eval_dataloader) - 1:
-                    # Last batch needs to be truncated on distributed systems as it contains additional samples
-                    predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
-                    references = references[: len(eval_dataloader.dataset) - samples_seen]
-                else:
-                    # Otherwise we add the number of samples seen
-                    samples_seen += references.shape[0]
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["label"]))
             accurate_preds = predictions == references
             accurate += accurate_preds.long().sum()
 
diff --git a/examples/complete_nlp_example.py b/examples/complete_nlp_example.py
index 000da6038..094f6cfb5 100644
--- a/examples/complete_nlp_example.py
+++ b/examples/complete_nlp_example.py
@@ -212,23 +212,13 @@ def collate_fn(examples):
                     accelerator.save_state(output_dir)
 
         model.eval()
-        samples_seen = 0
         for step, batch in enumerate(eval_dataloader):
             # We could avoid this line since we set the accelerator with `device_placement=True`.
             batch.to(accelerator.device)
             with torch.no_grad():
                 outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1)
-            # It is slightly faster to call this once, than multiple times
-            predictions, references = accelerator.gather((predictions, batch["labels"]))
-            if accelerator.use_distributed:
-                if step == len(eval_dataloader) - 1:
-                    # Last batch needs to be truncated on distributed systems as it contains additional samples
-                    predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
-                    references = references[: len(eval_dataloader.dataset) - samples_seen]
-                else:
-                    # Otherwise we add the number of samples seen
-                    samples_seen += references.shape[0]
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
             metric.add_batch(
                 predictions=predictions,
                 references=references,
diff --git a/examples/cv_example.py b/examples/cv_example.py
index 875b68398..1118a2f0e 100644
--- a/examples/cv_example.py
+++ b/examples/cv_example.py
@@ -166,7 +166,6 @@ def training_function(config, args):
         model.eval()
         accurate = 0
         num_elems = 0
-        samples_seen = 0
         for _, batch in enumerate(eval_dataloader):
             # We could avoid this line since we set the accelerator with `device_placement=True`.
             batch = {k: v.to(accelerator.device) for k, v in batch.items()}
@@ -174,16 +173,7 @@ def training_function(config, args):
             with torch.no_grad():
                 outputs = model(inputs)
             predictions = outputs.argmax(dim=-1)
-            # It is slightly faster to call this once, than multiple times
-            predictions, references = accelerator.gather((predictions, batch["label"]))
-            if accelerator.use_distributed:
-                if step == len(eval_dataloader) - 1:
-                    # Last batch needs to be truncated on distributed systems as it contains additional samples
-                    predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
-                    references = references[: len(eval_dataloader.dataset) - samples_seen]
-                else:
-                    # Otherwise we add the number of samples seen
-                    samples_seen += references.shape[0]
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["label"]))
             accurate_preds = predictions == references
             num_elems += accurate_preds.shape[0]
             accurate += accurate_preds.long().sum()
diff --git a/examples/nlp_example.py b/examples/nlp_example.py
index b1e7cba27..a126b5dd5 100644
--- a/examples/nlp_example.py
+++ b/examples/nlp_example.py
@@ -154,23 +154,13 @@ def training_function(config, args):
                 optimizer.zero_grad()
 
         model.eval()
-        samples_seen = 0
         for step, batch in enumerate(eval_dataloader):
             # We could avoid this line since we set the accelerator with `device_placement=True`.
             batch.to(accelerator.device)
             with torch.no_grad():
                 outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1)
-            # It is slightly faster to call this once, than multiple times
-            predictions, references = accelerator.gather((predictions, batch["labels"]))
-            if accelerator.use_distributed:
-                if step == len(eval_dataloader) - 1:
-                    # Last batch needs to be truncated on distributed systems as it contains additional samples
-                    predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
-                    references = references[: len(eval_dataloader.dataset) - samples_seen]
-                else:
-                    # Otherwise we add the number of samples seen
-                    samples_seen += references.shape[0]
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
             metric.add_batch(
                 predictions=predictions,
                 references=references,
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index bbd8db435..ef421a22b 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -956,9 +956,6 @@ def gather_for_metrics(self, tensor):
             tensor (`torch.Tensor`, or a nested tuple/list/dictionary of `torch.Tensor`):
                 The tensors for calculating metrics across all processes.
         """
-        raise NotImplementedError(
-            "Currently there are a number of bugs with this method. You should use `Accelerator.gather()` and drop the samples yourself for the time being."
-        )
         tensor = self.gather(tensor)
         if self.use_distributed:
             if self.gradient_state.remainder == -1:
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index 4cedd8f1a..b8152b302 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -470,11 +470,6 @@ def _fetch_batches(self, iterator):
 
     def __iter__(self):
         self.gradient_state._set_end_of_dataloader(False)
-        try:
-            length = getattr(self.dataset, "total_dataset_length", len(self.dataset))
-            self.gradient_state._set_remainder(length % self.batch_size)
-        except:
-            self.gradient_state._set_remainder(-1)
         main_iterator = None
         if self.state.process_index == 0:
             # We only iterate through the DataLoader on process 0.
@@ -501,7 +496,16 @@ def __iter__(self):
             observed_batch_size = find_batch_size(batch)
             batch_size = observed_batch_size // self.state.num_processes
 
-            if not self._drop_last and self._stop_iteration and observed_batch_size % self.state.num_processes != 0:
+            stop_iteration = self._stop_iteration
+            if not stop_iteration:
+                # We may still be at the end of the dataloader without knowing it yet: if there is nothing left in
+                # the dataloader since the number of batches is a round multiple of the number of processes.
+                next_batch, next_batch_info, next_skip = self._fetch_batches(main_iterator)
+                # next_batch_info[0] is None when there are no more batches, otherwise we still need to process them.
+                if self._stop_iteration and next_batch_info[0] is None:
+                    stop_iteration = True
+
+            if not self._drop_last and stop_iteration and observed_batch_size % self.state.num_processes != 0:
                 # If the last batch is not complete, let's add the first batch to it.
                 batch = concatenate([batch, first_batch], dim=0)
                 # Batch size computation above is wrong, it's off by 1 so we fix it.
@@ -510,16 +514,8 @@ def __iter__(self):
             data_slice = slice(self.state.process_index * batch_size, (self.state.process_index + 1) * batch_size)
             batch = slice_tensors(batch, data_slice)
 
-            stop_iteration = self._stop_iteration
-            if not stop_iteration:
-                # We may still be at the end of the dataloader without knowing it yet: if there is nothing left
-                # because by change the dataset had a round multiple of samples.
-                next_batch, next_batch_info, next_skip = self._fetch_batches(main_iterator)
-                # next_batch_info[0] is None when there are no more batches, otherwise we still need to process them.
-                if self._stop_iteration and next_batch_info[0] is None:
-                    stop_iteration = True
-
             if stop_iteration:
+                self.gradient_state._set_remainder(observed_batch_size)
                 self.gradient_state._set_end_of_dataloader(True)
             yield batch
 
diff --git a/src/accelerate/test_utils/__init__.py b/src/accelerate/test_utils/__init__.py
index 382bf81bf..43d5ed7ef 100644
--- a/src/accelerate/test_utils/__init__.py
+++ b/src/accelerate/test_utils/__init__.py
@@ -7,6 +7,7 @@
     execute_subprocess_async,
     require_cpu,
     require_cuda,
+    require_huggingface_suite,
     require_multi_gpu,
     require_single_gpu,
     require_tpu,
@@ -16,4 +17,4 @@
 from .training import RegressionDataset, RegressionModel
 
 
-from .scripts import test_metrics, test_script, test_sync  # isort:skip
+from .scripts import test_script, test_sync  # isort:skip
diff --git a/src/accelerate/test_utils/scripts/external_deps/__init__.py b/src/accelerate/test_utils/scripts/external_deps/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/src/accelerate/test_utils/scripts/external_deps/test_metrics.py b/src/accelerate/test_utils/scripts/external_deps/test_metrics.py
new file mode 100755
index 000000000..d234ee14a
--- /dev/null
+++ b/src/accelerate/test_utils/scripts/external_deps/test_metrics.py
@@ -0,0 +1,170 @@
+# 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 math
+from copy import deepcopy
+
+import torch
+from torch.utils.data import DataLoader
+
+import datasets
+import evaluate
+import transformers
+from accelerate import Accelerator
+from accelerate.test_utils import RegressionDataset, RegressionModel
+from accelerate.utils import is_tpu_available, set_seed
+from datasets import load_dataset
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+
+
+def get_basic_setup(accelerator, num_samples=82):
+    "Returns everything needed to perform basic training"
+    set_seed(42)
+    model = RegressionModel()
+    ddp_model = deepcopy(model)
+    dset = RegressionDataset(length=num_samples)
+    dataloader = DataLoader(dset, batch_size=16)
+    model.to(accelerator.device)
+    ddp_model, dataloader = accelerator.prepare(ddp_model, dataloader)
+    return model, ddp_model, dataloader
+
+
+def get_dataloader(accelerator: Accelerator, use_longest=False):
+    tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/mrpc-bert-base-cased")
+    dataset = load_dataset("glue", "mrpc", split="validation")
+
+    def tokenize_function(examples):
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    with accelerator.main_process_first():
+        tokenized_datasets = dataset.map(
+            tokenize_function,
+            batched=True,
+            remove_columns=["idx", "sentence1", "sentence2"],
+        )
+
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    def collate_fn(examples):
+        if use_longest:
+            return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+        return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+
+    return DataLoader(tokenized_datasets, shuffle=False, collate_fn=collate_fn, batch_size=16)
+
+
+def get_mrpc_setup(dispatch_batches, split_batches):
+    accelerator = Accelerator(dispatch_batches=dispatch_batches, split_batches=split_batches)
+    dataloader = get_dataloader(accelerator, not dispatch_batches)
+    model = AutoModelForSequenceClassification.from_pretrained(
+        "hf-internal-testing/mrpc-bert-base-cased", return_dict=True
+    )
+    ddp_model, ddp_dataloader = accelerator.prepare(model, dataloader)
+    return {"ddp": [ddp_model, ddp_dataloader, "cuda:0"], "no": [model, dataloader, accelerator.device]}, accelerator
+
+
+def generate_predictions(model, dataloader, accelerator):
+    logits_and_targets = []
+    for batch in dataloader:
+        input, target = batch.values()
+        with torch.no_grad():
+            logit = model(input)
+            logit, target = accelerator.gather_for_metrics((logit, target))
+            logits_and_targets.append((logit, target))
+    logits, targs = [], []
+    for (logit, targ) in logits_and_targets:
+        logits.append(logit)
+        targs.append(targ)
+    logits, targs = torch.cat(logits), torch.cat(targs)
+    return logits, targs
+
+
+def test_torch_metrics(accelerator: Accelerator, num_samples=82, dispatch_batches=False, split_batches=False):
+    model, ddp_model, dataloader = get_basic_setup(accelerator, num_samples)
+    logits, targs = generate_predictions(ddp_model, dataloader, accelerator)
+    assert (
+        len(logits) == num_samples
+    ), f"Unexpected number of inputs:\n    Expected: {num_samples}\n    Actual: {len(logits)}"
+
+
+def test_mrpc(dispatch_batches: bool = False, split_batches: bool = False):
+    metric = evaluate.load("glue", "mrpc")
+    setup, accelerator = get_mrpc_setup(dispatch_batches, split_batches)
+    # First do baseline
+    model, dataloader, device = setup["no"]
+    model.to(device)
+    model.eval()
+    for batch in dataloader:
+        batch.to(device)
+        with torch.inference_mode():
+            outputs = model(**batch)
+        preds = outputs.logits.argmax(dim=-1)
+        metric.add_batch(predictions=preds, references=batch["labels"])
+    baseline = metric.compute()
+
+    # Then do distributed
+    model, dataloader, device = setup["ddp"]
+    model.eval()
+    for batch in dataloader:
+        with torch.inference_mode():
+            outputs = model(**batch)
+        preds = outputs.logits.argmax(dim=-1)
+        references = batch["labels"]
+        preds, references = accelerator.gather_for_metrics((preds, references))
+        metric.add_batch(predictions=preds, references=references)
+    distributed = metric.compute()
+
+    for key in "accuracy f1".split():
+        assert math.isclose(
+            baseline[key], distributed[key]
+        ), f"Baseline and Distributed are not the same for key {key}:\n\tBaseline: {baseline[key]}\n\tDistributed: {distributed[key]}\n"
+
+
+def main():
+    accelerator = Accelerator(split_batches=False, dispatch_batches=False)
+    if accelerator.is_local_main_process:
+        datasets.utils.logging.set_verbosity_warning()
+        transformers.utils.logging.set_verbosity_warning()
+    else:
+        datasets.utils.logging.set_verbosity_error()
+        transformers.utils.logging.set_verbosity_error()
+    # These are a bit slower so they should only be ran on the GPU or TPU
+    if torch.cuda.is_available() or is_tpu_available():
+        if accelerator.is_local_main_process:
+            print("**Testing gather_for_metrics**")
+        for split_batches in [True, False]:
+            for dispatch_batches in [True, False]:
+                if accelerator.is_local_main_process:
+                    print(f"With: `split_batches={split_batches}`, `dispatch_batches={dispatch_batches}`")
+                test_mrpc(dispatch_batches, split_batches)
+                accelerator.state._reset_state()
+    if accelerator.is_local_main_process:
+        print("**Test torch metrics**")
+    for split_batches in [True, False]:
+        for dispatch_batches in [True, False]:
+            accelerator = Accelerator(split_batches=split_batches, dispatch_batches=dispatch_batches)
+            if accelerator.is_local_main_process:
+                print(f"With: `split_batches={split_batches}`, `dispatch_batches={dispatch_batches}`, length=99")
+            test_torch_metrics(accelerator, 99)
+            accelerator.state._reset_state()
+
+
+def _mp_fn(index):
+    # For xla_spawn (TPUs)
+    main()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/accelerate/test_utils/scripts/test_metrics.py b/src/accelerate/test_utils/scripts/test_metrics.py
deleted file mode 100644
index 8f057ccc5..000000000
--- a/src/accelerate/test_utils/scripts/test_metrics.py
+++ /dev/null
@@ -1,90 +0,0 @@
-# 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.
-
-from copy import deepcopy
-
-import torch
-from torch.utils.data import DataLoader
-
-from accelerate import Accelerator
-from accelerate.test_utils import RegressionDataset, RegressionModel
-from accelerate.utils import set_seed
-
-
-def get_setup(accelerator, num_samples=82):
-    "Returns everything needed to perform basic training"
-    set_seed(42)
-    model = RegressionModel()
-    ddp_model = deepcopy(model)
-    dset = RegressionDataset(length=num_samples)
-    dataloader = DataLoader(dset, batch_size=16)
-    model.to(accelerator.device)
-    ddp_model, dataloader = accelerator.prepare(ddp_model, dataloader)
-    return model, ddp_model, dataloader
-
-
-def generate_predictions(model, dataloader, accelerator):
-    logits_and_targets = []
-    for batch in dataloader:
-        input, target = batch.values()
-        with torch.no_grad():
-            logits = model(input)
-            logits, target = accelerator.gather_for_metrics((logits, target))
-            logits_and_targets.append((logits, target))
-    inps, targs = [], []
-    for (inp, targ) in logits_and_targets:
-        inps.append(inp)
-        targs.append(targ)
-    inps, targs = torch.cat(inps), torch.cat(targs)
-    return inps, targs
-
-
-def test_torch_metrics(accelerator: Accelerator, num_samples=82):
-    model, ddp_model, dataloader = get_setup(accelerator, num_samples)
-    inps, targs = generate_predictions(ddp_model, dataloader, accelerator)
-    assert (
-        len(inps) == num_samples
-    ), f"Unexpected number of inputs:\n    Expected: {num_samples}\n    Actual: {len(inps)}"
-
-
-def main():
-    accelerator = Accelerator(split_batches=False, dispatch_batches=False)
-    if accelerator.is_local_main_process:
-        print("**Test torch metrics**")
-        print("With: `split_batches=False`, `dispatch_batches=False`")
-    test_torch_metrics(accelerator)
-    accelerator.state._reset_state()
-    accelerator = Accelerator(split_batches=True, dispatch_batches=False)
-    if accelerator.is_local_main_process:
-        print("With: `split_batches=True`, `dispatch_batches=False`")
-    test_torch_metrics(accelerator)
-    accelerator.state._reset_state()
-    accelerator = Accelerator(split_batches=False, dispatch_batches=True)
-    if accelerator.is_local_main_process:
-        print("With: `split_batches=False`, `dispatch_batches=True`")
-    test_torch_metrics(accelerator)
-    accelerator.state._reset_state()
-    accelerator = Accelerator(split_batches=True, dispatch_batches=True)
-    if accelerator.is_local_main_process:
-        print("With: `split_batches=True`, `dispatch_batches=True`")
-    test_torch_metrics(accelerator)
-
-
-def _mp_fn(index):
-    # For xla_spawn (TPUs)
-    main()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/src/accelerate/test_utils/testing.py b/src/accelerate/test_utils/testing.py
index 145cc80b2..381f92c58 100644
--- a/src/accelerate/test_utils/testing.py
+++ b/src/accelerate/test_utils/testing.py
@@ -30,10 +30,12 @@
 from ..utils import (
     gather,
     is_comet_ml_available,
+    is_datasets_available,
     is_deepspeed_available,
     is_tensorboard_available,
     is_torch_version,
     is_tpu_available,
+    is_transformers_available,
     is_wandb_available,
 )
 
@@ -84,6 +86,15 @@ def require_cuda(test_case):
     return unittest.skipUnless(torch.cuda.is_available(), "test requires a GPU")(test_case)
 
 
+def require_huggingface_suite(test_case):
+    """
+    Decorator marking a test that requires transformers and datasets. These tests are skipped when they are not.
+    """
+    return unittest.skipUnless(
+        is_transformers_available() and is_datasets_available(), "test requires the Hugging Face suite"
+    )(test_case)
+
+
 def require_tpu(test_case):
     """
     Decorator marking a test that requires TPUs. These tests are skipped when there are no TPUs available.
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
index d1a015b93..db21afc3a 100644
--- a/src/accelerate/utils/__init__.py
+++ b/src/accelerate/utils/__init__.py
@@ -25,6 +25,7 @@
     is_boto3_available,
     is_ccl_available,
     is_comet_ml_available,
+    is_datasets_available,
     is_deepspeed_available,
     is_sagemaker_available,
     is_tensorboard_available,
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index bc4e3b4cd..074d02e4a 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -89,6 +89,10 @@ def is_transformers_available():
     return importlib.util.find_spec("transformers") is not None
 
 
+def is_datasets_available():
+    return importlib.util.find_spec("datasets") is not None
+
+
 def is_tensorboard_available():
     return importlib.util.find_spec("tensorboard") is not None or importlib.util.find_spec("tensorboardX") is not None
 
diff --git a/tests/test_metrics.py b/tests/test_metrics.py
index 7e42c793e..b26ef00b2 100644
--- a/tests/test_metrics.py
+++ b/tests/test_metrics.py
@@ -23,31 +23,36 @@
 from accelerate.test_utils import (
     execute_subprocess_async,
     require_cpu,
+    require_huggingface_suite,
     require_multi_gpu,
     require_single_gpu,
-    skip,
-    test_metrics,
 )
 from accelerate.utils import get_launch_prefix, patch_environment
 
 
-@skip
+@require_huggingface_suite
 class MetricTester(unittest.TestCase):
     def setUp(self):
         mod_file = inspect.getfile(accelerate.test_utils)
-        self.test_file_path = os.path.sep.join(mod_file.split(os.path.sep)[:-1] + ["scripts", "test_metrics.py"])
+        self.test_file_path = os.path.sep.join(
+            mod_file.split(os.path.sep)[:-1] + ["scripts", "external_deps", "test_metrics.py"]
+        )
+
+        from accelerate.test_utils.scripts.external_deps import test_metrics  # noqa: F401
+
+        self.test_metrics = test_metrics
 
     @require_cpu
     def test_metric_cpu_noop(self):
-        debug_launcher(test_metrics.main, num_processes=1)
+        debug_launcher(self.test_metrics.main, num_processes=1)
 
     @require_cpu
     def test_metric_cpu_multi(self):
-        debug_launcher(test_metrics.main)
+        debug_launcher(self.test_metrics.main)
 
     @require_single_gpu
     def test_metric_gpu(self):
-        test_metrics.main()
+        self.test_metrics.main()
 
     @require_multi_gpu
     def test_metric_gpu_multi(self):

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/590
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index a0278218b..806f0094d 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -1,6 +1,6 @@
 name: Run Tests
 
-# on: [pull_request]
+on: [pull_request]
 
 env:
   HF_HOME: ~/hf_cache

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/589
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index 531f8eff9..4cedd8f1a 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -400,7 +400,7 @@ class DataLoaderDispatcher(DataLoader):
         - **total_dataset_length** (`int`) -- Total length of the inner dataset across all processes.
     """
 
-    def __init__(self, dataset, split_batches: bool = False, **kwargs):
+    def __init__(self, dataset, split_batches: bool = False, _drop_last: bool = False, **kwargs):
         shuffle = False
         if is_torch_version(">=", "1.11.0"):
             from torch.utils.data.datapipes.iter.combinatorics import ShufflerIterDataPipe
@@ -419,6 +419,13 @@ def __init__(self, dataset, split_batches: bool = False, **kwargs):
 
         self.gradient_state = GradientState()
         self.state = AcceleratorState()
+        self._drop_last = _drop_last
+        try:
+            length = getattr(self.dataset, "total_dataset_length", len(self.dataset))
+            self.gradient_state._set_remainder(length % self.total_batch_size)
+        except Exception:
+            # We can safely pass because the default is -1
+            pass
 
     def _fetch_batches(self, iterator):
         batches, batch = None, None
@@ -448,7 +455,7 @@ def _fetch_batches(self, iterator):
         self._stop_iteration = batch_info[1]
         if self._stop_iteration:
             # If drop_last is False and split_batches is False, we may have a remainder to take care of.
-            if not self.split_batches and not self.drop_last:
+            if not self.split_batches and not self._drop_last:
                 if self.state.process_index == 0 and len(batches) > 0:
                     batch = concatenate(batches, dim=0)
                     batch_info = [get_data_structure(batch), False]
@@ -472,10 +479,12 @@ def __iter__(self):
         if self.state.process_index == 0:
             # We only iterate through the DataLoader on process 0.
             main_iterator = super().__iter__()
+        stop_iteration = False
         self._stop_iteration = False
         first_batch = None
-        batch, batch_info, skip = self._fetch_batches(main_iterator)
-        while True:
+        next_batch, next_batch_info, next_skip = self._fetch_batches(main_iterator)
+        while not stop_iteration:
+            batch, batch_info, skip = next_batch, next_batch_info, next_skip
             if skip:
                 continue
             if self.state.process_index != 0:
@@ -485,34 +494,40 @@ def __iter__(self):
             # Broadcast the batch before splitting it.
             batch = broadcast(batch, from_process=0)
 
-            if not self.drop_last and first_batch is None:
+            if not self._drop_last and first_batch is None:
                 # We keep at least num processes elements of the first batch to be able to complete the last batch
                 first_batch = slice_tensors(batch, slice(0, self.state.num_processes))
 
             observed_batch_size = find_batch_size(batch)
             batch_size = observed_batch_size // self.state.num_processes
 
-            if not self.drop_last and self._stop_iteration and observed_batch_size % self.state.num_processes != 0:
+            if not self._drop_last and self._stop_iteration and observed_batch_size % self.state.num_processes != 0:
                 # If the last batch is not complete, let's add the first batch to it.
                 batch = concatenate([batch, first_batch], dim=0)
+                # Batch size computation above is wrong, it's off by 1 so we fix it.
                 batch_size += 1
 
             data_slice = slice(self.state.process_index * batch_size, (self.state.process_index + 1) * batch_size)
-            next_batch, next_batch_info, next_skip = self._fetch_batches(main_iterator)
             batch = slice_tensors(batch, data_slice)
-            if not self._stop_iteration:
-                yield batch
-                batch, batch_info, skip = next_batch, next_batch_info, next_skip
-            else:
+
+            stop_iteration = self._stop_iteration
+            if not stop_iteration:
+                # We may still be at the end of the dataloader without knowing it yet: if there is nothing left
+                # because by change the dataset had a round multiple of samples.
+                next_batch, next_batch_info, next_skip = self._fetch_batches(main_iterator)
+                # next_batch_info[0] is None when there are no more batches, otherwise we still need to process them.
+                if self._stop_iteration and next_batch_info[0] is None:
+                    stop_iteration = True
+
+            if stop_iteration:
                 self.gradient_state._set_end_of_dataloader(True)
-                yield batch
-                break
+            yield batch
 
     def __len__(self):
         whole_length = super().__len__()
         if self.split_batches:
             return whole_length
-        elif self.drop_last:
+        elif self._drop_last:
             return whole_length // self.state.num_processes
         else:
             return math.ceil(whole_length / self.state.num_processes)
@@ -675,6 +690,7 @@ def prepare_data_loader(
             new_dataset,
             split_batches=split_batches,
             batch_sampler=new_batch_sampler,
+            _drop_last=dataloader.drop_last,
             **kwargs,
         )
     else:
diff --git a/src/accelerate/test_utils/scripts/test_sync.py b/src/accelerate/test_utils/scripts/test_sync.py
index ae5a2c65b..8cee87d3d 100644
--- a/src/accelerate/test_utils/scripts/test_sync.py
+++ b/src/accelerate/test_utils/scripts/test_sync.py
@@ -245,7 +245,7 @@ def main():
                         "**Test `accumulate` gradient accumulation, ",
                         f"`split_batches={split_batch}` and `dispatch_batches={dispatch_batches}`**",
                     )
-                test_gradient_accumulation(split_batch)
+                test_gradient_accumulation(split_batch, dispatch_batches)
     if state.local_process_index == 0:
         print(
             "**Test `accumulate` gradient accumulation with optimizer and scheduler, ",

PyTorch does not let us pass `drop_last`.
@@ -400,7 +400,7 @@ class DataLoaderDispatcher(DataLoader):
         - **total_dataset_length** (`int`) -- Total length of the inner dataset across all processes.
     """
 
-    def __init__(self, dataset, split_batches: bool = False, **kwargs):
+    def __init__(self, dataset, split_batches: bool = False, _drop_last: bool = False, **kwargs):
Removed this logic as we don't need to take the next batch to have an informed decision on whether we are the last batch or not. This was actually causing the dataloader to stop one batch too early.
@@ -485,34 +492,31 @@ def __iter__(self):
             # Broadcast the batch before splitting it.
             batch = broadcast(batch, from_process=0)
 
-            if not self.drop_last and first_batch is None:
+            if not self._drop_last and first_batch is None:
                 # We keep at least num processes elements of the first batch to be able to complete the last batch
                 first_batch = slice_tensors(batch, slice(0, self.state.num_processes))
 
             observed_batch_size = find_batch_size(batch)
             batch_size = observed_batch_size // self.state.num_processes
 
-            if not self.drop_last and self._stop_iteration and observed_batch_size % self.state.num_processes != 0:
+            if not self._drop_last and self._stop_iteration and observed_batch_size % self.state.num_processes != 0:
                 # If the last batch is not complete, let's add the first batch to it.
                 batch = concatenate([batch, first_batch], dim=0)
+                # Batch size computation above is wrong, it's off by 1 so we fix it.
                 batch_size += 1
 
             data_slice = slice(self.state.process_index * batch_size, (self.state.process_index + 1) * batch_size)
-            next_batch, next_batch_info, next_skip = self._fetch_batches(main_iterator)
             batch = slice_tensors(batch, data_slice)
-            if not self._stop_iteration:
-                yield batch
-                batch, batch_info, skip = next_batch, next_batch_info, next_skip
-            else:
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/588
diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
index 7100d2c88..188800527 100644
--- a/src/accelerate/utils/modeling.py
+++ b/src/accelerate/utils/modeling.py
@@ -405,17 +405,12 @@ def get_balanced_memory(
     leaves = [n for n in module_sizes if len([p for p in module_sizes if p.startswith(n) and len(p) > len(n)]) == 0]
     mean_leaves = int(sum([module_sizes[n] for n in leaves]) / len(leaves))
     buffer = int(1.25 * max(buffer, mean_leaves))
-    if low_zero:
-        per_gpu += buffer
-        gpu_zero = 0
-    else:
-        gpu_zero = per_gpu
-        per_gpu += buffer
+    per_gpu += buffer
 
     max_memory = get_max_memory(max_memory)
-    for i in range(num_devices):
-        # We still leave slightly more space on GPU 0 and only apply the buffer on the other devices.
-        max_memory[i] = min(gpu_zero if i == 0 else per_gpu, max_memory[i])
+    # The last device is left with max_memory just in case the buffer is not enough.
+    for i in range(num_devices - 1):
+        max_memory[i] = min(0 if low_zero and i == 0 else per_gpu, max_memory[i])
 
     if low_zero:
         min_zero = max(0, module_sizes[""] - sum([max_memory[i] for i in range(1, num_devices)]))
diff --git a/tests/test_modeling_utils.py b/tests/test_modeling_utils.py
index 799170257..f5c36ee3d 100644
--- a/tests/test_modeling_utils.py
+++ b/tests/test_modeling_utils.py
@@ -365,7 +365,11 @@ def test_get_balanced_memory(self):
         model = ModelForTest()
         # model has size 236: linear1 64, batchnorm 72, linear2 100
         max_memory = get_balanced_memory(model, max_memory={0: 200, 1: 200})
-        self.assertDictEqual({0: 118, 1: 200}, max_memory)
+        self.assertDictEqual({0: 200, 1: 200}, max_memory)
 
         max_memory = get_balanced_memory(model, max_memory={0: 300, 1: 300})
-        self.assertDictEqual({0: 118, 1: 215}, max_memory)
+        self.assertDictEqual({0: 215, 1: 300}, max_memory)
+
+        # Last device always get max memory to give more buffer and avoid accidental CPU offload
+        max_memory = get_balanced_memory(model, max_memory={0: 300, 1: 500})
+        self.assertDictEqual({0: 215, 1: 500}, max_memory)

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/583
diff --git a/examples/by_feature/checkpointing.py b/examples/by_feature/checkpointing.py
index 5d31dd8f4..1b462a1f4 100644
--- a/examples/by_feature/checkpointing.py
+++ b/examples/by_feature/checkpointing.py
@@ -72,11 +72,13 @@ def tokenize_function(examples):
         return outputs
 
     # Apply the method we just defined to all the examples in all the splits of the dataset
-    tokenized_datasets = datasets.map(
-        tokenize_function,
-        batched=True,
-        remove_columns=["idx", "sentence1", "sentence2"],
-    )
+    # starting with the main process first:
+    with accelerator.main_process_first():
+        tokenized_datasets = datasets.map(
+            tokenize_function,
+            batched=True,
+            remove_columns=["idx", "sentence1", "sentence2"],
+        )
 
     # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
     # transformers library
diff --git a/examples/by_feature/cross_validation.py b/examples/by_feature/cross_validation.py
index dc41b91e5..6dd61bbf8 100644
--- a/examples/by_feature/cross_validation.py
+++ b/examples/by_feature/cross_validation.py
@@ -92,11 +92,13 @@ def tokenize_function(examples):
         return outputs
 
     # Apply the method we just defined to all the examples in all the splits of the dataset
-    tokenized_datasets = datasets.map(
-        tokenize_function,
-        batched=True,
-        remove_columns=["idx", "sentence1", "sentence2"],
-    )
+    # starting with the main process first:
+    with accelerator.main_process_first():
+        tokenized_datasets = datasets.map(
+            tokenize_function,
+            batched=True,
+            remove_columns=["idx", "sentence1", "sentence2"],
+        )
 
     # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
     # transformers library
diff --git a/examples/by_feature/fsdp_with_peak_mem_tracking.py b/examples/by_feature/fsdp_with_peak_mem_tracking.py
index b95fb9908..a279e64eb 100644
--- a/examples/by_feature/fsdp_with_peak_mem_tracking.py
+++ b/examples/by_feature/fsdp_with_peak_mem_tracking.py
@@ -127,11 +127,13 @@ def tokenize_function(examples):
         return outputs
 
     # Apply the method we just defined to all the examples in all the splits of the dataset
-    tokenized_datasets = datasets.map(
-        tokenize_function,
-        batched=True,
-        remove_columns=["idx", "sentence1", "sentence2"],
-    )
+    # starting with the main process first:
+    with accelerator.main_process_first():
+        tokenized_datasets = datasets.map(
+            tokenize_function,
+            batched=True,
+            remove_columns=["idx", "sentence1", "sentence2"],
+        )
 
     # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
     # transformers library
diff --git a/examples/by_feature/gradient_accumulation.py b/examples/by_feature/gradient_accumulation.py
index 580d07a57..69856c97b 100644
--- a/examples/by_feature/gradient_accumulation.py
+++ b/examples/by_feature/gradient_accumulation.py
@@ -67,11 +67,13 @@ def tokenize_function(examples):
         return outputs
 
     # Apply the method we just defined to all the examples in all the splits of the dataset
-    tokenized_datasets = datasets.map(
-        tokenize_function,
-        batched=True,
-        remove_columns=["idx", "sentence1", "sentence2"],
-    )
+    # starting with the main process first:
+    with accelerator.main_process_first():
+        tokenized_datasets = datasets.map(
+            tokenize_function,
+            batched=True,
+            remove_columns=["idx", "sentence1", "sentence2"],
+        )
 
     # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
     # transformers library
diff --git a/examples/by_feature/memory.py b/examples/by_feature/memory.py
index d85319552..91f3e41d1 100644
--- a/examples/by_feature/memory.py
+++ b/examples/by_feature/memory.py
@@ -72,11 +72,13 @@ def tokenize_function(examples):
         return outputs
 
     # Apply the method we just defined to all the examples in all the splits of the dataset
-    tokenized_datasets = datasets.map(
-        tokenize_function,
-        batched=True,
-        remove_columns=["idx", "sentence1", "sentence2"],
-    )
+    # starting with the main process first:
+    with accelerator.main_process_first():
+        tokenized_datasets = datasets.map(
+            tokenize_function,
+            batched=True,
+            remove_columns=["idx", "sentence1", "sentence2"],
+        )
 
     # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
     # transformers library
diff --git a/examples/by_feature/multi_process_metrics.py b/examples/by_feature/multi_process_metrics.py
index cf581c73d..522cc571b 100644
--- a/examples/by_feature/multi_process_metrics.py
+++ b/examples/by_feature/multi_process_metrics.py
@@ -74,11 +74,13 @@ def tokenize_function(examples):
         return outputs
 
     # Apply the method we just defined to all the examples in all the splits of the dataset
-    tokenized_datasets = datasets.map(
-        tokenize_function,
-        batched=True,
-        remove_columns=["idx", "sentence1", "sentence2"],
-    )
+    # starting with the main process first:
+    with accelerator.main_process_first():
+        tokenized_datasets = datasets.map(
+            tokenize_function,
+            batched=True,
+            remove_columns=["idx", "sentence1", "sentence2"],
+        )
 
     # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
     # transformers library
diff --git a/examples/by_feature/tracking.py b/examples/by_feature/tracking.py
index 0da8c437a..d7248682d 100644
--- a/examples/by_feature/tracking.py
+++ b/examples/by_feature/tracking.py
@@ -72,11 +72,13 @@ def tokenize_function(examples):
         return outputs
 
     # Apply the method we just defined to all the examples in all the splits of the dataset
-    tokenized_datasets = datasets.map(
-        tokenize_function,
-        batched=True,
-        remove_columns=["idx", "sentence1", "sentence2"],
-    )
+    # starting with the main process first:
+    with accelerator.main_process_first():
+        tokenized_datasets = datasets.map(
+            tokenize_function,
+            batched=True,
+            remove_columns=["idx", "sentence1", "sentence2"],
+        )
 
     # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
     # transformers library
diff --git a/examples/complete_nlp_example.py b/examples/complete_nlp_example.py
index fc0fae90b..000da6038 100644
--- a/examples/complete_nlp_example.py
+++ b/examples/complete_nlp_example.py
@@ -91,11 +91,13 @@ def tokenize_function(examples):
         return outputs
 
     # Apply the method we just defined to all the examples in all the splits of the dataset
-    tokenized_datasets = datasets.map(
-        tokenize_function,
-        batched=True,
-        remove_columns=["idx", "sentence1", "sentence2"],
-    )
+    # starting with the main process first:
+    with accelerator.main_process_first():
+        tokenized_datasets = datasets.map(
+            tokenize_function,
+            batched=True,
+            remove_columns=["idx", "sentence1", "sentence2"],
+        )
 
     # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
     # transformers library
diff --git a/examples/nlp_example.py b/examples/nlp_example.py
index 33c0ed7aa..b1e7cba27 100644
--- a/examples/nlp_example.py
+++ b/examples/nlp_example.py
@@ -65,11 +65,13 @@ def tokenize_function(examples):
         return outputs
 
     # Apply the method we just defined to all the examples in all the splits of the dataset
-    tokenized_datasets = datasets.map(
-        tokenize_function,
-        batched=True,
-        remove_columns=["idx", "sentence1", "sentence2"],
-    )
+    # starting with the main process first:
+    with accelerator.main_process_first():
+        tokenized_datasets = datasets.map(
+            tokenize_function,
+            batched=True,
+            remove_columns=["idx", "sentence1", "sentence2"],
+        )
 
     # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
     # transformers library

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/581
diff --git a/examples/by_feature/checkpointing.py b/examples/by_feature/checkpointing.py
index 3860875e9..5d31dd8f4 100644
--- a/examples/by_feature/checkpointing.py
+++ b/examples/by_feature/checkpointing.py
@@ -234,6 +234,7 @@ def training_function(config, args):
                     accelerator.save_state(output_dir)
 
         model.eval()
+        samples_seen = 0
         for step, batch in enumerate(eval_dataloader):
             # We could avoid this line since we set the accelerator with `device_placement=True` (the default).
             batch.to(accelerator.device)
@@ -241,7 +242,15 @@ def training_function(config, args):
                 outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1)
             # It is slightly faster to call this once, than multiple times
-            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
+            predictions, references = accelerator.gather((predictions, batch["labels"]))
+            if accelerator.use_distributed:
+                if step == len(eval_dataloader) - 1:
+                    # Last batch needs to be truncated on distributed systems as it contains additional samples
+                    predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
+                    references = references[: len(eval_dataloader.dataset) - samples_seen]
+                else:
+                    # Otherwise we add the number of samples seen
+                    samples_seen += references.shape[0]
             metric.add_batch(
                 predictions=predictions,
                 references=references,
diff --git a/examples/by_feature/cross_validation.py b/examples/by_feature/cross_validation.py
index cd5b2fd02..dc41b91e5 100644
--- a/examples/by_feature/cross_validation.py
+++ b/examples/by_feature/cross_validation.py
@@ -201,13 +201,23 @@ def training_function(config, args):
                     optimizer.zero_grad()
 
             model.eval()
+            samples_seen = 0
             for step, batch in enumerate(eval_dataloader):
                 # We could avoid this line since we set the accelerator with `device_placement=True`.
                 batch.to(accelerator.device)
                 with torch.no_grad():
                     outputs = model(**batch)
                 predictions = outputs.logits.argmax(dim=-1)
-                predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
+                # It is slightly faster to call this once, than multiple times
+                predictions, references = accelerator.gather((predictions, batch["labels"]))
+                if accelerator.use_distributed:
+                    if step == len(eval_dataloader) - 1:
+                        # Last batch needs to be truncated on distributed systems as it contains additional samples
+                        predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
+                        references = references[: len(eval_dataloader.dataset) - samples_seen]
+                    else:
+                        # Otherwise we add the number of samples seen
+                        samples_seen += references.shape[0]
                 metric.add_batch(
                     predictions=predictions,
                     references=references,
diff --git a/examples/by_feature/fsdp_with_peak_mem_tracking.py b/examples/by_feature/fsdp_with_peak_mem_tracking.py
index 1426dd0cd..b95fb9908 100644
--- a/examples/by_feature/fsdp_with_peak_mem_tracking.py
+++ b/examples/by_feature/fsdp_with_peak_mem_tracking.py
@@ -272,6 +272,7 @@ def collate_fn(examples):
         # context manager to track the peak memory usage during the evaluation
         with TorchTracemalloc() as tracemalloc:
             model.eval()
+            samples_seen = 0
             for step, batch in enumerate(eval_dataloader):
                 # We could avoid this line since we set the accelerator with `device_placement=True`.
                 batch.to(accelerator.device)
@@ -279,7 +280,15 @@ def collate_fn(examples):
                     outputs = model(**batch)
                 predictions = outputs.logits.argmax(dim=-1)
                 # It is slightly faster to call this once, than multiple times
-                predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
+                predictions, references = accelerator.gather((predictions, batch["labels"]))
+                if accelerator.use_distributed:
+                    if step == len(eval_dataloader) - 1:
+                        # Last batch needs to be truncated on distributed systems as it contains additional samples
+                        predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
+                        references = references[: len(eval_dataloader.dataset) - samples_seen]
+                    else:
+                        # Otherwise we add the number of samples seen
+                        samples_seen += references.shape[0]
                 metric.add_batch(
                     predictions=predictions,
                     references=references,
diff --git a/examples/by_feature/gradient_accumulation.py b/examples/by_feature/gradient_accumulation.py
index 2e6de6a32..580d07a57 100644
--- a/examples/by_feature/gradient_accumulation.py
+++ b/examples/by_feature/gradient_accumulation.py
@@ -168,13 +168,23 @@ def training_function(config, args):
                 optimizer.zero_grad()
 
         model.eval()
+        samples_seen = 0
         for step, batch in enumerate(eval_dataloader):
             # We could avoid this line since we set the accelerator with `device_placement=True`.
             batch.to(accelerator.device)
             with torch.no_grad():
                 outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1)
-            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
+            # It is slightly faster to call this once, than multiple times
+            predictions, references = accelerator.gather((predictions, batch["labels"]))
+            if accelerator.use_distributed:
+                if step == len(eval_dataloader) - 1:
+                    # Last batch needs to be truncated on distributed systems as it contains additional samples
+                    predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
+                    references = references[: len(eval_dataloader.dataset) - samples_seen]
+                else:
+                    # Otherwise we add the number of samples seen
+                    samples_seen += references.shape[0]
             metric.add_batch(
                 predictions=predictions,
                 references=references,
diff --git a/examples/by_feature/memory.py b/examples/by_feature/memory.py
index 6f7cf00fd..d85319552 100644
--- a/examples/by_feature/memory.py
+++ b/examples/by_feature/memory.py
@@ -180,13 +180,23 @@ def inner_training_loop(batch_size):
                     optimizer.zero_grad()
 
             model.eval()
+            samples_seen = 0
             for step, batch in enumerate(eval_dataloader):
                 # We could avoid this line since we set the accelerator with `device_placement=True`.
                 batch.to(accelerator.device)
                 with torch.no_grad():
                     outputs = model(**batch)
                 predictions = outputs.logits.argmax(dim=-1)
-                predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
+                # It is slightly faster to call this once, than multiple times
+                predictions, references = accelerator.gather((predictions, batch["labels"]))
+                if accelerator.use_distributed:
+                    if step == len(eval_dataloader) - 1:
+                        # Last batch needs to be truncated on distributed systems as it contains additional samples
+                        predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
+                        references = references[: len(eval_dataloader.dataset) - samples_seen]
+                    else:
+                        # Otherwise we add the number of samples seen
+                        samples_seen += references.shape[0]
                 metric.add_batch(
                     predictions=predictions,
                     references=references,
diff --git a/examples/by_feature/multi_process_metrics.py b/examples/by_feature/multi_process_metrics.py
index f39b064d6..cf581c73d 100644
--- a/examples/by_feature/multi_process_metrics.py
+++ b/examples/by_feature/multi_process_metrics.py
@@ -190,8 +190,6 @@ def training_function(config, args):
                 else:
                     # Otherwise we add the number of samples seen
                     samples_seen += references.shape[0]
-            # All of this can be avoided if you use `Accelerator.gather_for_metrics` instead of `Accelerator.gather`:
-            # accelerator.gather_for_metrics((predictions, references))
             metric.add_batch(
                 predictions=predictions,
                 references=references,
diff --git a/examples/by_feature/tracking.py b/examples/by_feature/tracking.py
index d8c985e88..0da8c437a 100644
--- a/examples/by_feature/tracking.py
+++ b/examples/by_feature/tracking.py
@@ -192,6 +192,7 @@ def training_function(config, args):
                 optimizer.zero_grad()
 
         model.eval()
+        samples_seen = 0
         for step, batch in enumerate(eval_dataloader):
             # We could avoid this line since we set the accelerator with `device_placement=True` (the default).
             batch.to(accelerator.device)
@@ -199,7 +200,15 @@ def training_function(config, args):
                 outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1)
             # It is slightly faster to call this once, than multiple times
-            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
+            predictions, references = accelerator.gather((predictions, batch["labels"]))
+            if accelerator.use_distributed:
+                if step == len(eval_dataloader) - 1:
+                    # Last batch needs to be truncated on distributed systems as it contains additional samples
+                    predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
+                    references = references[: len(eval_dataloader.dataset) - samples_seen]
+                else:
+                    # Otherwise we add the number of samples seen
+                    samples_seen += references.shape[0]
             metric.add_batch(
                 predictions=predictions,
                 references=references,
diff --git a/examples/complete_cv_example.py b/examples/complete_cv_example.py
index ee051a0c5..f984406bf 100644
--- a/examples/complete_cv_example.py
+++ b/examples/complete_cv_example.py
@@ -232,6 +232,7 @@ def training_function(config, args):
                     accelerator.save_state(output_dir)
         model.eval()
         accurate = 0
+        samples_seen = 0
         for step, batch in enumerate(eval_dataloader):
             # We could avoid this line since we set the accelerator with `device_placement=True`.
             batch = {k: v.to(accelerator.device) for k, v in batch.items()}
@@ -239,8 +240,17 @@ def training_function(config, args):
             with torch.no_grad():
                 outputs = model(inputs)
             predictions = outputs.argmax(dim=-1)
-            predictions, labels = accelerator.gather_for_metrics((predictions, batch["label"]))
-            accurate_preds = predictions == labels
+            # It is slightly faster to call this once, than multiple times
+            predictions, references = accelerator.gather((predictions, batch["label"]))
+            if accelerator.use_distributed:
+                if step == len(eval_dataloader) - 1:
+                    # Last batch needs to be truncated on distributed systems as it contains additional samples
+                    predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
+                    references = references[: len(eval_dataloader.dataset) - samples_seen]
+                else:
+                    # Otherwise we add the number of samples seen
+                    samples_seen += references.shape[0]
+            accurate_preds = predictions == references
             accurate += accurate_preds.long().sum()
 
         eval_metric = accurate.item() / accelerator.gradient_state.samples_seen
diff --git a/examples/complete_nlp_example.py b/examples/complete_nlp_example.py
index 88d379c8e..fc0fae90b 100644
--- a/examples/complete_nlp_example.py
+++ b/examples/complete_nlp_example.py
@@ -210,6 +210,7 @@ def collate_fn(examples):
                     accelerator.save_state(output_dir)
 
         model.eval()
+        samples_seen = 0
         for step, batch in enumerate(eval_dataloader):
             # We could avoid this line since we set the accelerator with `device_placement=True`.
             batch.to(accelerator.device)
@@ -217,7 +218,15 @@ def collate_fn(examples):
                 outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1)
             # It is slightly faster to call this once, than multiple times
-            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
+            predictions, references = accelerator.gather((predictions, batch["labels"]))
+            if accelerator.use_distributed:
+                if step == len(eval_dataloader) - 1:
+                    # Last batch needs to be truncated on distributed systems as it contains additional samples
+                    predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
+                    references = references[: len(eval_dataloader.dataset) - samples_seen]
+                else:
+                    # Otherwise we add the number of samples seen
+                    samples_seen += references.shape[0]
             metric.add_batch(
                 predictions=predictions,
                 references=references,
diff --git a/examples/cv_example.py b/examples/cv_example.py
index 6663b2fe9..875b68398 100644
--- a/examples/cv_example.py
+++ b/examples/cv_example.py
@@ -166,6 +166,7 @@ def training_function(config, args):
         model.eval()
         accurate = 0
         num_elems = 0
+        samples_seen = 0
         for _, batch in enumerate(eval_dataloader):
             # We could avoid this line since we set the accelerator with `device_placement=True`.
             batch = {k: v.to(accelerator.device) for k, v in batch.items()}
@@ -173,8 +174,17 @@ def training_function(config, args):
             with torch.no_grad():
                 outputs = model(inputs)
             predictions = outputs.argmax(dim=-1)
-            predictions, labels = accelerator.gather_for_metrics((predictions, batch["label"]))
-            accurate_preds = predictions == labels
+            # It is slightly faster to call this once, than multiple times
+            predictions, references = accelerator.gather((predictions, batch["label"]))
+            if accelerator.use_distributed:
+                if step == len(eval_dataloader) - 1:
+                    # Last batch needs to be truncated on distributed systems as it contains additional samples
+                    predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
+                    references = references[: len(eval_dataloader.dataset) - samples_seen]
+                else:
+                    # Otherwise we add the number of samples seen
+                    samples_seen += references.shape[0]
+            accurate_preds = predictions == references
             num_elems += accurate_preds.shape[0]
             accurate += accurate_preds.long().sum()
 
diff --git a/examples/nlp_example.py b/examples/nlp_example.py
index 2ef18598b..33c0ed7aa 100644
--- a/examples/nlp_example.py
+++ b/examples/nlp_example.py
@@ -152,13 +152,23 @@ def training_function(config, args):
                 optimizer.zero_grad()
 
         model.eval()
+        samples_seen = 0
         for step, batch in enumerate(eval_dataloader):
             # We could avoid this line since we set the accelerator with `device_placement=True`.
             batch.to(accelerator.device)
             with torch.no_grad():
                 outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1)
-            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
+            # It is slightly faster to call this once, than multiple times
+            predictions, references = accelerator.gather((predictions, batch["labels"]))
+            if accelerator.use_distributed:
+                if step == len(eval_dataloader) - 1:
+                    # Last batch needs to be truncated on distributed systems as it contains additional samples
+                    predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
+                    references = references[: len(eval_dataloader.dataset) - samples_seen]
+                else:
+                    # Otherwise we add the number of samples seen
+                    samples_seen += references.shape[0]
             metric.add_batch(
                 predictions=predictions,
                 references=references,
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 977cd5d24..5a9a55c68 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -943,6 +943,9 @@ def gather_for_metrics(self, tensor):
             tensor (`torch.Tensor`, or a nested tuple/list/dictionary of `torch.Tensor`):
                 The tensors for calculating metrics across all processes.
         """
+        raise NotImplementedError(
+            "Currently there are a number of bugs with this method. You should use `Accelerator.gather()` and drop the samples yourself for the time being."
+        )
         tensor = self.gather(tensor)
         if self.use_distributed:
             if self.gradient_state.remainder == -1:
diff --git a/src/accelerate/test_utils/__init__.py b/src/accelerate/test_utils/__init__.py
index d65099602..382bf81bf 100644
--- a/src/accelerate/test_utils/__init__.py
+++ b/src/accelerate/test_utils/__init__.py
@@ -10,6 +10,7 @@
     require_multi_gpu,
     require_single_gpu,
     require_tpu,
+    skip,
     slow,
 )
 from .training import RegressionDataset, RegressionModel
diff --git a/src/accelerate/test_utils/testing.py b/src/accelerate/test_utils/testing.py
index b5cccc1ff..145cc80b2 100644
--- a/src/accelerate/test_utils/testing.py
+++ b/src/accelerate/test_utils/testing.py
@@ -57,6 +57,11 @@ def parse_flag_from_env(key, default=False):
 _run_slow_tests = parse_flag_from_env("RUN_SLOW", default=False)
 
 
+def skip(test_case):
+    "Decorator that skips a test unconditionally"
+    return unittest.skip("Test was skipped")(test_case)
+
+
 def slow(test_case):
     """
     Decorator marking a test as slow. Slow tests are skipped by default. Set the RUN_SLOW environment variable to a
diff --git a/tests/test_metrics.py b/tests/test_metrics.py
index c744a2fdb..7e42c793e 100644
--- a/tests/test_metrics.py
+++ b/tests/test_metrics.py
@@ -25,11 +25,13 @@
     require_cpu,
     require_multi_gpu,
     require_single_gpu,
+    skip,
     test_metrics,
 )
 from accelerate.utils import get_launch_prefix, patch_environment
 
 
+@skip
 class MetricTester(unittest.TestCase):
     def setUp(self):
         mod_file = inspect.getfile(accelerate.test_utils)

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/580
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 9ee840440..599e710f0 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -752,7 +752,7 @@ def launch_command(args):
     # Get the default from the config file.
     if args.config_file is not None or os.path.isfile(default_config_file) and not args.cpu:
         defaults = load_config_from_file(args.config_file)
-        if not args.multi_gpu and not args.tpu and not args.use_deepspeed:
+        if not args.multi_gpu and not args.tpu and not args.use_deepspeed and not args.use_fsdp:
             args.use_deepspeed = defaults.distributed_type == DistributedType.DEEPSPEED
             args.multi_gpu = defaults.distributed_type == DistributedType.MULTI_GPU
             args.tpu = defaults.distributed_type == DistributedType.TPU

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/579
diff --git a/examples/by_feature/checkpointing.py b/examples/by_feature/checkpointing.py
index 5e364491e..3860875e9 100644
--- a/examples/by_feature/checkpointing.py
+++ b/examples/by_feature/checkpointing.py
@@ -241,7 +241,7 @@ def training_function(config, args):
                 outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1)
             # It is slightly faster to call this once, than multiple times
-            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]), eval_dataloader)
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
             metric.add_batch(
                 predictions=predictions,
                 references=references,
diff --git a/examples/by_feature/cross_validation.py b/examples/by_feature/cross_validation.py
index 79a207fa8..cd5b2fd02 100644
--- a/examples/by_feature/cross_validation.py
+++ b/examples/by_feature/cross_validation.py
@@ -207,9 +207,7 @@ def training_function(config, args):
                 with torch.no_grad():
                     outputs = model(**batch)
                 predictions = outputs.logits.argmax(dim=-1)
-                predictions, references = accelerator.gather_for_metrics(
-                    (predictions, batch["labels"]), eval_dataloader
-                )
+                predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
                 metric.add_batch(
                     predictions=predictions,
                     references=references,
@@ -228,7 +226,7 @@ def training_function(config, args):
             with torch.no_grad():
                 outputs = model(**batch)
             predictions = outputs.logits
-            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]), test_dataloader)
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
             fold_predictions.append(predictions.cpu())
             if i == 0:
                 # We need all of the test predictions
diff --git a/examples/by_feature/fsdp_with_peak_mem_tracking.py b/examples/by_feature/fsdp_with_peak_mem_tracking.py
index 8046b4557..1426dd0cd 100644
--- a/examples/by_feature/fsdp_with_peak_mem_tracking.py
+++ b/examples/by_feature/fsdp_with_peak_mem_tracking.py
@@ -279,9 +279,7 @@ def collate_fn(examples):
                     outputs = model(**batch)
                 predictions = outputs.logits.argmax(dim=-1)
                 # It is slightly faster to call this once, than multiple times
-                predictions, references = accelerator.gather_for_metrics(
-                    (predictions, batch["labels"]), eval_dataloader
-                )
+                predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
                 metric.add_batch(
                     predictions=predictions,
                     references=references,
diff --git a/examples/by_feature/gradient_accumulation.py b/examples/by_feature/gradient_accumulation.py
index fe3d948f5..2e6de6a32 100644
--- a/examples/by_feature/gradient_accumulation.py
+++ b/examples/by_feature/gradient_accumulation.py
@@ -174,7 +174,7 @@ def training_function(config, args):
             with torch.no_grad():
                 outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1)
-            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]), eval_dataloader)
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
             metric.add_batch(
                 predictions=predictions,
                 references=references,
diff --git a/examples/by_feature/memory.py b/examples/by_feature/memory.py
index f71283561..6f7cf00fd 100644
--- a/examples/by_feature/memory.py
+++ b/examples/by_feature/memory.py
@@ -186,9 +186,7 @@ def inner_training_loop(batch_size):
                 with torch.no_grad():
                     outputs = model(**batch)
                 predictions = outputs.logits.argmax(dim=-1)
-                predictions, references = accelerator.gather_for_metrics(
-                    (predictions, batch["labels"]), eval_dataloader
-                )
+                predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
                 metric.add_batch(
                     predictions=predictions,
                     references=references,
diff --git a/examples/by_feature/multi_process_metrics.py b/examples/by_feature/multi_process_metrics.py
index d69c0bc90..f39b064d6 100644
--- a/examples/by_feature/multi_process_metrics.py
+++ b/examples/by_feature/multi_process_metrics.py
@@ -191,7 +191,7 @@ def training_function(config, args):
                     # Otherwise we add the number of samples seen
                     samples_seen += references.shape[0]
             # All of this can be avoided if you use `Accelerator.gather_for_metrics` instead of `Accelerator.gather`:
-            # accelerator.gather_for_metrics((predictions, references), eval_dataloader)
+            # accelerator.gather_for_metrics((predictions, references))
             metric.add_batch(
                 predictions=predictions,
                 references=references,
diff --git a/examples/by_feature/tracking.py b/examples/by_feature/tracking.py
index d453b770c..d8c985e88 100644
--- a/examples/by_feature/tracking.py
+++ b/examples/by_feature/tracking.py
@@ -199,7 +199,7 @@ def training_function(config, args):
                 outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1)
             # It is slightly faster to call this once, than multiple times
-            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]), eval_dataloader)
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
             metric.add_batch(
                 predictions=predictions,
                 references=references,
diff --git a/examples/complete_cv_example.py b/examples/complete_cv_example.py
index 17cc83170..ee051a0c5 100644
--- a/examples/complete_cv_example.py
+++ b/examples/complete_cv_example.py
@@ -239,7 +239,7 @@ def training_function(config, args):
             with torch.no_grad():
                 outputs = model(inputs)
             predictions = outputs.argmax(dim=-1)
-            predictions, labels = accelerator.gather_for_metrics((predictions, batch["label"]), eval_dataloader)
+            predictions, labels = accelerator.gather_for_metrics((predictions, batch["label"]))
             accurate_preds = predictions == labels
             accurate += accurate_preds.long().sum()
 
diff --git a/examples/complete_nlp_example.py b/examples/complete_nlp_example.py
index dc0cf43ed..88d379c8e 100644
--- a/examples/complete_nlp_example.py
+++ b/examples/complete_nlp_example.py
@@ -217,7 +217,7 @@ def collate_fn(examples):
                 outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1)
             # It is slightly faster to call this once, than multiple times
-            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]), eval_dataloader)
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
             metric.add_batch(
                 predictions=predictions,
                 references=references,
diff --git a/examples/cv_example.py b/examples/cv_example.py
index aa97ce2d7..6663b2fe9 100644
--- a/examples/cv_example.py
+++ b/examples/cv_example.py
@@ -173,7 +173,7 @@ def training_function(config, args):
             with torch.no_grad():
                 outputs = model(inputs)
             predictions = outputs.argmax(dim=-1)
-            predictions, labels = accelerator.gather_for_metrics((predictions, batch["label"]), eval_dataloader)
+            predictions, labels = accelerator.gather_for_metrics((predictions, batch["label"]))
             accurate_preds = predictions == labels
             num_elems += accurate_preds.shape[0]
             accurate += accurate_preds.long().sum()
diff --git a/examples/nlp_example.py b/examples/nlp_example.py
index d64ab3bb7..2ef18598b 100644
--- a/examples/nlp_example.py
+++ b/examples/nlp_example.py
@@ -158,7 +158,7 @@ def training_function(config, args):
             with torch.no_grad():
                 outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1)
-            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]), eval_dataloader)
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
             metric.add_batch(
                 predictions=predictions,
                 references=references,
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index acf960026..977cd5d24 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -934,7 +934,7 @@ def gather(self, tensor):
         """
         return gather(tensor)
 
-    def gather_for_metrics(self, tensor, dataloader):
+    def gather_for_metrics(self, tensor):
         """
         Gathers `tensor` and potentially drops duplicates in the last batch if on a distributed system. Should be used
         for gathering the inputs and targets for metric calculation.
@@ -942,17 +942,20 @@ def gather_for_metrics(self, tensor, dataloader):
         Args:
             tensor (`torch.Tensor`, or a nested tuple/list/dictionary of `torch.Tensor`):
                 The tensors for calculating metrics across all processes.
-            dataloader (`torch.utils.data.DataLoader`):
-                A dataloader prepared with `Accelerator.prepare`
         """
         tensor = self.gather(tensor)
         if self.use_distributed:
+            if self.gradient_state.remainder == -1:
+                logger.info(
+                    "The used dataset had no length, returning gathered tensors. You should drop the remainder yourself."
+                )
+                return tensor
             try:
                 # Then see if we're on the last batch of our eval dataloader
                 if self.gradient_state.end_of_dataloader:
                     # Last batch needs to be truncated on distributed systems as it contains additional samples
                     def _adjust_samples(tensor):
-                        return tensor[: dataloader.total_dataset_length - self.gradient_state.samples_seen]
+                        return tensor[: self.gradient_state.remainder]
 
                     return recursively_apply(_adjust_samples, tensor)
                 else:
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index ba3338546..531f8eff9 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -138,6 +138,10 @@ def __init__(
         self.batch_size = batch_sampler.batch_size
         self.drop_last = batch_sampler.drop_last
 
+    @property
+    def total_length(self):
+        return len(self.batch_sampler)
+
     def __len__(self):
         if self.split_batches:
             return len(self.batch_sampler)
@@ -332,7 +336,12 @@ def __iter__(self):
         if self.rng_types is not None:
             synchronize_rng_states(self.rng_types, self.generator)
         self.gradient_state._set_end_of_dataloader(False)
-        self.gradient_state._set_samples_seen(0)
+        try:
+            length = getattr(self.dataset, "total_dataset_length", len(self.dataset))
+            self.gradient_state._set_remainder(length % self.total_batch_size)
+        except Exception:
+            # We can safely pass because the default is -1
+            pass
         dataloader_iter = super().__iter__()
         # We iterate one batch ahead to check when we are at the end
         try:
@@ -346,12 +355,10 @@ def __iter__(self):
                 if self.device is not None:
                     current_batch = send_to_device(current_batch, self.device)
                 next_batch = next(dataloader_iter)
-                self.gradient_state._iterate_samples_seen(find_batch_size(current_batch))
                 yield current_batch
                 current_batch = next_batch
             except StopIteration:
                 self.gradient_state._set_end_of_dataloader(True)
-                self.gradient_state._iterate_samples_seen(find_batch_size(current_batch))
                 yield current_batch
                 break
 
@@ -365,7 +372,10 @@ def total_batch_size(self):
 
     @property
     def total_dataset_length(self):
-        return len(self.dataset)
+        if hasattr("total_length", self.dataset):
+            return self.dataset.total_length
+        else:
+            return len(self.dataset)
 
 
 class DataLoaderDispatcher(DataLoader):
@@ -453,7 +463,11 @@ def _fetch_batches(self, iterator):
 
     def __iter__(self):
         self.gradient_state._set_end_of_dataloader(False)
-        self.gradient_state._set_samples_seen(0)
+        try:
+            length = getattr(self.dataset, "total_dataset_length", len(self.dataset))
+            self.gradient_state._set_remainder(length % self.batch_size)
+        except:
+            self.gradient_state._set_remainder(-1)
         main_iterator = None
         if self.state.process_index == 0:
             # We only iterate through the DataLoader on process 0.
@@ -487,11 +501,9 @@ def __iter__(self):
             next_batch, next_batch_info, next_skip = self._fetch_batches(main_iterator)
             batch = slice_tensors(batch, data_slice)
             if not self._stop_iteration:
-                self.gradient_state._iterate_samples_seen(batch_size)
                 yield batch
                 batch, batch_info, skip = next_batch, next_batch_info, next_skip
             else:
-                self.gradient_state._iterate_samples_seen(batch_size)
                 self.gradient_state._set_end_of_dataloader(True)
                 yield batch
                 break
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index 06a25bc4e..27a627ae9 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -258,6 +258,7 @@ class GradientState:
 
         - **sync_gradients** (`bool`) -- Whether the gradients should be synced
         - **end_of_dataloader** (`bool`) -- Whether we have reached the end the current dataloader
+        - **remainder** (`int`) -- The number of extra samples that were added from padding the dataloader
     """
 
     _shared_state = {}
@@ -267,14 +268,14 @@ def __init__(self):
         if not getattr(self, "initialized", False):
             self.sync_gradients = True
             self.end_of_dataloader = False
-            self.samples_seen = 0
+            self.remainder = -1
         self.initialized = True
 
     def __repr__(self):
         return (
             f"Sync Gradients: {self.sync_gradients}\n"
             f"At end of current dataloader: {self.end_of_dataloader}\n"
-            f"Samples seen: {self.samples_seen}"
+            f"Extra samples added: {self.remainder}"
         )
 
     def _set_sync_gradients(self, sync_gradients):
@@ -285,10 +286,6 @@ def _set_end_of_dataloader(self, end_of_dataloader):
         "Private function that sets whether the end of the current dataloader has been reached. Users should not have to call this."
         self.end_of_dataloader = end_of_dataloader
 
-    def _set_samples_seen(self, samples_seen):
-        "Private function that sets the number of samples iterated over. Users should not have to call this."
-        self.samples_seen = samples_seen
-
-    def _iterate_samples_seen(self, iteration: int = 1):
-        "Private function that iterates the number of samples seen by an iteration. Users should not have to call this."
-        self._set_samples_seen(self.samples_seen + iteration)
+    def _set_remainder(self, remainder):
+        "Private function that sets the number of remaining samples at the end of the dataloader"
+        self.remainder = remainder
diff --git a/src/accelerate/test_utils/scripts/test_metrics.py b/src/accelerate/test_utils/scripts/test_metrics.py
index c426d1fc5..8f057ccc5 100644
--- a/src/accelerate/test_utils/scripts/test_metrics.py
+++ b/src/accelerate/test_utils/scripts/test_metrics.py
@@ -22,50 +22,63 @@
 from accelerate.utils import set_seed
 
 
-def get_setup(accelerator):
+def get_setup(accelerator, num_samples=82):
     "Returns everything needed to perform basic training"
     set_seed(42)
     model = RegressionModel()
     ddp_model = deepcopy(model)
-    dset = RegressionDataset(length=80)
+    dset = RegressionDataset(length=num_samples)
     dataloader = DataLoader(dset, batch_size=16)
     model.to(accelerator.device)
     ddp_model, dataloader = accelerator.prepare(ddp_model, dataloader)
     return model, ddp_model, dataloader
 
 
-def accuracy(predictions, labels) -> float:
-    """
-    Get the accuracy with respect to the most likely label
-    """
-    return (predictions == labels).float().mean()
-
-
-def test_torch_metrics():
-    accelerator = Accelerator()
-    model, ddp_model, dataloader = get_setup(accelerator)
+def generate_predictions(model, dataloader, accelerator):
+    logits_and_targets = []
     for batch in dataloader:
-        ddp_input, ddp_target = batch.values()
-        # First do single process
-        input, target = accelerator.gather((ddp_input, ddp_target))
-        input, target = input.to(accelerator.device), target.to(accelerator.device)
+        input, target = batch.values()
         with torch.no_grad():
             logits = model(input)
-            accuracy_single = accuracy(logits.argmax(dim=-1), target)
-        # Then do multiprocess
-        with torch.no_grad():
-            logits = ddp_model(ddp_input)
-            logits, target = accelerator.gather_for_metrics((logits, ddp_target), dataloader)
-            accuracy_multi = accuracy(logits.argmax(dim=-1), target)
-        assert torch.allclose(accuracy_single, accuracy_multi), "The two accuracies were not the same!"
+            logits, target = accelerator.gather_for_metrics((logits, target))
+            logits_and_targets.append((logits, target))
+    inps, targs = [], []
+    for (inp, targ) in logits_and_targets:
+        inps.append(inp)
+        targs.append(targ)
+    inps, targs = torch.cat(inps), torch.cat(targs)
+    return inps, targs
+
+
+def test_torch_metrics(accelerator: Accelerator, num_samples=82):
+    model, ddp_model, dataloader = get_setup(accelerator, num_samples)
+    inps, targs = generate_predictions(ddp_model, dataloader, accelerator)
+    assert (
+        len(inps) == num_samples
+    ), f"Unexpected number of inputs:\n    Expected: {num_samples}\n    Actual: {len(inps)}"
 
 
 def main():
-    accelerator = Accelerator()
-    state = accelerator.state
-    if state.local_process_index == 0:
+    accelerator = Accelerator(split_batches=False, dispatch_batches=False)
+    if accelerator.is_local_main_process:
         print("**Test torch metrics**")
-    test_torch_metrics()
+        print("With: `split_batches=False`, `dispatch_batches=False`")
+    test_torch_metrics(accelerator)
+    accelerator.state._reset_state()
+    accelerator = Accelerator(split_batches=True, dispatch_batches=False)
+    if accelerator.is_local_main_process:
+        print("With: `split_batches=True`, `dispatch_batches=False`")
+    test_torch_metrics(accelerator)
+    accelerator.state._reset_state()
+    accelerator = Accelerator(split_batches=False, dispatch_batches=True)
+    if accelerator.is_local_main_process:
+        print("With: `split_batches=False`, `dispatch_batches=True`")
+    test_torch_metrics(accelerator)
+    accelerator.state._reset_state()
+    accelerator = Accelerator(split_batches=True, dispatch_batches=True)
+    if accelerator.is_local_main_process:
+        print("With: `split_batches=True`, `dispatch_batches=True`")
+    test_torch_metrics(accelerator)
 
 
 def _mp_fn(index):

I think this can be an info, no need for a warning here.
@@ -934,25 +934,28 @@ def gather(self, tensor):
         """
         return gather(tensor)
 
-    def gather_for_metrics(self, tensor, dataloader):
+    def gather_for_metrics(self, tensor):
         """
         Gathers `tensor` and potentially drops duplicates in the last batch if on a distributed system. Should be used
         for gathering the inputs and targets for metric calculation.
 
         Args:
             tensor (`torch.Tensor`, or a nested tuple/list/dictionary of `torch.Tensor`):
                 The tensors for calculating metrics across all processes.
-            dataloader (`torch.utils.data.DataLoader`):
-                A dataloader prepared with `Accelerator.prepare`
         """
         tensor = self.gather(tensor)
         if self.use_distributed:
+            if self.gradient_state.remainder == -1:
+                logger.warn(
+                    "The used dataset had no length, returning gathered tensors. You should drop the remainder yourself."
+                )
There should be an exception ehre :-)
@@ -332,7 +336,11 @@ def __iter__(self):
         if self.rng_types is not None:
             synchronize_rng_states(self.rng_types, self.generator)
         self.gradient_state._set_end_of_dataloader(False)
-        self.gradient_state._set_samples_seen(0)
+        try:
+            length = getattr(self.dataset, "total_dataset_length", len(self.dataset))
+            self.gradient_state._set_remainder(length % self.batch_size)
+        except:
Nice new test!
@@ -22,50 +22,63 @@
 from accelerate.utils import set_seed
 
 
-def get_setup(accelerator):
+def get_setup(accelerator, num_samples=82):
     "Returns everything needed to perform basic training"
     set_seed(42)
     model = RegressionModel()
     ddp_model = deepcopy(model)
-    dset = RegressionDataset(length=80)
+    dset = RegressionDataset(length=num_samples)
     dataloader = DataLoader(dset, batch_size=16)
     model.to(accelerator.device)
     ddp_model, dataloader = accelerator.prepare(ddp_model, dataloader)
     return model, ddp_model, dataloader
 
 
-def accuracy(predictions, labels) -> float:
-    """
-    Get the accuracy with respect to the most likely label
-    """
-    return (predictions == labels).float().mean()
-
-
-def test_torch_metrics():
-    accelerator = Accelerator()
-    model, ddp_model, dataloader = get_setup(accelerator)
+def generate_predictions(model, dataloader, accelerator):
+    logits_and_targets = []
     for batch in dataloader:
-        ddp_input, ddp_target = batch.values()
-        # First do single process
-        input, target = accelerator.gather((ddp_input, ddp_target))
-        input, target = input.to(accelerator.device), target.to(accelerator.device)
+        input, target = batch.values()
         with torch.no_grad():
             logits = model(input)
-            accuracy_single = accuracy(logits.argmax(dim=-1), target)
-        # Then do multiprocess
-        with torch.no_grad():
-            logits = ddp_model(ddp_input)
-            logits, target = accelerator.gather_for_metrics((logits, ddp_target), dataloader)
-            accuracy_multi = accuracy(logits.argmax(dim=-1), target)
-        assert torch.allclose(accuracy_single, accuracy_multi), "The two accuracies were not the same!"
+            logits, target = accelerator.gather_for_metrics((logits, target))
+            logits_and_targets.append((logits, target))
+    inps, targs = [], []
+    for (inp, targ) in logits_and_targets:
+        inps.append(inp)
+        targs.append(targ)
+    inps, targs = torch.cat(inps), torch.cat(targs)
+    return inps, targs
+
+
+def test_torch_metrics(accelerator: Accelerator, num_samples=82):
+    model, ddp_model, dataloader = get_setup(accelerator, num_samples)
+    inps, targs = generate_predictions(ddp_model, dataloader, accelerator)
+    assert (
+        len(inps) == num_samples
+    ), f"Unexpected number of inputs:\n    Expected: {num_samples}\n    Actual: {len(inps)}"
```suggestion
self.gradient_state._set_remainder(length % self.total_batch_size)
```
@@ -453,7 +463,11 @@ def _fetch_batches(self, iterator):
 
     def __iter__(self):
         self.gradient_state._set_end_of_dataloader(False)
-        self.gradient_state._set_samples_seen(0)
+        try:
+            length = getattr(self.dataset, "total_dataset_length", len(self.dataset))
+            self.gradient_state._set_remainder(length % self.batch_size)
Irrespective of the above change, `test_metrics.py` fails when `dispatch_batches=True` when the `num_samples % num_processes != 0` , e.g., num_samples=97 and num_processes=2. Below is the modified test with print statements. I think, `IterableDatasetShard` should be used in conjunction with  `dispatch_batches=True` for expected behaviour or `DataLoaderDispatcher` needs to be fixed to handle normal dataloaders. 

![Screenshot 2022-07-28 at 10 54 03 AM](https://user-images.githubusercontent.com/13534540/181427059-38d54e57-894e-4ae9-bf6c-9add9549c1a5.png)

@@ -453,7 +463,11 @@ def _fetch_batches(self, iterator):
 
     def __iter__(self):
         self.gradient_state._set_end_of_dataloader(False)
-        self.gradient_state._set_samples_seen(0)
+        try:
+            length = getattr(self.dataset, "total_dataset_length", len(self.dataset))
+            self.gradient_state._set_remainder(length % self.batch_size)
Thanks, will change the test case slightly for this and resolve today, though a bit weird that 82 didn't catch this
@@ -453,7 +463,11 @@ def _fetch_batches(self, iterator):
 
     def __iter__(self):
         self.gradient_state._set_end_of_dataloader(False)
-        self.gradient_state._set_samples_seen(0)
+        try:
+            length = getattr(self.dataset, "total_dataset_length", len(self.dataset))
+            self.gradient_state._set_remainder(length % self.batch_size)
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/578
diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
index bd4b72682..7100d2c88 100644
--- a/src/accelerate/utils/modeling.py
+++ b/src/accelerate/utils/modeling.py
@@ -127,7 +127,7 @@ def set_module_tensor_to_device(
 
         if is_buffer:
             module._buffers[tensor_name] = new_value
-        elif torch.device(device) != module._parameters[tensor_name].device:
+        elif value is not None or torch.device(device) != module._parameters[tensor_name].device:
             param_cls = type(module._parameters[tensor_name])
             kwargs = module._parameters[tensor_name].__dict__
             new_value = param_cls(new_value, requires_grad=old_value.requires_grad, **kwargs).to(device)

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/577
diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
index 244bf89b3..bd4b72682 100644
--- a/src/accelerate/utils/modeling.py
+++ b/src/accelerate/utils/modeling.py
@@ -127,13 +127,11 @@ def set_module_tensor_to_device(
 
         if is_buffer:
             module._buffers[tensor_name] = new_value
-        elif module._parameters[tensor_name].device.type != "cuda":
+        elif torch.device(device) != module._parameters[tensor_name].device:
             param_cls = type(module._parameters[tensor_name])
             kwargs = module._parameters[tensor_name].__dict__
             new_value = param_cls(new_value, requires_grad=old_value.requires_grad, **kwargs).to(device)
             module._parameters[tensor_name] = new_value
-        else:
-            module._parameters[tensor_name] = module._parameters[tensor_name].to(device)
 
 
 def named_module_tensors(module: nn.Module, include_buffers: bool = True, recurse: bool = False):

```suggestion
        elif device != module._parameters[tensor_name].device.index:
```
The `device` param is an integer so to properly compare the devices we need to access the index of the `device` attribute of the parameter ;) 
@@ -127,13 +127,11 @@ def set_module_tensor_to_device(
 
         if is_buffer:
             module._buffers[tensor_name] = new_value
-        elif module._parameters[tensor_name].device.type != "cuda":
+        elif device != module._parameters[tensor_name].device:
Ah should be `torch.device(device)` on the left side to cover all cases!
@@ -127,13 +127,11 @@ def set_module_tensor_to_device(
 
         if is_buffer:
             module._buffers[tensor_name] = new_value
-        elif module._parameters[tensor_name].device.type != "cuda":
+        elif device != module._parameters[tensor_name].device:
Fixed [here](https://github.com/huggingface/accelerate/pull/576/commits/10334600357bdd014acd67bb255884ee10b54fe6) normally.
@@ -127,13 +127,11 @@ def set_module_tensor_to_device(
 
         if is_buffer:
             module._buffers[tensor_name] = new_value
-        elif module._parameters[tensor_name].device.type != "cuda":
+        elif device != module._parameters[tensor_name].device:
Can confirm this fixes the issue too !
@@ -127,13 +127,11 @@ def set_module_tensor_to_device(
 
         if is_buffer:
             module._buffers[tensor_name] = new_value
-        elif module._parameters[tensor_name].device.type != "cuda":
+        elif device != module._parameters[tensor_name].device:
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/576
diff --git a/.gitignore b/.gitignore
index 7f145cd5c..da99824aa 100644
--- a/.gitignore
+++ b/.gitignore
@@ -135,4 +135,7 @@ dmypy.json
 .idea
 
 # Mac .DS_Store
-.DS_Store
\ No newline at end of file
+.DS_Store
+
+# More test things
+wandb
\ No newline at end of file
diff --git a/tests/test_examples.py b/tests/test_examples.py
index 8983468e1..ac8c4de52 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -20,6 +20,8 @@
 import unittest
 from unittest import mock
 
+import torch
+
 from accelerate.test_utils.examples import compare_against_test
 from accelerate.test_utils.testing import TempDirTestCase, require_trackers, run_command, slow
 from accelerate.utils import write_basic_config
@@ -169,8 +171,16 @@ def test_load_states_by_steps(self):
         --resume_from_checkpoint {os.path.join(self.tmpdir, "step_2")}
         """.split()
         output = run_command(self._launch_args + testargs, return_stdout=True)
-        self.assertNotIn("epoch 0:", output)
-        self.assertIn("epoch 1:", output)
+        if torch.cuda.is_available():
+            num_processes = torch.cuda.device_count()
+        else:
+            num_processes = 1
+        if num_processes > 1:
+            self.assertNotIn("epoch 0:", output)
+            self.assertIn("epoch 1:", output)
+        else:
+            self.assertIn("epoch 0:", output)
+            self.assertIn("epoch 1:", output)
 
     @slow
     def test_cross_validation(self):

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/573
diff --git a/tests/test_examples.py b/tests/test_examples.py
index 361e19642..8983468e1 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -166,7 +166,7 @@ def test_load_states_by_epoch(self):
     def test_load_states_by_steps(self):
         testargs = f"""
         examples/by_feature/checkpointing.py
-        --resume_from_checkpoint {os.path.join(self.tmpdir, "step_5")}
+        --resume_from_checkpoint {os.path.join(self.tmpdir, "step_2")}
         """.split()
         output = run_command(self._launch_args + testargs, return_stdout=True)
         self.assertNotIn("epoch 0:", output)

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/572
diff --git a/examples/by_feature/checkpointing.py b/examples/by_feature/checkpointing.py
index 1b21ebed7..5e364491e 100644
--- a/examples/by_feature/checkpointing.py
+++ b/examples/by_feature/checkpointing.py
@@ -107,6 +107,9 @@ def collate_fn(examples):
 
 
 def training_function(config, args):
+    # For testing only
+    if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+        config["num_epochs"] = 2
     # Initialize accelerator
     accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
diff --git a/examples/by_feature/fsdp_with_peak_mem_tracking.py b/examples/by_feature/fsdp_with_peak_mem_tracking.py
index 09199154a..8046b4557 100644
--- a/examples/by_feature/fsdp_with_peak_mem_tracking.py
+++ b/examples/by_feature/fsdp_with_peak_mem_tracking.py
@@ -83,6 +83,9 @@ def __exit__(self, *exc):
 
 
 def training_function(config, args):
+    # For testing only
+    if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+        config["num_epochs"] = 2
     # Initialize accelerator
     if args.with_tracking:
         accelerator = Accelerator(
diff --git a/examples/by_feature/gradient_accumulation.py b/examples/by_feature/gradient_accumulation.py
index d95f9bef3..fe3d948f5 100644
--- a/examples/by_feature/gradient_accumulation.py
+++ b/examples/by_feature/gradient_accumulation.py
@@ -102,6 +102,9 @@ def collate_fn(examples):
 
 
 def training_function(config, args):
+    # For testing only
+    if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+        config["num_epochs"] = 2
     # New Code #
     gradient_accumulation_steps = int(args.gradient_accumulation_steps)
     # Initialize accelerator
diff --git a/examples/by_feature/memory.py b/examples/by_feature/memory.py
index b11c80fd0..f71283561 100644
--- a/examples/by_feature/memory.py
+++ b/examples/by_feature/memory.py
@@ -107,6 +107,9 @@ def collate_fn(examples):
 
 
 def training_function(config, args):
+    # For testing only
+    if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+        config["num_epochs"] = 2
     # Initialize accelerator
     accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
diff --git a/examples/by_feature/multi_process_metrics.py b/examples/by_feature/multi_process_metrics.py
index dc47f3db6..d69c0bc90 100644
--- a/examples/by_feature/multi_process_metrics.py
+++ b/examples/by_feature/multi_process_metrics.py
@@ -109,6 +109,9 @@ def collate_fn(examples):
 
 
 def training_function(config, args):
+    # For testing only
+    if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+        config["num_epochs"] = 2
     # Initialize accelerator
     accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
diff --git a/examples/by_feature/tracking.py b/examples/by_feature/tracking.py
index 00c8665dd..d453b770c 100644
--- a/examples/by_feature/tracking.py
+++ b/examples/by_feature/tracking.py
@@ -107,6 +107,9 @@ def collate_fn(examples):
 
 
 def training_function(config, args):
+    # For testing only
+    if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+        config["num_epochs"] = 2
     # Initialize Accelerator
 
     # New Code #
diff --git a/src/accelerate/test_utils/examples.py b/src/accelerate/test_utils/examples.py
index 4e4092c0e..f459e03c4 100644
--- a/src/accelerate/test_utils/examples.py
+++ b/src/accelerate/test_utils/examples.py
@@ -115,11 +115,17 @@ def compare_against_test(base_filename: str, feature_filename: str, parser_only:
     # Specific code in our script that differs from the full version, aka what is new
     new_feature_code = []
     passed_idxs = []  # We keep track of the idxs just in case it's a repeated statement
-    for i, line in enumerate(feature_file_func):
+    it = iter(feature_file_func)
+    for i in range(len(feature_file_func) - 1):
         if i not in passed_idxs:
+            line = next(it)
             if (line not in full_file_func) and (line.lstrip() != _dl_line):
-                new_feature_code.append(line)
-                passed_idxs.append(i)
+                if "TESTING_MOCKED_DATALOADERS" not in line:
+                    new_feature_code.append(line)
+                    passed_idxs.append(i)
+                else:
+                    # Skip over the `config['num_epochs'] = 2` statement
+                    _ = next(it)
 
     # Extract out just the new parts from the full_file_training_func
     new_full_example_parts = []
@@ -127,8 +133,9 @@ def compare_against_test(base_filename: str, feature_filename: str, parser_only:
     for i, line in enumerate(base_file_func):
         if i not in passed_idxs:
             if (line not in full_file_func) and (line.lstrip() != _dl_line):
-                new_full_example_parts.append(line)
-                passed_idxs.append(i)
+                if "TESTING_MOCKED_DATALOADERS" not in line:
+                    new_full_example_parts.append(line)
+                    passed_idxs.append(i)
 
     # Finally, get the overall diff
     diff_from_example = [line for line in new_feature_code if line not in new_full_example_parts]
diff --git a/tests/test_examples.py b/tests/test_examples.py
index 9e1854cd8..361e19642 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -20,8 +20,6 @@
 import unittest
 from unittest import mock
 
-import torch
-
 from accelerate.test_utils.examples import compare_against_test
 from accelerate.test_utils.testing import TempDirTestCase, require_trackers, run_command, slow
 from accelerate.utils import write_basic_config
@@ -145,7 +143,7 @@ def test_checkpointing_by_epoch(self):
         --output_dir {self.tmpdir}
         """.split()
         run_command(self._launch_args + testargs)
-        self.assertTrue(os.path.exists(os.path.join(self.tmpdir, "epoch_1")))
+        self.assertTrue(os.path.exists(os.path.join(self.tmpdir, "epoch_0")))
 
     def test_checkpointing_by_steps(self):
         testargs = f"""
@@ -154,17 +152,16 @@ def test_checkpointing_by_steps(self):
         --output_dir {self.tmpdir}
         """.split()
         _ = run_command(self._launch_args + testargs)
-        self.assertTrue(os.path.exists(os.path.join(self.tmpdir, "step_5")))
+        self.assertTrue(os.path.exists(os.path.join(self.tmpdir, "step_2")))
 
     def test_load_states_by_epoch(self):
         testargs = f"""
         examples/by_feature/checkpointing.py
-        --resume_from_checkpoint {os.path.join(self.tmpdir, "epoch_1")}
+        --resume_from_checkpoint {os.path.join(self.tmpdir, "epoch_0")}
         """.split()
         output = run_command(self._launch_args + testargs, return_stdout=True)
         self.assertNotIn("epoch 0:", output)
-        self.assertNotIn("epoch 1:", output)
-        self.assertIn("epoch 2:", output)
+        self.assertIn("epoch 1:", output)
 
     def test_load_states_by_steps(self):
         testargs = f"""
@@ -172,17 +169,8 @@ def test_load_states_by_steps(self):
         --resume_from_checkpoint {os.path.join(self.tmpdir, "step_5")}
         """.split()
         output = run_command(self._launch_args + testargs, return_stdout=True)
-        if torch.cuda.is_available():
-            num_processes = torch.cuda.device_count()
-        else:
-            num_processes = 1
-        if num_processes > 1:
-            self.assertNotIn("epoch 0:", output)
-            self.assertNotIn("epoch 1:", output)
-        else:
-            self.assertNotIn("epoch 0:", output)
-            self.assertIn("epoch 1:", output)
-        self.assertIn("epoch 2:", output)
+        self.assertNotIn("epoch 0:", output)
+        self.assertIn("epoch 1:", output)
 
     @slow
     def test_cross_validation(self):

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/571
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index a0f9fab40..b84a90c6e 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -29,10 +29,12 @@
 
 try:
     import torch_ccl  # noqa: F401
+
     _ccl_available = True
 except ImportError:
-    try: 
-        import oneccl_bindings_for_pytorch
+    try:
+        import oneccl_bindings_for_pytorch  # noqa: F401
+
         _ccl_available = True
     except ImportError:
         _ccl_available = False

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/569
diff --git a/.github/workflows/check_dependencies.yml b/.github/workflows/check_dependencies.yml
index 1e8a4e1bb..f8ac1f492 100644
--- a/.github/workflows/check_dependencies.yml
+++ b/.github/workflows/check_dependencies.yml
@@ -1,4 +1,4 @@
-name: Trigger docker images and run slow tests
+name: Trigger docker images and run tests
 
 on:
   push:
diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml
index 0e8f31407..c0c9c62f2 100644
--- a/.github/workflows/nightly.yml
+++ b/.github/workflows/nightly.yml
@@ -1,4 +1,4 @@
-name: Self-hosted runner (scheduled)
+name: Self-hosted runner with slow tests (scheduled)
 
 on:
   workflow_dispatch:

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/568
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index c846dd65d..06a25bc4e 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -17,7 +17,7 @@
 
 import torch
 
-from .utils import DistributedType, is_ccl_available, is_deepspeed_available, is_tpu_available
+from .utils import DistributedType, get_ccl_version, is_ccl_available, is_deepspeed_available, is_tpu_available
 from .utils.dataclasses import SageMakerDistributedType
 
 
@@ -166,6 +166,10 @@ def __init__(
             elif get_int_from_env(["PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE", "WORLD_SIZE"], 1) > 1:
                 self.distributed_type = DistributedType.MULTI_CPU
                 if is_ccl_available() and get_int_from_env(["CCL_WORKER_COUNT"], 0) > 0:
+                    if get_ccl_version() >= "1.12":
+                        import oneccl_bindings_for_pytorch  # noqa: F401
+                    else:
+                        import torch_ccl  # noqa: F401
                     backend = "ccl"
                 elif torch.distributed.is_mpi_available():
                     backend = "mpi"
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
index f6b83d2f9..0395a271d 100644
--- a/src/accelerate/utils/__init__.py
+++ b/src/accelerate/utils/__init__.py
@@ -19,6 +19,7 @@
     TensorInformation,
 )
 from .imports import (
+    get_ccl_version,
     is_apex_available,
     is_bf16_available,
     is_boto3_available,
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index b84a90c6e..da63d0aa1 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -27,19 +27,6 @@
     import importlib.metadata as importlib_metadata
 
 
-try:
-    import torch_ccl  # noqa: F401
-
-    _ccl_available = True
-except ImportError:
-    try:
-        import oneccl_bindings_for_pytorch  # noqa: F401
-
-        _ccl_available = True
-    except ImportError:
-        _ccl_available = False
-
-
 try:
     import torch_xla.core.xla_model as xm  # noqa: F401
 
@@ -49,7 +36,14 @@
 
 
 def is_ccl_available():
-    return _ccl_available
+    return (
+        importlib.util.find_spec("torch_ccl") is not None
+        or importlib.util.find_spec("oneccl_bindings_for_pytorch") is not None
+    )
+
+
+def get_ccl_version():
+    return importlib_metadata.version("oneccl_bind_pt")
 
 
 def is_apex_available():

@liangan1 @yao-matrix @sgugger please have a review
_The documentation is not available anymore as the PR was closed or merged._
Thanks for this! Looks great to me 😄 could you run `make style; make quality` to get past that failure please? And then we're good!
@muellerzr thanks for the suggestion, fix the "make quality" issue.
done for the rebase @sgugger 
https://api.github.com/repos/huggingface/accelerate/pulls/567
diff --git a/setup.py b/setup.py
index cad067ad5..a59d1ecaf 100644
--- a/setup.py
+++ b/setup.py
@@ -19,7 +19,7 @@
 extras["quality"] = ["black ~= 22.0", "isort >= 5.5.4", "flake8 >= 3.8.3", "hf-doc-builder >= 0.3.0"]
 extras["docs"] = []
 extras["test_prod"] = ["pytest", "pytest-xdist", "pytest-subtests", "parameterized"]
-extras["test_dev"] = ["datasets", "evaluate", "transformers", "scipy", "sklearn", "deepspeed"]
+extras["test_dev"] = ["datasets", "evaluate", "transformers", "scipy", "sklearn", "deepspeed", "tqdm"]
 extras["testing"] = extras["test_prod"] + extras["test_dev"]
 
 extras["test_trackers"] = ["wandb", "comet-ml", "tensorboard"]
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
index f6b83d2f9..af8ddd736 100644
--- a/src/accelerate/utils/__init__.py
+++ b/src/accelerate/utils/__init__.py
@@ -100,3 +100,4 @@
     write_basic_config,
 )
 from .random import set_seed, synchronize_rng_state, synchronize_rng_states
+from .tqdm import tqdm
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index b84a90c6e..4b796e62c 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -113,3 +113,7 @@ def is_boto3_available():
 
 def is_sagemaker_available():
     return importlib.util.find_spec("sagemaker") is not None
+
+
+def is_tqdm_available():
+    return importlib.util.find_spec("tqdm") is not None
diff --git a/src/accelerate/utils/tqdm.py b/src/accelerate/utils/tqdm.py
new file mode 100644
index 000000000..be489bd90
--- /dev/null
+++ b/src/accelerate/utils/tqdm.py
@@ -0,0 +1,37 @@
+# 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.
+
+from .imports import is_tqdm_available
+
+
+if is_tqdm_available():
+    import tqdm.auto as _tqdm
+
+from ..state import AcceleratorState
+
+
+def tqdm(main_process_only: bool = True, *args, **kwargs):
+    """
+    Wrapper around `tqdm.tqdm` that optionally displays only on the main process.
+
+    Args:
+        main_process_only (`bool`, *optional*):
+            Whether to display the progress bar only on the main process
+    """
+    if not is_tqdm_available():
+        raise ImportError("Accelerate's `tqdm` module requires `tqdm` to be installed. Please run `pip install tqdm`.")
+    disable = False
+    if main_process_only:
+        disable = AcceleratorState().local_process_index == 0
+    return _tqdm(*args, **kwargs, disable=disable)

Let's make tqdm a soft dependency.
@@ -50,7 +50,7 @@
         ]
     },
     python_requires=">=3.7.0",
-    install_requires=["numpy>=1.17", "packaging>=20.0", "psutil", "pyyaml", "torch>=1.4.0"],
+    install_requires=["numpy>=1.17", "packaging>=20.0", "psutil", "pyyaml", "torch>=1.4.0", "tqdm"],
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/564
diff --git a/setup.py b/setup.py
index 60b6647fd..cad067ad5 100644
--- a/setup.py
+++ b/setup.py
@@ -19,7 +19,7 @@
 extras["quality"] = ["black ~= 22.0", "isort >= 5.5.4", "flake8 >= 3.8.3", "hf-doc-builder >= 0.3.0"]
 extras["docs"] = []
 extras["test_prod"] = ["pytest", "pytest-xdist", "pytest-subtests", "parameterized"]
-extras["test_dev"] = ["datasets<=2.2.2", "evaluate", "transformers", "scipy", "sklearn", "deepspeed"]
+extras["test_dev"] = ["datasets", "evaluate", "transformers", "scipy", "sklearn", "deepspeed"]
 extras["testing"] = extras["test_prod"] + extras["test_dev"]
 
 extras["test_trackers"] = ["wandb", "comet-ml", "tensorboard"]

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/563
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 514e77d05..9ee840440 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -27,8 +27,10 @@
 
 import torch
 
+import psutil
 from accelerate.commands.config import default_config_file, load_config_from_file
 from accelerate.commands.config.config_args import SageMakerConfig
+from accelerate.state import get_int_from_env
 from accelerate.utils import (
     ComputeEnvironment,
     DistributedType,
@@ -273,7 +275,7 @@ def launch_command_parser(subparsers=None):
     parser.add_argument(
         "--num_cpu_threads_per_process",
         type=int,
-        default=1,
+        default=None,
         help="The number of CPU threads per process. Can be tuned for optimal performance.",
     )
     parser.add_argument(
@@ -337,6 +339,17 @@ def simple_launcher(args):
         mixed_precision = "fp16"
 
     current_env["MIXED_PRECISION"] = str(mixed_precision)
+    if args.num_cpu_threads_per_process is None:
+        local_size = get_int_from_env(
+            ["MPI_LOCALNRANKS", "OMPI_COMM_WORLD_LOCAL_SIZE", "MV2_COMM_WORLD_LOCAL_SIZE"], 1
+        )
+        args.num_cpu_threads_per_process = int(psutil.cpu_count(logical=False) / local_size)
+        if args.num_cpu_threads_per_process == 0:
+            args.num_cpu_threads_per_process = 1
+        logger.info(
+            f"num_cpu_threads_per_process unset, we set it at {args.num_cpu_threads_per_process} to improve oob performance."
+        )
+
     current_env["OMP_NUM_THREADS"] = str(args.num_cpu_threads_per_process)
 
     process = subprocess.Popen(cmd, env=current_env)
@@ -434,7 +447,9 @@ def multi_gpu_launcher(args):
             current_env["FSDP_BACKWARD_PREFETCH"] = str(args.fsdp_backward_prefetch_policy)
         if args.fsdp_state_dict_type is not None:
             current_env["FSDP_STATE_DICT_TYPE"] = str(args.fsdp_state_dict_type)
-
+    if args.num_cpu_threads_per_process is None:
+        args.num_cpu_threads_per_process = 1
+        logger.info(f"num_cpu_threads_per_process unset, we set it at {args.num_cpu_threads_per_process}.")
     current_env["OMP_NUM_THREADS"] = str(args.num_cpu_threads_per_process)
     process = subprocess.Popen(cmd, env=current_env)
     process.wait()

We shouldn't change this value if the user has set it to 1. Maybe change the default of `args.num_cpu_threads_per_process` to `None` and then set it to the right default if it's unset.
@@ -307,6 +309,20 @@ def simple_launcher(args):
         mixed_precision = "fp16"
 
     current_env["MIXED_PRECISION"] = str(mixed_precision)
+    if args.num_cpu_threads_per_process == 1:
Adapt the warning to the behavior suggested above (no default was provided, we picked xxx). Also use f-strings please.
@@ -307,6 +309,20 @@ def simple_launcher(args):
         mixed_precision = "fp16"
 
     current_env["MIXED_PRECISION"] = str(mixed_precision)
+    if args.num_cpu_threads_per_process == 1:
+        local_size = get_int_from_env(
+            ["MPI_LOCALNRANKS", "OMPI_COMM_WORLD_LOCAL_SIZE", "MV2_COMM_WORLD_LOCAL_SIZE"], 1
+        )
+        args.num_cpu_threads_per_process = int(psutil.cpu_count(logical=False) / local_size)
+        if args.num_cpu_threads_per_process == 0:
+            args.num_cpu_threads_per_process = 1
+        if args.num_cpu_threads_per_process != 1:
+            logger.warn(
+                'num_cpu_threads_per_process is reset to %d instead of default 1 to improve oob performance.\
+        \nyou could tun num_cpu_threads_per_process by pass --num_cpu_threads_per_process xxx to "accelerate launch "\n'
+                % args.num_cpu_threads_per_process
apply you suggestion @sgugger 
@@ -307,6 +309,20 @@ def simple_launcher(args):
         mixed_precision = "fp16"
 
     current_env["MIXED_PRECISION"] = str(mixed_precision)
+    if args.num_cpu_threads_per_process == 1:
+        local_size = get_int_from_env(
+            ["MPI_LOCALNRANKS", "OMPI_COMM_WORLD_LOCAL_SIZE", "MV2_COMM_WORLD_LOCAL_SIZE"], 1
+        )
+        args.num_cpu_threads_per_process = int(psutil.cpu_count(logical=False) / local_size)
+        if args.num_cpu_threads_per_process == 0:
+            args.num_cpu_threads_per_process = 1
+        if args.num_cpu_threads_per_process != 1:
+            logger.warn(
+                'num_cpu_threads_per_process is reset to %d instead of default 1 to improve oob performance.\
+        \nyou could tun num_cpu_threads_per_process by pass --num_cpu_threads_per_process xxx to "accelerate launch "\n'
+                % args.num_cpu_threads_per_process
```suggestion
            f"num_cpu_threads_per_process unset, we set it at {args.num_cpu_threads_per_process} to improve oob performance."
```
@@ -337,6 +339,17 @@ def simple_launcher(args):
         mixed_precision = "fp16"
 
     current_env["MIXED_PRECISION"] = str(mixed_precision)
+    if args.num_cpu_threads_per_process is None:
+        local_size = get_int_from_env(
+            ["MPI_LOCALNRANKS", "OMPI_COMM_WORLD_LOCAL_SIZE", "MV2_COMM_WORLD_LOCAL_SIZE"], 1
+        )
+        args.num_cpu_threads_per_process = int(psutil.cpu_count(logical=False) / local_size)
+        if args.num_cpu_threads_per_process == 0:
+            args.num_cpu_threads_per_process = 1
+        logger.warn(
+            f"num_cpu_threads_per_process unset, we pick as {args.num_cpu_threads_per_process} to improve oob performance"
```suggestion
        logger.info(
```
Maybe the info level is enough?
@@ -337,6 +339,17 @@ def simple_launcher(args):
         mixed_precision = "fp16"
 
     current_env["MIXED_PRECISION"] = str(mixed_precision)
+    if args.num_cpu_threads_per_process is None:
+        local_size = get_int_from_env(
+            ["MPI_LOCALNRANKS", "OMPI_COMM_WORLD_LOCAL_SIZE", "MV2_COMM_WORLD_LOCAL_SIZE"], 1
+        )
+        args.num_cpu_threads_per_process = int(psutil.cpu_count(logical=False) / local_size)
+        if args.num_cpu_threads_per_process == 0:
+            args.num_cpu_threads_per_process = 1
+        logger.warn(
```suggestion
        logger.info(f"num_cpu_threads_per_process unset, we set it at {args.num_cpu_threads_per_process}.")
```
@@ -434,7 +447,9 @@ def multi_gpu_launcher(args):
             current_env["FSDP_BACKWARD_PREFETCH"] = str(args.fsdp_backward_prefetch_policy)
         if args.fsdp_state_dict_type is not None:
             current_env["FSDP_STATE_DICT_TYPE"] = str(args.fsdp_state_dict_type)
-
+    if args.num_cpu_threads_per_process is None:
+        args.num_cpu_threads_per_process = 1
+        logger.warn(f"num_cpu_threads_per_process unset, we pick as {args.num_cpu_threads_per_process}")
change to info
@@ -337,6 +339,17 @@ def simple_launcher(args):
         mixed_precision = "fp16"
 
     current_env["MIXED_PRECISION"] = str(mixed_precision)
+    if args.num_cpu_threads_per_process is None:
+        local_size = get_int_from_env(
+            ["MPI_LOCALNRANKS", "OMPI_COMM_WORLD_LOCAL_SIZE", "MV2_COMM_WORLD_LOCAL_SIZE"], 1
+        )
+        args.num_cpu_threads_per_process = int(psutil.cpu_count(logical=False) / local_size)
+        if args.num_cpu_threads_per_process == 0:
+            args.num_cpu_threads_per_process = 1
+        logger.warn(
_The documentation is not available anymore as the PR was closed or merged._
@sgugger @yao-matrix @liangan1 please help review it
https://api.github.com/repos/huggingface/accelerate/pulls/562
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 35afd6c5d..5aea696b4 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -446,16 +446,13 @@ def _prepare_one(self, obj, first_pass=False):
             if isinstance(obj, torch.utils.data.DataLoader):
                 return self.prepare_data_loader(obj)
             elif isinstance(obj, torch.nn.Module):
-                self._models.append(obj)
                 return self.prepare_model(obj)
             elif isinstance(obj, torch.optim.Optimizer):
                 optimizer = self.prepare_optimizer(obj)
-                self._optimizers.append(optimizer)
                 return optimizer
         # Second pass of preparation: LR scheduler (which need the full list of optimizers)
         elif isinstance(obj, torch.optim.lr_scheduler._LRScheduler):
             scheduler = self.prepare_scheduler(obj)
-            self._schedulers.append(scheduler)
             return scheduler
         # Return the unprocessed object if previous criteria was not met
         return obj
@@ -570,6 +567,7 @@ def prepare(self, *args):
         return result if len(result) > 1 else result[0]
 
     def prepare_model(self, model):
+        self._models.append(model)
         if self.device_placement and self.distributed_type != DistributedType.FSDP:
             model = model.to(self.device)
         if self.distributed_type == DistributedType.MULTI_GPU:
@@ -782,7 +780,9 @@ def prepare_data_loader(self, data_loader):
         )
 
     def prepare_optimizer(self, optimizer):
-        return AcceleratedOptimizer(optimizer, device_placement=self.device_placement, scaler=self.scaler)
+        optimizer = AcceleratedOptimizer(optimizer, device_placement=self.device_placement, scaler=self.scaler)
+        self._optimizers.append(optimizer)
+        return optimizer
 
     def prepare_scheduler(self, scheduler):
         # We try to find the optimizer associated with `scheduler`, the default is the full list.
@@ -791,13 +791,14 @@ def prepare_scheduler(self, scheduler):
             if getattr(scheduler, "optimizer", None) == opt.optimizer:
                 optimizer = opt
                 break
-
-        return AcceleratedScheduler(
+        scheduler = AcceleratedScheduler(
             scheduler,
             optimizer,
             step_with_optimizer=self.step_scheduler_with_optimizer,
             split_batches=self.split_batches,
         )
+        self._schedulers.append(scheduler)
+        return scheduler
 
     def backward(self, loss, **kwargs):
         """

Thanks! If we want to go this route, it would be good to update the rest of them as well. I.e. optimizer and scheduler
_The documentation is not available anymore as the PR was closed or merged._
@muellerzr thanks for the comment, I updated the optimizer and scheduler as well. If any other changes are needed, please let me know.
https://api.github.com/repos/huggingface/accelerate/pulls/556
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index a660db858..8bffd8e99 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -29,10 +29,13 @@
 
 try:
     import torch_ccl  # noqa: F401
-
     _ccl_available = True
 except ImportError:
-    _ccl_available = False
+    try: 
+        import oneccl_bindings_for_pytorch
+        _ccl_available = True
+     except ImportError:
+        _ccl_available = False
 
 
 try:

https://api.github.com/repos/huggingface/accelerate/pulls/554
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 8d4138651..e6a23e46d 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -16,6 +16,7 @@
 
 import argparse
 import importlib
+import logging
 import os
 import subprocess
 import sys
@@ -24,6 +25,8 @@
 from pathlib import Path
 from typing import Dict, List
 
+import torch
+
 from accelerate.commands.config import default_config_file, load_config_from_file
 from accelerate.commands.config.config_args import SageMakerConfig
 from accelerate.utils import (
@@ -40,6 +43,9 @@
 from accelerate.utils.dataclasses import SageMakerDistributedType
 
 
+logger = logging.getLogger(__name__)
+
+
 def launch_command_parser(subparsers=None):
     if subparsers is not None:
         parser = subparsers.add_parser("launch")
@@ -659,6 +665,7 @@ def launch_command(args):
         raise ValueError("You can only pick one between `--multi_gpu`, `--use_deepspeed`, `--tpu`, `--use_fsdp`.")
 
     defaults = None
+    warned = []
     # Get the default from the config file.
     if args.config_file is not None or os.path.isfile(default_config_file) and not args.cpu:
         defaults = load_config_from_file(args.config_file)
@@ -684,7 +691,6 @@ def launch_command(args):
                     and getattr(args, name, None) is None
                 ):
                     setattr(args, name, attr)
-
         if not args.mixed_precision:
             if args.fp16:
                 args.mixed_precision = "fp16"
@@ -692,7 +698,32 @@ def launch_command(args):
                 args.mixed_precision = defaults.mixed_precision
     else:
         if args.num_processes is None:
+            warned.append("\t`--num_processes` was set to a value of `1`")
             args.num_processes = 1
+        if args.num_machines is None:
+            warned.append("\t`--num_machines` was set to a value of `1`")
+            args.num_machines = 1
+        if args.mixed_precision is None:
+            warned.append("\t`--mixed_precision` was set to a value of `'no'`")
+            args.mixed_precision = "no"
+        if not hasattr(args, "use_cpu"):
+            args.use_cpu = args.cpu
+    if args.multi_gpu and args.num_processes == 1:
+        args.num_processes = torch.cuda.device_count()
+        if not any("--num_processes" in warn for warn in warned):
+            warned.append(f"\t`--num_processes` was set to `{args.num_processes}`")
+        else:
+            for i, warn in enumerate(warned):
+                if "--num_processes" in warn:
+                    warned[i] = warn.replace("`1`", f"`{args.num_processes}`")
+
+    if any(warned):
+        message = "The following values were not passed to `accelerate launch` and had defaults used instead:\n"
+        message += "\n".join(warned)
+        message += (
+            "\nTo avoid this warning pass in values for each of the problematic parameters or run `accelerate config`."
+        )
+        logger.warn(message)
 
     # Use the proper launcher
     if args.use_deepspeed and not args.cpu:

```suggestion
            logger.warn("To avoid these warnings pass in values for each of the problematic parameters or run `accelerate config`.")
```
@@ -692,16 +696,26 @@ def launch_command(args):
             else:
                 args.mixed_precision = defaults.mixed_precision
     else:
+        warned = False
         if args.num_processes is None:
+            logger.warn("`--num_processes` was not set, using a value of `1`.")
+            warned = True
             args.num_processes = 1
         if args.num_machines is None:
+            warned = True
+            logger.warn("`--num_machines` was not set, using a value of `1`.")
             args.num_machines = 1
         if args.mixed_precision is None:
+            warned = True
+            logger.warn("`--mixed_precision` was not set, using a value of `'no'`.")
             args.mixed_precision = "no"
         if not hasattr(args, "use_cpu"):
-            args.use_cpu = torch.cuda.is_available()
+            args.use_cpu = args.cpu
+        if warned:
+            logger.warn("To avoid these warnings pass in values for each of the problematic parameters")
This one should have the text above to avoid warning (which you can put in a constant).
@@ -692,16 +696,26 @@ def launch_command(args):
             else:
                 args.mixed_precision = defaults.mixed_precision
     else:
+        warned = False
         if args.num_processes is None:
+            logger.warn("`--num_processes` was not set, using a value of `1`.")
+            warned = True
             args.num_processes = 1
         if args.num_machines is None:
+            warned = True
+            logger.warn("`--num_machines` was not set, using a value of `1`.")
             args.num_machines = 1
         if args.mixed_precision is None:
+            warned = True
+            logger.warn("`--mixed_precision` was not set, using a value of `'no'`.")
             args.mixed_precision = "no"
         if not hasattr(args, "use_cpu"):
-            args.use_cpu = torch.cuda.is_available()
+            args.use_cpu = args.cpu
+        if warned:
+            logger.warn("To avoid these warnings pass in values for each of the problematic parameters")
 
     if args.multi_gpu and args.num_processes == 1:
+        logger.warn("`--multi_gpu` was passed but `num_processes` was not set. Automatically using all available GPUs")
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/553
diff --git a/docs/source/accelerator.mdx b/docs/source/accelerator.mdx
index 5202aeb8c..268b09d1f 100644
--- a/docs/source/accelerator.mdx
+++ b/docs/source/accelerator.mdx
@@ -44,11 +44,12 @@ actually be performed, and auto-scale the loss:
 
 ```python
 accelerator = Accelerator(gradient_accumulation_steps=2)
+model, optimizer, training_dataloader = accelerator.prepare(model, optimizer, training_dataloader)
 
-for (input, label) in enumerate(training_dataloader):
+for input, label in training_dataloader:
     with accelerator.accumulate(model):
         predictions = model(input)
-        loss = loss_function(predictions, labels)
+        loss = loss_function(predictions, label)
         accelerator.backward(loss)
         optimizer.step()
         scheduler.step()

_The documentation is not available anymore as the PR was closed or merged._
@sgugger Thanks, I applied the change.
Thanks a lot!
https://api.github.com/repos/huggingface/accelerate/pulls/551
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index c0871b28c..fe44d5f61 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -12,7 +12,9 @@ jobs:
     strategy:
       matrix:
         test-kind: [
-          test,
+          test_prod,
+          test_core,
+          test_big_modeling,
           test_deepspeed,
           test_example_differences,
           test_checkpoint_step,
@@ -37,8 +39,9 @@ jobs:
     - name: Install the library
       run: |
         pip install --upgrade pip
-        pip install -e .[test,test_trackers]
-        if [ ${{ matrix.test-kind }} = test_rest ]; then pip uninstall comet_ml -y; fi
+        if [[ ${{ matrix.test-kind }} = test_prod ]]; then pip install -e .[test_prod]; fi
+        if [[ ${{ matrix.test-kind }} != test_prod ]]; then pip install -e .[test,test_trackers]; fi
+        if [[ ${{ matrix.test-kind }} = test_rest ]]; then pip uninstall comet_ml -y; fi
     
     - name: Run Tests
       run: |
diff --git a/Makefile b/Makefile
index 6dab6af11..4104f1bba 100644
--- a/Makefile
+++ b/Makefile
@@ -27,6 +27,12 @@ style:
 test:
 	python -m pytest -s -v ./tests/ --ignore=./tests/test_examples.py
 
+test_big_modeling:
+	python -m pytest -s -v ./tests/test_big_modeling.py
+
+test_core:
+	python -m pytest -s -v ./tests/ --ignore=./tests/test_examples.py --ignore=./tests/deepspeed --ignore=./tests/test_big_modeling.py
+
 test_deepspeed:
 	python -m pytest -s -v ./tests/deepspeed
 
@@ -43,5 +49,9 @@ test_checkpoint_epoch:
 test_checkpoint_step:
 	python -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k "by_step"
 
+# Same as test but used to install only the base dependencies
+test_prod:
+	$(MAKE) test_core
+
 test_rest:
 	python -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k "not by_step and not by_epoch"
\ No newline at end of file
diff --git a/setup.py b/setup.py
index 77679a51e..dedd077f3 100644
--- a/setup.py
+++ b/setup.py
@@ -18,18 +18,9 @@
 extras = {}
 extras["quality"] = ["black ~= 22.0", "isort >= 5.5.4", "flake8 >= 3.8.3", "hf-doc-builder >= 0.3.0"]
 extras["docs"] = []
-extras["test"] = [
-    "pytest",
-    "pytest-xdist",
-    "pytest-subtests",
-    "datasets<=2.2.2",
-    "evaluate",
-    "transformers",
-    "scipy",
-    "sklearn",
-    "parameterized",
-    "deepspeed",
-]
+extras["test_prod"] = ["pytest", "pytest-xdist", "pytest-subtests", "parameterized"]
+extras["test_dev"] = ["datasets<=2.2.2", "evaluate", "transformers", "scipy", "sklearn", "deepspeed"]
+extras["test"] = extras["test_prod"] + extras["test_dev"]
 
 extras["test_trackers"] = ["wandb", "comet-ml", "tensorboard"]
 extras["dev"] = extras["quality"] + extras["test"]
diff --git a/src/accelerate/test_utils/scripts/test_metrics.py b/src/accelerate/test_utils/scripts/test_metrics.py
index 442053213..c426d1fc5 100644
--- a/src/accelerate/test_utils/scripts/test_metrics.py
+++ b/src/accelerate/test_utils/scripts/test_metrics.py
@@ -17,7 +17,6 @@
 import torch
 from torch.utils.data import DataLoader
 
-import evaluate
 from accelerate import Accelerator
 from accelerate.test_utils import RegressionDataset, RegressionModel
 from accelerate.utils import set_seed
@@ -61,26 +60,6 @@ def test_torch_metrics():
         assert torch.allclose(accuracy_single, accuracy_multi), "The two accuracies were not the same!"
 
 
-def test_evaluate_metrics():
-    metric = evaluate.load("accuracy")
-    accelerator = Accelerator()
-    model, ddp_model, dataloader = get_setup(accelerator)
-    for batch in dataloader:
-        ddp_input, ddp_target = batch.values()
-        # First do single process
-        input, target = accelerator.gather((ddp_input, ddp_target))
-        input, target = input.to(accelerator.device), target.to(accelerator.device)
-        with torch.no_grad():
-            logits = model(input)
-            accuracy_single = metric.compute(logits, target)
-        # Then do multiprocess
-        with torch.no_grad():
-            logits = ddp_model(ddp_input)
-            logits, target = accelerator.gather_for_metrics((logits, ddp_target), dataloader)
-            accuracy_multi = metric.compute(logits, target)
-        assert torch.allclose(accuracy_single, accuracy_multi), "The two accuracies were not the same!"
-
-
 def main():
     accelerator = Accelerator()
     state = accelerator.state

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/547
diff --git a/.github/workflows/on-merge.yml b/.github/workflows/on-merge.yml
index 59677b34e..028b99d9a 100644
--- a/.github/workflows/on-merge.yml
+++ b/.github/workflows/on-merge.yml
@@ -25,7 +25,7 @@ jobs:
           source activate accelerate
           git config --global --add safe.directory '*'
           git fetch && git checkout ${{ github.sha }}
-          pip install -e .[test,test_trackers]
+          pip install -e .[testing,test_trackers]
 
       - name: Run test on GPUs
         run: |
@@ -52,7 +52,7 @@ jobs:
           source activate accelerate
           git config --global --add safe.directory '*'
           git fetch && git checkout ${{ github.sha }}
-          pip install -e .[test,test_trackers]
+          pip install -e .[testing,test_trackers]
 
       - name: Run test on GPUs
         run: |
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index fe44d5f61..09330eb63 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -40,7 +40,7 @@ jobs:
       run: |
         pip install --upgrade pip
         if [[ ${{ matrix.test-kind }} = test_prod ]]; then pip install -e .[test_prod]; fi
-        if [[ ${{ matrix.test-kind }} != test_prod ]]; then pip install -e .[test,test_trackers]; fi
+        if [[ ${{ matrix.test-kind }} != test_prod ]]; then pip install -e .[testing,test_trackers]; fi
         if [[ ${{ matrix.test-kind }} = test_rest ]]; then pip uninstall comet_ml -y; fi
     
     - name: Run Tests
diff --git a/docker/accelerate-cpu/Dockerfile b/docker/accelerate-cpu/Dockerfile
index 5d1286772..a872e6fb2 100644
--- a/docker/accelerate-cpu/Dockerfile
+++ b/docker/accelerate-cpu/Dockerfile
@@ -21,7 +21,7 @@ WORKDIR /workspace
 RUN python3 -m pip install --upgrade --no-cache-dir pip
 RUN python3 -m pip install --no-cache-dir \
     jupyter \
-    git+https://github.com/huggingface/accelerate#egg=accelerate[test,test_trackers] \
+    git+https://github.com/huggingface/accelerate#egg=accelerate[testing,test_trackers] \
     --extra-index-url https://download.pytorch.org/whl/cpu
     
 # Stage 2
diff --git a/docker/accelerate-gpu/Dockerfile b/docker/accelerate-gpu/Dockerfile
index 746e65350..44aa94b29 100644
--- a/docker/accelerate-gpu/Dockerfile
+++ b/docker/accelerate-gpu/Dockerfile
@@ -22,7 +22,7 @@ SHELL ["/bin/bash", "-c"]
 # Activate the conda env and install torch + accelerate
 RUN source activate accelerate && \
     python3 -m pip install --no-cache-dir \
-    git+https://github.com/huggingface/accelerate#egg=accelerate[test,test_trackers] \
+    git+https://github.com/huggingface/accelerate#egg=accelerate[testing,test_trackers] \
     --extra-index-url https://download.pytorch.org/whl/cu113
 
 # Stage 2
diff --git a/setup.py b/setup.py
index dedd077f3..60b6647fd 100644
--- a/setup.py
+++ b/setup.py
@@ -20,10 +20,10 @@
 extras["docs"] = []
 extras["test_prod"] = ["pytest", "pytest-xdist", "pytest-subtests", "parameterized"]
 extras["test_dev"] = ["datasets<=2.2.2", "evaluate", "transformers", "scipy", "sklearn", "deepspeed"]
-extras["test"] = extras["test_prod"] + extras["test_dev"]
+extras["testing"] = extras["test_prod"] + extras["test_dev"]
 
 extras["test_trackers"] = ["wandb", "comet-ml", "tensorboard"]
-extras["dev"] = extras["quality"] + extras["test"]
+extras["dev"] = extras["quality"] + extras["testing"]
 
 extras["sagemaker"] = [
     "sagemaker",  # boto3 is a required package in sagemaker

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/545
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index 6d3e7d3ff..c846dd65d 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -130,8 +130,16 @@ def __init__(
                 ), "DeepSpeed is not available => install it using `pip3 install deepspeed` or build it from source"
                 self.distributed_type = DistributedType.DEEPSPEED
                 if not torch.distributed.is_initialized():
-                    torch.distributed.init_process_group(backend="nccl", **kwargs)
+                    from .utils import compare_versions
+
                     self.backend = "nccl"
+                    if compare_versions("deepspeed", ">", "0.6.5"):
+                        from deepspeed import comm as dist
+
+                        dist.init_distributed(dist_backend=self.backend)
+                    else:
+                        torch.distributed.init_process_group(backend="nccl", **kwargs)
+
                 self.num_processes = torch.distributed.get_world_size()
                 self.process_index = torch.distributed.get_rank()
                 self.local_process_index = int(os.environ.get("LOCAL_RANK", -1))

```suggestion
    "deepspeed<=0.6.5",
```
I don't think we required 0.6.5 as a min version before, so let's not pin exactly.
@@ -28,7 +28,7 @@
     "scipy",
     "sklearn",
     "parameterized",
-    "deepspeed",
+    "deepspeed==0.6.5",
This can be imported at the top.
@@ -130,8 +130,16 @@ def __init__(
                 ), "DeepSpeed is not available => install it using `pip3 install deepspeed` or build it from source"
                 self.distributed_type = DistributedType.DEEPSPEED
                 if not torch.distributed.is_initialized():
-                    torch.distributed.init_process_group(backend="nccl", **kwargs)
+                    from .utils import compare_versions
that was leading to circular dependency error

@@ -130,8 +130,16 @@ def __init__(
                 ), "DeepSpeed is not available => install it using `pip3 install deepspeed` or build it from source"
                 self.distributed_type = DistributedType.DEEPSPEED
                 if not torch.distributed.is_initialized():
-                    torch.distributed.init_process_group(backend="nccl", **kwargs)
+                    from .utils import compare_versions
Ah yes, I can see that! Good for me then!
@@ -130,8 +130,16 @@ def __init__(
                 ), "DeepSpeed is not available => install it using `pip3 install deepspeed` or build it from source"
                 self.distributed_type = DistributedType.DEEPSPEED
                 if not torch.distributed.is_initialized():
-                    torch.distributed.init_process_group(backend="nccl", **kwargs)
+                    from .utils import compare_versions
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/544
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 3cc1059e1..cc4a67571 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -927,7 +927,7 @@ def wait_for_everyone(self):
         """
         wait_for_everyone()
 
-    def init_trackers(self, project_name: str, config: Optional[dict] = None):
+    def init_trackers(self, project_name: str, config: Optional[dict] = None, init_kwargs: Optional[dict] = {}):
         """
         Initializes a run for all trackers stored in `self.log_with`, potentially with starting configurations
 
@@ -936,6 +936,12 @@ def init_trackers(self, project_name: str, config: Optional[dict] = None):
                 The name of the project. All trackers will save their data based on this
             config (`dict`, *optional*):
                 Optional starting configuration to be logged.
+            init_kwargs (`dict`, *optional*):
+                A nested dictionary of kwargs to be passed to a specific tracker's `__init__` function. Should be
+                formatted like this:
+                ```python
+                {"wandb": {"tags": ["tag_a", "tag_b"]}}
+                ```
         """
         self.trackers = []
         for tracker in self.log_with:
@@ -946,14 +952,16 @@ def init_trackers(self, project_name: str, config: Optional[dict] = None):
                 tracker_init = LOGGER_TYPE_TO_CLASS[str(tracker)]
                 if getattr(tracker_init, "requires_logging_directory"):
                     # We can skip this check since it was done in `__init__`
-                    self.trackers.append(tracker_init(project_name, self.logging_dir))
+                    self.trackers.append(
+                        tracker_init(project_name, self.logging_dir, **init_kwargs.get(str(tracker), {}))
+                    )
                 else:
-                    self.trackers.append(tracker_init(project_name))
+                    self.trackers.append(tracker_init(project_name, **init_kwargs.get(str(tracker), {})))
         if config is not None:
             for tracker in self.trackers:
                 tracker.store_init_configuration(config)
 
-    def log(self, values: dict, step: Optional[int] = None):
+    def log(self, values: dict, step: Optional[int] = None, log_kwargs: Optional[dict] = {}):
         """
         Logs `values` to all stored trackers in `self.trackers`.
 
@@ -962,10 +970,16 @@ def log(self, values: dict, step: Optional[int] = None):
                 Values should be a dictionary-like object containing only types `int`, `float`, or `str`.
             step (`int`, *optional*):
                 The run step. If included, the log will be affiliated with this step.
+            log_kwargs (`dict`, *optional*):
+                A nested dictionary of kwargs to be passed to a specific tracker's `log` function. Should be formatted
+                like this:
+                ```python
+                {"wandb": {"tags": ["tag_a", "tag_b"]}}
+                ```
         """
         if self.is_main_process:
             for tracker in self.trackers:
-                tracker.log(values, step=step)
+                tracker.log(values, step=step, **log_kwargs.get(tracker.name, {}))
 
     def end_training(self):
         """
diff --git a/src/accelerate/tracking.py b/src/accelerate/tracking.py
index 3efff77c9..b89630a4e 100644
--- a/src/accelerate/tracking.py
+++ b/src/accelerate/tracking.py
@@ -52,8 +52,16 @@ def get_available_trackers():
 class GeneralTracker(object, metaclass=ABCMeta):
     """
     A base Tracker class to be used for all logging integration implementations.
+
+    Each function should take in `**kwargs` that will automatically be passed in from a base dictionary provided to
+    [`Accelerator`]
     """
 
+    @abstractproperty
+    def name(self):
+        "String representation of the python class name"
+        pass
+
     @abstractproperty
     def requires_logging_directory(self):
         """
@@ -75,7 +83,7 @@ def store_init_configuration(self, values: dict):
         pass
 
     @abstractmethod
-    def log(self, values: dict, step: Optional[int]):
+    def log(self, values: dict, step: Optional[int], **kwargs):
         """
         Logs `values` to the current run. Base `log` implementations of a tracking API should go in here, along with
         special behavior for the `step parameter.
@@ -105,14 +113,17 @@ class TensorBoardTracker(GeneralTracker):
             The name of the experiment run
         logging_dir (`str`, `os.PathLike`):
             Location for TensorBoard logs to be stored.
+        kwargs:
+            Additional key word arguments passed along to the `tensorboard.SummaryWriter.__init__` method.
     """
 
+    name = "tensorboard"
     requires_logging_directory = True
 
-    def __init__(self, run_name: str, logging_dir: Optional[Union[str, os.PathLike]]):
+    def __init__(self, run_name: str, logging_dir: Optional[Union[str, os.PathLike]], **kwargs):
         self.run_name = run_name
         self.logging_dir = os.path.join(logging_dir, run_name)
-        self.writer = tensorboard.SummaryWriter(self.logging_dir)
+        self.writer = tensorboard.SummaryWriter(self.logging_dir, **kwargs)
         logger.info(f"Initialized TensorBoard project {self.run_name} logging to {self.logging_dir}")
         logger.info(
             "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
@@ -131,7 +142,7 @@ def store_init_configuration(self, values: dict):
         self.writer.flush()
         logger.info("Stored initial configuration hyperparameters to TensorBoard")
 
-    def log(self, values: dict, step: Optional[int] = None):
+    def log(self, values: dict, step: Optional[int] = None, **kwargs):
         """
         Logs `values` to the current run.
 
@@ -141,14 +152,17 @@ def log(self, values: dict, step: Optional[int] = None):
                 `str` to `float`/`int`.
             step (`int`, *optional*):
                 The run step. If included, the log will be affiliated with this step.
+            kwargs:
+                Additional key word arguments passed along to either `SummaryWriter.add_scaler`,
+                `SummaryWriter.add_text`, or `SummaryWriter.add_scalers` method based on the contents of `values`.
         """
         for k, v in values.items():
             if isinstance(v, (int, float)):
-                self.writer.add_scalar(k, v, global_step=step)
+                self.writer.add_scalar(k, v, global_step=step, **kwargs)
             elif isinstance(v, str):
-                self.writer.add_text(k, v, global_step=step)
+                self.writer.add_text(k, v, global_step=step, **kwargs)
             elif isinstance(v, dict):
-                self.writer.add_scalars(k, v, global_step=step)
+                self.writer.add_scalars(k, v, global_step=step, **kwargs)
         self.writer.flush()
         logger.info("Successfully logged to TensorBoard")
 
@@ -167,13 +181,16 @@ class WandBTracker(GeneralTracker):
     Args:
         run_name (`str`):
             The name of the experiment run.
+        kwargs:
+            Additional key word arguments passed along to the `wandb.init` method.
     """
 
+    name = "wandb"
     requires_logging_directory = False
 
-    def __init__(self, run_name: str):
+    def __init__(self, run_name: str, **kwargs):
         self.run_name = run_name
-        self.run = wandb.init(project=self.run_name)
+        self.run = wandb.init(project=self.run_name, **kwargs)
         logger.info(f"Initialized WandB project {self.run_name}")
         logger.info(
             "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
@@ -191,7 +208,7 @@ def store_init_configuration(self, values: dict):
         wandb.config.update(values)
         logger.info("Stored initial configuration hyperparameters to WandB")
 
-    def log(self, values: dict, step: Optional[int] = None):
+    def log(self, values: dict, step: Optional[int] = None, **kwargs):
         """
         Logs `values` to the current run.
 
@@ -201,8 +218,10 @@ def log(self, values: dict, step: Optional[int] = None):
                 `str` to `float`/`int`.
             step (`int`, *optional*):
                 The run step. If included, the log will be affiliated with this step.
+            kwargs:
+                Additional key word arguments passed along to the `wandb.log` method.
         """
-        self.run.log(values, step=step)
+        self.run.log(values, step=step, **kwargs)
         logger.info("Successfully logged to WandB")
 
     def finish(self):
@@ -222,13 +241,16 @@ class CometMLTracker(GeneralTracker):
     Args:
         run_name (`str`):
             The name of the experiment run.
+        kwargs:
+            Additional key word arguments passed along to the `Experiment.__init__` method.
     """
 
+    name = "comet_ml"
     requires_logging_directory = False
 
-    def __init__(self, run_name: str):
+    def __init__(self, run_name: str, **kwargs):
         self.run_name = run_name
-        self.writer = Experiment(project_name=run_name)
+        self.writer = Experiment(project_name=run_name, **kwargs)
         logger.info(f"Initialized CometML project {self.run_name}")
         logger.info(
             "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
@@ -246,7 +268,7 @@ def store_init_configuration(self, values: dict):
         self.writer.log_parameters(values)
         logger.info("Stored initial configuration hyperparameters to CometML")
 
-    def log(self, values: dict, step: Optional[int] = None):
+    def log(self, values: dict, step: Optional[int] = None, **kwargs):
         """
         Logs `values` to the current run.
 
@@ -256,16 +278,19 @@ def log(self, values: dict, step: Optional[int] = None):
                 `str` to `float`/`int`.
             step (`int`, *optional*):
                 The run step. If included, the log will be affiliated with this step.
+            kwargs:
+                Additional key word arguments passed along to either `Experiment.log_metric`, `Experiment.log_other`,
+                or `Experiment.log_metrics` method based on the contents of `values`.
         """
         if step is not None:
             self.writer.set_step(step)
         for k, v in values.items():
             if isinstance(v, (int, float)):
-                self.writer.log_metric(k, v, step=step)
+                self.writer.log_metric(k, v, step=step, **kwargs)
             elif isinstance(v, str):
-                self.writer.log_other(k, v)
+                self.writer.log_other(k, v, **kwargs)
             elif isinstance(v, dict):
-                self.writer.log_metrics(v, step=step)
+                self.writer.log_metrics(v, step=step, **kwargs)
         logger.info("Successfully logged to CometML")
 
     def finish(self):
diff --git a/tests/test_tracking.py b/tests/test_tracking.py
index a37ae7128..c32502796 100644
--- a/tests/test_tracking.py
+++ b/tests/test_tracking.py
@@ -100,7 +100,8 @@ def test_init_trackers(self):
         project_name = "test_project_with_config"
         accelerator = Accelerator(log_with="wandb")
         config = {"num_iterations": 12, "learning_rate": 1e-2, "some_boolean": False, "some_string": "some_value"}
-        accelerator.init_trackers(project_name, config)
+        kwargs = {"wandb": {"tags": ["my_tag"]}}
+        accelerator.init_trackers(project_name, config, kwargs)
         accelerator.end_training()
         # The latest offline log is stored at wandb/latest-run/*.wandb
         for child in Path(f"{self.tmpdir}/wandb/latest-run").glob("*"):
@@ -116,6 +117,7 @@ def test_init_trackers(self):
         self.assertEqual(self.get_value_from_log("learning_rate", cleaned_log), "0.01")
         self.assertEqual(self.get_value_from_log("some_boolean", cleaned_log), "false")
         self.assertEqual(self.get_value_from_log("some_string", cleaned_log), "some_value")
+        self.assertIn("my_tag", cleaned_log)
 
     def test_log(self):
         project_name = "test_project_with_log"
@@ -214,6 +216,7 @@ class MyCustomTracker(GeneralTracker):
         "some_string",
     ]
 
+    name = "my_custom_tracker"
     requires_logging_directory = False
 
     def __init__(self, dir: str):

```suggestion
                formatted like this:
```
@@ -936,6 +936,12 @@ def init_trackers(self, project_name: str, config: Optional[dict] = None):
                 The name of the project. All trackers will save their data based on this
             config (`dict`, *optional*):
                 Optional starting configuration to be logged.
+            init_kwargs (`dict`, *optional*):
+                A nested dictionary of kwargs to be passed to a specific tracker's `__init__` function. Should be
+                formatted as {tracker_name:{kwarg_a:value_a}}, such as:
```suggestion
                like this:
```
@@ -962,10 +970,16 @@ def log(self, values: dict, step: Optional[int] = None):
                 Values should be a dictionary-like object containing only types `int`, `float`, or `str`.
             step (`int`, *optional*):
                 The run step. If included, the log will be affiliated with this step.
+            log_kwargs (`dict`, *optional*):
+                A nested dictionary of kwargs to be passed to a specific tracker's `log` function. Should be formatted
+                as {tracker_name:{kwarg_a:value_a}}, such as:
```suggestion
    Each function should take in `**kwargs` that will automatically be passed in from a base dictionary provided to [`Accelerator`]
```
@@ -52,8 +52,15 @@ def get_available_trackers():
 class GeneralTracker(object, metaclass=ABCMeta):
     """
     A base Tracker class to be used for all logging integration implementations.
+
+    Each function should take in `**kwargs` that will automatically be passed in from a base dictionary in Accelerator
```suggestion
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
```
@@ -8,7 +8,7 @@
 #
 # 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.
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY name, either express or implied.
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/542
diff --git a/docs/source/quicktour.mdx b/docs/source/quicktour.mdx
index 2e9cbe203..4b144ba05 100644
--- a/docs/source/quicktour.mdx
+++ b/docs/source/quicktour.mdx
@@ -120,14 +120,13 @@ validation_dataloader = accelerator.prepare(validation_dataloader)
 
 As for your training dataloader, it will mean that (should you run your script on multiple devices) each device will
 only see part of the evaluation data. This means you will need to group your predictions together. This is very easy to
-do with the [`~Accelerator.gather`] method.
+do with the [`~Accelerator.gather_for_metrics`] method.
 
 ```python
 for inputs, targets in validation_dataloader:
     predictions = model(inputs)
     # Gather all predictions and targets
-    all_predictions = accelerator.gather(predictions)
-    all_targets = accelerator.gather(targets)
+    all_predictions, all_targets = accelerator.gather_for_metrics((predictions, targets))
     # Example of use with a *Datasets.Metric*
     metric.add_batch(all_predictions, all_targets)
 ```
@@ -141,11 +140,17 @@ As for the training dataloader, passing your validation dataloader through
 Any instruction using your training dataloader length (for instance if you need the number of total training steps
 to create a learning rate scheduler) should go after the call to [`~Accelerator.prepare`].
 
+As some data at the end of the dataset may be duplicated so the batch can divide equally to all workers, metrics should be 
+calculated through the [`~Accelerator.gather_for_metrics`] method to automatically remove the duplicated data.
+
+If for some reason you don't wish to have this automatically done, [`~Accelerator.gather`] can be used instead to gather 
+the data across all processes and this can manually be done instead.
+
 </Tip>
 
 <Tip warning={true}>
 
-The [`~Accelerator.gather`] method requires the tensors to be all the same size on each process. If
+The [`~Accelerator.gather`] and [`~Accelerator.gather_for_metrics`] methods require the tensors to be all the same size on each process. If
 you have tensors of different sizes on each process (for instance when dynamically padding to the maximum length in
 a batch), you should use the [`~Accelerator.pad_across_processes`] method to pad you tensor to the
 biggest size across processes.
diff --git a/examples/by_feature/checkpointing.py b/examples/by_feature/checkpointing.py
index 6880f4626..1b21ebed7 100644
--- a/examples/by_feature/checkpointing.py
+++ b/examples/by_feature/checkpointing.py
@@ -238,7 +238,7 @@ def training_function(config, args):
                 outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1)
             # It is slightly faster to call this once, than multiple times
-            predictions, references = accelerator.gather((predictions, batch["labels"]))
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]), eval_dataloader)
             metric.add_batch(
                 predictions=predictions,
                 references=references,
diff --git a/examples/by_feature/cross_validation.py b/examples/by_feature/cross_validation.py
index 95af2a6c9..79a207fa8 100644
--- a/examples/by_feature/cross_validation.py
+++ b/examples/by_feature/cross_validation.py
@@ -207,7 +207,9 @@ def training_function(config, args):
                 with torch.no_grad():
                     outputs = model(**batch)
                 predictions = outputs.logits.argmax(dim=-1)
-                predictions, references = accelerator.gather((predictions, batch["labels"]))
+                predictions, references = accelerator.gather_for_metrics(
+                    (predictions, batch["labels"]), eval_dataloader
+                )
                 metric.add_batch(
                     predictions=predictions,
                     references=references,
@@ -226,7 +228,7 @@ def training_function(config, args):
             with torch.no_grad():
                 outputs = model(**batch)
             predictions = outputs.logits
-            predictions, references = accelerator.gather((predictions, batch["labels"]))
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]), test_dataloader)
             fold_predictions.append(predictions.cpu())
             if i == 0:
                 # We need all of the test predictions
diff --git a/examples/by_feature/fsdp_with_peak_mem_tracking.py b/examples/by_feature/fsdp_with_peak_mem_tracking.py
index 8da906ee9..8e3b9e130 100644
--- a/examples/by_feature/fsdp_with_peak_mem_tracking.py
+++ b/examples/by_feature/fsdp_with_peak_mem_tracking.py
@@ -270,7 +270,6 @@ def collate_fn(examples):
         # context manager to track the peak memory usage during the evaluation
         with TorchTracemalloc() as tracemalloc:
             model.eval()
-            samples_seen = 0
             for step, batch in enumerate(eval_dataloader):
                 # We could avoid this line since we set the accelerator with `device_placement=True`.
                 batch.to(accelerator.device)
@@ -278,15 +277,9 @@ def collate_fn(examples):
                     outputs = model(**batch)
                 predictions = outputs.logits.argmax(dim=-1)
                 # It is slightly faster to call this once, than multiple times
-                predictions, references = accelerator.gather(
-                    (predictions, batch["labels"])
-                )  # If we are in a multiprocess environment, the last batch has duplicates
-                if accelerator.use_distributed:
-                    if step == len(eval_dataloader) - 1:
-                        predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
-                        references = references[: len(eval_dataloader.dataset) - samples_seen]
-                    else:
-                        samples_seen += references.shape[0]
+                predictions, references = accelerator.gather_for_metrics(
+                    (predictions, batch["labels"]), eval_dataloader
+                )
                 metric.add_batch(
                     predictions=predictions,
                     references=references,
diff --git a/examples/by_feature/gradient_accumulation.py b/examples/by_feature/gradient_accumulation.py
index 403587f10..d95f9bef3 100644
--- a/examples/by_feature/gradient_accumulation.py
+++ b/examples/by_feature/gradient_accumulation.py
@@ -171,7 +171,7 @@ def training_function(config, args):
             with torch.no_grad():
                 outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1)
-            predictions, references = accelerator.gather((predictions, batch["labels"]))
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]), eval_dataloader)
             metric.add_batch(
                 predictions=predictions,
                 references=references,
diff --git a/examples/by_feature/memory.py b/examples/by_feature/memory.py
index 65feed51c..b11c80fd0 100644
--- a/examples/by_feature/memory.py
+++ b/examples/by_feature/memory.py
@@ -183,7 +183,9 @@ def inner_training_loop(batch_size):
                 with torch.no_grad():
                     outputs = model(**batch)
                 predictions = outputs.logits.argmax(dim=-1)
-                predictions, references = accelerator.gather((predictions, batch["labels"]))
+                predictions, references = accelerator.gather_for_metrics(
+                    (predictions, batch["labels"]), eval_dataloader
+                )
                 metric.add_batch(
                     predictions=predictions,
                     references=references,
diff --git a/examples/by_feature/multi_process_metrics.py b/examples/by_feature/multi_process_metrics.py
index 741e62b8c..dc47f3db6 100644
--- a/examples/by_feature/multi_process_metrics.py
+++ b/examples/by_feature/multi_process_metrics.py
@@ -187,6 +187,8 @@ def training_function(config, args):
                 else:
                     # Otherwise we add the number of samples seen
                     samples_seen += references.shape[0]
+            # All of this can be avoided if you use `Accelerator.gather_for_metrics` instead of `Accelerator.gather`:
+            # accelerator.gather_for_metrics((predictions, references), eval_dataloader)
             metric.add_batch(
                 predictions=predictions,
                 references=references,
diff --git a/examples/by_feature/tracking.py b/examples/by_feature/tracking.py
index 3278a913a..78d4af422 100644
--- a/examples/by_feature/tracking.py
+++ b/examples/by_feature/tracking.py
@@ -197,7 +197,7 @@ def training_function(config, args):
                 outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1)
             # It is slightly faster to call this once, than multiple times
-            predictions, references = accelerator.gather((predictions, batch["labels"]))
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]), eval_dataloader)
             metric.add_batch(
                 predictions=predictions,
                 references=references,
diff --git a/examples/complete_cv_example.py b/examples/complete_cv_example.py
index 6e414bc9c..8f893b7c4 100644
--- a/examples/complete_cv_example.py
+++ b/examples/complete_cv_example.py
@@ -233,7 +233,6 @@ def training_function(config, args):
                     accelerator.save_state(output_dir)
         model.eval()
         accurate = 0
-        samples_seen = 0
         for step, batch in enumerate(eval_dataloader):
             # We could avoid this line since we set the accelerator with `device_placement=True`.
             batch = {k: v.to(accelerator.device) for k, v in batch.items()}
@@ -241,19 +240,11 @@ def training_function(config, args):
             with torch.no_grad():
                 outputs = model(inputs)
             predictions = outputs.argmax(dim=-1)
-            predictions, references = accelerator.gather((predictions, batch["label"]))
-            if accelerator.use_distributed:
-                if step == len(eval_dataloader) - 1:
-                    predictions = predictions[: len(eval_dataloader) - samples_seen]
-                    references = references[: len(eval_dataloader) - samples_seen]
-                else:
-                    samples_seen += references.shape[0]
-            else:
-                samples_seen += references.shape[0]
-            accurate_preds = predictions == references
+            predictions, labels = accelerator.gather_for_metrics((predictions, batch["label"]), eval_dataloader)
+            accurate_preds = predictions == labels
             accurate += accurate_preds.long().sum()
 
-        eval_metric = accurate.item() / samples_seen
+        eval_metric = accurate.item() / accelerator.gradient_state.samples_seen
         # Use accelerator.print to print only on the main process.
         accelerator.print(f"epoch {epoch}: {100 * eval_metric:.2f}")
         if args.with_tracking:
diff --git a/examples/complete_nlp_example.py b/examples/complete_nlp_example.py
index 6e6aec0ff..572bc9a23 100644
--- a/examples/complete_nlp_example.py
+++ b/examples/complete_nlp_example.py
@@ -211,7 +211,6 @@ def collate_fn(examples):
                     accelerator.save_state(output_dir)
 
         model.eval()
-        samples_seen = 0
         for step, batch in enumerate(eval_dataloader):
             # We could avoid this line since we set the accelerator with `device_placement=True`.
             batch.to(accelerator.device)
@@ -219,15 +218,7 @@ def collate_fn(examples):
                 outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1)
             # It is slightly faster to call this once, than multiple times
-            predictions, references = accelerator.gather(
-                (predictions, batch["labels"])
-            )  # If we are in a multiprocess environment, the last batch has duplicates
-            if accelerator.use_distributed:
-                if step == len(eval_dataloader) - 1:
-                    predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
-                    references = references[: len(eval_dataloader.dataset) - samples_seen]
-                else:
-                    samples_seen += references.shape[0]
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]), eval_dataloader)
             metric.add_batch(
                 predictions=predictions,
                 references=references,
diff --git a/examples/cv_example.py b/examples/cv_example.py
index 483bddac3..aa97ce2d7 100644
--- a/examples/cv_example.py
+++ b/examples/cv_example.py
@@ -173,7 +173,8 @@ def training_function(config, args):
             with torch.no_grad():
                 outputs = model(inputs)
             predictions = outputs.argmax(dim=-1)
-            accurate_preds = accelerator.gather(predictions) == accelerator.gather(batch["label"])
+            predictions, labels = accelerator.gather_for_metrics((predictions, batch["label"]), eval_dataloader)
+            accurate_preds = predictions == labels
             num_elems += accurate_preds.shape[0]
             accurate += accurate_preds.long().sum()
 
diff --git a/examples/nlp_example.py b/examples/nlp_example.py
index 55dedae5e..d64ab3bb7 100644
--- a/examples/nlp_example.py
+++ b/examples/nlp_example.py
@@ -158,7 +158,7 @@ def training_function(config, args):
             with torch.no_grad():
                 outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1)
-            predictions, references = accelerator.gather((predictions, batch["labels"]))
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]), eval_dataloader)
             metric.add_batch(
                 predictions=predictions,
                 references=references,
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 3cc1059e1..6a2aa9703 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -51,6 +51,7 @@
     is_torch_version,
     is_tpu_available,
     pad_across_processes,
+    recursively_apply,
     reduce,
     save,
     wait_for_everyone,
@@ -874,6 +875,35 @@ def gather(self, tensor):
         """
         return gather(tensor)
 
+    def gather_for_metrics(self, tensor, dataloader):
+        """
+        Gathers `tensor` and potentially drops duplicates in the last batch if on a distributed system. Should be used
+        for gathering the inputs and targets for metric calculation.
+
+        Args:
+            tensor (`torch.Tensor`, or a nested tuple/list/dictionary of `torch.Tensor`):
+                The tensors for calculating metrics across all processes.
+            dataloader (`torch.utils.data.DataLoader`):
+                A dataloader prepared with `Accelerator.prepare`
+        """
+        tensor = self.gather(tensor)
+        if self.use_distributed:
+            try:
+                # Then see if we're on the last batch of our eval dataloader
+                if self.gradient_state.end_of_dataloader:
+                    # Last batch needs to be truncated on distributed systems as it contains additional samples
+                    def _adjust_samples(tensor):
+                        return tensor[: dataloader.total_dataset_length - self.gradient_state.samples_seen]
+
+                    return recursively_apply(_adjust_samples, tensor)
+                else:
+                    # Not at the end of the dataloader, no need to adjust the tensors
+                    return tensor
+            except:
+                # Dataset had no length or raised an error
+                return tensor
+        return tensor
+
     def reduce(self, tensor, reduction="sum"):
         """
         Reduce the values in *tensor* across all processes based on *reduction*.
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index fddfe8955..ba3338546 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -47,12 +47,18 @@ class MpDeviceLoaderWrapper(xpl.MpDeviceLoader):
         - **total_batch_size** (`int`) -- Total batch size of the dataloader across all processes.
             Equal to the original batch size when `split_batches=True`; otherwise the original batch size * the total
             number of processes
+
+        - **total_dataset_length** (`int`) -- Total length of the inner dataset across all processes.
         """
 
         @property
         def total_batch_size(self):
             return self._loader.total_batch_size
 
+        @property
+        def total_dataset_length(self):
+            return self._loader.total_dataset_length
+
 
 logger = get_logger(__name__)
 
@@ -311,6 +317,8 @@ class DataLoaderShard(DataLoader):
         - **total_batch_size** (`int`) -- Total batch size of the dataloader across all processes.
             Equal to the original batch size when `split_batches=True`; otherwise the original batch size * the total
             number of processes
+
+        - **total_dataset_length** (`int`) -- Total length of the inner dataset across all processes.
     """
 
     def __init__(self, dataset, device=None, rng_types=None, generator=None, **kwargs):
@@ -324,11 +332,13 @@ def __iter__(self):
         if self.rng_types is not None:
             synchronize_rng_states(self.rng_types, self.generator)
         self.gradient_state._set_end_of_dataloader(False)
+        self.gradient_state._set_samples_seen(0)
         dataloader_iter = super().__iter__()
         # We iterate one batch ahead to check when we are at the end
         try:
             current_batch = next(dataloader_iter)
         except StopIteration:
+            self.gradient_state._iterate_samples_seen(find_batch_size(current_batch))
             yield
         while True:
             try:
@@ -336,10 +346,12 @@ def __iter__(self):
                 if self.device is not None:
                     current_batch = send_to_device(current_batch, self.device)
                 next_batch = next(dataloader_iter)
+                self.gradient_state._iterate_samples_seen(find_batch_size(current_batch))
                 yield current_batch
                 current_batch = next_batch
             except StopIteration:
                 self.gradient_state._set_end_of_dataloader(True)
+                self.gradient_state._iterate_samples_seen(find_batch_size(current_batch))
                 yield current_batch
                 break
 
@@ -351,6 +363,10 @@ def total_batch_size(self):
             else (self.batch_sampler.batch_size * self.batch_sampler.num_processes)
         )
 
+    @property
+    def total_dataset_length(self):
+        return len(self.dataset)
+
 
 class DataLoaderDispatcher(DataLoader):
     """
@@ -370,6 +386,8 @@ class DataLoaderDispatcher(DataLoader):
         - **total_batch_size** (`int`) -- Total batch size of the dataloader across all processes.
             Equal to the original batch size when `split_batches=True`; otherwise the original batch size * the total
             number of processes
+
+        - **total_dataset_length** (`int`) -- Total length of the inner dataset across all processes.
     """
 
     def __init__(self, dataset, split_batches: bool = False, **kwargs):
@@ -435,6 +453,7 @@ def _fetch_batches(self, iterator):
 
     def __iter__(self):
         self.gradient_state._set_end_of_dataloader(False)
+        self.gradient_state._set_samples_seen(0)
         main_iterator = None
         if self.state.process_index == 0:
             # We only iterate through the DataLoader on process 0.
@@ -466,12 +485,15 @@ def __iter__(self):
 
             data_slice = slice(self.state.process_index * batch_size, (self.state.process_index + 1) * batch_size)
             next_batch, next_batch_info, next_skip = self._fetch_batches(main_iterator)
+            batch = slice_tensors(batch, data_slice)
             if not self._stop_iteration:
-                yield slice_tensors(batch, data_slice)
+                self.gradient_state._iterate_samples_seen(batch_size)
+                yield batch
                 batch, batch_info, skip = next_batch, next_batch_info, next_skip
             else:
+                self.gradient_state._iterate_samples_seen(batch_size)
                 self.gradient_state._set_end_of_dataloader(True)
-                yield slice_tensors(batch, data_slice)
+                yield batch
                 break
 
     def __len__(self):
@@ -489,6 +511,10 @@ def total_batch_size(self):
             self.dataset.batch_size if self.split_batches else (self.dataset.batch_size * self.dataset.num_processes)
         )
 
+    @property
+    def total_dataset_length(self):
+        return len(self.dataset)
+
 
 def prepare_data_loader(
     dataloader: DataLoader,
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index 3c06968aa..6d3e7d3ff 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -255,10 +255,15 @@ def __init__(self):
         if not getattr(self, "initialized", False):
             self.sync_gradients = True
             self.end_of_dataloader = False
+            self.samples_seen = 0
         self.initialized = True
 
     def __repr__(self):
-        return f"Sync Gradients: {self.sync_gradients}\n" f"At end of current dataloader: {self.end_of_dataloader}\n"
+        return (
+            f"Sync Gradients: {self.sync_gradients}\n"
+            f"At end of current dataloader: {self.end_of_dataloader}\n"
+            f"Samples seen: {self.samples_seen}"
+        )
 
     def _set_sync_gradients(self, sync_gradients):
         "Private function that sets whether gradients should be synchronized. Users should not have to call this."
@@ -267,3 +272,11 @@ def _set_sync_gradients(self, sync_gradients):
     def _set_end_of_dataloader(self, end_of_dataloader):
         "Private function that sets whether the end of the current dataloader has been reached. Users should not have to call this."
         self.end_of_dataloader = end_of_dataloader
+
+    def _set_samples_seen(self, samples_seen):
+        "Private function that sets the number of samples iterated over. Users should not have to call this."
+        self.samples_seen = samples_seen
+
+    def _iterate_samples_seen(self, iteration: int = 1):
+        "Private function that iterates the number of samples seen by an iteration. Users should not have to call this."
+        self._set_samples_seen(self.samples_seen + iteration)
diff --git a/src/accelerate/test_utils/__init__.py b/src/accelerate/test_utils/__init__.py
index 2d8527be2..d65099602 100644
--- a/src/accelerate/test_utils/__init__.py
+++ b/src/accelerate/test_utils/__init__.py
@@ -15,4 +15,4 @@
 from .training import RegressionDataset, RegressionModel
 
 
-from .scripts import test_script, test_sync  # isort:skip
+from .scripts import test_metrics, test_script, test_sync  # isort:skip
diff --git a/src/accelerate/test_utils/scripts/test_metrics.py b/src/accelerate/test_utils/scripts/test_metrics.py
new file mode 100644
index 000000000..442053213
--- /dev/null
+++ b/src/accelerate/test_utils/scripts/test_metrics.py
@@ -0,0 +1,98 @@
+# 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.
+
+from copy import deepcopy
+
+import torch
+from torch.utils.data import DataLoader
+
+import evaluate
+from accelerate import Accelerator
+from accelerate.test_utils import RegressionDataset, RegressionModel
+from accelerate.utils import set_seed
+
+
+def get_setup(accelerator):
+    "Returns everything needed to perform basic training"
+    set_seed(42)
+    model = RegressionModel()
+    ddp_model = deepcopy(model)
+    dset = RegressionDataset(length=80)
+    dataloader = DataLoader(dset, batch_size=16)
+    model.to(accelerator.device)
+    ddp_model, dataloader = accelerator.prepare(ddp_model, dataloader)
+    return model, ddp_model, dataloader
+
+
+def accuracy(predictions, labels) -> float:
+    """
+    Get the accuracy with respect to the most likely label
+    """
+    return (predictions == labels).float().mean()
+
+
+def test_torch_metrics():
+    accelerator = Accelerator()
+    model, ddp_model, dataloader = get_setup(accelerator)
+    for batch in dataloader:
+        ddp_input, ddp_target = batch.values()
+        # First do single process
+        input, target = accelerator.gather((ddp_input, ddp_target))
+        input, target = input.to(accelerator.device), target.to(accelerator.device)
+        with torch.no_grad():
+            logits = model(input)
+            accuracy_single = accuracy(logits.argmax(dim=-1), target)
+        # Then do multiprocess
+        with torch.no_grad():
+            logits = ddp_model(ddp_input)
+            logits, target = accelerator.gather_for_metrics((logits, ddp_target), dataloader)
+            accuracy_multi = accuracy(logits.argmax(dim=-1), target)
+        assert torch.allclose(accuracy_single, accuracy_multi), "The two accuracies were not the same!"
+
+
+def test_evaluate_metrics():
+    metric = evaluate.load("accuracy")
+    accelerator = Accelerator()
+    model, ddp_model, dataloader = get_setup(accelerator)
+    for batch in dataloader:
+        ddp_input, ddp_target = batch.values()
+        # First do single process
+        input, target = accelerator.gather((ddp_input, ddp_target))
+        input, target = input.to(accelerator.device), target.to(accelerator.device)
+        with torch.no_grad():
+            logits = model(input)
+            accuracy_single = metric.compute(logits, target)
+        # Then do multiprocess
+        with torch.no_grad():
+            logits = ddp_model(ddp_input)
+            logits, target = accelerator.gather_for_metrics((logits, ddp_target), dataloader)
+            accuracy_multi = metric.compute(logits, target)
+        assert torch.allclose(accuracy_single, accuracy_multi), "The two accuracies were not the same!"
+
+
+def main():
+    accelerator = Accelerator()
+    state = accelerator.state
+    if state.local_process_index == 0:
+        print("**Test torch metrics**")
+    test_torch_metrics()
+
+
+def _mp_fn(index):
+    # For xla_spawn (TPUs)
+    main()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/test_metrics.py b/tests/test_metrics.py
new file mode 100644
index 000000000..c744a2fdb
--- /dev/null
+++ b/tests/test_metrics.py
@@ -0,0 +1,55 @@
+# 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 inspect
+import os
+import unittest
+
+import torch
+
+import accelerate
+from accelerate import debug_launcher
+from accelerate.test_utils import (
+    execute_subprocess_async,
+    require_cpu,
+    require_multi_gpu,
+    require_single_gpu,
+    test_metrics,
+)
+from accelerate.utils import get_launch_prefix, patch_environment
+
+
+class MetricTester(unittest.TestCase):
+    def setUp(self):
+        mod_file = inspect.getfile(accelerate.test_utils)
+        self.test_file_path = os.path.sep.join(mod_file.split(os.path.sep)[:-1] + ["scripts", "test_metrics.py"])
+
+    @require_cpu
+    def test_metric_cpu_noop(self):
+        debug_launcher(test_metrics.main, num_processes=1)
+
+    @require_cpu
+    def test_metric_cpu_multi(self):
+        debug_launcher(test_metrics.main)
+
+    @require_single_gpu
+    def test_metric_gpu(self):
+        test_metrics.main()
+
+    @require_multi_gpu
+    def test_metric_gpu_multi(self):
+        print(f"Found {torch.cuda.device_count()} devices.")
+        cmd = get_launch_prefix() + [f"--nproc_per_node={torch.cuda.device_count()}", self.test_file_path]
+        with patch_environment(omp_num_threads=1):
+            execute_subprocess_async(cmd, env=os.environ.copy())

Maybe mention it that if the user doesn't want that behavior for some reason, they can just use `gather`?
@@ -141,11 +140,14 @@ As for the training dataloader, passing your validation dataloader through
 Any instruction using your training dataloader length (for instance if you need the number of total training steps
 to create a learning rate scheduler) should go after the call to [`~Accelerator.prepare`].
 
+As some data at the end of the dataset may be duplicated so the batch can divide equally to all workers, metrics should be 
+calculated through the [`~Accelerator.gather_for_metrics`] method to automatically remove the duplicated data.
_The documentation is not available anymore as the PR was closed or merged._
is this functionality in 0.11.0?
It is not, we suggest not using it for now and doing the check manually as shown in the metric example script as some bugs were discovered: https://github.com/huggingface/accelerate/issues/575
Cool, thank you, excited for this change when it happens
https://api.github.com/repos/huggingface/accelerate/pulls/540
diff --git a/src/accelerate/big_modeling.py b/src/accelerate/big_modeling.py
index 4014e3972..85b1f3237 100644
--- a/src/accelerate/big_modeling.py
+++ b/src/accelerate/big_modeling.py
@@ -65,7 +65,9 @@ def init_empty_weights(include_buffers: bool = False):
     def register_empty_parameter(module, name, param):
         old_register_parameter(module, name, param)
         if param is not None:
-            module._parameters[name] = nn.Parameter(module._parameters[name].to(torch.device("meta")))
+            param_cls = type(module._parameters[name])
+            kwargs = module._parameters[name].__dict__
+            module._parameters[name] = param_cls(module._parameters[name].to(torch.device("meta")), **kwargs)
 
     def register_empty_buffer(module, name, buffer):
         old_register_buffer(module, name, buffer)
diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
index 62c4f9761..2dc85ef70 100644
--- a/src/accelerate/utils/modeling.py
+++ b/src/accelerate/utils/modeling.py
@@ -124,11 +124,16 @@ def set_module_tensor_to_device(
             new_value = value.to(device)
         else:
             new_value = torch.tensor(value, device=device)
-    if is_buffer:
-        module._buffers[tensor_name] = new_value
-    else:
-        new_value = nn.Parameter(new_value, requires_grad=old_value.requires_grad)
-        module._parameters[tensor_name] = new_value
+
+        if is_buffer:
+            module._buffers[tensor_name] = new_value
+        elif module._parameters[tensor_name].device.type != "cuda":
+            param_cls = type(module._parameters[tensor_name])
+            kwargs = module._parameters[tensor_name].__dict__
+            new_value = param_cls(new_value, requires_grad=old_value.requires_grad, **kwargs).to(device)
+            module._parameters[tensor_name] = new_value
+        else:
+            module._parameters[tensor_name] = module._parameters[tensor_name].to(device)
 
 
 def named_module_tensors(module: nn.Module, include_buffers: bool = True, recurse: bool = False):

This should stay at the previous level of indentation. To make sure there is absolutely zero breaking change, how about doing an `else`  with just `module._parameters[tensor_name].to(device)` ?
@@ -124,11 +124,14 @@ def set_module_tensor_to_device(
             new_value = value.to(device)
         else:
             new_value = torch.tensor(value, device=device)
-    if is_buffer:
-        module._buffers[tensor_name] = new_value
-    else:
-        new_value = nn.Parameter(new_value, requires_grad=old_value.requires_grad)
-        module._parameters[tensor_name] = new_value
+
+        if is_buffer:
+            module._buffers[tensor_name] = new_value
+        elif module._parameters[tensor_name].device.type != "cuda":
+            param_cls = type(module._parameters[tensor_name])
+            kwargs = module._parameters[tensor_name].__dict__
+            new_value = param_cls(new_value, requires_grad=old_value.requires_grad, **kwargs).to(device)
+            module._parameters[tensor_name] = new_value
Just added it and seems to work fine! 🚀 
Can we wait a bit before merging the PR?
@@ -124,11 +124,14 @@ def set_module_tensor_to_device(
             new_value = value.to(device)
         else:
             new_value = torch.tensor(value, device=device)
-    if is_buffer:
-        module._buffers[tensor_name] = new_value
-    else:
-        new_value = nn.Parameter(new_value, requires_grad=old_value.requires_grad)
-        module._parameters[tensor_name] = new_value
+
+        if is_buffer:
+            module._buffers[tensor_name] = new_value
+        elif module._parameters[tensor_name].device.type != "cuda":
+            param_cls = type(module._parameters[tensor_name])
+            kwargs = module._parameters[tensor_name].__dict__
+            new_value = param_cls(new_value, requires_grad=old_value.requires_grad, **kwargs).to(device)
+            module._parameters[tensor_name] = new_value
Sure thing, let me know when it's good for you. Also don't forget to run `make style` on your branch.
@@ -124,11 +124,14 @@ def set_module_tensor_to_device(
             new_value = value.to(device)
         else:
             new_value = torch.tensor(value, device=device)
-    if is_buffer:
-        module._buffers[tensor_name] = new_value
-    else:
-        new_value = nn.Parameter(new_value, requires_grad=old_value.requires_grad)
-        module._parameters[tensor_name] = new_value
+
+        if is_buffer:
+            module._buffers[tensor_name] = new_value
+        elif module._parameters[tensor_name].device.type != "cuda":
+            param_cls = type(module._parameters[tensor_name])
+            kwargs = module._parameters[tensor_name].__dict__
+            new_value = param_cls(new_value, requires_grad=old_value.requires_grad, **kwargs).to(device)
+            module._parameters[tensor_name] = new_value
Great thanks a lot!!
@@ -124,11 +124,14 @@ def set_module_tensor_to_device(
             new_value = value.to(device)
         else:
             new_value = torch.tensor(value, device=device)
-    if is_buffer:
-        module._buffers[tensor_name] = new_value
-    else:
-        new_value = nn.Parameter(new_value, requires_grad=old_value.requires_grad)
-        module._parameters[tensor_name] = new_value
+
+        if is_buffer:
+            module._buffers[tensor_name] = new_value
+        elif module._parameters[tensor_name].device.type != "cuda":
+            param_cls = type(module._parameters[tensor_name])
+            kwargs = module._parameters[tensor_name].__dict__
+            new_value = param_cls(new_value, requires_grad=old_value.requires_grad, **kwargs).to(device)
+            module._parameters[tensor_name] = new_value
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/539
diff --git a/docs/source/big_modeling.mdx b/docs/source/big_modeling.mdx
index 051d1bda7..006c8bc4e 100644
--- a/docs/source/big_modeling.mdx
+++ b/docs/source/big_modeling.mdx
@@ -205,7 +205,7 @@ This only supports inference of your model, not training. Most of the computatio
 
 ## Designing a device map
 
-You can let 🤗 Accelerate handle the device map computation by setting `device_map="auto"` or create one yourself, if you want more control over where each layer should go.
+You can let 🤗 Accelerate handle the device map computation by setting `device_map` to one of the supported options (`"auto"`, `"balanced"`, `"balanced_low_0"`, `"sequential"`) or create one yourself, if you want more control over where each layer should go.
 
 <Tip>
 
@@ -213,6 +213,19 @@ You can derive all sizes of the model (and thus compute a `device_map`) on a mod
 
 </Tip>
 
+All the options will produce the same result when you don't have enough GPU memory to accomodate the whole model (which is to fit everything that can on the GPU, then offload weights on the CPU or even on the disk if there is not enough RAM). 
+
+When you have more GPU memory available than the model size, here the difference between each option:
+- `"auto"` and `"balanced"` evenly split the model on all available GPUs, making it possible for you to use a batch size greater than 1.
+- `"balanced_low_0"` evenly splits the model on all GPUs except the first one, and only puts on GPU 0 what does not fit on the others. This option is great when you need to use GPU 0 for some processing of the outputs, like when using the `generate` function for Transformers models
+- `"sequential"` will fit what it can on GPU 0, then move on GPU 1 and so forth (so won't use the last GPUs if it doesn't need to).
+
+<Tip>
+
+The options `"auto"` and `"balanced"` produce the same results for now, but the behavior of `"auto"` might change in the future if we find a strategy that makes more sense, while `"balanced"` will stay stable.
+
+</Tip>
+
 First note that you can limit the memory used on each GPU by using the `max_memory` argument (available in [`infer_auto_device_map`] and in all functions using it). When setting `max_memory`, you should pass along a dictionary containing the GPU identifiers (for instance `0`, `1` etc.) and the `"cpu"` key for the maximum RAM you want used for CPU offload. The values can either be an integer (in bytes) or a string representing a number with its unit, such as `"10GiB"` or `"10GB"`.
 
 Here is an example where we don't want to use more than 10GiB on each of two GPUs and no more than 30GiB of CPU RAM for the model weights:
diff --git a/src/accelerate/big_modeling.py b/src/accelerate/big_modeling.py
index 907dfdac4..311111ad4 100644
--- a/src/accelerate/big_modeling.py
+++ b/src/accelerate/big_modeling.py
@@ -24,6 +24,7 @@
     OffloadedWeightsLoader,
     check_device_map,
     extract_submodules_state_dict,
+    get_balanced_memory,
     infer_auto_device_map,
     load_checkpoint_in_model,
     offload_state_dict,
@@ -278,7 +279,8 @@ def load_checkpoint_and_dispatch(
             A map that specifies where each submodule should go. It doesn't need to be refined to each parameter/buffer
             name, once a given module name is inside, every submodule of it will be sent to the same device.
 
-            To have Accelerate compute the most optimized `device_map` automatically, set `device_map="auto"`.
+            To have Accelerate compute the most optimized `device_map` automatically, set `device_map="auto"`. For more
+            information about each option see [here](big_modeling#designing-a-device-map).
         max_memory (`Dict`, *optional*):
             A dictionary device identifier to maximum memory. Will default to the maximum memory available for each GPU
             and the available CPU RAM if unset.
@@ -302,7 +304,20 @@ def load_checkpoint_and_dispatch(
             called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
             `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
     """
-    if device_map == "auto":
+    if isinstance(device_map, str) and device_map not in ["auto", "balanced", "balanced_low_0", "sequential"]:
+        raise ValueError(
+            "If passing a string for `device_map`, please choose 'auto', 'balanced', 'balanced_low_0' or "
+            "'sequential'."
+        )
+    if device_map != "sequential":
+        max_memory = get_balanced_memory(
+            model,
+            max_memory=max_memory,
+            no_split_module_classes=no_split_module_classes,
+            dtype=dtype,
+            low_zero=(device_map == "balanced_low_0"),
+        )
+    if isinstance(device_map, str):
         device_map = infer_auto_device_map(
             model, max_memory=max_memory, no_split_module_classes=no_split_module_classes, dtype=dtype
         )
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
index 350efc899..f6b83d2f9 100644
--- a/src/accelerate/utils/__init__.py
+++ b/src/accelerate/utils/__init__.py
@@ -37,6 +37,7 @@
     convert_file_size_to_int,
     dtype_byte_size,
     find_tied_parameters,
+    get_balanced_memory,
     get_max_layer_size,
     get_max_memory,
     infer_auto_device_map,
diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
index 62c4f9761..e412ce0a0 100644
--- a/src/accelerate/utils/modeling.py
+++ b/src/accelerate/utils/modeling.py
@@ -326,6 +326,101 @@ def load_offloaded_weights(model, index, offload_folder):
         set_module_tensor_to_device(model, param_name, "cpu", value=torch.tensor(weight))
 
 
+def get_balanced_memory(
+    model: nn.Module,
+    max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None,
+    no_split_module_classes: Optional[List[str]] = None,
+    dtype: Optional[Union[str, torch.dtype]] = None,
+    low_zero: bool = False,
+):
+    """
+    Compute a `max_memory` dictionary for [`infer_auto_device_map`] that will balance the use of each available GPU.
+
+    <Tip>
+
+    All computation is done analyzing sizes and dtypes of the model parameters. As a result, the model can be on the
+    meta device (as it would if initialized within the `init_empty_weights` context manager).
+
+    </Tip>
+
+    Args:
+        model (`torch.nn.Module`): The model to analyze.
+        max_memory (`Dict`, *optional*):
+            A dictionary device identifier to maximum memory. Will default to the maximum memory available if unset.
+        no_split_module_classes (`List[str]`, *optional*):
+            A list of layer class names that should never be split across device (for instance any layer that has a
+            residual connection).
+        dtype (`str` or `torch.dtype`, *optional*):
+            If provided, the weights will be converted to that type when loaded.
+        low_zero (`bool`, *optional*):
+            Minimizes the number of weights on GPU 0, which is convenient when it's used for other operations (like the
+            Transformers generate function).
+    """
+    # Get default / clean up max_memory
+    max_memory = get_max_memory(max_memory)
+
+    if not torch.cuda.is_available():
+        return max_memory
+
+    num_devices = len([d for d in max_memory if torch.device(d).type == "cuda"])
+    module_sizes = compute_module_sizes(model, dtype=dtype)
+    per_gpu = module_sizes[""] // (num_devices - 1 if low_zero else num_devices)
+
+    # We can't just set the memory to model_size // num_devices as it will end being too small: each GPU will get
+    # slightly less layers and some layers will end up offload at the end. So this function computes a buffer size to
+    # add which is the biggest of:
+    # - the size of no split block (if applicable)
+    # - the mean of the layer sizes
+    if no_split_module_classes is None:
+        no_split_module_classes = []
+    elif not isinstance(no_split_module_classes, (list, tuple)):
+        no_split_module_classes = [no_split_module_classes]
+
+    # Identify the size of the no_split_block modules
+    if len(no_split_module_classes) > 0:
+        no_split_children = {}
+        for name, size in module_sizes.items():
+            if name == "":
+                continue
+            submodule = model
+            for submodule_name in name.split("."):
+                submodule = getattr(submodule, submodule_name)
+            class_name = submodule.__class__.__name__
+            if class_name in no_split_module_classes and class_name not in no_split_children:
+                no_split_children[class_name] = size
+
+            if set(no_split_children.keys()) == set(no_split_module_classes):
+                break
+        buffer = max(no_split_children.values()) if len(no_split_children) > 0 else 0
+    else:
+        buffer = 0
+
+    # Compute mean of final modules. In the first dict of module sizes, leaves are the parameters
+    leaves = [n for n in module_sizes if len([p for p in module_sizes if p.startswith(n) and len(p) > len(n)]) == 0]
+    module_sizes = {n: v for n, v in module_sizes.items() if n not in leaves}
+    # Once removed, leaves are the final modules.
+    leaves = [n for n in module_sizes if len([p for p in module_sizes if p.startswith(n) and len(p) > len(n)]) == 0]
+    mean_leaves = int(sum([module_sizes[n] for n in leaves]) / len(leaves))
+    buffer = int(1.25 * max(buffer, mean_leaves))
+    if low_zero:
+        per_gpu += buffer
+        gpu_zero = 0
+    else:
+        gpu_zero = per_gpu
+        per_gpu += buffer
+
+    max_memory = get_max_memory(max_memory)
+    for i in range(num_devices):
+        # We still leave slightly more space on GPU 0 and only apply the buffer on the other devices.
+        max_memory[i] = min(gpu_zero if i == 0 else per_gpu, max_memory[i])
+
+    if low_zero:
+        min_zero = max(0, module_sizes[""] - sum([max_memory[i] for i in range(1, num_devices)]))
+        max_memory[0] = min(min_zero, max_memory[0])
+
+    return max_memory
+
+
 def infer_auto_device_map(
     model: nn.Module,
     max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None,
diff --git a/tests/test_modeling_utils.py b/tests/test_modeling_utils.py
index 15e164218..799170257 100644
--- a/tests/test_modeling_utils.py
+++ b/tests/test_modeling_utils.py
@@ -26,6 +26,7 @@
     clean_device_map,
     compute_module_sizes,
     find_tied_parameters,
+    get_balanced_memory,
     infer_auto_device_map,
     load_checkpoint_in_model,
     named_module_tensors,
@@ -358,3 +359,13 @@ def test_infer_auto_device_map(self):
         device_map = infer_auto_device_map(model, max_memory={0: 400, 1: 500})
         expected = {"0": 0, "2.linear2": 0, "1": 1, "2.linear1": 1, "2.batchnorm": 1}
         self.assertDictEqual(device_map, expected)
+
+    @require_cuda
+    def test_get_balanced_memory(self):
+        model = ModelForTest()
+        # model has size 236: linear1 64, batchnorm 72, linear2 100
+        max_memory = get_balanced_memory(model, max_memory={0: 200, 1: 200})
+        self.assertDictEqual({0: 118, 1: 200}, max_memory)
+
+        max_memory = get_balanced_memory(model, max_memory={0: 300, 1: 300})
+        self.assertDictEqual({0: 118, 1: 215}, max_memory)

Docstring format nit
@@ -326,6 +326,95 @@ def load_offloaded_weights(model, index, offload_folder):
         set_module_tensor_to_device(model, param_name, "cpu", value=torch.tensor(weight))
 
 
+def get_balanced_memory(
+    model: nn.Module,
+    max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None,
+    no_split_module_classes: Optional[List[str]] = None,
+    dtype: Optional[Union[str, torch.dtype]] = None,
+    low_zero: bool = False,
+):
+    """
+    Compute a `max_memory` dictionary for [`infer_auto_device_map`] that will balance the use of each available GPU.
+
+    <Tip>
+
+    All computation is done analyzing sizes and dtypes of the model parameters. As a result, the model can be on the
+    meta device (as it would if initialized within the `init_empty_weights` context manager).
+
+    </Tip>
+
+    Args:
+        model (`torch.nn.Module`): The model to analyze.
I don't see any nit ;-) 
@@ -326,6 +326,95 @@ def load_offloaded_weights(model, index, offload_folder):
         set_module_tensor_to_device(model, param_name, "cpu", value=torch.tensor(weight))
 
 
+def get_balanced_memory(
+    model: nn.Module,
+    max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None,
+    no_split_module_classes: Optional[List[str]] = None,
+    dtype: Optional[Union[str, torch.dtype]] = None,
+    low_zero: bool = False,
+):
+    """
+    Compute a `max_memory` dictionary for [`infer_auto_device_map`] that will balance the use of each available GPU.
+
+    <Tip>
+
+    All computation is done analyzing sizes and dtypes of the model parameters. As a result, the model can be on the
+    meta device (as it would if initialized within the `init_empty_weights` context manager).
+
+    </Tip>
+
+    Args:
+        model (`torch.nn.Module`): The model to analyze.
If it resonates I would rename it to something more self-documenting? e.g. `minimize_gpu0_memory` or `minimize_first_gpu_memory` 
@@ -326,6 +326,95 @@ def load_offloaded_weights(model, index, offload_folder):
         set_module_tensor_to_device(model, param_name, "cpu", value=torch.tensor(weight))
 
 
+def get_balanced_memory(
+    model: nn.Module,
+    max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None,
+    no_split_module_classes: Optional[List[str]] = None,
+    dtype: Optional[Union[str, torch.dtype]] = None,
+    low_zero: bool = False,
```suggestion
All the options will produce the same result when you don't have enough GPU memory to accomodate the whole model (which is to fit everything that can on the GPU, then offload weights on the CPU or even on the disk if there is not enough RAM). 

When you have more GPU memory available than the model size, here the difference between each option:
```
Let's add a break here
@@ -205,14 +205,25 @@ This only supports inference of your model, not training. Most of the computatio
 
 ## Designing a device map
 
-You can let 🤗 Accelerate handle the device map computation by setting `device_map="auto"` or create one yourself, if you want more control over where each layer should go.
+You can let 🤗 Accelerate handle the device map computation by setting `device_map` to one of the supported options (`"auto"`, `"balanced"`, `"balanced_low_0"`, `"sequential"`) or create one yourself, if you want more control over where each layer should go.
 
 <Tip>
 
 You can derive all sizes of the model (and thus compute a `device_map`) on a model that is on the meta device.
 
 </Tip>
 
+All the options will produce the same result when you don't have enough GPU memory to accomodate the whole model (which is to fit everything that can on the GPU, then offload weights on the CPU or even on the disk if there is not enough RAM). When you have more GPU memory available than the model size, here the difference between each option:
_The documentation is not available anymore as the PR was closed or merged._
Yes I will update it once the other PR is merged :-) 
> This leads to equal allocation across all GPUs. Ideally it should be reworked to leave the gpu0 as close to empty as possible. So probably first trying to spread the weights across all but the first gpu, while leaving enough memory for activation calculation and temps. And only then assign any remaining weights to the first gpu.

That is exactly what the option `"balanced_low_0"` will do.
https://api.github.com/repos/huggingface/accelerate/pulls/534
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 09330eb63..806f0094d 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -16,6 +16,7 @@ jobs:
           test_core,
           test_big_modeling,
           test_deepspeed,
+          test_fsdp,
           test_example_differences,
           test_checkpoint_step,
           test_checkpoint_epoch,
diff --git a/Makefile b/Makefile
index 4104f1bba..9cae9d74d 100644
--- a/Makefile
+++ b/Makefile
@@ -31,11 +31,15 @@ test_big_modeling:
 	python -m pytest -s -v ./tests/test_big_modeling.py
 
 test_core:
-	python -m pytest -s -v ./tests/ --ignore=./tests/test_examples.py --ignore=./tests/deepspeed --ignore=./tests/test_big_modeling.py
+	python -m pytest -s -v ./tests/ --ignore=./tests/test_examples.py --ignore=./tests/deepspeed --ignore=./tests/test_big_modeling.py \
+	--ignore=./tests/fsdp
 
 test_deepspeed:
 	python -m pytest -s -v ./tests/deepspeed
 
+test_fsdp:
+	python -m pytest -s -v ./tests/fsdp
+
 test_examples:
 	python -m pytest -s -v ./tests/test_examples.py
 
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 35afd6c5d..c961f5be0 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -31,6 +31,7 @@
 from .state import AcceleratorState, GradientState
 from .tracking import LOGGER_TYPE_TO_CLASS, GeneralTracker, filter_trackers
 from .utils import (
+    MODEL_NAME,
     DeepSpeedPlugin,
     DistributedDataParallelKwargs,
     DistributedType,
@@ -520,7 +521,8 @@ def prepare(self, *args):
             if model_count > 1 and optimizer_present:
                 raise ValueError(
                     "For FSDP to work with multiple models (>1), "
-                    "prepare must be called for all the models before optimizers are created"
+                    "prepare must be called for all the models before optimizers are created. "
+                    "Then pass the optimizers to the prepare call in the same order as corresponding models."
                 )
             elif model_count == 1 and optimizer_present:
                 logger.warn(
@@ -596,6 +598,7 @@ def prepare_model(self, model):
                 )
                 if not fsdp_plugin.cpu_offload.offload_params:
                     model.to(self.device)
+            self._models[-1] = model
         elif self.distributed_type == DistributedType.MULTI_CPU:
             kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {}
             model = torch.nn.parallel.DistributedDataParallel(model, **kwargs)
@@ -1042,9 +1045,44 @@ def save_state(self, output_dir: str):
         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, unwrap=False) for m in self._models]
+
+        # Save the models taking care of FSDP and DeepSpeed nuances
+        weights = []
+        for i, model in enumerate(self._models):
+            if self.distributed_type == DistributedType.FSDP:
+                logger.info("Saving FSDP model")
+                self.state.fsdp_plugin.save_model(self, model, output_dir, i)
+                logger.info(f"FSDP Model saved to output dir {output_dir}")
+            elif self.distributed_type == DistributedType.DEEPSPEED:
+                logger.info("Saving DeepSpeed Model and Optimizer")
+                ckpt_id = f"{MODEL_NAME}" if i == 0 else f"{MODEL_NAME}_{i}"
+                model.save_checkpoint(output_dir, ckpt_id)
+                logger.info(f"DeepSpeed Model and Optimizer saved to output dir {os.path.join(output_dir, ckpt_id)}")
+            else:
+                weights.append(self.get_state_dict(model, unwrap=False))
+
+        # Save the optimizers taking care of FSDP and DeepSpeed nuances
+        optimizers = []
+        if self.distributed_type == DistributedType.FSDP:
+            for opt in self._optimizers:
+                logger.info("Saving FSDP Optimizer")
+                self.state.fsdp_plugin.save_optimizer(self, opt, self._models[i], output_dir, i)
+                logger.info(f"FSDP Optimizer saved to output dir {output_dir}")
+        elif self.distributed_type != DistributedType.DEEPSPEED:
+            optimizers = self._optimizers
+
+        # Save the lr schedulers taking care of DeepSpeed nuances
+        schedulers = []
+        if self.distributed_type == DistributedType.DEEPSPEED:
+            for i, scheduler in enumerate(self._schedulers):
+                if isinstance(scheduler, DeepSpeedSchedulerWrapper):
+                    continue
+                schedulers.append(scheduler)
+        else:
+            schedulers = self._schedulers
+
         save_location = save_accelerator_state(
-            output_dir, weights, self._optimizers, self._schedulers, self.state.process_index, self.scaler
+            output_dir, weights, optimizers, schedulers, self.state.process_index, self.scaler
         )
         for i, obj in enumerate(self._custom_objects):
             save_custom_state(obj, output_dir, i)
@@ -1063,9 +1101,43 @@ def load_state(self, input_dir: str):
         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._schedulers, self.state.process_index, self.scaler
-        )
+
+        # Load the models taking care of FSDP and DeepSpeed nuances
+        models = []
+        for i, model in enumerate(self._models):
+            if self.distributed_type == DistributedType.FSDP:
+                logger.info("Loading FSDP model")
+                self.state.fsdp_plugin.load_model(self, model, input_dir, i)
+                logger.info(f"FSDP Model loaded from input dir {input_dir}")
+            elif self.distributed_type == DistributedType.DEEPSPEED:
+                logger.info("Loading DeepSpeed Model and Optimizer")
+                ckpt_id = f"{MODEL_NAME}" if i == 0 else f"{MODEL_NAME}_{i}"
+                model.load_checkpoint(input_dir, ckpt_id)
+                logger.info(f"DeepSpeed Model and Optimizer loaded from input dir {os.path.join(input_dir, ckpt_id)}")
+            else:
+                models.append(model)
+
+        # Load the optimizers taking care of FSDP and DeepSpeed nuances
+        optimizers = []
+        if self.distributed_type == DistributedType.FSDP:
+            for i, opt in enumerate(self._optimizers):
+                logger.info("Loading FSDP Optimizer")
+                self.state.fsdp_plugin.load_optimizer(self, opt, self._models[i], input_dir, i)
+                logger.info(f"FSDP Optimizer loaded from input dir {input_dir}")
+        elif self.distributed_type != DistributedType.DEEPSPEED:
+            optimizers = self._optimizers
+
+        # Load the lr schedulers taking care of DeepSpeed nuances
+        schedulers = []
+        if self.distributed_type == DistributedType.DEEPSPEED:
+            for i, scheduler in enumerate(self._schedulers):
+                if isinstance(scheduler, DeepSpeedSchedulerWrapper):
+                    continue
+                schedulers.append(scheduler)
+        else:
+            schedulers = self._schedulers
+
+        load_accelerator_state(input_dir, models, optimizers, schedulers, self.state.process_index, self.scaler)
         custom_checkpoints = [f for f in os.listdir(input_dir) if "custom_checkpoint" in f]
         if len(custom_checkpoints) != len(self._custom_objects):
             err = "Warning! Number of found checkpoints does not match the number of registered objects:"
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index fd345fa1c..ba677e3eb 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -20,6 +20,7 @@
     FSDP_AUTO_WRAP_POLICY,
     FSDP_BACKWARD_PREFETCH,
     FSDP_SHARDING_STRATEGY,
+    FSDP_STATE_DICT_TYPE,
 )
 from .config_args import ClusterConfig
 from .config_utils import _ask_field, _convert_distributed_mode, _convert_yes_no_to_bool
@@ -210,12 +211,12 @@ def get_cluster_input():
             for i, strategy in enumerate(FSDP_SHARDING_STRATEGY):
                 sharding_strategy_query += f"[{i+1}] {strategy}, "
             sharding_strategy_query = sharding_strategy_query[:-2] + ")? [1]: "
-            fsdp_config["sharding_strategy"] = _ask_field(
+            fsdp_config["fsdp_sharding_strategy"] = _ask_field(
                 sharding_strategy_query,
                 lambda x: int(x),
                 default=1,
             )
-            fsdp_config["offload_params"] = _ask_field(
+            fsdp_config["fsdp_offload_params"] = _ask_field(
                 "Do you want to offload parameters and gradients to CPU? [yes/NO]: ",
                 _convert_yes_no_to_bool,
                 default=False,
@@ -228,15 +229,15 @@ def get_cluster_input():
             fsdp_config["fsdp_auto_wrap_policy"] = _ask_field(
                 fsdp_wrap_query,
                 lambda x: FSDP_AUTO_WRAP_POLICY[int(x)],
-                default=FSDP_AUTO_WRAP_POLICY[0],
+                default="TRANSFORMER_BASED_WRAP",
             )
             if fsdp_config["fsdp_auto_wrap_policy"] == FSDP_AUTO_WRAP_POLICY[0]:
-                fsdp_config["transformer_layer_cls_to_wrap"] = _ask_field(
+                fsdp_config["fsdp_transformer_layer_cls_to_wrap"] = _ask_field(
                     "What is the transformer layer class name (case-sensitive) to wrap ,e.g, `BertLayer`, `GPTJBlock`, `T5Block` ...? : ",
                     lambda x: str(x),
                 )
             elif fsdp_config["fsdp_auto_wrap_policy"] == FSDP_AUTO_WRAP_POLICY[1]:
-                fsdp_config["min_num_params"] = _ask_field(
+                fsdp_config["fsdp_min_num_params"] = _ask_field(
                     "What should be your FSDP's minimum number of parameters for Default Auto Wrapping Policy? [1e8]: ",
                     lambda x: int(x),
                     default=1e8,
@@ -248,7 +249,16 @@ def get_cluster_input():
             fsdp_config["fsdp_backward_prefetch_policy"] = _ask_field(
                 fsdp_backward_prefetch_query,
                 lambda x: FSDP_BACKWARD_PREFETCH[int(x)],
-                default=FSDP_BACKWARD_PREFETCH[0],
+                default="BACKWARD_PRE",
+            )
+            fsdp_state_dict_type_query = "What should be your FSDP's state dict type ("
+            for i, state_dict_type in enumerate(FSDP_STATE_DICT_TYPE):
+                fsdp_state_dict_type_query += f"[{i}] {state_dict_type}, "
+            fsdp_state_dict_type_query = fsdp_state_dict_type_query[:-2] + ")? [0]: "
+            fsdp_config["fsdp_state_dict_type"] = _ask_field(
+                fsdp_state_dict_type_query,
+                lambda x: FSDP_STATE_DICT_TYPE[int(x)],
+                default="FULL_STATE_DICT",
             )
 
     if distributed_type == DistributedType.TPU:
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index e6a23e46d..514e77d05 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -148,19 +148,19 @@ def launch_command_parser(subparsers=None):
         help="Whether to use fsdp.",
     )
     parser.add_argument(
-        "--offload_params",
+        "--fsdp_offload_params",
         default="false",
         type=str,
         help="Decides Whether (true|false) to offload parameters and gradients to CPU. (useful only when `use_fsdp` flag is passed).",
     )
     parser.add_argument(
-        "--min_num_params",
+        "--fsdp_min_num_params",
         type=int,
         default=1e8,
         help="FSDP's minimum number of parameters for Default Auto Wrapping. (useful only when `use_fsdp` flag is passed).",
     )
     parser.add_argument(
-        "--sharding_strategy",
+        "--fsdp_sharding_strategy",
         type=int,
         default=1,
         help="FSDP's Sharding Strategy. (useful only when `use_fsdp` flag is passed).",
@@ -172,7 +172,7 @@ def launch_command_parser(subparsers=None):
         help="FSDP's auto wrap policy. (useful only when `use_fsdp` flag is passed).",
     )
     parser.add_argument(
-        "--transformer_layer_cls_to_wrap",
+        "--fsdp_transformer_layer_cls_to_wrap",
         default=None,
         type=str,
         help="Transformer layer class name (case-sensitive) to wrap ,e.g, `BertLayer`, `GPTJBlock`, `T5Block` .... "
@@ -184,6 +184,36 @@ def launch_command_parser(subparsers=None):
         type=str,
         help="FSDP's backward prefetch policy. (useful only when `use_fsdp` flag is passed).",
     )
+    parser.add_argument(
+        "--fsdp_state_dict_type",
+        default=None,
+        type=str,
+        help="FSDP's state dict type. (useful only when `use_fsdp` flag is passed).",
+    )
+    parser.add_argument(
+        "--offload_params",
+        default=None,
+        type=str,
+        help="This argument is deprecated. Use `fsdp_offload_params` instead.",
+    )
+    parser.add_argument(
+        "--min_num_params",
+        type=int,
+        default=None,
+        help="This argument is deprecated. Use `fsdp_min_num_params` instead.",
+    )
+    parser.add_argument(
+        "--sharding_strategy",
+        type=int,
+        default=None,
+        help="This argument is deprecated. Use `fsdp_sharding_strategy` instead.",
+    )
+    parser.add_argument(
+        "--transformer_layer_cls_to_wrap",
+        default=None,
+        type=str,
+        help="This argument is deprecated. Use `fsdp_transformer_layer_cls_to_wrap` instead.",
+    )
     parser.add_argument(
         "--tpu", default=False, action="store_true", help="Whether or not this should launch a TPU training."
     )
@@ -360,13 +390,51 @@ def multi_gpu_launcher(args):
 
     current_env["MIXED_PRECISION"] = str(mixed_precision)
     if args.use_fsdp:
+        if args.sharding_strategy is not None:
+            warnings.warn(
+                "`sharding_strategy` is deprecated and will be removed in version 0.13.0 of 🤗 Accelerate. Use"
+                " `fsdp_sharding_strategy` instead",
+                FutureWarning,
+            )
+            args.fsdp_sharding_strategy = args.sharding_strategy
+
+        if args.offload_params is not None:
+            warnings.warn(
+                "`offload_params` is deprecated and will be removed in version 0.13.0 of 🤗 Accelerate. Use"
+                " `fsdp_offload_params` instead",
+                FutureWarning,
+            )
+            args.fsdp_offload_params = args.offload_params
+
+        if args.min_num_params is not None:
+            warnings.warn(
+                "`min_num_params` is deprecated and will be removed in version 0.13.0 of 🤗 Accelerate. Use"
+                " `fsdp_min_num_params` instead",
+                FutureWarning,
+            )
+            args.fsdp_min_num_params = args.min_num_params
+
+        if args.transformer_layer_cls_to_wrap is not None:
+            warnings.warn(
+                "`transformer_layer_cls_to_wrap` is deprecated and will be removed in version 0.13.0 of 🤗 Accelerate. Use"
+                " `fsdp_transformer_layer_cls_to_wrap` instead",
+                FutureWarning,
+            )
+            args.fsdp_transformer_layer_cls_to_wrap = args.transformer_layer_cls_to_wrap
+
         current_env["USE_FSDP"] = "true"
-        current_env["FSDP_AUTO_WRAP_POLICY"] = str(args.fsdp_auto_wrap_policy)
-        current_env["FSDP_TRANSFORMER_CLS_TO_WRAP"] = str(args.transformer_layer_cls_to_wrap)
-        current_env["FSDP_OFFLOAD_PARAMS"] = str(args.offload_params).lower()
-        current_env["FSDP_MIN_NUM_PARAMS"] = str(args.min_num_params)
-        current_env["FSDP_SHARDING_STRATEGY"] = str(args.sharding_strategy)
-        current_env["FSDP_BACKWARD_PREFETCH"] = str(args.fsdp_backward_prefetch_policy)
+        current_env["FSDP_SHARDING_STRATEGY"] = str(args.fsdp_sharding_strategy)
+        current_env["FSDP_OFFLOAD_PARAMS"] = str(args.fsdp_offload_params).lower()
+        current_env["FSDP_MIN_NUM_PARAMS"] = str(args.fsdp_min_num_params)
+        if args.fsdp_auto_wrap_policy is not None:
+            current_env["FSDP_AUTO_WRAP_POLICY"] = str(args.fsdp_auto_wrap_policy)
+        if args.fsdp_transformer_layer_cls_to_wrap is not None:
+            current_env["FSDP_TRANSFORMER_CLS_TO_WRAP"] = str(args.fsdp_transformer_layer_cls_to_wrap)
+        if args.fsdp_backward_prefetch_policy is not None:
+            current_env["FSDP_BACKWARD_PREFETCH"] = str(args.fsdp_backward_prefetch_policy)
+        if args.fsdp_state_dict_type is not None:
+            current_env["FSDP_STATE_DICT_TYPE"] = str(args.fsdp_state_dict_type)
+
     current_env["OMP_NUM_THREADS"] = str(args.num_cpu_threads_per_process)
     process = subprocess.Popen(cmd, env=current_env)
     process.wait()
@@ -682,7 +750,10 @@ def launch_command(args):
                         if getattr(args, k) is None:
                             setattr(args, k, defaults.deepspeed_config[k])
                     for k in defaults.fsdp_config:
-                        setattr(args, k, defaults.fsdp_config[k])
+                        arg_to_set = k
+                        if "fsdp" not in arg_to_set:
+                            arg_to_set = "fsdp_" + arg_to_set
+                        setattr(args, arg_to_set, defaults.fsdp_config[k])
                     continue
 
                 # Those args are handled separately
diff --git a/src/accelerate/test_utils/scripts/test_checkpointing.py b/src/accelerate/test_utils/scripts/test_checkpointing.py
new file mode 100644
index 000000000..cde602dfa
--- /dev/null
+++ b/src/accelerate/test_utils/scripts/test_checkpointing.py
@@ -0,0 +1,269 @@
+# coding=utf-8
+# Copyright 2022 The HuggingFace Inc. 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 json
+import os
+
+import torch
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+
+import evaluate
+from accelerate import Accelerator, DistributedType
+from accelerate.utils.deepspeed import DummyOptim, DummyScheduler
+from datasets import load_dataset
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
+
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+def get_dataloaders(accelerator: Accelerator, batch_size: int = 16, model_name: str = "bert-base-cased"):
+    """
+    Creates a set of `DataLoader`s for the `glue` dataset.
+
+    Args:
+        accelerator (`Accelerator`):
+            An `Accelerator` object
+        batch_size (`int`, *optional*):
+            The batch size for the train and validation DataLoaders.
+        model_name (`str`, *optional*):
+    """
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    datasets = load_dataset("glue", "mrpc")
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function, batched=True, remove_columns=["idx", "sentence1", "sentence2"], load_from_cache_file=False
+    )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    return train_dataloader, eval_dataloader
+
+
+def evaluation_loop(accelerator, model, eval_dataloader, metric):
+    model.eval()
+    samples_seen = 0
+    for step, batch in enumerate(eval_dataloader):
+        # We could avoid this line since we set the accelerator with `device_placement=True`.
+        batch.to(accelerator.device)
+        with torch.no_grad():
+            outputs = model(**batch)
+        predictions = outputs.logits.argmax(dim=-1)
+        # It is slightly faster to call this once, than multiple times
+        predictions, references = accelerator.gather(
+            (predictions, batch["labels"])
+        )  # If we are in a multiprocess environment, the last batch has duplicates
+        if accelerator.use_distributed:
+            if step == len(eval_dataloader) - 1:
+                predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
+                references = references[: len(eval_dataloader.dataset) - samples_seen]
+            else:
+                samples_seen += references.shape[0]
+        metric.add_batch(
+            predictions=predictions,
+            references=references,
+        )
+
+    eval_metric = metric.compute()
+    return eval_metric["accuracy"]
+
+
+def training_function(config, args):
+    # Initialize accelerator
+    accelerator = Accelerator()
+
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+    model_name = args.model_name_or_path
+
+    set_seed(seed)
+    train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size, model_name)
+
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained(model_name, return_dict=True)
+
+    # Instantiate optimizer
+    optimizer_cls = (
+        AdamW
+        if accelerator.state.deepspeed_plugin is None
+        or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config
+        else DummyOptim
+    )
+    optimizer = optimizer_cls(params=model.parameters(), lr=lr)
+
+    if accelerator.state.deepspeed_plugin is not None:
+        gradient_accumulation_steps = accelerator.state.deepspeed_plugin.deepspeed_config[
+            "gradient_accumulation_steps"
+        ]
+    else:
+        gradient_accumulation_steps = 1
+    max_training_steps = (len(train_dataloader) * num_epochs) // gradient_accumulation_steps
+
+    # Instantiate scheduler
+    if (
+        accelerator.state.deepspeed_plugin is None
+        or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
+    ):
+        lr_scheduler = get_linear_schedule_with_warmup(
+            optimizer=optimizer,
+            num_warmup_steps=0,
+            num_training_steps=max_training_steps,
+        )
+    else:
+        lr_scheduler = DummyScheduler(optimizer, total_num_steps=max_training_steps, warmup_num_steps=0)
+
+    # Prepare everything
+    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
+    # prepare method.
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    # We need to keep track of how many total steps we have iterated over
+    overall_step = 0
+    # We also need to keep track of the stating epoch so files are named properly
+    starting_epoch = 0
+    metric = evaluate.load("glue", "mrpc")
+    ending_epoch = num_epochs
+
+    if args.partial_train_epoch is not None:
+        ending_epoch = args.partial_train_epoch
+
+    if args.resume_from_checkpoint:
+        accelerator.load_state(args.resume_from_checkpoint)
+        epoch_string = args.resume_from_checkpoint.split("epoch_")[1]
+        state_epoch_num = ""
+        for char in epoch_string:
+            if char.isdigit():
+                state_epoch_num += char
+            else:
+                break
+        starting_epoch = int(state_epoch_num) + 1
+        accuracy = evaluation_loop(accelerator, model, eval_dataloader, metric)
+        accelerator.print("resumed checkpoint performance:", accuracy)
+        accelerator.print("resumed checkpoint's scheduler's lr:", lr_scheduler.get_lr()[0])
+        accelerator.print("resumed optimizers's lr:", optimizer.param_groups[0]["lr"])
+        with open(os.path.join(args.output_dir, f"state_{starting_epoch-1}.json"), "r") as f:
+            resumed_state = json.load(f)
+            assert resumed_state["accuracy"] == accuracy, "Accuracy mismatch, loading from checkpoint failed"
+            assert (
+                resumed_state["lr"] == lr_scheduler.get_lr()[0]
+            ), "Scheduler learning rate mismatch, loading from checkpoint failed"
+            assert (
+                resumed_state["optimizer_lr"] == optimizer.param_groups[0]["lr"]
+            ), "Optimizer learning rate mismatch, loading from checkpoint failed"
+            assert resumed_state["epoch"] == starting_epoch - 1, "Epoch mismatch, loading from checkpoint failed"
+            return
+
+    # Now we train the model
+    state = {}
+    for epoch in range(starting_epoch, ending_epoch):
+        model.train()
+        for step, batch in enumerate(train_dataloader):
+            outputs = model(**batch)
+            loss = outputs.loss
+            loss = loss / gradient_accumulation_steps
+            accelerator.backward(loss)
+            if step % gradient_accumulation_steps == 0:
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+
+            overall_step += 1
+        output_dir = f"epoch_{epoch}"
+        output_dir = os.path.join(args.output_dir, output_dir)
+        accelerator.save_state(output_dir)
+        accuracy = evaluation_loop(accelerator, model, eval_dataloader, metric)
+        state["accuracy"] = accuracy
+        state["lr"] = lr_scheduler.get_lr()[0]
+        state["optimizer_lr"] = optimizer.param_groups[0]["lr"]
+        state["epoch"] = epoch
+        state["step"] = overall_step
+        accelerator.print(f"epoch {epoch}:", state)
+
+        accelerator.wait_for_everyone()
+        if accelerator.is_main_process:
+            with open(os.path.join(args.output_dir, f"state_{epoch}.json"), "w") as f:
+                json.dump(state, f)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Simple example of training script tracking peak GPU memory usage.")
+    parser.add_argument(
+        "--model_name_or_path",
+        type=str,
+        default="bert-base-cased",
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+        required=False,
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=".",
+        help="Optional save directory where all checkpoint folders will be stored. Default is the current working directory.",
+    )
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        default=None,
+        help="If the training should continue from a checkpoint folder.",
+    )
+    parser.add_argument(
+        "--partial_train_epoch",
+        type=int,
+        default=None,
+        help="If passed, the training will stop after this number of epochs.",
+    )
+    parser.add_argument(
+        "--num_epochs",
+        type=int,
+        default=2,
+        help="Number of train epochs.",
+    )
+    args = parser.parse_args()
+    config = {"lr": 2e-5, "num_epochs": args.num_epochs, "seed": 42, "batch_size": 16}
+
+    training_function(config, args)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/accelerate/test_utils/scripts/test_peak_memory_usage.py b/src/accelerate/test_utils/scripts/test_peak_memory_usage.py
new file mode 100644
index 000000000..7bb5ca3bf
--- /dev/null
+++ b/src/accelerate/test_utils/scripts/test_peak_memory_usage.py
@@ -0,0 +1,258 @@
+# coding=utf-8
+# Copyright 2022 The HuggingFace Inc. 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 gc
+import json
+import os
+
+import torch
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+
+from accelerate import Accelerator, DistributedType
+from accelerate.utils.deepspeed import DummyOptim, DummyScheduler
+from datasets import load_dataset
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
+
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+# Converting Bytes to Megabytes
+def b2mb(x):
+    return int(x / 2**20)
+
+
+# This context manager is used to track the peak memory usage of the process
+class TorchTracemalloc:
+    def __enter__(self):
+        gc.collect()
+        torch.cuda.empty_cache()
+        torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
+        self.begin = torch.cuda.memory_allocated()
+        return self
+
+    def __exit__(self, *exc):
+        gc.collect()
+        torch.cuda.empty_cache()
+        self.end = torch.cuda.memory_allocated()
+        self.peak = torch.cuda.max_memory_allocated()
+        self.used = b2mb(self.end - self.begin)
+        self.peaked = b2mb(self.peak - self.begin)
+        # print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")
+
+
+def get_dataloaders(
+    accelerator: Accelerator,
+    batch_size: int = 16,
+    model_name: str = "bert-base-cased",
+    n_train: int = 320,
+    n_val: int = 160,
+):
+    """
+    Creates a set of `DataLoader`s for the `glue` dataset.
+
+    Args:
+        accelerator (`Accelerator`):
+            An `Accelerator` object
+        batch_size (`int`, *optional*):
+            The batch size for the train and validation DataLoaders.
+        model_name (`str`, *optional*):
+            The name of the model to use.
+        n_train (`int`, *optional*):
+            The number of training examples to use.
+        n_val (`int`, *optional*):
+            The number of validation examples to use.
+    """
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    datasets = load_dataset(
+        "glue", "mrpc", split={"train": f"train[:{n_train}]", "validation": f"validation[:{n_val}]"}
+    )
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function, batched=True, remove_columns=["idx", "sentence1", "sentence2"], load_from_cache_file=False
+    )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    return train_dataloader, eval_dataloader
+
+
+def training_function(config, args):
+    # Initialize accelerator
+    accelerator = Accelerator()
+
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+    model_name = args.model_name_or_path
+
+    set_seed(seed)
+    train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size, model_name, args.n_train, args.n_val)
+
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained(model_name, return_dict=True)
+
+    # Instantiate optimizer
+    optimizer_cls = (
+        AdamW
+        if accelerator.state.deepspeed_plugin is None
+        or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config
+        else DummyOptim
+    )
+    optimizer = optimizer_cls(params=model.parameters(), lr=lr)
+
+    if accelerator.state.deepspeed_plugin is not None:
+        gradient_accumulation_steps = accelerator.state.deepspeed_plugin.deepspeed_config[
+            "gradient_accumulation_steps"
+        ]
+    else:
+        gradient_accumulation_steps = 1
+    max_training_steps = (len(train_dataloader) * num_epochs) // gradient_accumulation_steps
+
+    # Instantiate scheduler
+    if (
+        accelerator.state.deepspeed_plugin is None
+        or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
+    ):
+        lr_scheduler = get_linear_schedule_with_warmup(
+            optimizer=optimizer,
+            num_warmup_steps=0,
+            num_training_steps=max_training_steps,
+        )
+    else:
+        lr_scheduler = DummyScheduler(optimizer, total_num_steps=max_training_steps, warmup_num_steps=0)
+
+    # Prepare everything
+    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
+    # prepare method.
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    # We need to keep track of how many total steps we have iterated over
+    overall_step = 0
+    # We also need to keep track of the stating epoch so files are named properly
+    starting_epoch = 0
+
+    # Now we train the model
+    train_total_peak_memory = {}
+    for epoch in range(starting_epoch, num_epochs):
+        with TorchTracemalloc() as tracemalloc:
+            model.train()
+            for step, batch in enumerate(train_dataloader):
+                outputs = model(**batch)
+                loss = outputs.loss
+                loss = loss / gradient_accumulation_steps
+                accelerator.backward(loss)
+                if step % gradient_accumulation_steps == 0:
+                    optimizer.step()
+                    lr_scheduler.step()
+                    optimizer.zero_grad()
+
+                overall_step += 1
+
+        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
+        accelerator.print("Memory before entering the train : {}".format(b2mb(tracemalloc.begin)))
+        accelerator.print("Memory consumed at the end of the train (end-begin): {}".format(tracemalloc.used))
+        accelerator.print("Peak Memory consumed during the train (max-begin): {}".format(tracemalloc.peaked))
+        accelerator.print(
+            "Total Peak Memory consumed during the train (max): {}".format(
+                tracemalloc.peaked + b2mb(tracemalloc.begin)
+            )
+        )
+        train_total_peak_memory[f"epoch-{epoch}"] = tracemalloc.peaked + b2mb(tracemalloc.begin)
+        if args.peak_memory_upper_bound is not None:
+            assert (
+                train_total_peak_memory[f"epoch-{epoch}"] <= args.peak_memory_upper_bound
+            ), "Peak memory usage exceeded the upper bound"
+
+    accelerator.wait_for_everyone()
+    if accelerator.is_main_process:
+        with open(os.path.join(args.output_dir, "peak_memory_utilization.json"), "w") as f:
+            json.dump(train_total_peak_memory, f)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Simple example of training script tracking peak GPU memory usage.")
+    parser.add_argument(
+        "--model_name_or_path",
+        type=str,
+        default="bert-base-cased",
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+        required=False,
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=".",
+        help="Optional save directory where all checkpoint folders will be stored. Default is the current working directory.",
+    )
+    parser.add_argument(
+        "--peak_memory_upper_bound",
+        type=float,
+        default=None,
+        help="The upper bound of peak memory usage in MB. If set, the training will throw an error if the peak memory usage exceeds this value.",
+    )
+    parser.add_argument(
+        "--n_train",
+        type=int,
+        default=320,
+        help="Number of training examples to use.",
+    )
+    parser.add_argument(
+        "--n_val",
+        type=int,
+        default=160,
+        help="Number of validation examples to use.",
+    )
+    parser.add_argument(
+        "--num_epochs",
+        type=int,
+        default=1,
+        help="Number of train epochs.",
+    )
+    args = parser.parse_args()
+    config = {"lr": 2e-5, "num_epochs": args.num_epochs, "seed": 42, "batch_size": 16}
+    training_function(config, args)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/accelerate/test_utils/scripts/test_performance.py b/src/accelerate/test_utils/scripts/test_performance.py
new file mode 100644
index 000000000..324a1854e
--- /dev/null
+++ b/src/accelerate/test_utils/scripts/test_performance.py
@@ -0,0 +1,231 @@
+# coding=utf-8
+# Copyright 2022 The HuggingFace Inc. 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 json
+import os
+
+import torch
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+
+import evaluate
+from accelerate import Accelerator, DistributedType
+from accelerate.utils.deepspeed import DummyOptim, DummyScheduler
+from datasets import load_dataset
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
+
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+def get_dataloaders(accelerator: Accelerator, batch_size: int = 16, model_name: str = "bert-base-cased"):
+    """
+    Creates a set of `DataLoader`s for the `glue` dataset.
+
+    Args:
+        accelerator (`Accelerator`):
+            An `Accelerator` object
+        batch_size (`int`, *optional*):
+            The batch size for the train and validation DataLoaders.
+        model_name (`str`, *optional*):
+    """
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    datasets = load_dataset("glue", "mrpc")
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function, batched=True, remove_columns=["idx", "sentence1", "sentence2"], load_from_cache_file=False
+    )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    return train_dataloader, eval_dataloader
+
+
+def training_function(config, args):
+    # Initialize accelerator
+    accelerator = Accelerator()
+
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+    model_name = args.model_name_or_path
+
+    set_seed(seed)
+    train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size, model_name)
+
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained(model_name, return_dict=True)
+
+    # Instantiate optimizer
+    optimizer_cls = (
+        AdamW
+        if accelerator.state.deepspeed_plugin is None
+        or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config
+        else DummyOptim
+    )
+    optimizer = optimizer_cls(params=model.parameters(), lr=lr)
+
+    if accelerator.state.deepspeed_plugin is not None:
+        gradient_accumulation_steps = accelerator.state.deepspeed_plugin.deepspeed_config[
+            "gradient_accumulation_steps"
+        ]
+    else:
+        gradient_accumulation_steps = 1
+    max_training_steps = (len(train_dataloader) * num_epochs) // gradient_accumulation_steps
+
+    # Instantiate scheduler
+    if (
+        accelerator.state.deepspeed_plugin is None
+        or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
+    ):
+        lr_scheduler = get_linear_schedule_with_warmup(
+            optimizer=optimizer,
+            num_warmup_steps=0,
+            num_training_steps=max_training_steps,
+        )
+    else:
+        lr_scheduler = DummyScheduler(optimizer, total_num_steps=max_training_steps, warmup_num_steps=0)
+
+    # Prepare everything
+    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
+    # prepare method.
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    # We need to keep track of how many total steps we have iterated over
+    overall_step = 0
+    # We also need to keep track of the stating epoch so files are named properly
+    starting_epoch = 0
+
+    # Now we train the model
+    metric = evaluate.load("glue", "mrpc")
+    best_performance = 0
+    performance_metric = {}
+    for epoch in range(starting_epoch, num_epochs):
+        model.train()
+        for step, batch in enumerate(train_dataloader):
+            outputs = model(**batch)
+            loss = outputs.loss
+            loss = loss / gradient_accumulation_steps
+            accelerator.backward(loss)
+            if step % gradient_accumulation_steps == 0:
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+
+            overall_step += 1
+
+        model.eval()
+        samples_seen = 0
+        for step, batch in enumerate(eval_dataloader):
+            # We could avoid this line since we set the accelerator with `device_placement=True`.
+            batch.to(accelerator.device)
+            with torch.no_grad():
+                outputs = model(**batch)
+            predictions = outputs.logits.argmax(dim=-1)
+            # It is slightly faster to call this once, than multiple times
+            predictions, references = accelerator.gather(
+                (predictions, batch["labels"])
+            )  # If we are in a multiprocess environment, the last batch has duplicates
+            if accelerator.use_distributed:
+                if step == len(eval_dataloader) - 1:
+                    predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
+                    references = references[: len(eval_dataloader.dataset) - samples_seen]
+                else:
+                    samples_seen += references.shape[0]
+            metric.add_batch(
+                predictions=predictions,
+                references=references,
+            )
+
+        eval_metric = metric.compute()
+        # Use accelerator.print to print only on the main process.
+        accelerator.print(f"epoch {epoch}:", eval_metric)
+        performance_metric[f"epoch-{epoch}"] = eval_metric["accuracy"]
+
+        if best_performance < eval_metric["accuracy"]:
+            best_performance = eval_metric["accuracy"]
+
+    if args.performance_lower_bound is not None:
+        assert (
+            args.performance_lower_bound <= best_performance
+        ), f"Best performance metric {best_performance} is lower than the lower bound {args.performance_lower_bound}"
+
+    accelerator.wait_for_everyone()
+    if accelerator.is_main_process:
+        with open(os.path.join(args.output_dir, "all_results.json"), "w") as f:
+            json.dump(performance_metric, f)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Simple example of training script tracking peak GPU memory usage.")
+    parser.add_argument(
+        "--model_name_or_path",
+        type=str,
+        default="bert-base-cased",
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+        required=False,
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=".",
+        help="Optional save directory where all checkpoint folders will be stored. Default is the current working directory.",
+    )
+    parser.add_argument(
+        "--performance_lower_bound",
+        type=float,
+        default=None,
+        help="Optional lower bound for the performance metric. If set, the training will throw error when the performance metric drops below this value.",
+    )
+    parser.add_argument(
+        "--num_epochs",
+        type=int,
+        default=3,
+        help="Number of train epochs.",
+    )
+    args = parser.parse_args()
+    config = {"lr": 2e-5, "num_epochs": args.num_epochs, "seed": 42, "batch_size": 16}
+    training_function(config, args)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/accelerate/test_utils/testing.py b/src/accelerate/test_utils/testing.py
index 0b3145070..b5cccc1ff 100644
--- a/src/accelerate/test_utils/testing.py
+++ b/src/accelerate/test_utils/testing.py
@@ -32,6 +32,7 @@
     is_comet_ml_available,
     is_deepspeed_available,
     is_tensorboard_available,
+    is_torch_version,
     is_tpu_available,
     is_wandb_available,
 )
@@ -108,6 +109,13 @@ def require_deepspeed(test_case):
     return unittest.skipUnless(is_deepspeed_available(), "test requires DeepSpeed")(test_case)
 
 
+def require_fsdp(test_case):
+    """
+    Decorator marking a test that requires FSDP installed. These tests are skipped when FSDP isn't installed
+    """
+    return unittest.skipUnless(is_torch_version(">=", "1.12.0"), "test requires torch version >= 1.12.0")(test_case)
+
+
 def require_tensorboard(test_case):
     """
     Decorator marking a test that requires tensorboard installed. These tests are skipped when tensorboard isn't
diff --git a/src/accelerate/utils/constants.py b/src/accelerate/utils/constants.py
index 00556cad6..4e7c71853 100644
--- a/src/accelerate/utils/constants.py
+++ b/src/accelerate/utils/constants.py
@@ -27,6 +27,7 @@
 FSDP_SHARDING_STRATEGY = ["FULL_SHARD", "SHARD_GRAD_OP", "NO_SHARD"]
 FSDP_AUTO_WRAP_POLICY = ["TRANSFORMER_BASED_WRAP", "SIZE_BASED_WRAP", "NO_WRAP"]
 FSDP_BACKWARD_PREFETCH = ["BACKWARD_PRE", "BACKWARD_POST", "NO_PREFETCH"]
+FSDP_STATE_DICT_TYPE = ["FULL_STATE_DICT", "LOCAL_STATE_DICT", "SHARDED_STATE_DICT"]
 DEEPSPEED_MULTINODE_LAUNCHERS = ["pdsh", "standard", "openmpi", "mvapich"]
 
 STR_OPERATION_TO_FUNC = {">": op.gt, ">=": op.ge, "==": op.eq, "!=": op.ne, "<=": op.le, "<": op.lt}
diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py
index 8c8bdf106..350d95320 100644
--- a/src/accelerate/utils/dataclasses.py
+++ b/src/accelerate/utils/dataclasses.py
@@ -28,7 +28,7 @@
 
 import torch
 
-from .constants import FSDP_AUTO_WRAP_POLICY, FSDP_BACKWARD_PREFETCH
+from .constants import FSDP_AUTO_WRAP_POLICY, FSDP_BACKWARD_PREFETCH, FSDP_STATE_DICT_TYPE, MODEL_NAME, OPTIMIZER_NAME
 
 
 class KwargsHandler:
@@ -455,8 +455,28 @@ class FullyShardedDataParallelPlugin:
         metadata={"help": "A list of modules to ignore for FSDP."},
     )
 
+    state_dict_type: "typing.Any" = field(
+        default=None,
+        metadata={
+            "help": "FSDP State Dict Type of type `torch.distributed.fsdp.fully_sharded_data_parallel.StateDictType`"
+        },
+    )
+
+    state_dict_config: "typing.Any" = field(
+        default=None,
+        metadata={
+            "help": "FSDP State Dict Config of type `torch.distributed.fsdp.fully_sharded_data_parallel.StateDictConfig`"
+        },
+    )
+
     def __post_init__(self):
-        from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch, CPUOffload, ShardingStrategy
+        from torch.distributed.fsdp.fully_sharded_data_parallel import (
+            BackwardPrefetch,
+            CPUOffload,
+            ShardingStrategy,
+            StateDictType,
+            _state_dict_type_to_config,
+        )
 
         if self.sharding_strategy is None:
             self.sharding_strategy = ShardingStrategy(int(os.environ.get("FSDP_SHARDING_STRATEGY", 1)))
@@ -468,10 +488,21 @@ def __post_init__(self):
                 self.cpu_offload = CPUOffload(offload_params=False)
 
         if self.backward_prefetch is None:
-            prefetch_policy = os.environ.get("FSDP_BACKWARD_PREFETCH", FSDP_BACKWARD_PREFETCH[-1])
+            prefetch_policy = os.environ.get("FSDP_BACKWARD_PREFETCH", "NO_PREFETCH")
             if prefetch_policy != FSDP_BACKWARD_PREFETCH[-1]:
                 self.backward_prefetch = BackwardPrefetch(FSDP_BACKWARD_PREFETCH.index(prefetch_policy) + 1)
 
+        if self.state_dict_type is None:
+            state_dict_type_policy = os.environ.get("FSDP_STATE_DICT_TYPE", "FULL_STATE_DICT")
+            self.state_dict_type = StateDictType(FSDP_STATE_DICT_TYPE.index(state_dict_type_policy) + 1)
+
+            if self.state_dict_type == StateDictType.FULL_STATE_DICT:
+                self.state_dict_config = _state_dict_type_to_config[self.state_dict_type](
+                    offload_to_cpu=True, rank0_only=True
+                )
+            else:
+                self.state_dict_config = _state_dict_type_to_config[self.state_dict_type]()
+
     @staticmethod
     def get_module_class_from_name(module, name):
         """
@@ -496,7 +527,7 @@ def set_auto_wrap_policy(self, model):
         from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy, transformer_auto_wrap_policy
 
         if self.auto_wrap_policy is None:
-            auto_wrap_policy = os.environ.get("FSDP_AUTO_WRAP_POLICY", FSDP_AUTO_WRAP_POLICY[-1])
+            auto_wrap_policy = os.environ.get("FSDP_AUTO_WRAP_POLICY", "NO_WRAP")
             if auto_wrap_policy == FSDP_AUTO_WRAP_POLICY[0]:
                 transformer_cls_to_wrap = os.environ.get("FSDP_TRANSFORMER_CLS_TO_WRAP", "")
                 transformer_cls_to_wrap = FullyShardedDataParallelPlugin.get_module_class_from_name(
@@ -527,3 +558,84 @@ def set_mixed_precision(self, mixed_precision):
 
         if self.mixed_precision_policy is None:
             self.mixed_precision_policy = MixedPrecision(param_dtype=dtype, reduce_dtype=dtype, buffer_dtype=dtype)
+
+    def save_model(self, accelerator, model, output_dir, model_index=0):
+        from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP
+        from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType
+
+        if self.state_dict_type == StateDictType.FULL_STATE_DICT:
+            with FSDP.state_dict_type(model, self.state_dict_type, self.state_dict_config):
+                state_dict = model.state_dict()
+            weights_name = f"{MODEL_NAME}.bin" if model_index == 0 else f"{MODEL_NAME}_{model_index}.bin"
+            output_model_file = os.path.join(output_dir, weights_name)
+            if accelerator.process_index == 0:
+                print(f"Saving model to {output_model_file}")
+                torch.save(state_dict, output_model_file)
+                print(f"Model saved to {output_model_file}")
+        else:
+            with FSDP.state_dict_type(model, self.state_dict_type, self.state_dict_config):
+                state_dict = model.state_dict()
+            weights_name = (
+                f"{MODEL_NAME}_rank{accelerator.process_index}.bin"
+                if model_index == 0
+                else f"{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin"
+            )
+            output_model_file = os.path.join(output_dir, weights_name)
+            print(f"Saving model to {output_model_file}")
+            torch.save(state_dict, output_model_file)
+            print(f"Model saved to {output_model_file}")
+
+    def load_model(self, accelerator, model, input_dir, model_index=0):
+        from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP
+        from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType
+
+        accelerator.wait_for_everyone()
+
+        if self.state_dict_type == StateDictType.FULL_STATE_DICT:
+            weights_name = f"{MODEL_NAME}.bin" if model_index == 0 else f"{MODEL_NAME}_{model_index}.bin"
+            input_model_file = os.path.join(input_dir, weights_name)
+            accelerator.print(f"Loading model from {input_model_file}")
+            state_dict = torch.load(input_model_file)
+            accelerator.print(f"Model loaded from {input_model_file}")
+        else:
+            weights_name = (
+                f"{MODEL_NAME}_rank{accelerator.process_index}.bin"
+                if model_index == 0
+                else f"{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin"
+            )
+            input_model_file = os.path.join(input_dir, weights_name)
+            print(f"Loading model from {input_model_file}")
+            state_dict = torch.load(input_model_file)
+            print(f"Model loaded from {input_model_file}")
+        with FSDP.state_dict_type(model, self.state_dict_type, self.state_dict_config):
+            model.load_state_dict(state_dict)
+
+    def save_optimizer(self, accelerator, optimizer, model, output_dir, optimizer_index=0, optim_input=None):
+        from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP
+
+        optim_state = FSDP.full_optim_state_dict(model, optimizer, optim_input=optim_input)
+        if accelerator.process_index == 0:
+            optim_state_name = (
+                f"{OPTIMIZER_NAME}.bin" if optimizer_index == 0 else f"{OPTIMIZER_NAME}_{optimizer_index}.bin"
+            )
+            output_optimizer_file = os.path.join(output_dir, optim_state_name)
+            print(f"Saving Optimizer state to {output_optimizer_file}")
+            torch.save(optim_state, output_optimizer_file)
+            print(f"Optimizer state saved in {output_optimizer_file}")
+
+    def load_optimizer(self, accelerator, optimizer, model, input_dir, optimizer_index=0):
+        from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP
+
+        accelerator.wait_for_everyone()
+        full_osd = None
+        if accelerator.process_index == 0:
+            optimizer_name = (
+                f"{OPTIMIZER_NAME}.bin" if optimizer_index == 0 else f"{OPTIMIZER_NAME}_{optimizer_index}.bin"
+            )
+            input_optimizer_file = os.path.join(input_dir, optimizer_name)
+            print(f"Loading Optimizer state from {input_optimizer_file}")
+            full_osd = torch.load(input_optimizer_file)
+            print(f"Optimizer state loaded from {input_optimizer_file}")
+        # called from all ranks, though only rank0 has a valid param for full_osd
+        sharded_osd = FSDP.scatter_full_optim_state_dict(full_osd, model)
+        optimizer.load_state_dict(sharded_osd)
diff --git a/src/accelerate/utils/operations.py b/src/accelerate/utils/operations.py
index 7b3c8de5a..42868a0a5 100644
--- a/src/accelerate/utils/operations.py
+++ b/src/accelerate/utils/operations.py
@@ -254,7 +254,11 @@ def gather_object(object: Any):
     """
     if AcceleratorState().distributed_type == DistributedType.TPU:
         raise NotImplementedError("gather objects in TPU is not supported")
-    elif AcceleratorState().distributed_type in [DistributedType.DEEPSPEED, DistributedType.MULTI_GPU]:
+    elif AcceleratorState().distributed_type in [
+        DistributedType.DEEPSPEED,
+        DistributedType.MULTI_GPU,
+        DistributedType.FSDP,
+    ]:
         return _gpu_gather_object(object)
     elif AcceleratorState().distributed_type == DistributedType.MULTI_CPU:
         return _cpu_gather_object(object)
@@ -293,7 +297,11 @@ def broadcast(tensor, from_process: int = 0):
     """
     if AcceleratorState().distributed_type == DistributedType.TPU:
         return _tpu_broadcast(tensor, src=from_process, name="accelerate.utils.broadcast")
-    elif AcceleratorState().distributed_type in [DistributedType.DEEPSPEED, DistributedType.MULTI_GPU]:
+    elif AcceleratorState().distributed_type in [
+        DistributedType.DEEPSPEED,
+        DistributedType.MULTI_GPU,
+        DistributedType.FSDP,
+    ]:
         return _gpu_broadcast(tensor, src=from_process)
     elif AcceleratorState().distributed_type == DistributedType.MULTI_CPU:
         return _gpu_broadcast(tensor, src=from_process)
@@ -317,7 +325,11 @@ def broadcast_object_list(object_list, from_process: int = 0):
     if AcceleratorState().distributed_type == DistributedType.TPU:
         for i, obj in enumerate(object_list):
             object_list[i] = xm.mesh_reduce("accelerate.utils.broadcast_object_list", obj, lambda x: x[from_process])
-    elif AcceleratorState().distributed_type in [DistributedType.DEEPSPEED, DistributedType.MULTI_GPU]:
+    elif AcceleratorState().distributed_type in [
+        DistributedType.DEEPSPEED,
+        DistributedType.MULTI_GPU,
+        DistributedType.FSDP,
+    ]:
         torch.distributed.broadcast_object_list(object_list, src=from_process)
     elif AcceleratorState().distributed_type == DistributedType.MULTI_CPU:
         torch.distributed.broadcast_object_list(object_list, src=from_process)
@@ -433,7 +445,11 @@ def _reduce_across_processes(tensor, reduction="mean"):
         if state.distributed_type == DistributedType.TPU:
             xm.all_reduce("sum", cloned_tensor)
             return cloned_tensor
-        elif state.distributed_type in [DistributedType.DEEPSPEED, DistributedType.MULTI_GPU]:
+        elif state.distributed_type in [
+            DistributedType.DEEPSPEED,
+            DistributedType.MULTI_GPU,
+            DistributedType.FSDP,
+        ]:
             torch.distributed.all_reduce(cloned_tensor, ReduceOp.SUM)
             return cloned_tensor
         else:
diff --git a/src/accelerate/utils/other.py b/src/accelerate/utils/other.py
index 206c70589..ff360038d 100644
--- a/src/accelerate/utils/other.py
+++ b/src/accelerate/utils/other.py
@@ -65,6 +65,7 @@ def wait_for_everyone():
         AcceleratorState().distributed_type == DistributedType.MULTI_GPU
         or AcceleratorState().distributed_type == DistributedType.MULTI_CPU
         or AcceleratorState().distributed_type == DistributedType.DEEPSPEED
+        or AcceleratorState().distributed_type == DistributedType.FSDP
     ):
         torch.distributed.barrier()
     elif AcceleratorState().distributed_type == DistributedType.TPU:
diff --git a/tests/deepspeed/test_deepspeed.py b/tests/deepspeed/test_deepspeed.py
index 12cca415c..6b37eb93e 100644
--- a/tests/deepspeed/test_deepspeed.py
+++ b/tests/deepspeed/test_deepspeed.py
@@ -25,10 +25,18 @@
 import torch
 from torch.utils.data import DataLoader
 
+import accelerate
 from accelerate.accelerator import Accelerator
 from accelerate.scheduler import AcceleratedScheduler
 from accelerate.state import AcceleratorState
-from accelerate.test_utils.testing import require_cuda, require_deepspeed
+from accelerate.test_utils.testing import (
+    TempDirTestCase,
+    execute_subprocess_async,
+    require_cuda,
+    require_deepspeed,
+    require_multi_gpu,
+    slow,
+)
 from accelerate.test_utils.training import RegressionDataset
 from accelerate.utils.dataclasses import DeepSpeedPlugin
 from accelerate.utils.deepspeed import (
@@ -38,6 +46,7 @@
     DummyOptim,
     DummyScheduler,
 )
+from accelerate.utils.other import patch_environment
 from parameterized import parameterized
 from transformers import AutoModel, AutoModelForCausalLM, get_scheduler
 from transformers.testing_utils import mockenv_context
@@ -118,6 +127,10 @@ def setUp(self):
             WORLD_SIZE="1",
         )
 
+    def tearDown(self):
+        super().tearDown()
+        AcceleratorState._reset_state()
+
     def get_config_dict(self, stage):
         # As some tests modify the dict, always make a copy
         return deepcopy(self.ds_config_dict[stage])
@@ -260,11 +273,10 @@ def test_init_zero3(self):
         )
 
         with mockenv_context(**self.dist_env):
-            accelerator = Accelerator(deepspeed_plugin=deepspeed_plugin)
+            accelerator = Accelerator(deepspeed_plugin=deepspeed_plugin)  # noqa: F841
             from transformers.deepspeed import is_deepspeed_zero3_enabled
 
             self.assertTrue(is_deepspeed_zero3_enabled())
-            accelerator.state.initialized = False
 
     @parameterized.expand(optim_scheduler_params, name_func=parameterized_custom_name_func)
     def test_prepare_deepspeed(self, optim_type, scheduler_type):
@@ -479,7 +491,6 @@ def test_prepare_deepspeed(self, optim_type, scheduler_type):
                     "You can only specify `accelerate.utils.DummyScheduler` in the code when using `accelerate.utils.DummyOptim`."
                     in str(cm.exception)
                 )
-        accelerator.state.initialized = False
 
     def test_save_checkpoints(self):
         deepspeed_plugin = DeepSpeedPlugin(
@@ -533,7 +544,6 @@ def test_save_checkpoints(self):
                 "To save the full checkpoint, run `model.save_checkpoint(save_dir)` and use `zero_to_fp32.py` to recover weights."
             )
             self.assertTrue(msg in str(cm.exception))
-        accelerator.state.initialized = False
 
     def test_autofill_dsconfig(self):
         deepspeed_plugin = DeepSpeedPlugin(
@@ -581,4 +591,213 @@ def test_autofill_dsconfig(self):
             self.assertFalse(
                 accelerator.deepspeed_config["zero_optimization"]["stage3_gather_16bit_weights_on_model_save"]
             )
-        accelerator.state.initialized = False
+
+    def test_basic_run(self):
+        mod_file = inspect.getfile(accelerate.test_utils)
+        test_file_path = os.path.sep.join(mod_file.split(os.path.sep)[:-1] + ["scripts", "test_performance.py"])
+        with tempfile.TemporaryDirectory() as dirpath:
+            cmd = [
+                "accelerate",
+                "launch",
+                "--num_processes=1",
+                "--num_machines=1",
+                "--machine_rank=0",
+                "--mixed_precision=fp16",
+                "--use_deepspeed",
+                "--gradient_accumulation_steps=1",
+                "--zero_stage=2",
+                "--offload_optimizer_device=none",
+                "--offload_param_device=none",
+                test_file_path,
+                "--model_name_or_path=distilbert-base-uncased",
+                "--num_epochs=1",
+                f"--output_dir={dirpath}",
+            ]
+            with patch_environment(omp_num_threads=1):
+                execute_subprocess_async(cmd, env=os.environ.copy())
+
+
+@require_deepspeed
+@require_multi_gpu
+@slow
+class DeepSpeedIntegrationTest(TempDirTestCase):
+    def setUp(self):
+        super().setUp()
+        self._test_file_path = inspect.getfile(self.__class__)
+        path = Path(self._test_file_path).resolve()
+        self.test_file_dir_str = str(path.parents[0])
+
+        self.ds_config_file = dict(
+            zero2=f"{self.test_file_dir_str}/ds_config_zero2.json",
+            zero3=f"{self.test_file_dir_str}/ds_config_zero3.json",
+        )
+
+        self.stages = [1, 2, 3]
+        self.zero3_offload_config = False
+        self.performance_lower_bound = 0.83
+        self.peak_memory_usage_upper_bound = {
+            "multi_gpu_fp16": 3200,
+            "deepspeed_stage_1_fp16": 1600,
+            "deepspeed_stage_2_fp16": 2500,
+            "deepspeed_stage_3_zero_init_fp16": 2800,
+            "deepspeed_stage_3_cpu_offload_fp16": 1900,
+        }
+        self.n_train = 160
+        self.n_val = 160
+
+        mod_file = inspect.getfile(accelerate.test_utils)
+        self.test_scripts_folder = os.path.sep.join(mod_file.split(os.path.sep)[:-1] + ["scripts"])
+
+    def test_performance(self):
+        self.test_file_path = os.path.join(self.test_scripts_folder, "test_performance.py")
+        cmd = [
+            "accelerate",
+            "launch",
+            "--num_processes=2",
+            "--num_machines=1",
+            "--machine_rank=0",
+            "--mixed_precision=fp16",
+            "--use_deepspeed",
+            "--gradient_accumulation_steps=1",
+            "--gradient_clipping=1",
+            "--zero3_init_flag=True",
+            "--zero3_save_16bit_model=True",
+        ]
+        for stage in self.stages:
+            if stage == 1:
+                continue
+            cmd_stage = cmd.copy()
+            cmd_stage.extend([f"--zero_stage={stage}"])
+            cmd_stage.extend(["--offload_optimizer_device=none", "--offload_param_device=none"])
+            if self.zero3_offload_config:
+                with io.open(self.ds_config_file[ZERO3], "r", encoding="utf-8") as f:
+                    ds_config = json.load(f)
+                    del ds_config["bf16"]
+                    del ds_config["optimizer"]["params"]["torch_adam"]
+                    del ds_config["optimizer"]["params"]["adam_w_mode"]
+                    ds_config["fp16"]["enabled"] = True
+                    ds_config_path = os.path.join(self.tmpdir, "ds_config.json")
+                    with open(ds_config_path, "w") as out_file:
+                        json.dump(ds_config, out_file)
+
+                cmd_stage.extend([f"--deepspeed_config_file={ds_config_path}"])
+
+            cmd_stage.extend(
+                [
+                    self.test_file_path,
+                    f"--output_dir={self.tmpdir}",
+                    f"--performance_lower_bound={self.performance_lower_bound}",
+                ]
+            )
+            with patch_environment(omp_num_threads=1):
+                execute_subprocess_async(cmd_stage, env=os.environ.copy())
+
+    def test_checkpointing(self):
+        self.test_file_path = os.path.join(self.test_scripts_folder, "test_checkpointing.py")
+        cmd = [
+            "accelerate",
+            "launch",
+            "--num_processes=2",
+            "--num_machines=1",
+            "--machine_rank=0",
+            "--mixed_precision=fp16",
+            "--use_deepspeed",
+            "--gradient_accumulation_steps=1",
+            "--gradient_clipping=1",
+            "--zero3_init_flag=True",
+            "--zero3_save_16bit_model=True",
+        ]
+        for stage in self.stages:
+            if stage == 1:
+                continue
+            cmd_stage = cmd.copy()
+            cmd_stage.extend([f"--zero_stage={stage}"])
+            cmd_stage.extend(["--offload_optimizer_device=none", "--offload_param_device=none"])
+            if self.zero3_offload_config:
+                with io.open(self.ds_config_file[ZERO3], "r", encoding="utf-8") as f:
+                    ds_config = json.load(f)
+                    del ds_config["bf16"]
+                    del ds_config["optimizer"]["params"]["torch_adam"]
+                    del ds_config["optimizer"]["params"]["adam_w_mode"]
+                    ds_config["fp16"]["enabled"] = True
+                    ds_config_path = os.path.join(self.tmpdir, "ds_config.json")
+                    with open(ds_config_path, "w") as out_file:
+                        json.dump(ds_config, out_file)
+
+                cmd_stage.extend([f"--deepspeed_config_file={ds_config_path}"])
+
+            cmd_stage.extend(
+                [
+                    self.test_file_path,
+                    f"--output_dir={self.tmpdir}",
+                    "--partial_train_epoch=1",
+                ]
+            )
+            with patch_environment(omp_num_threads=1):
+                execute_subprocess_async(cmd_stage, env=os.environ.copy())
+
+            cmd_stage = cmd_stage[:-1]
+            resume_from_checkpoint = os.path.join(self.tmpdir, "epoch_0")
+            cmd_stage.extend(
+                [
+                    f"--resume_from_checkpoint={resume_from_checkpoint}",
+                ]
+            )
+            with patch_environment(omp_num_threads=1):
+                execute_subprocess_async(cmd_stage, env=os.environ.copy())
+
+    def test_peak_memory_usage(self):
+        self.test_file_path = os.path.join(self.test_scripts_folder, "test_peak_memory_usage.py")
+        cmd = [
+            "accelerate",
+            "launch",
+            "--num_processes=2",
+            "--num_machines=1",
+            "--machine_rank=0",
+        ]
+        for spec, peak_mem_upper_bound in self.peak_memory_usage_upper_bound.items():
+            cmd_stage = cmd.copy()
+            if "fp16" in spec:
+                cmd_stage.extend(["--mixed_precision=fp16"])
+
+            if "multi_gpu" in spec:
+                continue
+            else:
+                cmd_stage.extend(
+                    [
+                        "--use_deepspeed",
+                        "--gradient_accumulation_steps=1",
+                        "--gradient_clipping=1",
+                        "--zero3_init_flag=True",
+                        "--zero3_save_16bit_model=True",
+                    ]
+                )
+                for i in range(3):
+                    if f"stage_{i+1}" in spec:
+                        cmd_stage.extend([f"--zero_stage={i+1}"])
+                        break
+                cmd_stage.extend(["--offload_optimizer_device=none", "--offload_param_device=none"])
+                if "cpu_offload" in spec:
+                    with io.open(self.ds_config_file[ZERO3], "r", encoding="utf-8") as f:
+                        ds_config = json.load(f)
+                        del ds_config["bf16"]
+                        del ds_config["fp16"]
+                        del ds_config["optimizer"]["params"]["torch_adam"]
+                        del ds_config["optimizer"]["params"]["adam_w_mode"]
+                        ds_config_path = os.path.join(self.tmpdir, "ds_config.json")
+                        with open(ds_config_path, "w") as out_file:
+                            json.dump(ds_config, out_file)
+
+                    cmd_stage.extend([f"--deepspeed_config_file={ds_config_path}"])
+
+            cmd_stage.extend(
+                [
+                    self.test_file_path,
+                    f"--output_dir={self.tmpdir}",
+                    f"--peak_memory_upper_bound={peak_mem_upper_bound}",
+                    f"--n_train={self.n_train}",
+                    f"--n_val={self.n_val}",
+                ]
+            )
+            with patch_environment(omp_num_threads=1):
+                execute_subprocess_async(cmd_stage, env=os.environ.copy())
diff --git a/tests/fsdp/test_fsdp.py b/tests/fsdp/test_fsdp.py
new file mode 100644
index 000000000..8ad088c04
--- /dev/null
+++ b/tests/fsdp/test_fsdp.py
@@ -0,0 +1,332 @@
+# 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 inspect
+import os
+import unittest
+
+import torch
+
+import accelerate
+from accelerate.accelerator import Accelerator
+from accelerate.state import AcceleratorState
+from accelerate.test_utils.testing import (
+    TempDirTestCase,
+    execute_subprocess_async,
+    require_cuda,
+    require_fsdp,
+    require_multi_gpu,
+    slow,
+)
+from accelerate.utils.constants import (
+    FSDP_AUTO_WRAP_POLICY,
+    FSDP_BACKWARD_PREFETCH,
+    FSDP_SHARDING_STRATEGY,
+    FSDP_STATE_DICT_TYPE,
+)
+from accelerate.utils.dataclasses import FullyShardedDataParallelPlugin
+from accelerate.utils.other import patch_environment
+from transformers import AutoModel
+from transformers.testing_utils import mockenv_context
+from transformers.trainer_utils import set_seed
+
+
+set_seed(42)
+
+BERT_BASE_CASED = "bert-base-cased"
+FP16 = "fp16"
+BF16 = "bf16"
+dtypes = [FP16, BF16]
+
+
+@require_fsdp
+@require_cuda
+class FSDPPluginIntegration(unittest.TestCase):
+    def setUp(self):
+        super().setUp()
+
+        self.dist_env = dict(
+            USE_FSDP="true",
+            MASTER_ADDR="localhost",
+            MASTER_PORT="10999",
+            RANK="0",
+            LOCAL_RANK="0",
+            WORLD_SIZE="1",
+        )
+
+    def tearDown(self):
+        super().tearDown()
+        AcceleratorState._reset_state()
+
+    def test_sharding_strategy(self):
+        from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy
+
+        for i, strategy in enumerate(FSDP_SHARDING_STRATEGY):
+            env = self.dist_env.copy()
+            env["FSDP_SHARDING_STRATEGY"] = f"{i + 1}"
+            env["FSDP_SHARDING_STRATEGY_NAME"] = strategy
+            with mockenv_context(**env):
+                fsdp_plugin = FullyShardedDataParallelPlugin()
+                self.assertEqual(fsdp_plugin.sharding_strategy, ShardingStrategy(i + 1))
+
+    def test_backward_prefetch(self):
+        from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch
+
+        for i, prefetch_policy in enumerate(FSDP_BACKWARD_PREFETCH):
+            env = self.dist_env.copy()
+            env["FSDP_BACKWARD_PREFETCH"] = prefetch_policy
+            with mockenv_context(**env):
+                fsdp_plugin = FullyShardedDataParallelPlugin()
+                if prefetch_policy == "NO_PREFETCH":
+                    self.assertIsNone(fsdp_plugin.backward_prefetch)
+                else:
+                    self.assertEqual(fsdp_plugin.backward_prefetch, BackwardPrefetch(i + 1))
+
+    def test_state_dict_type(self):
+        from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType, _state_dict_type_to_config
+
+        for i, state_dict_type in enumerate(FSDP_STATE_DICT_TYPE):
+            env = self.dist_env.copy()
+            env["FSDP_STATE_DICT_TYPE"] = state_dict_type
+            with mockenv_context(**env):
+                fsdp_plugin = FullyShardedDataParallelPlugin()
+                self.assertEqual(fsdp_plugin.state_dict_type, StateDictType(i + 1))
+                self.assertEqual(
+                    type(fsdp_plugin.state_dict_config), type(_state_dict_type_to_config[StateDictType(i + 1)]())
+                )
+                if state_dict_type == "FULL_STATE_DICT":
+                    self.assertTrue(fsdp_plugin.state_dict_config.offload_to_cpu)
+                    self.assertTrue(fsdp_plugin.state_dict_config.rank0_only)
+
+    def test_auto_wrap_policy(self):
+        model = AutoModel.from_pretrained(BERT_BASE_CASED)
+        for policy in FSDP_AUTO_WRAP_POLICY:
+            env = self.dist_env.copy()
+            env["FSDP_AUTO_WRAP_POLICY"] = policy
+            if policy == "TRANSFORMER_BASED_WRAP":
+                env["FSDP_TRANSFORMER_CLS_TO_WRAP"] = "BertLayer"
+            elif policy == "SIZE_BASED_WRAP":
+                env["FSDP_MIN_NUM_PARAMS"] = "2000"
+            with mockenv_context(**env):
+                fsdp_plugin = FullyShardedDataParallelPlugin()
+                fsdp_plugin.set_auto_wrap_policy(model)
+                if policy == "NO_WRAP":
+                    self.assertIsNone(fsdp_plugin.auto_wrap_policy)
+                else:
+                    self.assertIsNotNone(fsdp_plugin.auto_wrap_policy)
+
+        env = self.dist_env.copy()
+        env["FSDP_AUTO_WRAP_POLICY"] = "TRANSFORMER_BASED_WRAP"
+        env["FSDP_TRANSFORMER_CLS_TO_WRAP"] = "T5Layer"
+        with mockenv_context(**env):
+            fsdp_plugin = FullyShardedDataParallelPlugin()
+            with self.assertRaises(Exception) as cm:
+                fsdp_plugin.set_auto_wrap_policy(model)
+            self.assertTrue("Could not find the transformer layer class to wrap in the model." in str(cm.exception))
+
+        env = self.dist_env.copy()
+        env["FSDP_AUTO_WRAP_POLICY"] = "SIZE_BASED_WRAP"
+        env["FSDP_MIN_NUM_PARAMS"] = "0"
+        with mockenv_context(**env):
+            fsdp_plugin = FullyShardedDataParallelPlugin()
+            fsdp_plugin.set_auto_wrap_policy(model)
+            self.assertIsNone(fsdp_plugin.auto_wrap_policy)
+
+    def test_mixed_precision(self):
+        from torch.distributed.fsdp.fully_sharded_data_parallel import MixedPrecision
+        from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler
+
+        for mp_dtype in dtypes:
+            env = self.dist_env.copy()
+            env["MIXED_PRECISION"] = mp_dtype
+            with mockenv_context(**env):
+                accelerator = Accelerator()
+                if mp_dtype == "fp16":
+                    dtype = torch.float16
+                elif mp_dtype == "bf16":
+                    dtype = torch.bfloat16
+                mp_policy = MixedPrecision(param_dtype=dtype, reduce_dtype=dtype, buffer_dtype=dtype)
+                self.assertEqual(accelerator.state.fsdp_plugin.mixed_precision_policy, mp_policy)
+                if mp_dtype == FP16:
+                    self.assertTrue(isinstance(accelerator.scaler, ShardedGradScaler))
+                elif mp_dtype == BF16:
+                    self.assertIsNone(accelerator.scaler)
+                AcceleratorState._reset_state()
+
+    def test_cpu_offload(self):
+        from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload
+
+        for flag in [True, False]:
+            env = self.dist_env.copy()
+            env["FSDP_OFFLOAD_PARAMS"] = str(flag).lower()
+            with mockenv_context(**env):
+                fsdp_plugin = FullyShardedDataParallelPlugin()
+                self.assertEqual(fsdp_plugin.cpu_offload, CPUOffload(offload_params=flag))
+
+
+@require_fsdp
+@require_multi_gpu
+@slow
+class FSDPIntegrationTest(TempDirTestCase):
+    def setUp(self):
+        super().setUp()
+        self.performance_lower_bound = 0.83
+        self.performance_configs = [
+            "fsdp_shard_grad_op_transformer_based_wrap",
+            "fsdp_full_shard_transformer_based_wrap",
+        ]
+        self.peak_memory_usage_upper_bound = {
+            "multi_gpu_fp16": 3200,
+            "fsdp_shard_grad_op_transformer_based_wrap_fp16": 2000,
+            "fsdp_full_shard_transformer_based_wrap_fp16": 1900,
+            "fsdp_full_shard_cpu_offload_transformer_based_wrap_fp32": 1500,  # fp16 was leading to indefinite hang
+        }
+        self.n_train = 160
+        self.n_val = 160
+
+        mod_file = inspect.getfile(accelerate.test_utils)
+        self.test_scripts_folder = os.path.sep.join(mod_file.split(os.path.sep)[:-1] + ["scripts"])
+
+    def test_performance(self):
+        self.test_file_path = os.path.join(self.test_scripts_folder, "test_performance.py")
+        cmd = ["accelerate", "launch", "--num_processes=2", "--num_machines=1", "--machine_rank=0", "--use_fsdp"]
+        for config in self.performance_configs:
+            cmd_config = cmd.copy()
+            for i, strategy in enumerate(FSDP_SHARDING_STRATEGY):
+                if strategy.lower() in config:
+                    cmd_config.append(f"--fsdp_sharding_strategy={i+1}")
+                    break
+
+            if "fp32" in config:
+                cmd_config.append("--mixed_precision=no")
+            else:
+                cmd_config.append("--mixed_precision=fp16")
+
+            if "cpu_offload" in config:
+                cmd_config.append("--fsdp_offload_params=True")
+
+            for policy in FSDP_AUTO_WRAP_POLICY:
+                if policy.lower() in config:
+                    cmd_config.append(f"--fsdp_auto_wrap_policy={policy}")
+                    break
+
+            if policy == "TRANSFORMER_BASED_WRAP":
+                cmd_config.append("--fsdp_transformer_layer_cls_to_wrap=BertLayer")
+            elif policy == "SIZE_BASED_WRAP":
+                cmd_config.append("--fsdp_min_num_params=2000")
+
+            cmd_config.extend(
+                [
+                    self.test_file_path,
+                    f"--output_dir={self.tmpdir}",
+                    f"--performance_lower_bound={self.performance_lower_bound}",
+                ]
+            )
+            with patch_environment(omp_num_threads=1):
+                execute_subprocess_async(cmd_config, env=os.environ.copy())
+
+    def test_checkpointing(self):
+        self.test_file_path = os.path.join(self.test_scripts_folder, "test_checkpointing.py")
+        cmd = [
+            "accelerate",
+            "launch",
+            "--num_processes=2",
+            "--num_machines=1",
+            "--machine_rank=0",
+            "--use_fsdp",
+            "--mixed_precision=fp16",
+            "--fsdp_transformer_layer_cls_to_wrap=BertLayer",
+        ]
+
+        for i, strategy in enumerate(FSDP_SHARDING_STRATEGY):
+            cmd_config = cmd.copy()
+            cmd_config.append(f"--fsdp_sharding_strategy={i+1}")
+            if strategy != "FULL_SHARD":
+                continue
+            state_dict_config_index = len(cmd_config)
+            for state_dict_type in FSDP_STATE_DICT_TYPE:
+                cmd_config = cmd_config[:state_dict_config_index]
+                if state_dict_type == "SHARDED_STATE_DICT":
+                    continue
+                cmd_config.append(f"--fsdp_state_dict_type={state_dict_type}")
+                cmd_config.extend(
+                    [
+                        self.test_file_path,
+                        f"--output_dir={self.tmpdir}",
+                        "--partial_train_epoch=1",
+                    ]
+                )
+                with patch_environment(omp_num_threads=1):
+                    execute_subprocess_async(cmd_config, env=os.environ.copy())
+
+                cmd_config = cmd_config[:-1]
+                resume_from_checkpoint = os.path.join(self.tmpdir, "epoch_0")
+                cmd_config.extend(
+                    [
+                        f"--resume_from_checkpoint={resume_from_checkpoint}",
+                    ]
+                )
+                with patch_environment(omp_num_threads=1):
+                    execute_subprocess_async(cmd_config, env=os.environ.copy())
+
+    def test_peak_memory_usage(self):
+        self.test_file_path = os.path.join(self.test_scripts_folder, "test_peak_memory_usage.py")
+        cmd = [
+            "accelerate",
+            "launch",
+            "--num_processes=2",
+            "--num_machines=1",
+            "--machine_rank=0",
+        ]
+        for spec, peak_mem_upper_bound in self.peak_memory_usage_upper_bound.items():
+            cmd_config = cmd.copy()
+            if "fp16" in spec:
+                cmd_config.extend(["--mixed_precision=fp16"])
+            else:
+                cmd_config.extend(["--mixed_precision=no"])
+
+            if "multi_gpu" in spec:
+                continue
+            else:
+                cmd_config.extend(["--use_fsdp"])
+                for i, strategy in enumerate(FSDP_SHARDING_STRATEGY):
+                    if strategy.lower() in spec:
+                        cmd_config.append(f"--fsdp_sharding_strategy={i+1}")
+                        break
+
+                if "cpu_offload" in spec:
+                    cmd_config.append("--fsdp_offload_params=True")
+
+                for policy in FSDP_AUTO_WRAP_POLICY:
+                    if policy.lower() in spec:
+                        cmd_config.append(f"--fsdp_auto_wrap_policy={policy}")
+                        break
+
+                if policy == "TRANSFORMER_BASED_WRAP":
+                    cmd_config.append("--fsdp_transformer_layer_cls_to_wrap=BertLayer")
+                elif policy == "SIZE_BASED_WRAP":
+                    cmd_config.append("--fsdp_min_num_params=2000")
+
+            cmd_config.extend(
+                [
+                    self.test_file_path,
+                    f"--output_dir={self.tmpdir}",
+                    f"--peak_memory_upper_bound={peak_mem_upper_bound}",
+                    f"--n_train={self.n_train}",
+                    f"--n_val={self.n_val}",
+                ]
+            )
+            with patch_environment(omp_num_threads=1):
+                execute_subprocess_async(cmd_config, env=os.environ.copy())

We can't change the API of this function, it's public and has been released now.
@@ -43,8 +44,9 @@
 
 
 def save_accelerator_state(
+    accelerator,
     output_dir: str,
-    model_states: List[dict],
+    models: list,
The logic specific to the Accelerator should go in the `Accelerator.save_state` method.
@@ -68,20 +72,40 @@ def save_accelerator_state(
             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}")
+    for i, model in enumerate(models):
+        if accelerator.distributed_type == DistributedType.FSDP:
+            logger.info("Saving FSDP model")
+            accelerator.state.fsdp_plugin.save_model(accelerator, model, output_dir, i)
+            logger.info(f"FSDP Model saved to output dir {output_dir}")
+        elif accelerator.distributed_type == DistributedType.DEEPSPEED:
+            logger.info("Saving DeepSpeed Model and Optimizer")
+            ckpt_id = f"{MODEL_NAME}" if i == 0 else f"{MODEL_NAME}_{i}"
+            model.save_checkpoint(output_dir, ckpt_id)
+            logger.info(f"DeepSpeed Model and Optimizer saved to output dir {os.path.join(output_dir, ckpt_id)}")
+        else:
+            state = accelerator.get_state_dict(model, unwrap=False)
+            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}")
Same here, we can't change the API and the logic specific to the Accelerator should go in the Accelerator.
@@ -109,11 +133,21 @@ def save_accelerator_state(
     return output_dir
 
 
-def load_accelerator_state(input_dir, models, optimizers, schedulers, process_index, scaler=None):
+def load_accelerator_state(
+    accelerator,
+    input_dir: str,
+    models: list,
+    optimizers: list,
+    schedulers: list,
+    process_index: int,
+    scaler: GradScaler = None,
+):
Chaing the name will break existing configs, and also it's already in something called `FSDP` so it seems redundant to add it here.
@@ -210,12 +211,12 @@ def get_cluster_input():
             for i, strategy in enumerate(FSDP_SHARDING_STRATEGY):
                 sharding_strategy_query += f"[{i+1}] {strategy}, "
             sharding_strategy_query = sharding_strategy_query[:-2] + ")? [1]: "
-            fsdp_config["sharding_strategy"] = _ask_field(
+            fsdp_config["fsdp_sharding_strategy"] = _ask_field(
Let's not rely on this constant and manually put lists and defaults here. We can always add more in the future.
@@ -250,6 +251,15 @@ def get_cluster_input():
                 lambda x: FSDP_BACKWARD_PREFETCH[int(x)],
                 default=FSDP_BACKWARD_PREFETCH[0],
             )
+            fsdp_state_dict_type_query = "What should be your FSDP's state dict type ("
+            for i, state_dict_type in enumerate(FSDP_STATE_DICT_TYPE):
+                fsdp_state_dict_type_query += f"[{i}] {state_dict_type}, "
+            fsdp_state_dict_type_query = fsdp_state_dict_type_query[:-2] + ")? [0]: "
+            fsdp_config["fsdp_state_dict_type"] = _ask_field(
+                fsdp_state_dict_type_query,
+                lambda x: FSDP_STATE_DICT_TYPE[int(x)],
+                default=FSDP_STATE_DICT_TYPE[0],
+            )
This and all below is breaking. We should at least have one deprecation cycle with the old names.
@@ -141,19 +141,19 @@ def launch_command_parser(subparsers=None):
         help="Whether to use fsdp.",
     )
     parser.add_argument(
-        "--offload_params",
+        "--fsdp_offload_params",
Let's write a default here to not rely on the order of this constant.
@@ -472,6 +492,17 @@ def __post_init__(self):
             if prefetch_policy != FSDP_BACKWARD_PREFETCH[-1]:
                 self.backward_prefetch = BackwardPrefetch(FSDP_BACKWARD_PREFETCH.index(prefetch_policy) + 1)
 
+        if self.state_dict_type is None:
+            state_dict_type_policy = os.environ.get("FSDP_STATE_DICT_TYPE", FSDP_STATE_DICT_TYPE[0])
Hello, had to do this so that launcher has arguments with `fsdp` prefix. I thought this was fair game as we specified FSDP integration was experimental given that FSDP itself has changed a lot. This was more around readability/grouping for folks to easily understand all the variables necessary for FSDP as there are a lot more now. If you think we can do away with it, I will revert the variable name changes.
@@ -210,12 +211,12 @@ def get_cluster_input():
             for i, strategy in enumerate(FSDP_SHARDING_STRATEGY):
                 sharding_strategy_query += f"[{i+1}] {strategy}, "
             sharding_strategy_query = sharding_strategy_query[:-2] + ")? [1]: "
-            fsdp_config["sharding_strategy"] = _ask_field(
+            fsdp_config["fsdp_sharding_strategy"] = _ask_field(
Hello, we maintain that constant in our repo at `accelerate.utils.constants.py` and hence I don't think this would be an issue. 
@@ -250,6 +251,15 @@ def get_cluster_input():
                 lambda x: FSDP_BACKWARD_PREFETCH[int(x)],
                 default=FSDP_BACKWARD_PREFETCH[0],
             )
+            fsdp_state_dict_type_query = "What should be your FSDP's state dict type ("
+            for i, state_dict_type in enumerate(FSDP_STATE_DICT_TYPE):
+                fsdp_state_dict_type_query += f"[{i}] {state_dict_type}, "
+            fsdp_state_dict_type_query = fsdp_state_dict_type_query[:-2] + ")? [0]: "
+            fsdp_config["fsdp_state_dict_type"] = _ask_field(
+                fsdp_state_dict_type_query,
+                lambda x: FSDP_STATE_DICT_TYPE[int(x)],
+                default=FSDP_STATE_DICT_TYPE[0],
+            )
It would be more readable to just have a value.
@@ -250,6 +251,15 @@ def get_cluster_input():
                 lambda x: FSDP_BACKWARD_PREFETCH[int(x)],
                 default=FSDP_BACKWARD_PREFETCH[0],
             )
+            fsdp_state_dict_type_query = "What should be your FSDP's state dict type ("
+            for i, state_dict_type in enumerate(FSDP_STATE_DICT_TYPE):
+                fsdp_state_dict_type_query += f"[{i}] {state_dict_type}, "
+            fsdp_state_dict_type_query = fsdp_state_dict_type_query[:-2] + ")? [0]: "
+            fsdp_config["fsdp_state_dict_type"] = _ask_field(
+                fsdp_state_dict_type_query,
+                lambda x: FSDP_STATE_DICT_TYPE[int(x)],
+                default=FSDP_STATE_DICT_TYPE[0],
+            )
As mentioned above, we are separately maintaining it in our repo and hence this should not be an issue. 
@@ -472,6 +492,17 @@ def __post_init__(self):
             if prefetch_policy != FSDP_BACKWARD_PREFETCH[-1]:
                 self.backward_prefetch = BackwardPrefetch(FSDP_BACKWARD_PREFETCH.index(prefetch_policy) + 1)
 
+        if self.state_dict_type is None:
+            state_dict_type_policy = os.environ.get("FSDP_STATE_DICT_TYPE", FSDP_STATE_DICT_TYPE[0])
Changing the launcher is fair game, but existing configs breaking when updating Accelerate is a bit too annoying to the user I think.
@@ -210,12 +211,12 @@ def get_cluster_input():
             for i, strategy in enumerate(FSDP_SHARDING_STRATEGY):
                 sharding_strategy_query += f"[{i+1}] {strategy}, "
             sharding_strategy_query = sharding_strategy_query[:-2] + ")? [1]: "
-            fsdp_config["sharding_strategy"] = _ask_field(
+            fsdp_config["fsdp_sharding_strategy"] = _ask_field(
As mentioned above, I thought this was fair game as we specified FSDP integration was experimental given that FSDP itself has changed a lot. This was more around readability/grouping for folks to easily understand all the variables necessary for FSDP as there are a lot more now. If you think we can do away with it at the expense of readability/clarity, I will revert the variable name changes.
@@ -141,19 +141,19 @@ def launch_command_parser(subparsers=None):
         help="Whether to use fsdp.",
     )
     parser.add_argument(
-        "--offload_params",
+        "--fsdp_offload_params",
You mean for the `default` or entire query?
@@ -250,6 +251,15 @@ def get_cluster_input():
                 lambda x: FSDP_BACKWARD_PREFETCH[int(x)],
                 default=FSDP_BACKWARD_PREFETCH[0],
             )
+            fsdp_state_dict_type_query = "What should be your FSDP's state dict type ("
+            for i, state_dict_type in enumerate(FSDP_STATE_DICT_TYPE):
+                fsdp_state_dict_type_query += f"[{i}] {state_dict_type}, "
+            fsdp_state_dict_type_query = fsdp_state_dict_type_query[:-2] + ")? [0]: "
+            fsdp_config["fsdp_state_dict_type"] = _ask_field(
+                fsdp_state_dict_type_query,
+                lambda x: FSDP_STATE_DICT_TYPE[int(x)],
+                default=FSDP_STATE_DICT_TYPE[0],
+            )
But it's not readable.
@@ -472,6 +492,17 @@ def __post_init__(self):
             if prefetch_policy != FSDP_BACKWARD_PREFETCH[-1]:
                 self.backward_prefetch = BackwardPrefetch(FSDP_BACKWARD_PREFETCH.index(prefetch_policy) + 1)
 
+        if self.state_dict_type is None:
+            state_dict_type_policy = os.environ.get("FSDP_STATE_DICT_TYPE", FSDP_STATE_DICT_TYPE[0])
Done.
@@ -472,6 +492,17 @@ def __post_init__(self):
             if prefetch_policy != FSDP_BACKWARD_PREFETCH[-1]:
                 self.backward_prefetch = BackwardPrefetch(FSDP_BACKWARD_PREFETCH.index(prefetch_policy) + 1)
 
+        if self.state_dict_type is None:
+            state_dict_type_policy = os.environ.get("FSDP_STATE_DICT_TYPE", FSDP_STATE_DICT_TYPE[0])
Done. Added deprecation args with the warnings.
@@ -141,19 +141,19 @@ def launch_command_parser(subparsers=None):
         help="Whether to use fsdp.",
     )
     parser.add_argument(
-        "--offload_params",
+        "--fsdp_offload_params",
Understood. Added code to handle this wherein old configs should work without any issues. 
@@ -210,12 +211,12 @@ def get_cluster_input():
             for i, strategy in enumerate(FSDP_SHARDING_STRATEGY):
                 sharding_strategy_query += f"[{i+1}] {strategy}, "
             sharding_strategy_query = sharding_strategy_query[:-2] + ")? [1]: "
-            fsdp_config["sharding_strategy"] = _ask_field(
+            fsdp_config["fsdp_sharding_strategy"] = _ask_field(
Done.
@@ -109,11 +133,21 @@ def save_accelerator_state(
     return output_dir
 
 
-def load_accelerator_state(input_dir, models, optimizers, schedulers, process_index, scaler=None):
+def load_accelerator_state(
+    accelerator,
+    input_dir: str,
+    models: list,
+    optimizers: list,
+    schedulers: list,
+    process_index: int,
+    scaler: GradScaler = None,
+):
Done.
@@ -68,20 +72,40 @@ def save_accelerator_state(
             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}")
+    for i, model in enumerate(models):
+        if accelerator.distributed_type == DistributedType.FSDP:
+            logger.info("Saving FSDP model")
+            accelerator.state.fsdp_plugin.save_model(accelerator, model, output_dir, i)
+            logger.info(f"FSDP Model saved to output dir {output_dir}")
+        elif accelerator.distributed_type == DistributedType.DEEPSPEED:
+            logger.info("Saving DeepSpeed Model and Optimizer")
+            ckpt_id = f"{MODEL_NAME}" if i == 0 else f"{MODEL_NAME}_{i}"
+            model.save_checkpoint(output_dir, ckpt_id)
+            logger.info(f"DeepSpeed Model and Optimizer saved to output dir {os.path.join(output_dir, ckpt_id)}")
+        else:
+            state = accelerator.get_state_dict(model, unwrap=False)
+            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}")
Done.
@@ -43,8 +44,9 @@
 
 
 def save_accelerator_state(
+    accelerator,
     output_dir: str,
-    model_states: List[dict],
+    models: list,
For the default
@@ -250,6 +251,15 @@ def get_cluster_input():
                 lambda x: FSDP_BACKWARD_PREFETCH[int(x)],
                 default=FSDP_BACKWARD_PREFETCH[0],
             )
+            fsdp_state_dict_type_query = "What should be your FSDP's state dict type ("
+            for i, state_dict_type in enumerate(FSDP_STATE_DICT_TYPE):
+                fsdp_state_dict_type_query += f"[{i}] {state_dict_type}, "
+            fsdp_state_dict_type_query = fsdp_state_dict_type_query[:-2] + ")? [0]: "
+            fsdp_config["fsdp_state_dict_type"] = _ask_field(
+                fsdp_state_dict_type_query,
+                lambda x: FSDP_STATE_DICT_TYPE[int(x)],
+                default=FSDP_STATE_DICT_TYPE[0],
+            )
Done.
@@ -250,6 +251,15 @@ def get_cluster_input():
                 lambda x: FSDP_BACKWARD_PREFETCH[int(x)],
                 default=FSDP_BACKWARD_PREFETCH[0],
             )
+            fsdp_state_dict_type_query = "What should be your FSDP's state dict type ("
+            for i, state_dict_type in enumerate(FSDP_STATE_DICT_TYPE):
+                fsdp_state_dict_type_query += f"[{i}] {state_dict_type}, "
+            fsdp_state_dict_type_query = fsdp_state_dict_type_query[:-2] + ")? [0]: "
+            fsdp_config["fsdp_state_dict_type"] = _ask_field(
+                fsdp_state_dict_type_query,
+                lambda x: FSDP_STATE_DICT_TYPE[int(x)],
+                default=FSDP_STATE_DICT_TYPE[0],
+            )
_The documentation is not available anymore as the PR was closed or merged._
`make test_deepspeed` results below (~6 minutes) :

![Screenshot 2022-07-26 at 1 11 22 PM](https://user-images.githubusercontent.com/13534540/180951182-b4bd178b-2d90-49a6-b6f1-9cfcf3d59ffd.png)

`make test_fsdp` results below (~4 minutes):

![Screenshot 2022-07-26 at 1 25 05 PM](https://user-images.githubusercontent.com/13534540/180953795-7e7937fb-723d-4848-82da-051328bb1566.png)


https://api.github.com/repos/huggingface/accelerate/pulls/532
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 5af38afaf..7f36121c1 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -434,22 +434,23 @@ def print(self, *args, **kwargs):
 
     def _prepare_one(self, obj, first_pass=False):
         # First pass of preparation: DataLoader, model, optimizer
-        if isinstance(obj, torch.utils.data.DataLoader) and first_pass:
-            return self.prepare_data_loader(obj)
-        elif isinstance(obj, torch.nn.Module) and first_pass:
-            self._models.append(obj)
-            return self.prepare_model(obj)
-        elif isinstance(obj, torch.optim.Optimizer) and first_pass:
-            optimizer = self.prepare_optimizer(obj)
-            self._optimizers.append(optimizer)
-            return optimizer
+        if first_pass:
+            if isinstance(obj, torch.utils.data.DataLoader):
+                return self.prepare_data_loader(obj)
+            elif isinstance(obj, torch.nn.Module):
+                self._models.append(obj)
+                return self.prepare_model(obj)
+            elif isinstance(obj, torch.optim.Optimizer):
+                optimizer = self.prepare_optimizer(obj)
+                self._optimizers.append(optimizer)
+                return optimizer
         # Second pass of preparation: LR scheduler (which need the full list of optimizers)
-        elif isinstance(obj, torch.optim.lr_scheduler._LRScheduler) and not first_pass:
+        elif isinstance(obj, torch.optim.lr_scheduler._LRScheduler):
             scheduler = self.prepare_scheduler(obj)
             self._schedulers.append(scheduler)
             return scheduler
-        else:
-            return obj
+        # Return the unprocessed object if previous criteria was not met
+        return obj
 
     def _prepare_fsdp(self, *args):
         result = []

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/531
diff --git a/docs/source/big_modeling.mdx b/docs/source/big_modeling.mdx
index a5cb3d4ed..051d1bda7 100644
--- a/docs/source/big_modeling.mdx
+++ b/docs/source/big_modeling.mdx
@@ -203,6 +203,65 @@ This only supports inference of your model, not training. Most of the computatio
 
 </Tip>
 
+## Designing a device map
+
+You can let 🤗 Accelerate handle the device map computation by setting `device_map="auto"` or create one yourself, if you want more control over where each layer should go.
+
+<Tip>
+
+You can derive all sizes of the model (and thus compute a `device_map`) on a model that is on the meta device.
+
+</Tip>
+
+First note that you can limit the memory used on each GPU by using the `max_memory` argument (available in [`infer_auto_device_map`] and in all functions using it). When setting `max_memory`, you should pass along a dictionary containing the GPU identifiers (for instance `0`, `1` etc.) and the `"cpu"` key for the maximum RAM you want used for CPU offload. The values can either be an integer (in bytes) or a string representing a number with its unit, such as `"10GiB"` or `"10GB"`.
+
+Here is an example where we don't want to use more than 10GiB on each of two GPUs and no more than 30GiB of CPU RAM for the model weights:
+
+```python
+from accelerate import infer_auto_device_map
+
+device_map = infer_auto_device_map(my_model, max_memory={0: "10GiB", 1: "10GiB", "cpu": "30GiB"})
+```
+
+<Tip warning={true}>
+
+When a first allocation happens in PyTorch, it loads CUDA kernels which take about 1-2GB of memory depending on the GPU. Therefore you always have less usable memory than the actual size of the GPU. To see how much memory is actually used do `torch.ones(1).cuda()` and look at the memory usage.
+
+Therefore when you create memory maps with `max_memory` make sure to adjust the avaialble memory accordingly to avoid out-of-memory errors.
+
+</Tip>
+
+Additionally, if you do some additional operations with your outputs without placing them back on the CPU (for instance inside the `generate` method of Transformers) and if you placed your inputs on a GPU, that GPU will consume more memory than the others (Accelerate always place the output back to the device of the input). Therefore if you would like to optimize the maximum batch size and you have many GPUs, give the first GPU less memory. For example, with BLOOM-176B on 8x80 A100 setup the close to ideal map is:
+
+```python
+max_memory = {0: "30GIB", 1: "46GIB", 2: "46GIB", 3: "46GIB", 4: "46GIB", 5: "46GIB", 6: "46GIB", 7: "46GIB"}
+```
+as you can see we gave the remaining 7 GPUs ~50% more memory than GPU 0.
+
+If you opt to fully design the `device_map` yourself, it should be a dictionary with keys being module names of your model and values being a valid device identifier (for instance an integer for the GPUs) or `"cpu"` for CPU offload, `"disk"` for disk offload. The keys need to cover the whole model, you can then define your device map as you wish: for instance if your model has two blocks (let's say `block1` and `block2`) which each contain three linear layers (let's say `linear1`, `linear2` and `linear3`), a valid device map can be:
+
+```python
+device_map = {"block1": 0, "block2": 1}
+```
+
+another one that is valid could be:
+
+```python
+device_map = {"block1": 0, "block2.linear1": 0, "block2.linear2": 1, "block2.linear3": 1}
+```
+
+On the other hand, this one is not valid as it does not cover every parameter of the model:
+
+```python
+device_map = {"block1": 0, "block2.linear1": 1, "block2.linear2": 1}
+```
+
+<Tip>
+
+To be the most efficient, make sure your device map puts the parameters on the GPUs in a sequential manner (e.g. don't put one of the first weights on GPU 0, then weights on GPU 1 and the last weight back to GPU 0) to avoid making many transfers of data between the GPUs.
+
+</Tip>
+
 ## Limits and further development
 
 We are aware of the current limitations in the API:

I think including the ticks here may have broken the right navbar in the docs?
@@ -203,6 +203,50 @@ This only supports inference of your model, not training. Most of the computatio
 
 </Tip>
 
+## Designing a `device_map`
Arg!
@@ -203,6 +203,50 @@ This only supports inference of your model, not training. Most of the computatio
 
 </Tip>
 
+## Designing a `device_map`
I think here you may want to say "want to use more than ... for model weights". The addition is to be specific that it's just "model weights". I don't think accelerate controls activation and temp memory, does it?

So should it say - don't forget to leave memory for activations?

We really need to try to come up with the math that converts batch size into how much memory it'd need.

One way is to do it empirically by simply measuring with a given bs and reporting back to user a breakdown - this is how much was used for weights and how much for the activations/temps. But may not succeed on OOM.
@@ -203,6 +203,50 @@ This only supports inference of your model, not training. Most of the computatio
 
 </Tip>
 
+## Designing a device map
+
+You can let 🤗 Accelerate handle the device map computation by setting `device_map="auto"` or create one yourself, if you want more control over where each layer should go.
+
+<Tip>
+
+You can derive all sizes of the model (and thus compute a `device_map`) on a model that is on the meta device.
+
+</Tip>
+
+First note that you can limit the memory used on each GPU by using the `max_memory` argument (available in [`infer_auto_device_map`] and in all functions using it). When setting `max_memory`, you should pass along a dictionary containing the GPU identifiers (for instance `0`, `1` etc.) and the `"cpu"` key for the maximum RAM you want used for CPU offload. The values can either be an integer (in bytes) or a string representing a number with its unit, such as `"10GiB"` or `"10GB"`.
+
+Here is an example where we don't want to use more than 10GiB on each of two GPUs and no more than 30GiB of CPU RAM:
```suggestion
When a first allocation happens in PyTorch, it loads CUDA kernels which take about 1-2GB of memory depending on the GPU. Therefore you always have less usable memory than the actual size of the GPU. To see how much memory is actually used do `torch.ones(1).cuda()` and look at the memory usage.

Therefore when you create memory maps with `max_memory` make sure to adjust the avaialble memory accordingly to avoid out-of-memory errors.
```

if you're open to a more detailed explanation of what's going on and how to measure it.
@@ -203,6 +203,50 @@ This only supports inference of your model, not training. Most of the computatio
 
 </Tip>
 
+## Designing a device map
+
+You can let 🤗 Accelerate handle the device map computation by setting `device_map="auto"` or create one yourself, if you want more control over where each layer should go.
+
+<Tip>
+
+You can derive all sizes of the model (and thus compute a `device_map`) on a model that is on the meta device.
+
+</Tip>
+
+First note that you can limit the memory used on each GPU by using the `max_memory` argument (available in [`infer_auto_device_map`] and in all functions using it). When setting `max_memory`, you should pass along a dictionary containing the GPU identifiers (for instance `0`, `1` etc.) and the `"cpu"` key for the maximum RAM you want used for CPU offload. The values can either be an integer (in bytes) or a string representing a number with its unit, such as `"10GiB"` or `"10GB"`.
+
+Here is an example where we don't want to use more than 10GiB on each of two GPUs and no more than 30GiB of CPU RAM:
+
+```python
+from accelerate import infer_auto_device_map
+
+device_map = infer_auto_device_map(my_model, max_memory={0: "10GiB", 1: "10GiB", "cpu": "30GiB"})
+```
+
+<Tip warning={true}>
+
+PyTorch takes some space on a GPU, which can be up to 2GB. You should account for this when passing along `max_memory` to avoid out-of-memory errors.
Should it recommend that one puts the params in a sequential order and not manually organize them so that they fit better? e.g. putting a huge 176 embedding somewhere with a few tiny weights would balance things nicer, but it'd be detrimental to performance.
@@ -203,6 +203,50 @@ This only supports inference of your model, not training. Most of the computatio
 
 </Tip>
 
+## Designing a device map
+
+You can let 🤗 Accelerate handle the device map computation by setting `device_map="auto"` or create one yourself, if you want more control over where each layer should go.
+
+<Tip>
+
+You can derive all sizes of the model (and thus compute a `device_map`) on a model that is on the meta device.
+
+</Tip>
+
+First note that you can limit the memory used on each GPU by using the `max_memory` argument (available in [`infer_auto_device_map`] and in all functions using it). When setting `max_memory`, you should pass along a dictionary containing the GPU identifiers (for instance `0`, `1` etc.) and the `"cpu"` key for the maximum RAM you want used for CPU offload. The values can either be an integer (in bytes) or a string representing a number with its unit, such as `"10GiB"` or `"10GB"`.
+
+Here is an example where we don't want to use more than 10GiB on each of two GPUs and no more than 30GiB of CPU RAM:
+
+```python
+from accelerate import infer_auto_device_map
+
+device_map = infer_auto_device_map(my_model, max_memory={0: "10GiB", 1: "10GiB", "cpu": "30GiB"})
+```
+
+<Tip warning={true}>
+
+PyTorch takes some space on a GPU, which can be up to 2GB. You should account for this when passing along `max_memory` to avoid out-of-memory errors.
+
+</Tip>
+
+If you opt to fully design the `device_map` yourself, it should be a dictionary with keys being module names of your model and values being a valid device identifier (for instance an integer for the GPUs) or `"cpu"` for CPU offload, `"disk"` for disk offload. The keys need to cover the whole model, you can then define your device map as you wish: for instance if your model has two blocks (let's say `block1` and `block2`) which each contain three linear layers (let's say `linear1`, `linear2` and `linear3`), a valid device map can be:
+
+```python
+device_map = {"block1": 0, "block2": 1}
+```
+
+another one that is valid could be:
+
+```python
+device_map = {"block1": 0, "block2.linear1": 0, "block2.linear2": 1, "block2.linear3": 1}
+```
+
+On the other hand, this one is not valid as it does not cover every parameter of the model:
+
+```python
+device_map = {"block1": 0, "block2.linear1": 1, "block2.linear2": 1}
+```
+
````suggestion

Additionally, it's important to know that the first GPU consumes more memory than the rest of the GPUs for managing activations, therefore if you would like to optimize the maximum batch size and you have many GPUs. Give the first GPU less memory. For example, with BLOOM-176B on 8x80 A100 setup the close to ideal map is:
```
    return {0: '30GIB', 1: '46GIB', 2: '46GIB', 3: '46GIB', 4: '46GIB', 5: '46GIB', 6: '46GIB', 7: '46GIB'}
```
as you can see we gave the remaining 7 GPUs ~50% more memory than GPU 0.
````
@@ -203,6 +203,50 @@ This only supports inference of your model, not training. Most of the computatio
 
 </Tip>
 
+## Designing a device map
+
+You can let 🤗 Accelerate handle the device map computation by setting `device_map="auto"` or create one yourself, if you want more control over where each layer should go.
+
+<Tip>
+
+You can derive all sizes of the model (and thus compute a `device_map`) on a model that is on the meta device.
+
+</Tip>
+
+First note that you can limit the memory used on each GPU by using the `max_memory` argument (available in [`infer_auto_device_map`] and in all functions using it). When setting `max_memory`, you should pass along a dictionary containing the GPU identifiers (for instance `0`, `1` etc.) and the `"cpu"` key for the maximum RAM you want used for CPU offload. The values can either be an integer (in bytes) or a string representing a number with its unit, such as `"10GiB"` or `"10GB"`.
+
+Here is an example where we don't want to use more than 10GiB on each of two GPUs and no more than 30GiB of CPU RAM:
+
+```python
+from accelerate import infer_auto_device_map
+
+device_map = infer_auto_device_map(my_model, max_memory={0: "10GiB", 1: "10GiB", "cpu": "30GiB"})
+```
+
+<Tip warning={true}>
+
+PyTorch takes some space on a GPU, which can be up to 2GB. You should account for this when passing along `max_memory` to avoid out-of-memory errors.
+
+</Tip>
+
This part is only for `generate`. Will rephrase a bit after accepting your suggestion.
@@ -203,6 +203,50 @@ This only supports inference of your model, not training. Most of the computatio
 
 </Tip>
 
+## Designing a device map
+
+You can let 🤗 Accelerate handle the device map computation by setting `device_map="auto"` or create one yourself, if you want more control over where each layer should go.
+
+<Tip>
+
+You can derive all sizes of the model (and thus compute a `device_map`) on a model that is on the meta device.
+
+</Tip>
+
+First note that you can limit the memory used on each GPU by using the `max_memory` argument (available in [`infer_auto_device_map`] and in all functions using it). When setting `max_memory`, you should pass along a dictionary containing the GPU identifiers (for instance `0`, `1` etc.) and the `"cpu"` key for the maximum RAM you want used for CPU offload. The values can either be an integer (in bytes) or a string representing a number with its unit, such as `"10GiB"` or `"10GB"`.
+
+Here is an example where we don't want to use more than 10GiB on each of two GPUs and no more than 30GiB of CPU RAM:
+
+```python
+from accelerate import infer_auto_device_map
+
+device_map = infer_auto_device_map(my_model, max_memory={0: "10GiB", 1: "10GiB", "cpu": "30GiB"})
+```
+
+<Tip warning={true}>
+
+PyTorch takes some space on a GPU, which can be up to 2GB. You should account for this when passing along `max_memory` to avoid out-of-memory errors.
+
+</Tip>
+
Good point!
@@ -203,6 +203,50 @@ This only supports inference of your model, not training. Most of the computatio
 
 </Tip>
 
+## Designing a device map
+
+You can let 🤗 Accelerate handle the device map computation by setting `device_map="auto"` or create one yourself, if you want more control over where each layer should go.
+
+<Tip>
+
+You can derive all sizes of the model (and thus compute a `device_map`) on a model that is on the meta device.
+
+</Tip>
+
+First note that you can limit the memory used on each GPU by using the `max_memory` argument (available in [`infer_auto_device_map`] and in all functions using it). When setting `max_memory`, you should pass along a dictionary containing the GPU identifiers (for instance `0`, `1` etc.) and the `"cpu"` key for the maximum RAM you want used for CPU offload. The values can either be an integer (in bytes) or a string representing a number with its unit, such as `"10GiB"` or `"10GB"`.
+
+Here is an example where we don't want to use more than 10GiB on each of two GPUs and no more than 30GiB of CPU RAM:
+
+```python
+from accelerate import infer_auto_device_map
+
+device_map = infer_auto_device_map(my_model, max_memory={0: "10GiB", 1: "10GiB", "cpu": "30GiB"})
+```
+
+<Tip warning={true}>
+
+PyTorch takes some space on a GPU, which can be up to 2GB. You should account for this when passing along `max_memory` to avoid out-of-memory errors.
+
+</Tip>
+
+If you opt to fully design the `device_map` yourself, it should be a dictionary with keys being module names of your model and values being a valid device identifier (for instance an integer for the GPUs) or `"cpu"` for CPU offload, `"disk"` for disk offload. The keys need to cover the whole model, you can then define your device map as you wish: for instance if your model has two blocks (let's say `block1` and `block2`) which each contain three linear layers (let's say `linear1`, `linear2` and `linear3`), a valid device map can be:
+
+```python
+device_map = {"block1": 0, "block2": 1}
+```
+
+another one that is valid could be:
+
+```python
+device_map = {"block1": 0, "block2.linear1": 0, "block2.linear2": 1, "block2.linear3": 1}
+```
+
+On the other hand, this one is not valid as it does not cover every parameter of the model:
+
+```python
+device_map = {"block1": 0, "block2.linear1": 1, "block2.linear2": 1}
+```
+
_The documentation is not available anymore as the PR was closed or merged._
Should have addressed all your comments @stas00, let me know if all looks right to you and I'll merge :-)
https://api.github.com/repos/huggingface/accelerate/pulls/530
diff --git a/src/accelerate/big_modeling.py b/src/accelerate/big_modeling.py
index 85b1f3237..907dfdac4 100644
--- a/src/accelerate/big_modeling.py
+++ b/src/accelerate/big_modeling.py
@@ -260,7 +260,7 @@ def load_checkpoint_and_dispatch(
     offload_folder: Optional[Union[str, os.PathLike]] = None,
     offload_buffers: bool = False,
     dtype: Optional[Union[str, torch.dtype]] = None,
-    offload_state_dict: bool = False,
+    offload_state_dict: Optional[bool] = None,
     preload_module_classes: Optional[List[str]] = None,
 ):
     """
@@ -292,9 +292,10 @@ def load_checkpoint_and_dispatch(
             well as the parameters.
         dtype (`str` or `torch.dtype`, *optional*):
             If provided, the weights will be converted to that type when loaded.
-        offload_state_dict (`bool`, *optional*, defaults to `False`):
+        offload_state_dict (`bool`, *optional*):
             If `True`, will temporarily offload the CPU state dict on the hard drive to avoig getting out of CPU RAM if
-            the weight of the CPU state dict + the biggest shard does not fit.
+            the weight of the CPU state dict + the biggest shard does not fit. Will default to `True` if the device map
+            picked contains `"disk"` values.
         preload_module_classes (`List[str]`, *optional*):
             A list of classes whose instances should load all their weights (even in the submodules) at the beginning
             of the forward. This should only be used for classes that have submodules which are registered but not
@@ -305,6 +306,8 @@ def load_checkpoint_and_dispatch(
         device_map = infer_auto_device_map(
             model, max_memory=max_memory, no_split_module_classes=no_split_module_classes, dtype=dtype
         )
+    if offload_state_dict is None and "disk" in device_map.values():
+        offload_state_dict = True
     load_checkpoint_in_model(
         model,
         checkpoint,

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/529
diff --git a/src/accelerate/commands/test.py b/src/accelerate/commands/test.py
index 379233679..41da75596 100644
--- a/src/accelerate/commands/test.py
+++ b/src/accelerate/commands/test.py
@@ -43,7 +43,7 @@ def test_command_parser(subparsers=None):
 
 
 def test_command(args):
-    script_name = os.path.sep.join(__file__.split(os.path.sep)[:-2] + ["test_utils", "test_script.py"])
+    script_name = os.path.sep.join(__file__.split(os.path.sep)[:-2] + ["test_utils", "scripts", "test_script.py"])
 
     test_args = f"""
         --config_file={args.config_file} {script_name}

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/528
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 9b6e7f069..e781a5cea 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -272,7 +272,7 @@ def __init__(
                 raise ValueError(err.format(mode="bf16", requirement="PyTorch >= 1.10 and a supported device."))
 
             # Only on the GPU do we care about scaling the gradients
-            if torch.cuda.is_available():
+            if torch.cuda.is_available() and self.device.type != "cpu":
                 kwargs = self.scaler_handler.to_kwargs() if self.scaler_handler is not None else {}
                 self.scaler = torch.cuda.amp.GradScaler(**kwargs)
 
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index cffe9990b..bda9e6b38 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -256,7 +256,14 @@ def simple_launcher(args):
     cmd.extend(args.training_script_args)
 
     current_env = os.environ.copy()
-    current_env["USE_CPU"] = str(args.cpu)
+    current_env["USE_CPU"] = str(args.cpu or args.use_cpu)
+    if args.num_machines > 1:
+        current_env["MASTER_ADDR"] = args.main_process_ip
+        current_env["MASTER_PORT"] = str(args.main_process_port)
+    elif args.num_processes > 1:
+        current_env["MASTER_ADDR"] = args.main_process_ip if args.main_process_ip is not None else "127.0.0.1"
+        current_env["MASTER_PORT"] = str(args.main_process_port) if args.main_process_port is not None else "29500"
+
     try:
         mixed_precision = PrecisionType(args.mixed_precision.lower())
     except ValueError:
@@ -269,6 +276,7 @@ def simple_launcher(args):
         mixed_precision = "fp16"
 
     current_env["MIXED_PRECISION"] = str(mixed_precision)
+    current_env["OMP_NUM_THREADS"] = str(args.num_cpu_threads_per_process)
 
     process = subprocess.Popen(cmd, env=current_env)
     process.wait()

Should this be under and `if args.num_machines > 1`? It seems related to multinode training.
@@ -257,6 +257,8 @@ def simple_launcher(args):
 
     current_env = os.environ.copy()
     current_env["USE_CPU"] = str(args.cpu)
+    current_env["MASTER_ADDR"] = args.main_process_ip
+    current_env["MASTER_PORT"] = str(args.main_process_port)
thanks for the suggestion, a patch is uploaded to fix it
@@ -257,6 +257,8 @@ def simple_launcher(args):
 
     current_env = os.environ.copy()
     current_env["USE_CPU"] = str(args.cpu)
+    current_env["MASTER_ADDR"] = args.main_process_ip
+    current_env["MASTER_PORT"] = str(args.main_process_port)
If they are needed, I would rather use the ones set par the user than hardcoding a default, which may create hard-to-debug issues for users.
@@ -257,6 +257,12 @@ def simple_launcher(args):
 
     current_env = os.environ.copy()
     current_env["USE_CPU"] = str(args.cpu)
+    if args.num_machines > 1:
+        current_env["MASTER_ADDR"] = args.main_process_ip
+        current_env["MASTER_PORT"] = str(args.main_process_port)
+    elif args.num_processes > 1:
+        current_env["MASTER_ADDR"] = "127.0.0.1"
+        current_env["MASTER_PORT"] = "29500"
hi, sgugger, now in the accelerate config, if num_machines == 1, the master ip and port is not set. and if num_process > 1 and DDP is used in one machine, we still need the master ip and port
@@ -257,6 +257,12 @@ def simple_launcher(args):
 
     current_env = os.environ.copy()
     current_env["USE_CPU"] = str(args.cpu)
+    if args.num_machines > 1:
+        current_env["MASTER_ADDR"] = args.main_process_ip
+        current_env["MASTER_PORT"] = str(args.main_process_port)
+    elif args.num_processes > 1:
+        current_env["MASTER_ADDR"] = "127.0.0.1"
+        current_env["MASTER_PORT"] = "29500"
But they could be set in the args passed, so maybe only use this default if the args is empty/None?
@@ -257,6 +257,12 @@ def simple_launcher(args):
 
     current_env = os.environ.copy()
     current_env["USE_CPU"] = str(args.cpu)
+    if args.num_machines > 1:
+        current_env["MASTER_ADDR"] = args.main_process_ip
+        current_env["MASTER_PORT"] = str(args.main_process_port)
+    elif args.num_processes > 1:
+        current_env["MASTER_ADDR"] = "127.0.0.1"
+        current_env["MASTER_PORT"] = "29500"
done, thanks for the suggestion
@@ -257,6 +257,12 @@ def simple_launcher(args):
 
     current_env = os.environ.copy()
     current_env["USE_CPU"] = str(args.cpu)
+    if args.num_machines > 1:
+        current_env["MASTER_ADDR"] = args.main_process_ip
+        current_env["MASTER_PORT"] = str(args.main_process_port)
+    elif args.num_processes > 1:
+        current_env["MASTER_ADDR"] = "127.0.0.1"
+        current_env["MASTER_PORT"] = "29500"
@sgugger @yao-matrix please help review
_The documentation is not available anymore as the PR was closed or merged._
Thansk again for your contribution!
https://api.github.com/repos/huggingface/accelerate/pulls/527
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 2dec26a98..9b6e7f069 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -161,7 +161,10 @@ def __init__(
         kwargs_handlers: Optional[List[KwargsHandler]] = None,
     ):
         self.logging_dir = logging_dir
-        self.log_with = filter_trackers(log_with, self.logging_dir)
+        trackers = filter_trackers(log_with, self.logging_dir)
+        if len(trackers) < 1 and log_with is not None:
+            warnings.warn(f"`log_with={log_with}` was passed but no supported trackers are currently installed.")
+        self.log_with = trackers
 
         if mixed_precision is not None:
             mixed_precision = str(mixed_precision)

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/524
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 2dec26a98..596998f1e 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -235,6 +235,13 @@ def __init__(
             **kwargs,
         )
 
+        if (
+            (mixed_precision != "bf16")
+            and getattr(self.state, "downcast_bfloat", False)
+            and (self.state.distributedType != DistributedType.TPU)
+        ):
+            raise ValueError("Can only use `downcast_bf16` when using `mixed_precision='bf16'` and on a TPU")
+
         if gradient_accumulation_steps > 1:
             if self.state.distributed_type == DistributedType.TPU:
                 raise NotImplementedError(
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index 0051d899a..8a4f38f9e 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -260,11 +260,18 @@ def get_cluster_input():
     else:
         mixed_precision = "no"
 
+    downcast_bf16 = "no"
+    if distributed_type == DistributedType.TPU and mixed_precision == "bf16":
+        downcast_bf16 = _ask_field(
+            "Should `torch.float` be cast as `bfloat16` and `torch.double` remain `float32` on TPUs?", default="no"
+        )
+
     return ClusterConfig(
         compute_environment=ComputeEnvironment.LOCAL_MACHINE,
         distributed_type=distributed_type,
         num_processes=num_processes,
         mixed_precision=mixed_precision,
+        downcast_bf16=downcast_bf16,
         machine_rank=machine_rank,
         num_machines=num_machines,
         main_process_ip=main_process_ip,
diff --git a/src/accelerate/commands/config/config_args.py b/src/accelerate/commands/config/config_args.py
index 1ef68200b..11ca0ca90 100644
--- a/src/accelerate/commands/config/config_args.py
+++ b/src/accelerate/commands/config/config_args.py
@@ -143,6 +143,8 @@ class ClusterConfig(BaseConfig):
     deepspeed_config: dict = None
     # args for fsdp
     fsdp_config: dict = None
+    # args for TPU
+    downcast_bf16: bool = False
 
     def __post_init__(self):
         if self.deepspeed_config is None:
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index cffe9990b..75c6c7b1c 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -34,6 +34,7 @@
     get_launch_prefix,
     is_deepspeed_available,
     is_sagemaker_available,
+    patch_environment,
 )
 from accelerate.utils.constants import DEEPSPEED_MULTINODE_LAUNCHERS
 from accelerate.utils.dataclasses import SageMakerDistributedType
@@ -198,6 +199,11 @@ def launch_command_parser(subparsers=None):
         default=None,
         help="The name of the main function to be executed in your script (only for TPU training).",
     )
+    parser.add_argument(
+        "--downcast_bf16",
+        action="store_true",
+        help="Whether when using bf16 precision on TPUs if both float and double tensors are cast to bfloat16 or if double tensors remain as float32",
+    )
     parser.add_argument(
         "-m",
         "--module",
@@ -425,9 +431,19 @@ def deepspeed_launcher(args):
 def tpu_launcher(args):
     import torch_xla.distributed.xla_multiprocessing as xmp
 
+    current_env = {}
+
     if args.no_python:
         raise ValueError("--no_python cannot be used with TPU launcher")
 
+    if args.mixed_precision == "bf16":
+        if args.downcast_bf16:
+            current_env["XLA_USE_BF16"] = "0"
+            current_env["XLA_DOWNCAST_BF16"] = "1"
+        else:
+            current_env["XLA_USE_BF16"] = "1"
+            current_env["XLA_DOWNCAST_BF16"] = "0"
+
     if args.module:
         mod_name = args.training_script
     else:
@@ -447,7 +463,8 @@ def tpu_launcher(args):
     sys.argv = [mod.__file__] + args.training_script_args
 
     main_function = getattr(mod, args.main_training_function)
-    xmp.spawn(PrepareForLaunch(main_function), args=(), nprocs=args.num_processes)
+    with patch_environment(**current_env):
+        xmp.spawn(PrepareForLaunch(main_function), args=(), nprocs=args.num_processes)
 
 
 def _convert_nargs_to_dict(nargs: List[str]) -> Dict[str, str]:
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index 4482341f3..3fcf7d2e2 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -115,7 +115,14 @@ def __init__(
                 self.local_process_index = xm.get_local_ordinal()
                 self.device = xm.xla_device()
                 if mixed_precision == "bf16":
-                    os.environ["XLA_USE_BF16"] = str(1)
+                    if os.environ.get("DOWNCAST_BF16"):
+                        os.environ["XLA_USE_BF16"] = str(0)
+                        os.environ["XLA_DOWNCAST_BF16"] = str(1)
+                        self.downcast_bfloat = True
+                    else:
+                        os.environ["XLA_USE_BF16"] = str(1)
+                        os.environ["XLA_DOWNCAST_BF16"] = str(0)
+                        self.downcast_bfloat = False
                 self.mixed_precision = mixed_precision
             elif os.environ.get("USE_DEEPSPEED", "false") == "true" and not cpu:
                 assert (

Very much against adding a new Accelerator parameter for such a niche feature. It should be in the config only.
@@ -126,6 +126,9 @@ class Accelerator:
             accept implementations of `GeneralTracker` for custom trackers, and can be combined with `"all"`.
         logging_dir (`str`, `os.PathLike`, *optional*):
             A path to a directory for storing logs of locally-compatible loggers.
+        downcast_bf16 (`bool`, *optional*):
+            If set to `True` and using `mixed_precision="bf16"` then `torch.float` will become `bfloat16` and
+            `torch.double` will remain `float32` on TPUs.
Doesn't seem to be able to be None, so I'm confused about this change in typing.
@@ -156,8 +159,9 @@ def __init__(
         rng_types: Optional[List[Union[str, RNGType]]] = None,
         log_with: Optional[List[Union[str, LoggerType, GeneralTracker]]] = None,
         logging_dir: Optional[Union[str, os.PathLike]] = None,
+        downcast_bf16: Optional[bool] = False,
         dispatch_batches: Optional[bool] = None,
-        step_scheduler_with_optimizer: bool = True,
+        step_scheduler_with_optimizer: Optional[bool] = True,
Okay, will adjust on Monday!

(This should be in the release after this next one anyways)
@@ -126,6 +126,9 @@ class Accelerator:
             accept implementations of `GeneralTracker` for custom trackers, and can be combined with `"all"`.
         logging_dir (`str`, `os.PathLike`, *optional*):
             A path to a directory for storing logs of locally-compatible loggers.
+        downcast_bf16 (`bool`, *optional*):
+            If set to `True` and using `mixed_precision="bf16"` then `torch.float` will become `bfloat16` and
+            `torch.double` will remain `float32` on TPUs.
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/523
diff --git a/examples/by_feature/fsdp_with_peak_mem_tracking.py b/examples/by_feature/fsdp_with_peak_mem_tracking.py
index eedbf3392..8da906ee9 100644
--- a/examples/by_feature/fsdp_with_peak_mem_tracking.py
+++ b/examples/by_feature/fsdp_with_peak_mem_tracking.py
@@ -162,6 +162,7 @@ def collate_fn(examples):
     # New Code #
     # For FSDP feature, it is highly recommended and efficient to prepare the model before creating optimizer
     model = accelerator.prepare(model)
+    accelerator.print(model)
 
     # Instantiate optimizer
     # New Code #
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 2dec26a98..c3182e441 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -262,8 +262,13 @@ def __init__(
             if not torch.cuda.is_available():
                 raise ValueError(err.format(mode="fp16", requirement="a GPU"))
             kwargs = self.scaler_handler.to_kwargs() if self.scaler_handler is not None else {}
-            self.scaler = torch.cuda.amp.GradScaler(**kwargs)
-        elif self.state.mixed_precision == "bf16":
+            if self.distributed_type == DistributedType.FSDP:
+                from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler
+
+                self.scaler = ShardedGradScaler(**kwargs)
+            else:
+                self.scaler = torch.cuda.amp.GradScaler(**kwargs)
+        elif self.state.mixed_precision == "bf16" and self.distributed_type != DistributedType.FSDP:
             self.native_amp = is_bf16_available(True)
             if mixed_precision == "bf16" and not self.native_amp and not is_tpu_available():
                 raise ValueError(err.format(mode="bf16", requirement="PyTorch >= 1.10 and a supported device."))
@@ -566,6 +571,7 @@ def prepare_model(self, model):
             # Check if the model is already a FSDP model due to `Manual Wrapping` and if so,
             # don't wrap it again
             if type(model) != FSDP:
+                self.state.fsdp_plugin.set_auto_wrap_policy(model)
                 fsdp_plugin = self.state.fsdp_plugin
                 model = FSDP(
                     model,
@@ -573,6 +579,7 @@ def prepare_model(self, model):
                     cpu_offload=fsdp_plugin.cpu_offload,
                     auto_wrap_policy=fsdp_plugin.auto_wrap_policy,
                     backward_prefetch=fsdp_plugin.backward_prefetch,
+                    mixed_precision=fsdp_plugin.mixed_precision_policy,
                     ignored_modules=fsdp_plugin.ignored_modules,
                 )
                 if not fsdp_plugin.cpu_offload.offload_params:
@@ -817,6 +824,7 @@ def clip_grad_norm_(self, parameters, max_norm, norm_type=2):
         Should be used in place of `torch.nn.utils.clip_grad_norm_`.
         """
         if self.distributed_type == DistributedType.FSDP:
+            self.unscale_gradients()
             parameters = [p for p in parameters]
             for model in self._models:
                 if parameters == [p for p in model.parameters()]:
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index 0051d899a..d9dd80094 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -15,7 +15,12 @@
 # limitations under the License.
 
 from ...utils import ComputeEnvironment, DistributedType, is_deepspeed_available, is_transformers_available
-from ...utils.constants import DEEPSPEED_MULTINODE_LAUNCHERS
+from ...utils.constants import (
+    DEEPSPEED_MULTINODE_LAUNCHERS,
+    FSDP_AUTO_WRAP_POLICY,
+    FSDP_BACKWARD_PREFETCH,
+    FSDP_SHARDING_STRATEGY,
+)
 from .config_args import ClusterConfig
 from .config_utils import _ask_field, _convert_distributed_mode, _convert_yes_no_to_bool
 
@@ -201,8 +206,12 @@ def get_cluster_input():
         if use_fsdp:
             distributed_type = DistributedType.FSDP
         if distributed_type == DistributedType.FSDP:
+            sharding_strategy_query = "What should be your sharding strategy ("
+            for i, strategy in enumerate(FSDP_SHARDING_STRATEGY):
+                sharding_strategy_query += f"[{i+1}] {strategy}, "
+            sharding_strategy_query = sharding_strategy_query[:-2] + ")? [1]: "
             fsdp_config["sharding_strategy"] = _ask_field(
-                "What should be your sharding strategy ([1] FULL_SHARD, [2] SHARD_GRAD_OP)? [1]: ",
+                sharding_strategy_query,
                 lambda x: int(x),
                 default=1,
             )
@@ -212,10 +221,34 @@ def get_cluster_input():
                 default=False,
                 error_message="Please enter yes or no.",
             )
-            fsdp_config["min_num_params"] = _ask_field(
-                "What should be your FSDP's minimum number of parameters for Default Auto Wrapping Policy? [1e8]: ",
-                lambda x: int(x),
-                default=1e8,
+            fsdp_wrap_query = "What should be your auto wrap policy ("
+            for i, wrap_policy in enumerate(FSDP_AUTO_WRAP_POLICY):
+                fsdp_wrap_query += f"[{i}] {wrap_policy}, "
+            fsdp_wrap_query = fsdp_wrap_query[:-2] + ")? [0]: "
+            fsdp_config["fsdp_auto_wrap_policy"] = _ask_field(
+                fsdp_wrap_query,
+                lambda x: FSDP_AUTO_WRAP_POLICY[int(x)],
+                default=FSDP_AUTO_WRAP_POLICY[0],
+            )
+            if fsdp_config["fsdp_auto_wrap_policy"] == FSDP_AUTO_WRAP_POLICY[0]:
+                fsdp_config["transformer_layer_cls_to_wrap"] = _ask_field(
+                    "What is the transformer layer class name (case-sensitive) to wrap ,e.g, `BertLayer`, `GPTJBlock`, `T5Block` ...? : ",
+                    lambda x: str(x),
+                )
+            elif fsdp_config["fsdp_auto_wrap_policy"] == FSDP_AUTO_WRAP_POLICY[1]:
+                fsdp_config["min_num_params"] = _ask_field(
+                    "What should be your FSDP's minimum number of parameters for Default Auto Wrapping Policy? [1e8]: ",
+                    lambda x: int(x),
+                    default=1e8,
+                )
+            fsdp_backward_prefetch_query = "What should be your FSDP's backward prefetch policy ("
+            for i, backward_prefetch_policy in enumerate(FSDP_BACKWARD_PREFETCH):
+                fsdp_backward_prefetch_query += f"[{i}] {backward_prefetch_policy}, "
+            fsdp_backward_prefetch_query = fsdp_backward_prefetch_query[:-2] + ")? [0]: "
+            fsdp_config["fsdp_backward_prefetch_policy"] = _ask_field(
+                fsdp_backward_prefetch_query,
+                lambda x: FSDP_BACKWARD_PREFETCH[int(x)],
+                default=FSDP_BACKWARD_PREFETCH[0],
             )
 
     if distributed_type == DistributedType.TPU:
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index cffe9990b..5bfd9894d 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -158,6 +158,25 @@ def launch_command_parser(subparsers=None):
         default=1,
         help="FSDP's Sharding Strategy. (useful only when `use_fsdp` flag is passed).",
     )
+    parser.add_argument(
+        "--fsdp_auto_wrap_policy",
+        type=str,
+        default=None,
+        help="FSDP's auto wrap policy. (useful only when `use_fsdp` flag is passed).",
+    )
+    parser.add_argument(
+        "--transformer_layer_cls_to_wrap",
+        default=None,
+        type=str,
+        help="Transformer layer class name (case-sensitive) to wrap ,e.g, `BertLayer`, `GPTJBlock`, `T5Block` .... "
+        "(useful only when `use_fsdp` flag is passed).",
+    )
+    parser.add_argument(
+        "--fsdp_backward_prefetch_policy",
+        default=None,
+        type=str,
+        help="FSDP's backward prefetch policy. (useful only when `use_fsdp` flag is passed).",
+    )
     parser.add_argument(
         "--tpu", default=False, action="store_true", help="Whether or not this should launch a TPU training."
     )
@@ -322,9 +341,12 @@ def multi_gpu_launcher(args):
     current_env["MIXED_PRECISION"] = str(mixed_precision)
     if args.use_fsdp:
         current_env["USE_FSDP"] = "true"
+        current_env["FSDP_AUTO_WRAP_POLICY"] = str(args.fsdp_auto_wrap_policy)
+        current_env["FSDP_TRANSFORMER_CLS_TO_WRAP"] = str(args.transformer_layer_cls_to_wrap)
         current_env["FSDP_OFFLOAD_PARAMS"] = str(args.offload_params).lower()
         current_env["FSDP_MIN_NUM_PARAMS"] = str(args.min_num_params)
         current_env["FSDP_SHARDING_STRATEGY"] = str(args.sharding_strategy)
+        current_env["FSDP_BACKWARD_PREFETCH"] = str(args.fsdp_backward_prefetch_policy)
     current_env["OMP_NUM_THREADS"] = str(args.num_cpu_threads_per_process)
     process = subprocess.Popen(cmd, env=current_env)
     process.wait()
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index 4482341f3..38cd40bc7 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -146,9 +146,7 @@ def __init__(
                 if os.environ.get("USE_FSDP", "false") == "true":
                     self.distributed_type = DistributedType.FSDP
                     if self.mixed_precision != "no":
-                        raise ValueError(
-                            "Mixed precision is currently not supported for FSDP. Please set `mixed_precision` to `no`."
-                        )
+                        fsdp_plugin.set_mixed_precision(self.mixed_precision)
                     self.fsdp_plugin = fsdp_plugin
             elif get_int_from_env(["PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE", "WORLD_SIZE"], 1) > 1:
                 self.distributed_type = DistributedType.MULTI_CPU
diff --git a/src/accelerate/utils/constants.py b/src/accelerate/utils/constants.py
index 5e31a96c9..00556cad6 100644
--- a/src/accelerate/utils/constants.py
+++ b/src/accelerate/utils/constants.py
@@ -24,6 +24,9 @@
 SAGEMAKER_PYTHON_VERSION = "py38"
 SAGEMAKER_TRANSFORMERS_VERSION = "4.17.0"
 SAGEMAKER_PARALLEL_EC2_INSTANCES = ["ml.p3.16xlarge", "ml.p3dn.24xlarge", "ml.p4dn.24xlarge"]
+FSDP_SHARDING_STRATEGY = ["FULL_SHARD", "SHARD_GRAD_OP", "NO_SHARD"]
+FSDP_AUTO_WRAP_POLICY = ["TRANSFORMER_BASED_WRAP", "SIZE_BASED_WRAP", "NO_WRAP"]
+FSDP_BACKWARD_PREFETCH = ["BACKWARD_PRE", "BACKWARD_POST", "NO_PREFETCH"]
 DEEPSPEED_MULTINODE_LAUNCHERS = ["pdsh", "standard", "openmpi", "mvapich"]
 
 STR_OPERATION_TO_FUNC = {">": op.gt, ">=": op.ge, "==": op.eq, "!=": op.ne, "<=": op.le, "<": op.lt}
diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py
index 28fc3826d..8c8bdf106 100644
--- a/src/accelerate/utils/dataclasses.py
+++ b/src/accelerate/utils/dataclasses.py
@@ -28,6 +28,8 @@
 
 import torch
 
+from .constants import FSDP_AUTO_WRAP_POLICY, FSDP_BACKWARD_PREFETCH
+
 
 class KwargsHandler:
     """
@@ -418,22 +420,35 @@ class FullyShardedDataParallelPlugin:
 
     sharding_strategy: "typing.Any" = field(
         default=None,
-        metadata={"help": "Possible options are [1] FULL_SHARD, [2] SHARD_GRAD_OP"},
+        metadata={
+            "help": "FSDP Sharding Strategy of type `torch.distributed.fsdp.fully_sharded_data_parallel.ShardingStrategy`"
+        },
     )
     backward_prefetch: "typing.Any" = field(
         default=None,
-        metadata={"help": "Possible options are [1] BACKWARD_PRE, [2] BACKWARD_POST"},
+        metadata={
+            "help": "FSDP Backward Prefetch of type `torch.distributed.fsdp.fully_sharded_data_parallel.BackwardPrefetch`"
+        },
     )
-    auto_wrap_policy: "typing.Any" = field(
+    mixed_precision_policy: "typing.Any" = field(
         default=None,
-        metadata={"help": "A callable specifying a policy to recursively wrap layers with FSDP"},
+        metadata={
+            "help": "A config to enable mixed precision training with FullyShardedDataParallel. "
+            "The 3 flags that are set are `param_dtype`, `reduce_dtype`, `buffer_dtype`. "
+            "Each flag expects `torch.dtype` as the value. "
+            "It is of type `torch.distributed.fsdp.fully_sharded_data_parallel.MixedPrecision`."
+        },
     )
-    cpu_offload: Optional[Callable] = field(
+    auto_wrap_policy: Optional[Callable] = field(
         default=None,
-        metadata={"help": "Decides Whether to offload parameters and gradients to CPU."},
+        metadata={"help": "A callable specifying a policy to recursively wrap layers with FSDP"},
     )
-    min_num_params: int = field(
-        default=None, metadata={"help": "FSDP's minimum number of parameters for Default Auto Wrapping."}
+    cpu_offload: "typing.Any" = field(
+        default=None,
+        metadata={
+            "help": "Decides Whether to offload parameters and gradients to CPU. "
+            "It is of type `torch.distributed.fsdp.fully_sharded_data_parallel.CPUOffload`."
+        },
     )
     ignored_modules: Optional[Iterable[torch.nn.Module]] = field(
         default=None,
@@ -441,8 +456,7 @@ class FullyShardedDataParallelPlugin:
     )
 
     def __post_init__(self):
-        from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload, ShardingStrategy
-        from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy
+        from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch, CPUOffload, ShardingStrategy
 
         if self.sharding_strategy is None:
             self.sharding_strategy = ShardingStrategy(int(os.environ.get("FSDP_SHARDING_STRATEGY", 1)))
@@ -453,11 +467,63 @@ def __post_init__(self):
             else:
                 self.cpu_offload = CPUOffload(offload_params=False)
 
-        if self.min_num_params is None:
-            self.min_num_params = int(os.environ.get("FSDP_MIN_NUM_PARAMS", 0))
+        if self.backward_prefetch is None:
+            prefetch_policy = os.environ.get("FSDP_BACKWARD_PREFETCH", FSDP_BACKWARD_PREFETCH[-1])
+            if prefetch_policy != FSDP_BACKWARD_PREFETCH[-1]:
+                self.backward_prefetch = BackwardPrefetch(FSDP_BACKWARD_PREFETCH.index(prefetch_policy) + 1)
+
+    @staticmethod
+    def get_module_class_from_name(module, name):
+        """
+        Gets a class from a module by its name.
+
+        Args:
+            module (`torch.nn.Module`): The module to get the class from.
+            name (`str`): The name of the class.
+        """
+        modules_children = list(module.children())
+        if module.__class__.__name__ == name:
+            return module.__class__
+        elif len(modules_children) == 0:
+            return
+        else:
+            for child_module in modules_children:
+                module_class = FullyShardedDataParallelPlugin.get_module_class_from_name(child_module, name)
+                if module_class is not None:
+                    return module_class
+
+    def set_auto_wrap_policy(self, model):
+        from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy, transformer_auto_wrap_policy
 
         if self.auto_wrap_policy is None:
-            if self.min_num_params > 0:
+            auto_wrap_policy = os.environ.get("FSDP_AUTO_WRAP_POLICY", FSDP_AUTO_WRAP_POLICY[-1])
+            if auto_wrap_policy == FSDP_AUTO_WRAP_POLICY[0]:
+                transformer_cls_to_wrap = os.environ.get("FSDP_TRANSFORMER_CLS_TO_WRAP", "")
+                transformer_cls_to_wrap = FullyShardedDataParallelPlugin.get_module_class_from_name(
+                    model, transformer_cls_to_wrap
+                )
+                if transformer_cls_to_wrap is None:
+                    raise Exception("Could not find the transformer layer class to wrap in the model.")
                 self.auto_wrap_policy = functools.partial(
-                    size_based_auto_wrap_policy, min_num_params=self.min_num_params
+                    transformer_auto_wrap_policy,
+                    # Transformer layer class to wrap
+                    transformer_layer_cls={transformer_cls_to_wrap},
                 )
+            elif auto_wrap_policy == FSDP_AUTO_WRAP_POLICY[1]:
+                min_num_params = int(os.environ.get("FSDP_MIN_NUM_PARAMS", 0))
+                if min_num_params > 0:
+                    self.auto_wrap_policy = functools.partial(
+                        size_based_auto_wrap_policy, min_num_params=min_num_params
+                    )
+
+    def set_mixed_precision(self, mixed_precision):
+        if mixed_precision == "fp16":
+            dtype = torch.float16
+        elif mixed_precision == "bf16":
+            dtype = torch.bfloat16
+        else:
+            raise ValueError(f"Unknown mixed precision value: {mixed_precision}")
+        from torch.distributed.fsdp.fully_sharded_data_parallel import MixedPrecision
+
+        if self.mixed_precision_policy is None:
+            self.mixed_precision_policy = MixedPrecision(param_dtype=dtype, reduce_dtype=dtype, buffer_dtype=dtype)
diff --git a/src/accelerate/utils/operations.py b/src/accelerate/utils/operations.py
index 7505c04ce..7b3c8de5a 100644
--- a/src/accelerate/utils/operations.py
+++ b/src/accelerate/utils/operations.py
@@ -217,7 +217,11 @@ def gather(tensor):
     """
     if AcceleratorState().distributed_type == DistributedType.TPU:
         return _tpu_gather(tensor, name="accelerate.utils.gather")
-    elif AcceleratorState().distributed_type in [DistributedType.DEEPSPEED, DistributedType.MULTI_GPU]:
+    elif AcceleratorState().distributed_type in [
+        DistributedType.DEEPSPEED,
+        DistributedType.MULTI_GPU,
+        DistributedType.FSDP,
+    ]:
         return _gpu_gather(tensor)
     elif AcceleratorState().distributed_type == DistributedType.MULTI_CPU:
         return _cpu_gather(tensor)

Is this supposed to stay here?
@@ -162,6 +162,7 @@ def collate_fn(examples):
     # New Code #
     # For FSDP feature, it is highly recommended and efficient to prepare the model before creating optimizer
     model = accelerator.prepare(model)
+    accelerator.print(model)
Just to be sure, this is in 1.12?
@@ -146,9 +146,7 @@ def __init__(
                 if os.environ.get("USE_FSDP", "false") == "true":
                     self.distributed_type = DistributedType.FSDP
                     if self.mixed_precision != "no":
-                        raise ValueError(
-                            "Mixed precision is currently not supported for FSDP. Please set `mixed_precision` to `no`."
-                        )
+                        fsdp_plugin.set_mixed_precision(self.mixed_precision)
Put the comment above so that black leaves everything on one line.
@@ -453,11 +467,57 @@ def __post_init__(self):
             else:
                 self.cpu_offload = CPUOffload(offload_params=False)
 
-        if self.min_num_params is None:
-            self.min_num_params = int(os.environ.get("FSDP_MIN_NUM_PARAMS", 0))
+    @staticmethod
+    def get_module_class_from_name(module, name):
+        """
+        Gets a class from a module by its name.
+
+        Args:
+            module (`torch.nn.Module`): The module to get the class from.
+            name (`str`): The name of the class.
+        """
+        modules_children = list(module.children())
+        if module.__class__.__name__ == name:
+            return module.__class__
+        elif len(modules_children) == 0:
+            return
+        else:
+            for child_module in modules_children:
+                module_class = FullyShardedDataParallelPlugin.get_module_class_from_name(child_module, name)
+                if module_class is not None:
+                    return module_class
+
+    def set_auto_wrap_policy(self, model):
+        from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy, transformer_auto_wrap_policy
 
         if self.auto_wrap_policy is None:
-            if self.min_num_params > 0:
+            auto_wrap_policy = os.environ.get("FSDP_AUTO_WRAP_POLICY", FSDP_AUTO_WRAP_POLICY[-1])
+            if auto_wrap_policy == FSDP_AUTO_WRAP_POLICY[0]:
+                transformer_cls_to_wrap = os.environ.get("FSDP_TRANSFORMER_CLS_TO_WRAP", "")
+                transformer_cls_to_wrap = FullyShardedDataParallelPlugin.get_module_class_from_name(
+                    model, transformer_cls_to_wrap
+                )
                 self.auto_wrap_policy = functools.partial(
-                    size_based_auto_wrap_policy, min_num_params=self.min_num_params
+                    transformer_auto_wrap_policy,
+                    transformer_layer_cls={
+                        transformer_cls_to_wrap,  # < ---- Transformer layer class to wrap
+                    },
Yes, to help users observe which layers have been auto-wrapped in FSDP modules.
@@ -162,6 +162,7 @@ def collate_fn(examples):
     # New Code #
     # For FSDP feature, it is highly recommended and efficient to prepare the model before creating optimizer
     model = accelerator.prepare(model)
+    accelerator.print(model)
Yes, the discussion is happening in this issue https://github.com/pytorch/pytorch/issues/79605#issuecomment-1184410231
@@ -146,9 +146,7 @@ def __init__(
                 if os.environ.get("USE_FSDP", "false") == "true":
                     self.distributed_type = DistributedType.FSDP
                     if self.mixed_precision != "no":
-                        raise ValueError(
-                            "Mixed precision is currently not supported for FSDP. Please set `mixed_precision` to `no`."
-                        )
+                        fsdp_plugin.set_mixed_precision(self.mixed_precision)
Done.
@@ -453,11 +467,57 @@ def __post_init__(self):
             else:
                 self.cpu_offload = CPUOffload(offload_params=False)
 
-        if self.min_num_params is None:
-            self.min_num_params = int(os.environ.get("FSDP_MIN_NUM_PARAMS", 0))
+    @staticmethod
+    def get_module_class_from_name(module, name):
+        """
+        Gets a class from a module by its name.
+
+        Args:
+            module (`torch.nn.Module`): The module to get the class from.
+            name (`str`): The name of the class.
+        """
+        modules_children = list(module.children())
+        if module.__class__.__name__ == name:
+            return module.__class__
+        elif len(modules_children) == 0:
+            return
+        else:
+            for child_module in modules_children:
+                module_class = FullyShardedDataParallelPlugin.get_module_class_from_name(child_module, name)
+                if module_class is not None:
+                    return module_class
+
+    def set_auto_wrap_policy(self, model):
+        from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy, transformer_auto_wrap_policy
 
         if self.auto_wrap_policy is None:
-            if self.min_num_params > 0:
+            auto_wrap_policy = os.environ.get("FSDP_AUTO_WRAP_POLICY", FSDP_AUTO_WRAP_POLICY[-1])
+            if auto_wrap_policy == FSDP_AUTO_WRAP_POLICY[0]:
+                transformer_cls_to_wrap = os.environ.get("FSDP_TRANSFORMER_CLS_TO_WRAP", "")
+                transformer_cls_to_wrap = FullyShardedDataParallelPlugin.get_module_class_from_name(
+                    model, transformer_cls_to_wrap
+                )
                 self.auto_wrap_policy = functools.partial(
-                    size_based_auto_wrap_policy, min_num_params=self.min_num_params
+                    transformer_auto_wrap_policy,
+                    transformer_layer_cls={
+                        transformer_cls_to_wrap,  # < ---- Transformer layer class to wrap
+                    },
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/522
diff --git a/src/accelerate/big_modeling.py b/src/accelerate/big_modeling.py
index 4014e3972..85b1f3237 100644
--- a/src/accelerate/big_modeling.py
+++ b/src/accelerate/big_modeling.py
@@ -65,7 +65,9 @@ def init_empty_weights(include_buffers: bool = False):
     def register_empty_parameter(module, name, param):
         old_register_parameter(module, name, param)
         if param is not None:
-            module._parameters[name] = nn.Parameter(module._parameters[name].to(torch.device("meta")))
+            param_cls = type(module._parameters[name])
+            kwargs = module._parameters[name].__dict__
+            module._parameters[name] = param_cls(module._parameters[name].to(torch.device("meta")), **kwargs)
 
     def register_empty_buffer(module, name, buffer):
         old_register_buffer(module, name, buffer)

Leaves those poor lines alone :-)
@@ -39,9 +39,7 @@ def init_empty_weights(include_buffers: bool = False):
     Args:
         include_buffers (`bool`, *optional*, defaults to `False`):
             Whether or not to also put all buffers on the meta device while initializing.
-
     Example:
-
🤣 yeah don't know what happened there
@@ -39,9 +39,7 @@ def init_empty_weights(include_buffers: bool = False):
     Args:
         include_buffers (`bool`, *optional*, defaults to `False`):
             Whether or not to also put all buffers on the meta device while initializing.
-
     Example:
-
Fixed in aafea4a1e18ac6562bc678f3e607e0d2c5f89b38 ;) 
@@ -39,9 +39,7 @@ def init_empty_weights(include_buffers: bool = False):
     Args:
         include_buffers (`bool`, *optional*, defaults to `False`):
             Whether or not to also put all buffers on the meta device while initializing.
-
     Example:
-
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/519
diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml
index d52cf5e1e..0e8f31407 100644
--- a/.github/workflows/nightly.yml
+++ b/.github/workflows/nightly.yml
@@ -35,6 +35,7 @@ jobs:
       - name: Run examples on GPUs
         run: |
           source activate accelerate
+          pip uninstall comet_ml -y
           make test_examples
 
   run_all_tests_multi_gpu:
@@ -64,4 +65,5 @@ jobs:
       - name: Run examples on GPUs
         run: |
           source activate accelerate
+          pip uninstall comet_ml -y
           make test_examples
\ No newline at end of file
diff --git a/.github/workflows/on-merge.yml b/.github/workflows/on-merge.yml
index b95474758..59677b34e 100644
--- a/.github/workflows/on-merge.yml
+++ b/.github/workflows/on-merge.yml
@@ -34,6 +34,7 @@ jobs:
       - name: Run examples on GPUs
         run: |
           source activate accelerate
+          pip uninstall comet_ml -y
           make test_examples
 
   run_all_tests_multi_gpu:
@@ -61,4 +62,5 @@ jobs:
       - name: Run examples on GPUs
         run: |
           source activate accelerate
+          pip uninstall comet_ml -y
           make test_examples
\ No newline at end of file
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index c585e742a..c0871b28c 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -38,6 +38,7 @@ jobs:
       run: |
         pip install --upgrade pip
         pip install -e .[test,test_trackers]
+        if [ ${{ matrix.test-kind }} = test_rest ]; then pip uninstall comet_ml -y; fi
     
     - name: Run Tests
       run: |
diff --git a/examples/by_feature/cross_validation.py b/examples/by_feature/cross_validation.py
index e380d2d85..95af2a6c9 100644
--- a/examples/by_feature/cross_validation.py
+++ b/examples/by_feature/cross_validation.py
@@ -125,7 +125,6 @@ def collate_fn(examples):
 
 def training_function(config, args):
     # New Code #
-    test_labels = None
     test_predictions = []
     # Download the dataset
     datasets = load_dataset("glue", "mrpc")
@@ -152,17 +151,15 @@ def training_function(config, args):
     # New Code #
     # Create our folds:
     folds = kfold.split(np.zeros(datasets["train"].num_rows), datasets["train"]["label"])
-
+    test_references = []
     # Iterate over them
-    for train_idxs, valid_idxs in folds:
+    for i, (train_idxs, valid_idxs) in enumerate(folds):
         train_dataloader, eval_dataloader, test_dataloader = get_fold_dataloaders(
             accelerator,
             datasets,
             train_idxs,
             valid_idxs,
         )
-        if test_labels is None:
-            test_labels = datasets["validation"]["label"]
         # Instantiate the model (we build the model here so that the seed also control new weights initialization)
         model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
 
@@ -231,19 +228,18 @@ def training_function(config, args):
             predictions = outputs.logits
             predictions, references = accelerator.gather((predictions, batch["labels"]))
             fold_predictions.append(predictions.cpu())
-            metric.add_batch(
-                predictions=predictions.argmax(dim=-1),
-                references=references,
-            )
-        test_metric = metric.compute()
+            if i == 0:
+                # We need all of the test predictions
+                test_references.append(references.cpu())
         # Use accelerator.print to print only on the main process.
         test_predictions.append(torch.cat(fold_predictions, dim=0))
         # We now need to release all our memory and get rid of the current model, optimizer, etc
         accelerator.free_memory()
     # New Code #
     # Finally we check the accuracy of our folded results:
-    preds = torch.stack(test_predictions, dim=0).sum(dim=0).div(int(config["n_splits"])).argmax(dim=-1)
-    test_metric = metric.compute(predictions=preds, references=test_labels)
+    test_references = torch.cat(test_references, dim=0)
+    preds = torch.stack(test_predictions, dim=0).sum(dim=0).div(int(args.num_folds)).argmax(dim=-1)
+    test_metric = metric.compute(predictions=preds, references=test_references)
     accelerator.print("Average test metrics from all folds:", test_metric)
 
 
diff --git a/examples/by_feature/tracking.py b/examples/by_feature/tracking.py
index ef2b1388b..3278a913a 100644
--- a/examples/by_feature/tracking.py
+++ b/examples/by_feature/tracking.py
@@ -165,8 +165,6 @@ def training_function(config, args):
     if args.with_tracking:
         if accelerator.is_main_process:
             run = os.path.split(__file__)[-1].split(".")[0]
-            if args.logging_dir:
-                run = os.path.join(args.logging_dir, run)
             accelerator.init_trackers(run, config)
 
     # Now we train the model
diff --git a/src/accelerate/test_utils/testing.py b/src/accelerate/test_utils/testing.py
index 017f66d06..0b3145070 100644
--- a/src/accelerate/test_utils/testing.py
+++ b/src/accelerate/test_utils/testing.py
@@ -15,6 +15,7 @@
 import asyncio
 import os
 import shutil
+import subprocess
 import sys
 import tempfile
 import unittest
@@ -129,6 +130,22 @@ def require_comet_ml(test_case):
     return unittest.skipUnless(is_comet_ml_available(), "test requires comet_ml")(test_case)
 
 
+_atleast_one_tracker_available = (
+    any([is_wandb_available(), is_tensorboard_available()]) and not is_comet_ml_available()
+)
+
+
+def require_trackers(test_case):
+    """
+    Decorator marking that a test requires at least one tracking library installed. These tests are skipped when none
+    are installed
+    """
+    return unittest.skipUnless(
+        _atleast_one_tracker_available,
+        "test requires at least one tracker to be available and for `comet_ml` to not be installed",
+    )(test_case)
+
+
 class TempDirTestCase(unittest.TestCase):
     """
     A TestCase class that keeps a single `tempfile.TemporaryDirectory` open for the duration of the class, wipes its
@@ -279,3 +296,24 @@ def execute_subprocess_async(cmd, env=None, stdin=None, timeout=180, quiet=False
         )
 
     return result
+
+
+class SubprocessCallException(Exception):
+    pass
+
+
+def run_command(command: List[str], return_stdout=False):
+    """
+    Runs `command` with `subprocess.check_output` and will potentially return the `stdout`. Will also properly capture
+    if an error occured while running `command`
+    """
+    try:
+        output = subprocess.check_output(command, stderr=subprocess.STDOUT)
+        if return_stdout:
+            if hasattr(output, "decode"):
+                output = output.decode("utf-8")
+            return output
+    except subprocess.CalledProcessError as e:
+        raise SubprocessCallException(
+            f"Command `{' '.join(command)}` failed with the following error:\n\n{e.output.decode()}"
+        ) from e
diff --git a/tests/test_examples.py b/tests/test_examples.py
index 4b298d253..9e1854cd8 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -16,7 +16,6 @@
 import os
 import re
 import shutil
-import subprocess
 import tempfile
 import unittest
 from unittest import mock
@@ -24,7 +23,7 @@
 import torch
 
 from accelerate.test_utils.examples import compare_against_test
-from accelerate.test_utils.testing import TempDirTestCase, slow
+from accelerate.test_utils.testing import TempDirTestCase, require_trackers, run_command, slow
 from accelerate.utils import write_basic_config
 
 
@@ -145,7 +144,7 @@ def test_checkpointing_by_epoch(self):
         --checkpointing_steps epoch
         --output_dir {self.tmpdir}
         """.split()
-        _ = subprocess.run(self._launch_args + testargs, stdout=subprocess.PIPE)
+        run_command(self._launch_args + testargs)
         self.assertTrue(os.path.exists(os.path.join(self.tmpdir, "epoch_1")))
 
     def test_checkpointing_by_steps(self):
@@ -154,7 +153,7 @@ def test_checkpointing_by_steps(self):
         --checkpointing_steps 1
         --output_dir {self.tmpdir}
         """.split()
-        _ = subprocess.run(self._launch_args + testargs, stdout=subprocess.PIPE, env=os.environ)
+        _ = run_command(self._launch_args + testargs)
         self.assertTrue(os.path.exists(os.path.join(self.tmpdir, "step_5")))
 
     def test_load_states_by_epoch(self):
@@ -162,9 +161,7 @@ def test_load_states_by_epoch(self):
         examples/by_feature/checkpointing.py
         --resume_from_checkpoint {os.path.join(self.tmpdir, "epoch_1")}
         """.split()
-        output = subprocess.run(
-            self._launch_args + testargs, universal_newlines=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE
-        ).stdout
+        output = run_command(self._launch_args + testargs, return_stdout=True)
         self.assertNotIn("epoch 0:", output)
         self.assertNotIn("epoch 1:", output)
         self.assertIn("epoch 2:", output)
@@ -174,9 +171,7 @@ def test_load_states_by_steps(self):
         examples/by_feature/checkpointing.py
         --resume_from_checkpoint {os.path.join(self.tmpdir, "step_5")}
         """.split()
-        output = subprocess.run(
-            self._launch_args + testargs, universal_newlines=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE
-        ).stdout
+        output = run_command(self._launch_args + testargs, return_stdout=True)
         if torch.cuda.is_available():
             num_processes = torch.cuda.device_count()
         else:
@@ -184,11 +179,10 @@ def test_load_states_by_steps(self):
         if num_processes > 1:
             self.assertNotIn("epoch 0:", output)
             self.assertNotIn("epoch 1:", output)
-            self.assertIn("epoch 2:", output)
         else:
             self.assertNotIn("epoch 0:", output)
             self.assertIn("epoch 1:", output)
-            self.assertIn("epoch 2:", output)
+        self.assertIn("epoch 2:", output)
 
     @slow
     def test_cross_validation(self):
@@ -197,16 +191,16 @@ def test_cross_validation(self):
         --num_folds 2
         """.split()
         with mock.patch.dict(os.environ, {"TESTING_MOCKED_DATALOADERS": "0"}):
-            output = subprocess.run(
-                self._launch_args + testargs, universal_newlines=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE
-            ).stdout
+            output = run_command(self._launch_args + testargs, return_stdout=True)
             results = ast.literal_eval(re.findall("({.+})", output)[-1])
             self.assertGreaterEqual(results["accuracy"], 0.75)
 
     def test_multi_process_metrics(self):
         testargs = ["examples/by_feature/multi_process_metrics.py"]
-        _ = subprocess.run(self._launch_args + testargs, stdout=subprocess.PIPE)
+        run_command(self._launch_args + testargs)
 
+    @require_trackers
+    @mock.patch.dict(os.environ, {"WANDB_MODE": "offline"})
     def test_tracking(self):
         with tempfile.TemporaryDirectory() as tmpdir:
             testargs = f"""
@@ -214,9 +208,9 @@ def test_tracking(self):
             --with_tracking
             --logging_dir {tmpdir}
             """.split()
-            _ = subprocess.run(self._launch_args + testargs, stdout=subprocess.PIPE)
+            run_command(self._launch_args + testargs)
             self.assertTrue(os.path.exists(os.path.join(tmpdir, "tracking")))
 
     def test_gradient_accumulation(self):
         testargs = ["examples/by_feature/gradient_accumulation.py"]
-        _ = subprocess.run(self._launch_args + testargs, stdout=subprocess.PIPE)
+        run_command(self._launch_args + testargs)

This will stop working on distributed setup. The predictions will need to be gathered (maybe not the labels).
@@ -229,20 +229,14 @@ def training_function(config, args):
             with torch.no_grad():
                 outputs = model(**batch)
             predictions = outputs.logits
-            predictions, references = accelerator.gather((predictions, batch["labels"]))
Should be fixed now
@@ -229,20 +229,14 @@ def training_function(config, args):
             with torch.no_grad():
                 outputs = model(**batch)
             predictions = outputs.logits
-            predictions, references = accelerator.gather((predictions, batch["labels"]))
_The documentation is not available anymore as the PR was closed or merged._
All slow ci is now passing here: https://github.com/huggingface/accelerate/runs/7358153325
https://api.github.com/repos/huggingface/accelerate/pulls/517
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index 907122036..0051d899a 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -15,6 +15,7 @@
 # limitations under the License.
 
 from ...utils import ComputeEnvironment, DistributedType, is_deepspeed_available, is_transformers_available
+from ...utils.constants import DEEPSPEED_MULTINODE_LAUNCHERS
 from .config_args import ClusterConfig
 from .config_utils import _ask_field, _convert_distributed_mode, _convert_yes_no_to_bool
 
@@ -144,6 +145,51 @@ def get_cluster_input():
                         "Please run `pip3 install transformers`."
                     )
 
+            if num_machines > 1:
+                launcher_query = "Which Type of launcher do you want to use "
+                for i, launcher in enumerate(DEEPSPEED_MULTINODE_LAUNCHERS):
+                    launcher_query += f"[{i}] {launcher}, "
+                launcher_query = launcher_query[:-2] + ")? [0]: "
+                deepspeed_config["deepspeed_multinode_launcher"] = _ask_field(
+                    launcher_query,
+                    lambda x: DEEPSPEED_MULTINODE_LAUNCHERS[int(x)],
+                    default=DEEPSPEED_MULTINODE_LAUNCHERS[0],
+                )
+
+                if deepspeed_config["deepspeed_multinode_launcher"] != DEEPSPEED_MULTINODE_LAUNCHERS[1]:
+                    deepspeed_config["deepspeed_hostfile"] = _ask_field(
+                        "DeepSpeed configures multi-node compute resources with hostfile. "
+                        "Each row is of the format `hostname slots=[num_gpus]`, e.g., `localhost slots=2`; "
+                        "for more information please refer official [documentation]"
+                        "(https://www.deepspeed.ai/getting-started/#resource-configuration-multi-node). "
+                        "Please specify the location of hostfile: ",
+                        lambda x: str(x),
+                    )
+
+                    is_exclusion_filter = _ask_field(
+                        "Do you want to specify exclusion filter string? [yes/NO]: ",
+                        _convert_yes_no_to_bool,
+                        default=False,
+                        error_message="Please enter yes or no.",
+                    )
+                    if is_exclusion_filter:
+                        deepspeed_config["deepspeed_exclusion_filter"] = _ask_field(
+                            "DeepSpeed exclusion filter string: ",
+                            lambda x: str(x),
+                        )
+
+                    is_inclusion_filter = _ask_field(
+                        "Do you want to specify inclusion filter string? [yes/NO]: ",
+                        _convert_yes_no_to_bool,
+                        default=False,
+                        error_message="Please enter yes or no.",
+                    )
+                    if is_inclusion_filter:
+                        deepspeed_config["deepspeed_inclusion_filter"] = _ask_field(
+                            "DeepSpeed inclusion filter string: ",
+                            lambda x: str(x),
+                        )
+
     fsdp_config = {}
     if distributed_type in [DistributedType.MULTI_GPU]:
         use_fsdp = _ask_field(
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index ac91ad1e9..cffe9990b 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -35,6 +35,7 @@
     is_deepspeed_available,
     is_sagemaker_available,
 )
+from accelerate.utils.constants import DEEPSPEED_MULTINODE_LAUNCHERS
 from accelerate.utils.dataclasses import SageMakerDistributedType
 
 
@@ -109,6 +110,30 @@ def launch_command_parser(subparsers=None):
         help="Decides Whether (true|false) to save 16-bit model weights when using ZeRO Stage-3. "
         "Only applicable with DeepSpeed ZeRO Stage-3.",
     )
+    parser.add_argument(
+        "--deepspeed_hostfile",
+        default=None,
+        type=str,
+        help="DeepSpeed hostfile for configuring multi-node compute resources.",
+    )
+    parser.add_argument(
+        "--deepspeed_exclusion_filter",
+        default=None,
+        type=str,
+        help="DeepSpeed exclusion filter string when using mutli-node setup.",
+    )
+    parser.add_argument(
+        "--deepspeed_inclusion_filter",
+        default=None,
+        type=str,
+        help="DeepSpeed inclusion filter string when using mutli-node setup.",
+    )
+    parser.add_argument(
+        "--deepspeed_multinode_launcher",
+        default=None,
+        type=str,
+        help="DeepSpeed multi-node launcher to use.",
+    )
     parser.add_argument(
         "--use_fsdp",
         default=False,
@@ -312,20 +337,42 @@ def deepspeed_launcher(args):
         raise ImportError("DeepSpeed is not installed => run `pip3 install deepspeed` or build it from source.")
     cmd = ["deepspeed", "--no_local_rank"]
     if args.num_machines > 1:
-        cmd.extend(
-            [
-                "--num_gpus",
-                str(args.num_processes // args.num_machines),
-                "--num_nodes",
-                str(args.num_machines),
-                "--node_rank",
-                str(args.machine_rank),
-                "--master_addr",
-                args.main_process_ip,
-                "--master_port",
-                str(args.main_process_port),
-            ]
-        )
+        if args.deepspeed_multinode_launcher == DEEPSPEED_MULTINODE_LAUNCHERS[1]:
+            cmd = get_launch_prefix()
+            cmd.extend(
+                [
+                    "--nproc_per_node",
+                    str(args.num_processes // args.num_machines),
+                    "--nnodes",
+                    str(args.num_machines),
+                    "--node_rank",
+                    str(args.machine_rank),
+                    "--master_addr",
+                    args.main_process_ip,
+                    "--master_port",
+                    str(args.main_process_port),
+                ]
+            )
+        else:
+            cmd.extend(
+                ["--hostfile", str(args.deepspeed_hostfile), "--launcher", str(args.deepspeed_multinode_launcher)]
+            )
+            if args.deepspeed_exclusion_filter is not None:
+                cmd.extend(
+                    [
+                        "--exclude",
+                        str(args.deepspeed_exclusion_filter),
+                    ]
+                )
+            elif args.deepspeed_inclusion_filter is not None:
+                cmd.extend(
+                    [
+                        "--include",
+                        str(args.deepspeed_inclusion_filter),
+                    ]
+                )
+            else:
+                cmd.extend(["--num_gpus", str(args.num_processes // args.num_machines)])
     else:
         cmd.extend(["--num_gpus", str(args.num_processes)])
 
@@ -350,6 +397,7 @@ def deepspeed_launcher(args):
         warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', DeprecationWarning)
         mixed_precision = "fp16"
 
+    current_env["PYTHONPATH"] = sys.executable
     current_env["MIXED_PRECISION"] = str(mixed_precision)
     current_env["USE_DEEPSPEED"] = "true"
     current_env["DEEPSPEED_ZERO_STAGE"] = str(args.zero_stage)
@@ -361,6 +409,13 @@ def deepspeed_launcher(args):
     current_env["DEEPSPEED_ZERO3_SAVE_16BIT_MODEL"] = str(args.zero3_save_16bit_model).lower()
     current_env["DEEPSPEED_CONFIG_FILE"] = str(args.deepspeed_config_file).lower()
 
+    if args.num_machines > 1 and args.deepspeed_multinode_launcher != DEEPSPEED_MULTINODE_LAUNCHERS[1]:
+        with open(".deepspeed_env", "a") as f:
+            for key, value in current_env.items():
+                if ";" in value or " " in value:
+                    continue
+                f.write(f"{key}={value}\n")
+
     process = subprocess.Popen(cmd, env=current_env)
     process.wait()
     if process.returncode != 0:
diff --git a/src/accelerate/utils/constants.py b/src/accelerate/utils/constants.py
index 173cfa790..5e31a96c9 100644
--- a/src/accelerate/utils/constants.py
+++ b/src/accelerate/utils/constants.py
@@ -24,5 +24,6 @@
 SAGEMAKER_PYTHON_VERSION = "py38"
 SAGEMAKER_TRANSFORMERS_VERSION = "4.17.0"
 SAGEMAKER_PARALLEL_EC2_INSTANCES = ["ml.p3.16xlarge", "ml.p3dn.24xlarge", "ml.p4dn.24xlarge"]
+DEEPSPEED_MULTINODE_LAUNCHERS = ["pdsh", "standard", "openmpi", "mvapich"]
 
 STR_OPERATION_TO_FUNC = {">": op.gt, ">=": op.ge, "==": op.eq, "!=": op.ne, "<=": op.le, "<": op.lt}

Could we add some documentation on the format of the hostfile? Then we can link it from here.
@@ -144,6 +145,46 @@ def get_cluster_input():
                         "Please run `pip3 install transformers`."
                     )
 
+            if num_machines > 1:
+                deepspeed_config["deepspeed_hostfile"] = _ask_field(
+                    "DeepSpeed configures multi-node compute resources with hostfile, please specify the location of hostfile: ",
Maybe start with this question and ask if users want to use the traditional PyTorch DDP launcher here, in which case the other questions are moot.
@@ -144,6 +145,46 @@ def get_cluster_input():
                         "Please run `pip3 install transformers`."
                     )
 
+            if num_machines > 1:
+                deepspeed_config["deepspeed_hostfile"] = _ask_field(
+                    "DeepSpeed configures multi-node compute resources with hostfile, please specify the location of hostfile: ",
+                    lambda x: str(x),
+                )
+
+                is_exclusion_filter = _ask_field(
+                    "Do you want to specify exclusion filter string? [yes/NO]: ",
+                    _convert_yes_no_to_bool,
+                    default=False,
+                    error_message="Please enter yes or no.",
+                )
+                if is_exclusion_filter:
+                    deepspeed_config["deepspeed_exclusion_filter"] = _ask_field(
+                        "DeepSpeed exclusion filter string: ",
+                        lambda x: str(x),
+                    )
+
+                is_inclusion_filter = _ask_field(
+                    "Do you want to specify inclusion filter string? [yes/NO]: ",
+                    _convert_yes_no_to_bool,
+                    default=False,
+                    error_message="Please enter yes or no.",
+                )
+                if is_inclusion_filter:
+                    deepspeed_config["deepspeed_inclusion_filter"] = _ask_field(
+                        "DeepSpeed inclusion filter string: ",
+                        lambda x: str(x),
+                    )
+
+                launcher_query = "Which Type of launcher do you want to use "
+                for i, launcher in enumerate(DEEPSPEED_MULTINODE_LAUNCHERS):
+                    launcher_query += f"[{i}] {launcher}, "
+                launcher_query = launcher_query[:-2] + ")? [0]: "
+                deepspeed_config["deepspeed_multinode_launcher"] = _ask_field(
Done.
@@ -144,6 +145,46 @@ def get_cluster_input():
                         "Please run `pip3 install transformers`."
                     )
 
+            if num_machines > 1:
+                deepspeed_config["deepspeed_hostfile"] = _ask_field(
+                    "DeepSpeed configures multi-node compute resources with hostfile, please specify the location of hostfile: ",
this is a problem, @pacman100 - how can a user then use `PYTHONPATH` if it's being hijacked by the launcher?

If you must manipulate `PYTHONPATH` you need to extend it and not override it.
@@ -350,6 +397,7 @@ def deepspeed_launcher(args):
         warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', DeprecationWarning)
         mixed_precision = "fp16"
 
+    current_env["PYTHONPATH"] = sys.executable
but also why would you set something that should be `path1:...:pathn` to find python libraries to a path to an executable which is not a directory.
@@ -350,6 +397,7 @@ def deepspeed_launcher(args):
         warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', DeprecationWarning)
         mixed_precision = "fp16"
 
+    current_env["PYTHONPATH"] = sys.executable
Proposed fix https://github.com/huggingface/accelerate/pull/741

@@ -350,6 +397,7 @@ def deepspeed_launcher(args):
         warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', DeprecationWarning)
         mixed_precision = "fp16"
 
+    current_env["PYTHONPATH"] = sys.executable
I don't remember but it was needed to know the python env on other nodes to run the code. If I don't do that, deepspeed code will still hijack it: https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/launcher/runner.py#L480-L484.

Below is the run with this line commented and you can notice that PYTHONPATH is still overridden by the above deepspeed code:

<img width="1502" alt="Screenshot 2022-10-05 at 12 35 26 PM" src="https://user-images.githubusercontent.com/13534540/194000694-8d5aa407-7ff3-42a0-9f84-908e8ae67713.png">
  
@@ -350,6 +397,7 @@ def deepspeed_launcher(args):
         warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', DeprecationWarning)
         mixed_precision = "fp16"
 
+    current_env["PYTHONPATH"] = sys.executable
The code you linked to:
```
        curr_path = os.path.abspath('.')
        if 'PYTHONPATH' in env:
            env['PYTHONPATH'] = curr_path + ":" + env['PYTHONPATH']
        else:
            env['PYTHONPATH'] = curr_path
```

1. it doesn't override the setting, it appends to it:
2. it doesn't add `sys.executable` but the current working directory. I think sort of pushing ``pwd`` of the launched-from-dir.
@@ -350,6 +397,7 @@ def deepspeed_launcher(args):
         warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', DeprecationWarning)
         mixed_precision = "fp16"
 
+    current_env["PYTHONPATH"] = sys.executable
Hello Stas, yeah, you are correct, my bad. After spending quite some time diving deep and setting up the multi-node setup to test things out, I have found the reason.

In Accelerate, all the DeepSpeed config from `accelerate config` command is passed as env variables, as shown towards the end of the above screenshot. DeepSpeed uses `.deepspeed_env` file for all these 
additional env variables instead of using the passed env variables as mentioned here: https://www.deepspeed.ai/getting-started/#multi-node-environment-variables

So, I was adding all the env variables to `.deepspeed_env` file as seen here: https://github.com/huggingface/accelerate/blob/main/src/accelerate/commands/launch.py#L590-L594
. The DeepSpeed snippet from the above message would extend/write to PYTHONPATH. However, just after that things will be read from `.deepspeed_env` and env variables will be reset or newly added. Here, the `pwd` extension would be reset and the original PYTHONPATH if sent from accelerate will be used instead of extended one. I overlooked extension usecase and confused `pwd` for python env path and just added that directly to the launch env variables leading to unintended behaviour.   
@@ -350,6 +397,7 @@ def deepspeed_launcher(args):
         warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', DeprecationWarning)
         mixed_precision = "fp16"
 
+    current_env["PYTHONPATH"] = sys.executable
Thank you for the extended commentary that gives more context to the problem, @pacman100 

And it looks like your suggestion in the PR resolving this issue. which has been merged now, so resolving this conversation. 

Thank you!
@@ -350,6 +397,7 @@ def deepspeed_launcher(args):
         warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', DeprecationWarning)
         mixed_precision = "fp16"
 
+    current_env["PYTHONPATH"] = sys.executable
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/514
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 354089c38..e8ff21884 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -855,17 +855,23 @@ def gather(self, tensor):
         """
         return gather(tensor)
 
-    def reduce(self, tensor: torch.Tensor, reduction="sum"):
+    def reduce(self, tensor, reduction="sum"):
         """
         Reduce the values in *tensor* across all processes based on *reduction*.
 
+        Note:
+            All processes get the reduced value.
+
         Args:
-            tensor (`torch.Tensor`):
+            tensor (`torch.Tensor`, or a nested tuple/list/dictionary of `torch.Tensor`):
                 The tensors to reduce across all processes.
             reduction (`str`, *optional*, defaults to "sum"):
                 A reduction type, can be one of 'sum', 'mean', or 'none'. If 'none', will not perform any operation.
+
+        Returns:
+            `torch.Tensor`, or a nested tuple/list/dictionary of `torch.Tensor`: The reduced tensor(s).
         """
-        reduce(tensor, reduction)
+        return reduce(tensor, reduction)
 
     def pad_across_processes(self, tensor, dim=0, pad_index=0, pad_first=False):
         """
diff --git a/src/accelerate/utils/operations.py b/src/accelerate/utils/operations.py
index 46753513e..7505c04ce 100644
--- a/src/accelerate/utils/operations.py
+++ b/src/accelerate/utils/operations.py
@@ -430,7 +430,7 @@ def _reduce_across_processes(tensor, reduction="mean"):
             xm.all_reduce("sum", cloned_tensor)
             return cloned_tensor
         elif state.distributed_type in [DistributedType.DEEPSPEED, DistributedType.MULTI_GPU]:
-            torch.distributed.reduce(cloned_tensor, ReduceOp.SUM)
+            torch.distributed.all_reduce(cloned_tensor, ReduceOp.SUM)
             return cloned_tensor
         else:
             if reduction == "sum":

Fix a bug when using reduce at a tensor.

Before:
```
a = torch.tensor(accelerator.process_index, device=accelerator.device)

accelerator.reduce(a, reduction='sum') # don't change a
# or
a = accelerator.reduce(a, reduction='sum') # a is None
```

After:
```
a = torch.tensor(accelerator.process_index, device=accelerator.device)
a = accelerator.reduce(a, reduction='sum') # a is correctly reduced
```

Thanks for the fix! 😄
_The documentation is not available anymore as the PR was closed or merged._
When I update the doc for this pr, I find that the `accelerate.utils.reduce` seems have inconsistent behavior on different type of devices.

In detail, it will reduce value on all processes, i.e. `all_reduce` when device is TPU. But only reduce to the main process, i.e. `reduce`, when device is GPU.

Could I change the behavior to `all_reduce` when device is GPU ?


For more infomation, see the doc stuff [here](https://github.com/wwhio/accelerate/blob/96a17aae1164da5511a50d7138bd6f8ae5303267/src/accelerate/accelerator.py#L862-L864) and the `accelerate.utils.reduce` code [here](https://github.com/wwhio/accelerate/blob/96a17aae1164da5511a50d7138bd6f8ae5303267/src/accelerate/utils/operations.py#L411-L441).

Yes, I think it should also be `all_reduce` on GPU to avoid bugs when not all processes have the same tensor later in the code.
https://api.github.com/repos/huggingface/accelerate/pulls/513
diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml
index 04571e50e..c2fd77282 100644
--- a/docs/source/_toctree.yml
+++ b/docs/source/_toctree.yml
@@ -9,6 +9,8 @@
 - sections:
   - local: big_modeling
     title: Handling big models
+  - local: gradient_accumulation
+    title: Gradient accumulation
   - local: sagemaker
     title: Amazon SageMaker
   title: Guides
diff --git a/docs/source/gradient_accumulation.mdx b/docs/source/gradient_accumulation.mdx
new file mode 100644
index 000000000..f6ab857d8
--- /dev/null
+++ b/docs/source/gradient_accumulation.mdx
@@ -0,0 +1,126 @@
+<!--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.
+-->
+
+# Performing gradient accumulation with 🤗 Accelerate
+
+Gradient accumulation is a technique where you can train on bigger batch sizes than 
+your machine would normally be able to fit into memory. This is done by accumulating gradients over
+several batches, and only stepping the optimizer after a certain number of batches have been performed.
+
+While technically standard gradient accumulation code would work fine in a distributed setup, it is not the most efficient
+method for doing so and you may experience considerable slowdowns!
+
+In this tutorial you will see how to quickly setup gradient accumulation and perform it with the utilities provided in 🤗 Accelerate,
+which can total to adding just one new line of code!
+
+This example will use a very simplistic PyTorch training loop that performs gradient accumulation every two batches:
+
+```python
+device = "cuda"
+model.to(device)
+
+gradient_accumulation_steps = 2
+
+for index, batch in enumerate(training_dataloader):
+    optimizer.zero_grad()
+    inputs, targets = batch
+    inputs = inputs.to(device)
+    targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    loss = loss / gradient_accumulation_steps
+    loss.backward()
+    if (index + 1) % gradient_accumulation_steps == 0:
+        optimizer.step()
+        scheduler.step()
+```
+
+## Converting it to 🤗 Accelerate
+
+First the code shown earlier will be converted to utilize 🤗 Accelerate without the special gradient accumulation helper:
+
+```diff
++ from accelerate import Accelerator
++ accelerator = Accelerator()
+
++ model, optimizer, training_dataloader, scheduler = accelerator.prepare(
++     model, optimizer, training_dataloader, scheduler
++ )
+
+  for index, batch in enumerate(training_dataloader):
+      optimizer.zero_grad()
+      inputs, targets = batch
+-     inputs = inputs.to(device)
+-     targets = targets.to(device)
+      outputs = model(inputs)
+      loss = loss_function(outputs, targets)
+      loss = loss / gradient_accumulation_steps
++     accelerator.backward(loss)
+      if (index+1) % gradient_accumulation_steps == 0:
+          optimizer.step()
+          scheduler.step()
+```
+
+<Tip warning={true}>
+In its current state, this code is not going to perform gradient accumulation efficiently due to a process called gradient synchronization.
+</Tip>
+
+## Letting 🤗 Accelerate handle gradient accumulation
+
+All that is left now is to let 🤗 Accelerate handle the gradient accumulation for us. To do so you should pass in a `gradient_accumulation_steps` parameter to [`Accelerator`], dictating the number 
+of steps to perform before each call to `step()` and how to automatically adjust the loss during the call to [`Accelerator.backward`]:
+
+```diff
+  from accelerate import Accelerator
+- accelerator = Accelerator()
++ accelerator = Accelerator(gradient_accumulation_steps=2)
+```
+
+From here you can use the [`Accelerator.accumulate`] context manager from inside your training loop to automatically perform the gradient accumulation for you!
+You just wrap it around the entire training part of your code: 
+
+```diff
+- for index, batch in enumerate(training_dataloader):
++ for batch in training_dataloader:
++    with accelerator.accumulate(model):
+          optimizer.zero_grad()
+          inputs, targets = batch
+          outputs = model(inputs)
+```
+
+and you can remove all the special checks for the step number and the loss adjustment:
+
+```diff
+- loss = loss / gradient_accumulation_steps
+  accelerator.backward(loss)
+- if (index+1) % gradient_accumulation_steps == 0:
+  optimizer.step()
+  scheduler.step()
+```
+
+As you can see the [`Accelerator`] is able to keep track of the batch number you are on and it will automatically know whether to step through the prepared optimizer and how to adjust the loss. 
+
+## The finished code
+
+Below is the finished implementation for performing gradient accumulation with 🤗 Accelerate
+
+```python
+for batch in training_dataloader:
+    with accelerator.accumulate(model):
+        optimizer.zero_grad()
+        inputs, targets = batch
+        outputs = model(inputs)
+        loss = loss_function(outputs, targets)
+        accelerator.backward(loss)
+        optimizer.step()
+        scheduler.step()
+```

Make sure to add a consistent indent level of 2 for the base, so that you can have `+ ` or `- ` without changing indentation on any line.
@@ -0,0 +1,126 @@
+<!--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.
+-->
+
+# Performing gradient accumulation with 🤗 Accelerate
+
+Gradient accumulation is a technique where you can train on bigger batch sizes than 
+your machine would normally be able to fit into memory. This is done by accumulating gradients over
+several batches, and only stepping the optimizer after a certain number of batches have been performed.
+
+While technically standard gradient accumulation code would work fine in a distributed setup, it is not the most efficient
+method for doing so and you may experience considerable slowdowns!
+
+In this tutorial you will see how to quickly setup gradient accumulation and perform it with the utilities provided in 🤗 Accelerate,
+which can total to adding just one new line of code!
+
+This example will use a very simplistic PyTorch training loop that performs gradient accumulation every two batches:
+
+```
+device = 'cuda'
+model.to(device)
+
+gradient_accumulation_steps = 2
+
+for index, batch in enumerate(training_dataloader):
+    optimizer.zero_grad()
+    inputs, targets = batch
+    inputs = inputs.to(device)
+    targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    loss = loss / gradient_accumulation_steps
+    loss.backward()
+    if (index+1) % gradient_accumulation_steps == 0:
+        optimizer.step()
+        scheduler.step()
+```
+
+## Converting it to 🤗 Accelerate
+
+First the code shown earlier will be converted to utilize 🤗 Accelerate without the special gradient accumulation helper:
+
+```diff
++from accelerate import Accelerator
++accelerator = Accelerator()
+
++model, optimizer, training_dataloader, scheduler = accelerator.prepare(
++    model, optimizer, training_dataloader, scheduler
++)
+
+for index, batch in enumerate(training_dataloader):
+    optimizer.zero_grad()
+    inputs, targets = batch
+-    inputs = inputs.to(device)
+-    targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    loss = loss / gradient_accumulation_steps
++    accelerator.backward(loss)
+    if (index+1) % gradient_accumulation_steps == 0:
+        optimizer.step()
+        scheduler.step()
Same here.
@@ -0,0 +1,126 @@
+<!--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.
+-->
+
+# Performing gradient accumulation with 🤗 Accelerate
+
+Gradient accumulation is a technique where you can train on bigger batch sizes than 
+your machine would normally be able to fit into memory. This is done by accumulating gradients over
+several batches, and only stepping the optimizer after a certain number of batches have been performed.
+
+While technically standard gradient accumulation code would work fine in a distributed setup, it is not the most efficient
+method for doing so and you may experience considerable slowdowns!
+
+In this tutorial you will see how to quickly setup gradient accumulation and perform it with the utilities provided in 🤗 Accelerate,
+which can total to adding just one new line of code!
+
+This example will use a very simplistic PyTorch training loop that performs gradient accumulation every two batches:
+
+```
+device = 'cuda'
+model.to(device)
+
+gradient_accumulation_steps = 2
+
+for index, batch in enumerate(training_dataloader):
+    optimizer.zero_grad()
+    inputs, targets = batch
+    inputs = inputs.to(device)
+    targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    loss = loss / gradient_accumulation_steps
+    loss.backward()
+    if (index+1) % gradient_accumulation_steps == 0:
+        optimizer.step()
+        scheduler.step()
+```
+
+## Converting it to 🤗 Accelerate
+
+First the code shown earlier will be converted to utilize 🤗 Accelerate without the special gradient accumulation helper:
+
+```diff
++from accelerate import Accelerator
++accelerator = Accelerator()
+
++model, optimizer, training_dataloader, scheduler = accelerator.prepare(
++    model, optimizer, training_dataloader, scheduler
++)
+
+for index, batch in enumerate(training_dataloader):
+    optimizer.zero_grad()
+    inputs, targets = batch
+-    inputs = inputs.to(device)
+-    targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    loss = loss / gradient_accumulation_steps
++    accelerator.backward(loss)
+    if (index+1) % gradient_accumulation_steps == 0:
+        optimizer.step()
+        scheduler.step()
+```
+
+<Tip warning={true}>
+In its current state, this code is not going to perform gradient accumulation efficiently due to a process called gradient synchronization. Read more about that in the Concepts tutorial!
+</Tip>
+
+## Letting 🤗 Accelerate handle gradient accumulation
+
+All that is left now is to let 🤗 Accelerate handle the gradient accumulation for us. To do so you should pass in a `gradient_accumulation_steps` parameter to `Accelerator`, dictating the number 
+of steps to perform before each call to `step()` and how to automatically adjust the loss during the call to [`Accelerator.backward`]:
+
+```diff
+from accelerate import Accelerator
+-accelerator = Accelerator()
++accelerator = Accelerator(gradient_accumulation_steps=2)
+```
+
+From here you can use the [`Accelerator.accumulate`] context manager from inside your training loop to automatically perform the gradient accumulation for you!
+You just wrap it around the entire training part of our code: 
+
+```diff
+-for index, batch in enumerate(training_dataloader):
++for batch in training_dataloader:
++   with accelerator.accumulate(model):
+        optimizer.zero_grad()
+        inputs, targets = batch
+        outputs = model(inputs)
+```
```suggestion
All that is left now is to let 🤗 Accelerate handle the gradient accumulation for us. To do so you should pass in a `gradient_accumulation_steps` parameter to [`Accelerator`], dictating the number 
```
@@ -0,0 +1,126 @@
+<!--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.
+-->
+
+# Performing gradient accumulation with 🤗 Accelerate
+
+Gradient accumulation is a technique where you can train on bigger batch sizes than 
+your machine would normally be able to fit into memory. This is done by accumulating gradients over
+several batches, and only stepping the optimizer after a certain number of batches have been performed.
+
+While technically standard gradient accumulation code would work fine in a distributed setup, it is not the most efficient
+method for doing so and you may experience considerable slowdowns!
+
+In this tutorial you will see how to quickly setup gradient accumulation and perform it with the utilities provided in 🤗 Accelerate,
+which can total to adding just one new line of code!
+
+This example will use a very simplistic PyTorch training loop that performs gradient accumulation every two batches:
+
+```
+device = 'cuda'
+model.to(device)
+
+gradient_accumulation_steps = 2
+
+for index, batch in enumerate(training_dataloader):
+    optimizer.zero_grad()
+    inputs, targets = batch
+    inputs = inputs.to(device)
+    targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    loss = loss / gradient_accumulation_steps
+    loss.backward()
+    if (index+1) % gradient_accumulation_steps == 0:
+        optimizer.step()
+        scheduler.step()
+```
+
+## Converting it to 🤗 Accelerate
+
+First the code shown earlier will be converted to utilize 🤗 Accelerate without the special gradient accumulation helper:
+
+```diff
++from accelerate import Accelerator
++accelerator = Accelerator()
+
++model, optimizer, training_dataloader, scheduler = accelerator.prepare(
++    model, optimizer, training_dataloader, scheduler
++)
+
+for index, batch in enumerate(training_dataloader):
+    optimizer.zero_grad()
+    inputs, targets = batch
+-    inputs = inputs.to(device)
+-    targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    loss = loss / gradient_accumulation_steps
++    accelerator.backward(loss)
+    if (index+1) % gradient_accumulation_steps == 0:
+        optimizer.step()
+        scheduler.step()
+```
+
+<Tip warning={true}>
+In its current state, this code is not going to perform gradient accumulation efficiently due to a process called gradient synchronization. Read more about that in the Concepts tutorial!
+</Tip>
+
+## Letting 🤗 Accelerate handle gradient accumulation
+
+All that is left now is to let 🤗 Accelerate handle the gradient accumulation for us. To do so you should pass in a `gradient_accumulation_steps` parameter to `Accelerator`, dictating the number 
Same there + the step lines change since they lose and indent level.
@@ -0,0 +1,126 @@
+<!--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.
+-->
+
+# Performing gradient accumulation with 🤗 Accelerate
+
+Gradient accumulation is a technique where you can train on bigger batch sizes than 
+your machine would normally be able to fit into memory. This is done by accumulating gradients over
+several batches, and only stepping the optimizer after a certain number of batches have been performed.
+
+While technically standard gradient accumulation code would work fine in a distributed setup, it is not the most efficient
+method for doing so and you may experience considerable slowdowns!
+
+In this tutorial you will see how to quickly setup gradient accumulation and perform it with the utilities provided in 🤗 Accelerate,
+which can total to adding just one new line of code!
+
+This example will use a very simplistic PyTorch training loop that performs gradient accumulation every two batches:
+
+```
+device = 'cuda'
+model.to(device)
+
+gradient_accumulation_steps = 2
+
+for index, batch in enumerate(training_dataloader):
+    optimizer.zero_grad()
+    inputs, targets = batch
+    inputs = inputs.to(device)
+    targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    loss = loss / gradient_accumulation_steps
+    loss.backward()
+    if (index+1) % gradient_accumulation_steps == 0:
+        optimizer.step()
+        scheduler.step()
+```
+
+## Converting it to 🤗 Accelerate
+
+First the code shown earlier will be converted to utilize 🤗 Accelerate without the special gradient accumulation helper:
+
+```diff
++from accelerate import Accelerator
++accelerator = Accelerator()
+
++model, optimizer, training_dataloader, scheduler = accelerator.prepare(
++    model, optimizer, training_dataloader, scheduler
++)
+
+for index, batch in enumerate(training_dataloader):
+    optimizer.zero_grad()
+    inputs, targets = batch
+-    inputs = inputs.to(device)
+-    targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    loss = loss / gradient_accumulation_steps
++    accelerator.backward(loss)
+    if (index+1) % gradient_accumulation_steps == 0:
+        optimizer.step()
+        scheduler.step()
+```
+
+<Tip warning={true}>
+In its current state, this code is not going to perform gradient accumulation efficiently due to a process called gradient synchronization. Read more about that in the Concepts tutorial!
+</Tip>
+
+## Letting 🤗 Accelerate handle gradient accumulation
+
+All that is left now is to let 🤗 Accelerate handle the gradient accumulation for us. To do so you should pass in a `gradient_accumulation_steps` parameter to `Accelerator`, dictating the number 
+of steps to perform before each call to `step()` and how to automatically adjust the loss during the call to [`Accelerator.backward`]:
+
+```diff
+from accelerate import Accelerator
+-accelerator = Accelerator()
++accelerator = Accelerator(gradient_accumulation_steps=2)
+```
+
+From here you can use the [`Accelerator.accumulate`] context manager from inside your training loop to automatically perform the gradient accumulation for you!
+You just wrap it around the entire training part of our code: 
+
+```diff
+-for index, batch in enumerate(training_dataloader):
++for batch in training_dataloader:
++   with accelerator.accumulate(model):
+        optimizer.zero_grad()
+        inputs, targets = batch
+        outputs = model(inputs)
+```
+
+and you can remove all the special checks for the step number and the loss adjustment:
+
+```diff
+-loss = loss / gradient_accumulation_steps
+accelerator.backward(loss)
+-if (index+1) % gradient_accumulation_steps == 0:
+optimizer.step()
+scheduler.step()
+```
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/511
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 7e99c76ef..354089c38 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -240,10 +240,6 @@ def __init__(
                 raise NotImplementedError(
                     "Gradient accumulation on TPU is not supported. Pass in `gradient_accumulation_steps=1`"
                 )
-            if dispatch_batches:
-                raise NotImplementedError(
-                    "Gradient accumulation with dispatched dataloaders is not supported. Pass in `gradient_accumulation_steps=1` or `dispatch_batches=False`"
-                )
 
         self.gradient_accumulation_steps = gradient_accumulation_steps
         self.device_placement = device_placement
@@ -397,7 +393,7 @@ def _do_sync(self):
             self.gradient_state._set_sync_gradients(True)
         else:
             self.step += 1
-        self.gradient_state._set_sync_gradients((self.step % self.gradient_accumulation_steps) == 0)
+            self.gradient_state._set_sync_gradients((self.step % self.gradient_accumulation_steps) == 0)
 
     @property
     def sync_gradients(self):
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index a2c7a4a9b..fddfe8955 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -392,74 +392,87 @@ def __init__(self, dataset, split_batches: bool = False, **kwargs):
         self.gradient_state = GradientState()
         self.state = AcceleratorState()
 
-    def __iter__(self):
-        state = AcceleratorState()
-        if state.process_index == 0:
-            # We only iterate through the DataLoader on process 0.
-            main_iterator = super().__iter__()
-        stop_iteration = False
-        first_batch = None
-        while not stop_iteration:
-            # On process 0, we gather the batch to dispatch.
-            if state.process_index == 0:
-                try:
-                    if self.split_batches:
-                        # One batch of the main iterator is dispatched and split.
-                        batch = next(main_iterator)
-                    else:
-                        # num_processes batches of the main iterator are concatenated then dispatched and split.
-                        # We add the batches one by one so we have the remainder available when drop_last=False.
-                        batches = []
-                        for _ in range(state.num_processes):
-                            batches.append(next(main_iterator))
-                        batch = concatenate(batches, dim=0)
-                    # In both cases, we need to get the structure of the batch that we will broadcast on other
-                    # processes to initialize the tensors with the right shape.
-                    # data_structure, stop_iteration
+    def _fetch_batches(self, iterator):
+        batches, batch = None, None
+        # On process 0, we gather the batch to dispatch.
+        if self.state.process_index == 0:
+            try:
+                if self.split_batches:
+                    # One batch of the main iterator is dispatched and split.
+                    batch = next(iterator)
+                else:
+                    # num_processes batches of the main iterator are concatenated then dispatched and split.
+                    # We add the batches one by one so we have the remainder available when drop_last=False.
+                    batches = []
+                    for _ in range(self.state.num_processes):
+                        batches.append(next(iterator))
+                    batch = concatenate(batches, dim=0)
+                # In both cases, we need to get the structure of the batch that we will broadcast on other
+                # processes to initialize the tensors with the right shape.
+                # data_structure, stop_iteration
+                batch_info = [get_data_structure(batch), False]
+            except StopIteration:
+                batch_info = [None, True]
+        else:
+            batch_info = [None, self._stop_iteration]
+        # This is inplace, so after this instruction, every process has the same `batch_info` as process 0.
+        broadcast_object_list(batch_info)
+        self._stop_iteration = batch_info[1]
+        if self._stop_iteration:
+            # If drop_last is False and split_batches is False, we may have a remainder to take care of.
+            if not self.split_batches and not self.drop_last:
+                if self.state.process_index == 0 and len(batches) > 0:
+                    batch = concatenate(batches, dim=0)
                     batch_info = [get_data_structure(batch), False]
-                except StopIteration:
+                else:
                     batch_info = [None, True]
+                broadcast_object_list(batch_info)
+                if batch_info[1]:
+                    return batch, batch_info, True
             else:
-                batch_info = [None, stop_iteration]
-
-            # This is inplace, so after this instruction, every process has the same `batch_info` as process 0.
-            broadcast_object_list(batch_info)
-            stop_iteration = batch_info[1]
-            if stop_iteration:
-                # If drop_last is False and split_batches is False, we may have a remainder to take care of.
-                if not self.split_batches and not self.drop_last:
-                    if state.process_index == 0 and len(batches) > 0:
-                        batch = concatenate(batches, dim=0)
-                        batch_info = [get_data_structure(batch), False]
-                    else:
-                        batch_info = [None, True]
-                    broadcast_object_list(batch_info)
-                    if batch_info[1]:
-                        continue
-                else:
-                    continue
+                return batch, batch_info, True
+        return batch, batch_info, False
 
-            if state.process_index != 0:
+    def __iter__(self):
+        self.gradient_state._set_end_of_dataloader(False)
+        main_iterator = None
+        if self.state.process_index == 0:
+            # We only iterate through the DataLoader on process 0.
+            main_iterator = super().__iter__()
+        self._stop_iteration = False
+        first_batch = None
+        batch, batch_info, skip = self._fetch_batches(main_iterator)
+        while True:
+            if skip:
+                continue
+            if self.state.process_index != 0:
                 # Initialize tensors on other processes than process 0.
                 batch = initialize_tensors(batch_info[0])
-            batch = send_to_device(batch, state.device)
+            batch = send_to_device(batch, self.state.device)
             # Broadcast the batch before splitting it.
             batch = broadcast(batch, from_process=0)
 
             if not self.drop_last and first_batch is None:
                 # We keep at least num processes elements of the first batch to be able to complete the last batch
-                first_batch = slice_tensors(batch, slice(0, state.num_processes))
+                first_batch = slice_tensors(batch, slice(0, self.state.num_processes))
 
             observed_batch_size = find_batch_size(batch)
-            batch_size = observed_batch_size // state.num_processes
+            batch_size = observed_batch_size // self.state.num_processes
 
-            if not self.drop_last and stop_iteration and observed_batch_size % state.num_processes != 0:
+            if not self.drop_last and self._stop_iteration and observed_batch_size % self.state.num_processes != 0:
                 # If the last batch is not complete, let's add the first batch to it.
                 batch = concatenate([batch, first_batch], dim=0)
                 batch_size += 1
 
-            data_slice = slice(state.process_index * batch_size, (state.process_index + 1) * batch_size)
-            yield slice_tensors(batch, data_slice)
+            data_slice = slice(self.state.process_index * batch_size, (self.state.process_index + 1) * batch_size)
+            next_batch, next_batch_info, next_skip = self._fetch_batches(main_iterator)
+            if not self._stop_iteration:
+                yield slice_tensors(batch, data_slice)
+                batch, batch_info, skip = next_batch, next_batch_info, next_skip
+            else:
+                self.gradient_state._set_end_of_dataloader(True)
+                yield slice_tensors(batch, data_slice)
+                break
 
     def __len__(self):
         whole_length = super().__len__()
diff --git a/src/accelerate/test_utils/scripts/test_sync.py b/src/accelerate/test_utils/scripts/test_sync.py
index 0b0dc7abd..ae5a2c65b 100644
--- a/src/accelerate/test_utils/scripts/test_sync.py
+++ b/src/accelerate/test_utils/scripts/test_sync.py
@@ -25,7 +25,7 @@
 from accelerate.utils import DistributedType, set_seed
 
 
-def check_model_parameters(model_a, model_b, did_step):
+def check_model_parameters(model_a, model_b, did_step, iteration):
     for param, grad_param in zip(model_a.parameters(), model_b.parameters()):
         if not param.requires_grad:
             continue
@@ -33,12 +33,12 @@ def check_model_parameters(model_a, model_b, did_step):
             # Grads should not be in sync
             assert (
                 torch.allclose(param.grad, grad_param.grad) is False
-            ), f"Gradients in sync when they should not be:\nmodel_a grad ({param.grad}) == model_b grad ({grad_param.grad})"
+            ), f"Gradients in sync when they should not be at iteration {iteration}:\nmodel_a grad ({param.grad}) == model_b grad ({grad_param.grad})"
         else:
             # Grads should be in sync
             assert (
                 torch.allclose(param.grad, grad_param.grad) is True
-            ), f"Gradients not in sync when they should be:\nmodel_a grad ({param.grad}) != model_b grad ({grad_param.grad})"
+            ), f"Gradients not in sync when they should be at iteration {iteration}:\nmodel_a grad ({param.grad}) != model_b grad ({grad_param.grad})"
 
 
 def step_model(model, input, target, accelerator, do_backward=True):
@@ -96,7 +96,7 @@ def test_noop_sync(accelerator):
             step_model(ddp_model, ddp_input, ddp_target, accelerator)
 
         # Since `no_sync` is a noop, `ddp_model` and `model` grads should always be in sync
-        check_model_parameters(model, ddp_model, True)
+        check_model_parameters(model, ddp_model, True, iteration)
         for param, ddp_param in zip(model.parameters(), ddp_model.parameters()):
             if not param.requires_grad:
                 continue
@@ -106,7 +106,7 @@ def test_noop_sync(accelerator):
 
         # Shuffle ddp_input on each iteration
         torch.manual_seed(1337 + iteration)
-        ddp_input = ddp_input[torch.randperm(16)]
+        ddp_input = ddp_input[torch.randperm(len(ddp_input))]
 
 
 def test_distributed_sync(accelerator):
@@ -146,11 +146,13 @@ def test_distributed_sync(accelerator):
 
         # Shuffle ddp_input on each iteration
         torch.manual_seed(1337 + iteration)
-        ddp_input = ddp_input[torch.randperm(16)]
+        ddp_input = ddp_input[torch.randperm(len(ddp_input))]
 
 
-def test_gradient_accumulation():
-    accelerator = Accelerator(gradient_accumulation_steps=2)
+def test_gradient_accumulation(split_batches=False, dispatch_batches=False):
+    accelerator = Accelerator(
+        gradient_accumulation_steps=2, split_batches=split_batches, dispatch_batches=dispatch_batches
+    )
     # Test that context manager behaves properly
     model, ddp_model, dataloader = get_training_setup(accelerator)
     for iteration, batch in enumerate(dataloader):
@@ -181,11 +183,13 @@ def test_gradient_accumulation():
 
         # Shuffle ddp_input on each iteration
         torch.manual_seed(1337 + iteration)
-        ddp_input = ddp_input[torch.randperm(16)]
+        ddp_input = ddp_input[torch.randperm(len(ddp_input))]
 
 
-def test_gradient_accumulation_with_opt_and_scheduler():
-    accelerator = Accelerator(gradient_accumulation_steps=2)
+def test_gradient_accumulation_with_opt_and_scheduler(split_batches=False, dispatch_batches=False):
+    accelerator = Accelerator(
+        gradient_accumulation_steps=2, split_batches=split_batches, dispatch_batches=dispatch_batches
+    )
     # Test that context manager behaves properly
     model, opt, sched, dataloader, ddp_model, ddp_opt, ddp_sched = get_training_setup(accelerator, True)
     for iteration, batch in enumerate(dataloader):
@@ -198,8 +202,11 @@ def test_gradient_accumulation_with_opt_and_scheduler():
         ddp_model.train()
         step_model(model, input, target, accelerator, False)
         opt.step()
-        for _ in range(accelerator.num_processes):
+        if split_batches:
             sched.step()
+        else:
+            for _ in range(accelerator.num_processes):
+                sched.step()
         opt.zero_grad()
         # Perform gradient accumulation under wrapper
         with accelerator.accumulate(ddp_model):
@@ -209,10 +216,12 @@ def test_gradient_accumulation_with_opt_and_scheduler():
             ddp_opt.zero_grad()
 
         # Learning rates should be the same
-        assert opt.param_groups[0]["lr"] == ddp_opt.param_groups[0]["lr"]
-        did_step = (((iteration + 1) % 2) == 0) or (iteration == (len(dataloader) - 1))
+        assert (
+            opt.param_groups[0]["lr"] == ddp_opt.param_groups[0]["lr"]
+        ), f'Learning rates found in each optimizer did not align\nopt: {opt.param_groups[0]["lr"]}\nDDP opt: {ddp_opt.param_groups[0]["lr"]}\n'
+        did_step = (((iteration + 1) % 2) == 0) or ((iteration + 1) == len(dataloader))
         if accelerator.num_processes > 1:
-            check_model_parameters(model, ddp_model, did_step)
+            check_model_parameters(model, ddp_model, did_step, iteration)
         # Shuffle ddp_input on each iteration
         torch.manual_seed(1337 + iteration)
 
@@ -229,12 +238,31 @@ def main():
             print("**Test Distributed `no_sync` context manager**")
         test_distributed_sync(accelerator)
     if state.distributed_type == DistributedType.MULTI_GPU:
-        if state.local_process_index == 0:
-            print("**Test `accumulate` gradient accumulation**")
-        test_gradient_accumulation()
+        for split_batch in [True, False]:
+            for dispatch_batches in [True, False]:
+                if state.local_process_index == 0:
+                    print(
+                        "**Test `accumulate` gradient accumulation, ",
+                        f"`split_batches={split_batch}` and `dispatch_batches={dispatch_batches}`**",
+                    )
+                test_gradient_accumulation(split_batch)
     if state.local_process_index == 0:
-        print("**Test `accumulate` gradient accumulation with optimizer and scheduler**")
+        print(
+            "**Test `accumulate` gradient accumulation with optimizer and scheduler, ",
+            "`split_batches=False`, `dispatch_batches=False`**",
+        )
     test_gradient_accumulation_with_opt_and_scheduler()
+    if state.distributed_type == DistributedType.MULTI_GPU:
+        for split_batch in [True, False]:
+            for dispatch_batches in [True, False]:
+                if not split_batch and not dispatch_batches:
+                    continue
+                if state.local_process_index == 0:
+                    print(
+                        "**Test `accumulate` gradient accumulation with optimizer and scheduler, ",
+                        f"`split_batches={split_batch}` and `dispatch_batches={dispatch_batches}`**",
+                    )
+                test_gradient_accumulation_with_opt_and_scheduler(split_batch, dispatch_batches)
 
 
 def _mp_fn(index):

Should be split in several lines to stay under the max char len of 119.
@@ -229,12 +238,28 @@ def main():
             print("**Test Distributed `no_sync` context manager**")
         test_distributed_sync(accelerator)
     if state.distributed_type == DistributedType.MULTI_GPU:
-        if state.local_process_index == 0:
-            print("**Test `accumulate` gradient accumulation**")
-        test_gradient_accumulation()
+        for split_batch in [True, False]:
+            for dispatch_batches in [True, False]:
+                if state.local_process_index == 0:
+                    print(
+                        f"**Test `accumulate` gradient accumulation, `split_batches={split_batch}` and `dispatch_batches={dispatch_batches}`**"
+                    )
+                test_gradient_accumulation(split_batch)
     if state.local_process_index == 0:
-        print("**Test `accumulate` gradient accumulation with optimizer and scheduler**")
+        print(
+            "**Test `accumulate` gradient accumulation with optimizer and scheduler, `split_batches=False`, `dispatch_batches=False`**"
+        )
     test_gradient_accumulation_with_opt_and_scheduler()
+    if state.distributed_type == DistributedType.MULTI_GPU:
+        for split_batch in [True, False]:
+            for dispatch_batches in [True, False]:
+                if not split_batch and not dispatch_batches:
+                    continue
+                if state.local_process_index == 0:
+                    print(
+                        f"**Test `accumulate` gradient accumulation with optimizer and scheduler, `split_batches={split_batch}` and `dispatch_batches={dispatch_batches}`**"
+                    )
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/510
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index f58752e53..a2c7a4a9b 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -463,7 +463,9 @@ def __iter__(self):
 
     def __len__(self):
         whole_length = super().__len__()
-        if self.drop_last:
+        if self.split_batches:
+            return whole_length
+        elif self.drop_last:
             return whole_length // self.state.num_processes
         else:
             return math.ceil(whole_length / self.state.num_processes)

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/509
diff --git a/Makefile b/Makefile
index 58cb89ab9..6dab6af11 100644
--- a/Makefile
+++ b/Makefile
@@ -1,6 +1,6 @@
 .PHONY: quality style test docs
 
-check_dirs := tests src examples
+check_dirs := tests src examples benchmarks
 
 # Check that source code meets quality standards
 
diff --git a/benchmarks/README.md b/benchmarks/README.md
new file mode 100644
index 000000000..243e9df58
--- /dev/null
+++ b/benchmarks/README.md
@@ -0,0 +1,46 @@
+# Big model inference benchmarks
+
+Running inference with Accelerate on big models.
+
+## Setup
+
+These benchmarks use the `transformers` library:
+
+```bash
+pip install transformers
+```
+
+To reproduce or test a new setup, run
+
+```py
+python inference_acc.py model_name
+```
+
+This script supports `gpt-j-6b`, `gpt-neox`, `opt` (30B version) and `T0pp` out of the box, but you can specify any valid checkpoint for `model_name`.
+
+To force a different `torch_dtype` than the one in the config: `--torch_dtype xxx`.
+
+If you get an error linked to disk offload, you need to add the option `--disk-offload`
+
+## Results
+
+On a setup with two Titan RTXs (24GB of RAM) and 32GB of RAM, we get the following benchmarks (T0pp does not run in float16, which is why it's not included).
+
+| Model | Model load time | Generation time | dtype | GPU 0 use | GPU 1 use | CPU use | Disk offload |
+|:-----:|:---------------:|:---------------:|:-----:|:---------:|:---------:|:-------:|:------------:|
+| GPT-J-6B | 8.7s | 0.05s per token | float16 | 11.7GB | 0GB | 0GB | no |
+| GPT-J-6B | 12.4s | 0.06s per token | float32 | 21.9GB | 1.5GB | 0GB | no |
+| GPT-Neo-X-20B | 30.9s | 0.08s per token | float16 | 21.5GB | 18GB | 0GB | no |
+| GPT-Neo-X-20B | 78.2s | 10.72s per token | float32 | 20.3GB | 22.7 GB | 24.4GB | yes |
+| T0pp (11B) | 29.4s | 0.05s per token | float32 | 21.1GB | 21.3GB | 0GB | no |
+| OPT-30B | 34.5s | 2.37s per token | float16 | 20.7GB | 22.3GB | 14.1GB | no |
+| OPT-30B | 112.3s | 33.9s per token | float32 | 20.2GB | 21.2GB | 23.5GB | yes |
+
+Note on the results:
+- using two GPUs instead of one does not slow down generation
+- using CPU offload slows down a bit (see OPT-30b)
+- using disk offload slows down a lot (need to implement prefetching)
+
+You will also note that Accelerate does not use anymore GPU and CPU RAM than necessary:
+- peak GPU memory is exactly the size of the model put on a given GPU
+- peak CPU memory is either the size of the biggest checkpoint shard or the part of the model offloaded on CPU, whichever is bigger.
\ No newline at end of file
diff --git a/benchmarks/big_model_inference.py b/benchmarks/big_model_inference.py
new file mode 100644
index 000000000..cb832d128
--- /dev/null
+++ b/benchmarks/big_model_inference.py
@@ -0,0 +1,143 @@
+# 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 time
+
+import torch
+
+import transformers
+from accelerate.utils import compute_module_sizes
+from measures_util import end_measure, log_measures, start_measure
+from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer
+
+
+DEFAULT_MODELS = {
+    "gpt-j-6b": {"is_causal": True, "model": "sgugger/sharded-gpt-j-6B", "tokenizer": "EleutherAI/gpt-j-6B"},
+    "gpt-neox": {"is_causal": True, "model": "EleutherAI/gpt-neox-20b"},
+    "opt": {"is_causal": True, "model": "facebook/opt-30b"},
+    "T0pp": {"is_causal": False, "model": "bigscience/T0pp", "model_revision": "sharded"},
+}
+
+PROMPTS = [
+    "Hello, my name is",
+    "Are unicorns real? Unicorns are",
+    "For the first time in several years,",
+    "My name is Julien and I am",
+    "The goal of life is",
+    "Whenever I'm sad, I like to",
+]
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Run and time generations on a big model using Accelerate.")
+    parser.add_argument("model_name", type=str, default=None, help="The name of the model to try.")
+    parser.add_argument(
+        "--tokenizer_name", type=str, default=None, help="The name of the tokenizer (if different from the model."
+    )
+    parser.add_argument("--is_causal", type=bool, default=None, help="Whether or not the model is causal.")
+    parser.add_argument(
+        "--model_revision", type=str, default=None, help="The revision to use for the model checkpoint."
+    )
+    parser.add_argument("--torch_dtype", type=str, default=None, help="The dtype for the model.")
+    parser.add_argument("--disk_offload", action="store_true")
+
+    args = parser.parse_args()
+
+    # Sanitize args
+    if args.model_name in DEFAULT_MODELS:
+        defaults = DEFAULT_MODELS[args.model_name]
+        args.model_name = defaults["model"]
+        if args.tokenizer_name is None:
+            args.tokenizer_name = defaults.get("tokenizer", args.model_name)
+        if args.is_causal is None:
+            args.is_causal = defaults["is_causal"]
+        if args.model_revision is None:
+            args.model_revision = defaults.get("model_revision", "main")
+
+    if args.is_causal is None:
+        raise ValueError("Could not infer the default for `--is_causal`, pass either True or False for it.")
+    if args.tokenizer_name is None:
+        args.tokenizer_name = args.model_name
+    if args.model_revision is None:
+        args.model_revision = "main"
+
+    return args
+
+
+def main():
+    transformers.utils.logging.set_verbosity_error()
+    args = parse_args()
+
+    if args.torch_dtype is None:
+        config = AutoConfig.from_pretrained(args.model_name)
+        torch_dtype = getattr(config, "torch_dtype", torch.float32)
+    else:
+        torch_dtype = getattr(torch, args.torch_dtype)
+    model_cls = AutoModelForCausalLM if args.is_causal else AutoModelForSeq2SeqLM
+    kwargs = {
+        "torch_dtype": torch_dtype,
+        "revision": args.model_revision,
+    }
+    if args.disk_offload:
+        kwargs["offload_folder"] = "tmp_offload"
+        kwargs["offload_state_dict"] = True
+
+    start_measures = start_measure()
+    model = model_cls.from_pretrained(args.model_name, device_map="auto", **kwargs)
+    end_measures = end_measure(start_measures)
+    log_measures(end_measures, "Model loading")
+
+    module_sizes = compute_module_sizes(model)
+    device_size = {v: 0 for v in model.hf_device_map.values()}
+    for module, device in model.hf_device_map.items():
+        device_size[device] += module_sizes[module]
+    message = "\n".join([f"- {device}: {size // 2**20}MiB" for device, size in device_size.items()])
+    print(f"\nTheoretical use:\n{message}")
+
+    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name)
+
+    start_measures = start_measure()
+    generation_times = []
+    gen_tokens = []
+    texts_outs = []
+    for prompt in PROMPTS:
+        inputs = tokenizer(prompt, return_tensors="pt").to(0)
+        tokens = inputs["input_ids"][0].tolist()
+        before_generate = time.time()
+        outputs = model.generate(inputs["input_ids"])
+        after_generate = time.time()
+        outputs = outputs[0].tolist()
+        num_gen_tokens = len(outputs) if outputs[: len(tokens)] != tokens else len(outputs) - len(tokens)
+        generation_time = after_generate - before_generate
+
+        text_out = tokenizer.decode(outputs, skip_special_tokens=True)
+        texts_outs.append(text_out)
+        generation_times.append(generation_time)
+        gen_tokens.append(num_gen_tokens)
+        print(f"Prompt: {prompt}\nGeneration {text_out}\nIn {generation_time:.2f}s for {num_gen_tokens} tokens\n")
+
+    end_measures = end_measure(start_measures)
+    log_measures(end_measures, "Model generation")
+
+    generation_times_per_token = [gen / tok for gen, tok in zip(generation_times, gen_tokens)]
+    avg_gen = sum(generation_times_per_token) / len(generation_times)
+    print(f"Average time of generation per token: {avg_gen:.2f}s")
+    print(f"First generation (avg time per token): {generation_times_per_token[0]:.2f}s")
+    avg_gen = sum(generation_times_per_token[1:]) / (len(generation_times_per_token) - 1)
+    print(f"Average time of generation per token (excluding the first): {avg_gen:.2f}s")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/benchmarks/measures_util.py b/benchmarks/measures_util.py
new file mode 100644
index 000000000..b6ac76b54
--- /dev/null
+++ b/benchmarks/measures_util.py
@@ -0,0 +1,86 @@
+import gc
+import threading
+import time
+
+import torch
+
+import psutil
+
+
+class PeakCPUMemory:
+    def __init__(self):
+        self.process = psutil.Process()
+        self.peak_monitoring = False
+
+    def peak_monitor(self):
+        self.cpu_memory_peak = -1
+
+        while True:
+            self.cpu_memory_peak = max(self.process.memory_info().rss, self.cpu_memory_peak)
+
+            # can't sleep or will not catch the peak right (this comment is here on purpose)
+            if not self.peak_monitoring:
+                break
+
+    def start(self):
+        self.peak_monitoring = True
+        self.thread = threading.Thread(target=self.peak_monitor)
+        self.thread.daemon = True
+        self.thread.start()
+
+    def stop(self):
+        self.peak_monitoring = False
+        self.thread.join()
+        return self.cpu_memory_peak
+
+
+cpu_peak_tracker = PeakCPUMemory()
+
+
+def start_measure():
+    # Time
+    measures = {"time": time.time()}
+
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    # CPU mem
+    measures["cpu"] = psutil.Process().memory_info().rss
+    cpu_peak_tracker.start()
+
+    # GPU mem
+    for i in range(torch.cuda.device_count()):
+        measures[str(i)] = torch.cuda.memory_allocated(i)
+    torch.cuda.reset_peak_memory_stats()
+
+    return measures
+
+
+def end_measure(start_measures):
+    # Time
+    measures = {"time": time.time() - start_measures["time"]}
+
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    # CPU mem
+    measures["cpu"] = (psutil.Process().memory_info().rss - start_measures["cpu"]) / 2**20
+    measures["cpu-peak"] = (cpu_peak_tracker.stop() - start_measures["cpu"]) / 2**20
+
+    # GPU mem
+    for i in range(torch.cuda.device_count()):
+        measures[str(i)] = (torch.cuda.memory_allocated(i) - start_measures[str(i)]) / 2**20
+        measures[f"{i}-peak"] = (torch.cuda.max_memory_allocated(i) - start_measures[str(i)]) / 2**20
+
+    return measures
+
+
+def log_measures(measures, description):
+    print(f"{description}:")
+    print(f"- Time: {measures['time']:.2f}s")
+    for i in range(torch.cuda.device_count()):
+        print(f"- GPU {i} allocated: {measures[str(i)]:.2f}MiB")
+        peak = measures[f"{i}-peak"]
+        print(f"- GPU {i} peak: {peak:.2f}MiB")
+    print(f"- CPU RAM allocated: {measures['cpu']:.2f}MiB")
+    print(f"- CPU RAM peak: {measures['cpu-peak']:.2f}MiB")

```suggestion
These benchmarks use the `transformers` library:
```
@@ -0,0 +1,46 @@
+# Big model inference benchmarks
+
+Running inference with Accelerate on big models.
+
+## Setup
+
+These benchmarks use the `Transformers` library:
Do they not also need `measures_util`?

Maybe add a `pip install accelerate[bench]`?
@@ -0,0 +1,46 @@
+# Big model inference benchmarks
+
+Running inference with Accelerate on big models.
+
+## Setup
+
+These benchmarks use the `Transformers` library:
+
+```bash
+pip install transformers
+```
Forgot one file :grimacing: 
@@ -0,0 +1,46 @@
+# Big model inference benchmarks
+
+Running inference with Accelerate on big models.
+
+## Setup
+
+These benchmarks use the `Transformers` library:
+
+```bash
+pip install transformers
+```
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_506). All of your documentation changes will be reflected on that endpoint.
https://api.github.com/repos/huggingface/accelerate/pulls/506
diff --git a/src/accelerate/commands/config/config_args.py b/src/accelerate/commands/config/config_args.py
index fe3d05bf3..1ef68200b 100644
--- a/src/accelerate/commands/config/config_args.py
+++ b/src/accelerate/commands/config/config_args.py
@@ -156,6 +156,7 @@ def __post_init__(self):
 class SageMakerConfig(BaseConfig):
     ec2_instance_type: str
     iam_role_name: str
+    image_uri: str
     profile: Optional[str] = None
     region: str = "us-east-1"
     num_machines: int = 1
@@ -163,3 +164,5 @@ class SageMakerConfig(BaseConfig):
     pytorch_version: str = SAGEMAKER_PYTORCH_VERSION
     transformers_version: str = SAGEMAKER_TRANSFORMERS_VERSION
     py_version: str = SAGEMAKER_PYTHON_VERSION
+    sagemaker_inputs_file: str = None
+    sagemaker_metrics_file: str = None
diff --git a/src/accelerate/commands/config/sagemaker.py b/src/accelerate/commands/config/sagemaker.py
index c0b17bc88..786881edf 100644
--- a/src/accelerate/commands/config/sagemaker.py
+++ b/src/accelerate/commands/config/sagemaker.py
@@ -20,7 +20,7 @@
 from ...utils.dataclasses import ComputeEnvironment, SageMakerDistributedType
 from ...utils.imports import is_boto3_available
 from .config_args import SageMakerConfig
-from .config_utils import _ask_field, _convert_sagemaker_distributed_mode
+from .config_utils import _ask_field, _convert_sagemaker_distributed_mode, _convert_yes_no_to_bool
 
 
 if is_boto3_available():
@@ -120,6 +120,42 @@ def get_sagemaker_input():
         print(f'Accelerate will create an iam role "{iam_role_name}" using the provided credentials')
         _create_iam_role_for_sagemaker(iam_role_name)
 
+    is_custom_docker_image = _ask_field(
+        "Do you want to use custom Docker image? [yes/NO]: ",
+        _convert_yes_no_to_bool,
+        default=False,
+        error_message="Please enter yes or no.",
+    )
+    docker_image = None
+    if is_custom_docker_image:
+        docker_image = _ask_field("Enter your Docker image: ", lambda x: str(x).lower())
+
+    is_sagemaker_inputs_enabled = _ask_field(
+        "Do you want to provide SageMaker input channels with data locations? [yes/NO]: ",
+        _convert_yes_no_to_bool,
+        default=False,
+        error_message="Please enter yes or no.",
+    )
+    sagemaker_inputs_file = None
+    if is_sagemaker_inputs_enabled:
+        sagemaker_inputs_file = _ask_field(
+            "Enter the path to the SageMaker inputs TSV file with columns (channel_name, data_location): ",
+            lambda x: str(x).lower(),
+        )
+
+    is_sagemaker_metrics_enabled = _ask_field(
+        "Do you want to enable SageMaker metrics? [yes/NO]: ",
+        _convert_yes_no_to_bool,
+        default=False,
+        error_message="Please enter yes or no.",
+    )
+    sagemaker_metrics_file = None
+    if is_sagemaker_metrics_enabled:
+        sagemaker_metrics_file = _ask_field(
+            "Enter the path to the SageMaker metrics TSV file with columns (metric_name, metric_regex): ",
+            lambda x: str(x).lower(),
+        )
+
     distributed_type = _ask_field(
         "Which type of machine are you using? ([0] No distributed training, [1] data parallelism): ",
         _convert_sagemaker_distributed_mode,
@@ -155,6 +191,7 @@ def get_sagemaker_input():
     )
 
     return SageMakerConfig(
+        image_uri=docker_image,
         compute_environment=ComputeEnvironment.AMAZON_SAGEMAKER,
         distributed_type=distributed_type,
         use_cpu=False,
@@ -164,4 +201,6 @@ def get_sagemaker_input():
         iam_role_name=iam_role_name,
         mixed_precision=mixed_precision,
         num_machines=num_machines,
+        sagemaker_inputs_file=sagemaker_inputs_file,
+        sagemaker_metrics_file=sagemaker_metrics_file,
     )
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 421bfe064..ac91ad1e9 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -493,13 +493,44 @@ def sagemaker_launcher(sagemaker_config: SageMakerConfig, args):
         "SAGEMAKER_DISTRIBUTED_TYPE": sagemaker_config.distributed_type.value,
     }
     # configure distribution set up
-    distribution = None  # TODO: not yet implemented
+    distribution = None
     if sagemaker_config.distributed_type == SageMakerDistributedType.DATA_PARALLEL:
         distribution = {"smdistributed": {"dataparallel": {"enabled": True}}}
 
+    # configure sagemaker inputs
+    sagemaker_inputs = None
+    if sagemaker_config.sagemaker_inputs_file is not None:
+        print(f"Loading SageMaker Inputs from {sagemaker_config.sagemaker_inputs_file} file")
+        sagemaker_inputs = {}
+        with open(sagemaker_config.sagemaker_inputs_file) as file:
+            for i, line in enumerate(file):
+                if i == 0:
+                    continue
+                l = line.split("\t")
+                sagemaker_inputs[l[0]] = l[1].strip()
+        print(f"Loaded SageMaker Inputs: {sagemaker_inputs}")
+
+    # configure sagemaker metrics
+    sagemaker_metrics = None
+    if sagemaker_config.sagemaker_metrics_file is not None:
+        print(f"Loading SageMaker Metrics from {sagemaker_config.sagemaker_metrics_file} file")
+        sagemaker_metrics = []
+        with open(sagemaker_config.sagemaker_metrics_file) as file:
+            for i, line in enumerate(file):
+                if i == 0:
+                    continue
+                l = line.split("\t")
+                metric_dict = {
+                    "Name": l[0],
+                    "Regex": l[1].strip(),
+                }
+                sagemaker_metrics.append(metric_dict)
+        print(f"Loaded SageMaker Metrics: {sagemaker_metrics}")
+
     # configure session
     print("Creating Estimator")
     huggingface_estimator = HuggingFace(
+        image_uri=sagemaker_config.image_uri,
         entry_point=entry_point,
         source_dir=source_dir,
         role=sagemaker_config.iam_role_name,
@@ -513,9 +544,10 @@ def sagemaker_launcher(sagemaker_config: SageMakerConfig, args):
         distribution=distribution,
         hyperparameters=hyperparameters,
         environment=environment,
+        metric_definitions=sagemaker_metrics,
     )
 
-    huggingface_estimator.fit()
+    huggingface_estimator.fit(inputs=sagemaker_inputs)
     print(f"You can find your model data at: {huggingface_estimator.model_data}")
 
 

I think we should default to the HF Image? 
@@ -156,10 +156,13 @@ def __post_init__(self):
 class SageMakerConfig(BaseConfig):
     ec2_instance_type: str
     iam_role_name: str
+    image_uri: str
Does accelerate output some default metrics, e.g. `loss` then we could already define them here? Or is there a "logging" schema, that way we could define a whole bunch of metrics already. 
@@ -156,10 +156,13 @@ def __post_init__(self):
 class SageMakerConfig(BaseConfig):
     ec2_instance_type: str
     iam_role_name: str
+    image_uri: str
     profile: Optional[str] = None
     region: str = "us-east-1"
     num_machines: int = 1
     base_job_name: str = f"accelerate-sagemaker-{num_machines}"
     pytorch_version: str = SAGEMAKER_PYTORCH_VERSION
     transformers_version: str = SAGEMAKER_TRANSFORMERS_VERSION
     py_version: str = SAGEMAKER_PYTHON_VERSION
+    sagemaker_inputs_file: str = None
+    sagemaker_metrics_file: str = None
`Accelerate` doesn't log any metrics by default, it enables running the same PyTorch code on any setup and as such user is in full control of what is being logged (flexibility).

@@ -156,10 +156,13 @@ def __post_init__(self):
 class SageMakerConfig(BaseConfig):
     ec2_instance_type: str
     iam_role_name: str
+    image_uri: str
     profile: Optional[str] = None
     region: str = "us-east-1"
     num_machines: int = 1
     base_job_name: str = f"accelerate-sagemaker-{num_machines}"
     pytorch_version: str = SAGEMAKER_PYTORCH_VERSION
     transformers_version: str = SAGEMAKER_TRANSFORMERS_VERSION
     py_version: str = SAGEMAKER_PYTHON_VERSION
+    sagemaker_inputs_file: str = None
+    sagemaker_metrics_file: str = None
Hello, HF image is used by default if `image_uri` isn't specified and it is based on the `Python`, `PyTorch` and `Transformers` versions.
@@ -156,10 +156,13 @@ def __post_init__(self):
 class SageMakerConfig(BaseConfig):
     ec2_instance_type: str
     iam_role_name: str
+    image_uri: str
_The documentation is not available anymore as the PR was closed or merged._
Wow @pacman100 you responded to that request quickly. Thanks everyone, the Accelerate team is awesome.
https://api.github.com/repos/huggingface/accelerate/pulls/504
diff --git a/examples/by_feature/gradient_accumulation.py b/examples/by_feature/gradient_accumulation.py
index d3b253bce..403587f10 100644
--- a/examples/by_feature/gradient_accumulation.py
+++ b/examples/by_feature/gradient_accumulation.py
@@ -137,7 +137,7 @@ def training_function(config, args):
     lr_scheduler = get_linear_schedule_with_warmup(
         optimizer=optimizer,
         num_warmup_steps=100,
-        num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,
+        num_training_steps=(len(train_dataloader) * num_epochs),
     )
 
     # Prepare everything
diff --git a/src/accelerate/optimizer.py b/src/accelerate/optimizer.py
index c5ed66cd4..4fad12c72 100644
--- a/src/accelerate/optimizer.py
+++ b/src/accelerate/optimizer.py
@@ -39,6 +39,9 @@ class AcceleratedOptimizer(torch.optim.Optimizer):
     """
     Internal wrapper around a torch optimizer.
 
+    Conditionally will perform `step` and `zero_grad` if gradients should be synchronized when performing gradient
+    accumulation.
+
     Args:
         optimizer (`torch.optim.optimizer.Optimizer`):
             The optimizer to wrap.
diff --git a/src/accelerate/scheduler.py b/src/accelerate/scheduler.py
index cfd6d719f..75695091e 100644
--- a/src/accelerate/scheduler.py
+++ b/src/accelerate/scheduler.py
@@ -16,7 +16,7 @@
 
 import warnings
 
-from .state import AcceleratorState, GradientState
+from .state import AcceleratorState
 
 
 warnings.filterwarnings("ignore", category=UserWarning, module="torch.optim.lr_scheduler")
@@ -25,10 +25,11 @@
 class AcceleratedScheduler:
     """
     A wrapper around a learning rate scheduler that will only step when the optimizer(s) have a training step. Useful
-    to avoid making a scheduler step too fast when:
+    to avoid making a scheduler step too fast when gradients went overflow and there was no training step (in mixed
+    precision training)
 
-    - gradients went overflow and there was no training step (in mixed precision training)
-    - step was skipped because of gradient accumulation
+    When performing gradient accumulation scheduler lengths should not be changed accordingly, accelerate will always
+    step the scheduler to account for it.
 
     Args:
         scheduler (`torch.optim.lr_scheduler._LRScheduler`):
@@ -47,7 +48,6 @@ def __init__(self, scheduler, optimizers, step_with_optimizer: bool = True, spli
         self.scheduler = scheduler
         self.optimizers = optimizers if isinstance(optimizers, (list, tuple)) else [optimizers]
         self.split_batches = split_batches
-        self.gradient_state = GradientState()
         self.step_with_optimizer = step_with_optimizer
 
     def step(self, *args, **kwargs):

```suggestion
    When performing gradient accumulation scheduler lengths should not be changed accordingly, accelerate will
    always step the scheduler to account for it.
```
@@ -25,10 +25,11 @@
 class AcceleratedScheduler:
     """
     A wrapper around a learning rate scheduler that will only step when the optimizer(s) have a training step. Useful
-    to avoid making a scheduler step too fast when:
+    to avoid making a scheduler step too fast when gradients went overflow and there was no training step (in mixed
+    precision training)
 
-    - gradients went overflow and there was no training step (in mixed precision training)
-    - step was skipped because of gradient accumulation
+    When performing gradient accumulation scheduler lengths should not be changed accordingly, accelerate will handle
+    the right times to step.
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/500
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 806f0094d..a0278218b 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -1,6 +1,6 @@
 name: Run Tests
 
-on: [pull_request]
+# on: [pull_request]
 
 env:
   HF_HOME: ~/hf_cache
diff --git a/Makefile b/Makefile
index 9cae9d74d..100a1484a 100644
--- a/Makefile
+++ b/Makefile
@@ -58,4 +58,4 @@ test_prod:
 	$(MAKE) test_core
 
 test_rest:
-	python -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k "not by_step and not by_epoch"
\ No newline at end of file
+	python -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k "not by_step and not by_epoch"
diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml
index c2fd77282..0da843ec2 100644
--- a/docs/source/_toctree.yml
+++ b/docs/source/_toctree.yml
@@ -1,38 +1,68 @@
-- sections: 
+- sections:
   - local: index
     title: 🤗 Accelerate
-  - local: quicktour
-    title: Quick tour
-  - local: installation
+  - local: basic_tutorials/install
     title: Installation
-  title: Get started
+  - local: quicktour
+    title: Quicktour
+  title: Getting started
+- sections:
+  - local: basic_tutorials/overview
+    title: Overview
+  - local: basic_tutorials/migration
+    title: Migrating to 🤗 Accelerate
+  - local: basic_tutorials/launch
+    title: Launching distributed code
+  - local: basic_tutorials/notebook
+    title: Launching distributed training from Jupyter Notebooks
+  title: Tutorials
 - sections:
-  - local: big_modeling
-    title: Handling big models
-  - local: gradient_accumulation
-    title: Gradient accumulation
-  - local: sagemaker
-    title: Amazon SageMaker
-  title: Guides
+  - local: usage_guides/gradient_accumulation
+    title: Performing gradient accumulation
+  - local: usage_guides/fsdp
+    title: Fully Sharded Data Parallelism
+  - local: usage_guides/checkpoint
+    title: Saving and loading training states
+  - local: usage_guides/deepspeed
+    title: How to use DeepSpeed
+  - local: usage_guides/tracking
+    title: Using experiment trackers
+  - local: usage_guides/big_modeling
+    title: How to use large models with small resources
+  - local: usage_guides/memory
+    title: How to avoid CUDA Out-of-Memory
+  - local: usage_guides/sagemaker
+    title: Using Accelerate on SageMaker
+  title: How-To Guides
 - sections:
-  - local: accelerator
-    title: Accelerator
-  - local: launcher
-    title: Notebook Launcher
-  - local: kwargs
-    title: Kwargs Handlers
-  - local: checkpoint
-    title: Checkpointing
-  - local: internal
-    title: Internals
-  - local: tracking
-    title: Experiment Tracking
-  - local: fsdp
-    title: Fully Sharded Data Parallel
-  - local: memory
-    title: Memory Utilities
-  - local: deepspeed
-    title: DeepSpeed
-  - local: utilities
-    title: General Utilities
-  title: API Reference
+  - local: concept_guides/gradient_synchronization
+    title: Gradient synchronization
+  - local: concept_guides/deferring_execution
+    title: Executing and deferring jobs
+  - local: concept_guides/training_tpu
+    title: TPU best practices
+  title: Concepts and fundamentals
+- sections: 
+  - local: package_reference/accelerator
+    title: Main Accelerator class
+  - local: package_reference/state
+    title: Stateful configuration classes
+  - local: package_reference/cli
+    title: The Command Line
+  - local: package_reference/torch_wrappers
+    title: Torch wrapper classes
+  - local: package_reference/tracking
+    title: Experiment trackers
+  - local: package_reference/launchers
+    title: Distributed launchers
+  - local: package_reference/deepspeed
+    title: DeepSpeed utilities
+  - local: package_reference/logging
+    title: Logging
+  - local: package_reference/big_modeling
+    title: Working with large models
+  - local: package_reference/kwargs
+    title: Kwargs handlers
+  - local: package_reference/utilities
+    title: Utility functions and classes
+  title: "Reference"
\ No newline at end of file
diff --git a/docs/source/accelerator.mdx b/docs/source/accelerator.mdx
deleted file mode 100644
index 268b09d1f..000000000
--- a/docs/source/accelerator.mdx
+++ /dev/null
@@ -1,65 +0,0 @@
-<!--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.
--->
-
-# Accelerator
-
-The [`Accelerator`] is the main class provided by 🤗 Accelerate. It serves at the main entrypoint for
-the API. To quickly adapt your script to work on any kind of setup with 🤗 Accelerate just:
-
-1. Initialize an [`Accelerator`] object (that we will call `accelerator` in the rest of this
-   page) as early as possible in your script.
-2. Pass along your model(s), optimizer(s), dataloader(s) to the [`~Accelerator.prepare`] method.
-3. (Optional but best practice) Remove all the `.cuda()` or `.to(device)` in your code and let the
-   `accelerator` handle device placement for you.
-4. Replace the `loss.backward()` in your code by `accelerator.backward(loss)`.
-5. (Optional, when using distributed evaluation) Gather your predictions and labels before storing them or using
-   them for metric computation using [`~Accelerator.gather`].
-
-This is all that is needed in most cases. For more advanced cases or a nicer experience here are the functions you
-should search for and replace by the corresponding methods of your `accelerator`:
-
-- `print` statements should be replaced by [`~Accelerator.print`] to be only printed once per
-  process.
-- Use [`~Accelerator.is_local_main_process`] for statements that should be executed once per server.
-- Use [`~Accelerator.is_main_process`] for statements that should be executed once only.
-- Use [`~Accelerator.wait_for_everyone`] to make sure all processes join that point before continuing
-  (useful before a model save for instance).
-- Use [`~Accelerator.unwrap_model`] to unwrap your model before saving it.
-- Use [`~Accelerator.save`] instead of `torch.save`.
-- Use [`~Accelerator.clip_grad_norm_`] instead of `torch.nn.utils.clip_grad_norm_` and
-  [`~Accelerator.clip_grad_value_`] instead of `torch.nn.utils.clip_grad_value_`.
-
-To perform gradient accumulation use [`~Accelerator.accumulate`] and specify a `gradient_accumulation_steps`. 
-This will also automatically ensure the gradients are synced or unsynced when on multi-device training, check if the step should
-actually be performed, and auto-scale the loss:
-
-```python
-accelerator = Accelerator(gradient_accumulation_steps=2)
-model, optimizer, training_dataloader = accelerator.prepare(model, optimizer, training_dataloader)
-
-for input, label in training_dataloader:
-    with accelerator.accumulate(model):
-        predictions = model(input)
-        loss = loss_function(predictions, label)
-        accelerator.backward(loss)
-        optimizer.step()
-        scheduler.step()
-        optimizer.zero_grad()
-```
-
-<Tip warning={true}>
-
-Using this with `dispatch_batches=True`  (which is the default for iterable datasets) is currently not supported.
-
-</Tip>
-
-[[autodoc]] Accelerator
diff --git a/docs/source/basic_tutorials/install.mdx b/docs/source/basic_tutorials/install.mdx
new file mode 100644
index 000000000..19630f575
--- /dev/null
+++ b/docs/source/basic_tutorials/install.mdx
@@ -0,0 +1,99 @@
+<!--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.
+-->
+
+# Installation and Configuration
+
+Before you start, you will need to setup your environment, install the appropriate packages, and configure 🤗 Accelerate. 🤗 Accelerate is tested on **Python 3.7+**.
+
+## Installing 🤗 Accelerate
+
+🤗 Accelerate is available on pypi and conda, as well as on GitHub. Details to install from each are below:
+
+### pip 
+
+To install 🤗 Accelerate from pypi, perform:
+
+```bash
+pip install accelerate
+```
+
+### conda
+
+🤗 Accelerate can also be installed with conda with:
+
+```bash
+conda install -c conda-forge accelerate
+```
+
+### Source
+
+New features are added every day that haven't been released yet. To try them out yourself, install
+from the GitHub repository:
+
+```bash
+pip install git+https://github.com/huggingface/accelerate
+```
+
+If you're working on contributing to the library or wish to play with the source code and see live 
+results as you run the code, an editable version can be installed from a locally-cloned version of the 
+repository:
+
+```bash
+git clone https://github.com/huggingface/accelerate
+cd accelerate
+pip install -e .
+```
+
+## Configuring 🤗 Accelerate
+
+After installing, you need to configure 🤗 Accelerate for how the current system is setup for training. 
+To do so run the following and answer the questions prompted to you:
+
+```bash
+accelerate config
+```
+
+To write a barebones configuration that doesn't include options such as DeepSpeed configuration or running on TPUs, you can quickly run:
+
+```bash
+python -c "from accelerate.utils import write_basic_config; write_basic_config(mixed_precision='fp16')"
+```
+🤗 Accelerate will automatically utilize the maximum number of GPUs available and set the mixed precision mode.
+
+To check that your configuration looks fine, run:
+
+```bash
+accelerate env
+```
+
+An example output is shown below, which describes two GPUs on a single machine with no mixed precision being used:
+
+```bash
+- `Accelerate` version: 0.11.0.dev0
+- Platform: Linux-5.10.0-15-cloud-amd64-x86_64-with-debian-11.3
+- Python version: 3.7.12
+- Numpy version: 1.19.5
+- PyTorch version (GPU?): 1.12.0+cu102 (True)
+- `Accelerate` default config:
+        - compute_environment: LOCAL_MACHINE
+        - distributed_type: MULTI_GPU
+        - mixed_precision: no
+        - use_cpu: False
+        - num_processes: 2
+        - machine_rank: 0
+        - num_machines: 1
+        - main_process_ip: None
+        - main_process_port: None
+        - main_training_function: main
+        - deepspeed_config: {}
+        - fsdp_config: {}
+```
\ No newline at end of file
diff --git a/docs/source/basic_tutorials/launch.mdx b/docs/source/basic_tutorials/launch.mdx
new file mode 100644
index 000000000..741920f26
--- /dev/null
+++ b/docs/source/basic_tutorials/launch.mdx
@@ -0,0 +1,178 @@
+<!--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.
+-->
+
+# Launching your 🤗 Accelerate scripts
+
+In the previous tutorial, you were introduced to how to modify your current training script to use 🤗 Accelerate.
+The final version of that code is shown below:
+
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    accelerator.backward(loss)
+    optimizer.step()
+    scheduler.step()
+```
+
+But how do you run this code and have it utilize the special hardware available to it?
+
+First you should rewrite the above code into a function, and make it callable as a script. For example:
+
+```diff
+  from accelerate import Accelerator
+  
++ def main():
+      accelerator = Accelerator()
+
+      model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+          model, optimizer, training_dataloader, scheduler
+      )
+
+      for batch in training_dataloader:
+          optimizer.zero_grad()
+          inputs, targets = batch
+          outputs = model(inputs)
+          loss = loss_function(outputs, targets)
+          accelerator.backward(loss)
+          optimizer.step()
+          scheduler.step()
+
++ if __name__ == "__main__":
++     main()
+```
+
+Next you need to launch it with `accelerate launch`. 
+
+<Tip warning={true}>
+
+  It's recommended you run `accelerate config` before using `accelerate launch` to configure your environment to your liking. 
+  Otherwise 🤗 Accelerate will use very basic defaults depending on your system setup.
+
+</Tip>
+
+
+## Using accelerate launch
+
+🤗 Accelerate has a special CLI command to help you launch your code in your system through `accelerate launch`.
+This command wraps around all of the different commands needed to launch your script on various platforms, without you having to remember what each of them are.
+
+<Tip>
+
+  If you are familiar with launching scripts in PyTorch yourself such as with `torchrun`, you can still do this. It is not required to use `accelerate launch`.
+
+</Tip>
+
+You can launch your script quickly by using:
+
+```bash
+accelerate launch {script_name.py} --arg1 --arg2 ...
+```
+
+Just put `accelerate launch` at the start of your command, and pass in additional arguments and parameters to your script afterwards like normal!
+
+Since this runs the various torch spawn methods, all of the expected environment variables can be modified here as well.
+For example, here is how to use `accelerate launch` with a single GPU:
+
+```bash
+CUDA_VISIBLE_DEVICES="0" accelerate launch {script_name.py} --arg1 --arg2 ...
+```
+
+You can also use `accelerate launch` without performing `accelerate config` first, but you may need to manually pass in the right configuration parameters.
+In this case, 🤗 Accelerate will make some hyperparameter decisions for you, e.g., if GPUs are available, it will use all of them by default without the mixed precision.
+Here is how you would use all GPUs and train with mixed precision disabled:
+
+```bash
+accelerate launch --multi_gpu {script_name.py} {--arg1} {--arg2} ...
+```
+
+To get more specific you should pass in the needed parameters yourself. For instance, here is how you 
+would also launch that same script on two GPUs using mixed precision while avoiding all of the warnings: 
+
+```bash
+accelerate launch --multi_gpu --mixed_precision=fp16 --num_processes=2 {script_name.py} {--arg1} {--arg2} ...
+```
+
+For a complete list of parameters you can pass in, run:
+
+```bash
+accelerate launch -h
+```
+
+<Tip>
+
+  Even if you are not using 🤗 Accelerate in your code, you can still use the launcher for starting your scripts!
+
+</Tip>
+
+For a visualization of this difference, that earlier `accelerate launch` on multi-gpu would look something like so with `torchrun`:
+
+```bash
+MIXED_PRECISION="fp16" torchrun --nproc_per_node=2 --num_machines=1 {script_name.py} {--arg1} {--arg2} ...
+```
+
+## Why you should always use `accelerate config`
+
+Why is it useful to the point you should **always** run `accelerate config`? 
+
+Remember that earlier call to `accelerate launch` as well as `torchrun`?
+Post configuration, to run that script with the needed parts you just need to use `accelerate launch` outright, without passing anything else in:
+
+```bash
+accelerate launch {script_name.py} {--arg1} {--arg2} ...
+```
+
+
+## Custom Configurations
+
+As briefly mentioned earlier, `accelerate launch` should be mostly used through combining set configurations 
+made with the `accelerate config` command. These configs are saved to a `default_config.yaml` file in your cache folder for 🤗 Accelerate. 
+This cache folder is located at (with decreasing order of priority):
+
+- The content of your environment variable `HF_HOME` suffixed with `accelerate`.
+- If it does not exist, the content of your environment variable `XDG_CACHE_HOME` suffixed with
+  `huggingface/accelerate`.
+- If this does not exist either, the folder `~/.cache/huggingface/accelerate`.
+
+To have multiple configurations, the flag `--config_file` can be passed to the `accelerate launch` command paired 
+with the location of the custom yaml. 
+
+An example yaml may look something like the following for two GPUs on a single machine using `fp16` for mixed precision:
+```yaml
+compute_environment: LOCAL_MACHINE
+deepspeed_config: {}
+distributed_type: MULTI_GPU
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 2
+use_cpu: false
+```
+
+Launching a script from the location of that custom yaml file looks like the following:
+```bash
+accelerate launch --config_file {path/to/config/my_config_file.yaml} {script_name.py} {--arg1} {--arg2} ...
+```
\ No newline at end of file
diff --git a/docs/source/basic_tutorials/migration.mdx b/docs/source/basic_tutorials/migration.mdx
new file mode 100644
index 000000000..ab703c960
--- /dev/null
+++ b/docs/source/basic_tutorials/migration.mdx
@@ -0,0 +1,123 @@
+<!--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.
+-->
+
+# Migrating your code to 🤗 Accelerate
+
+This tutorial will detail how to easily convert existing PyTorch code to use 🤗 Accelerate!
+You'll see that by just changing a few lines of code, 🤗 Accelerate can perform its magic and get you on 
+your way towards running your code on distributed systems with ease!
+
+## The base training loop
+
+To begin, write out a very basic PyTorch training loop. 
+
+<Tip>
+
+    We are under the presumption that `training_dataloader`, `model`, `optimizer`, `scheduler`, and `loss_function` have been defined beforehand.
+
+</Tip>
+
+```python
+device = "cuda"
+model.to(device)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    inputs = inputs.to(device)
+    targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    loss.backward()
+    optimizer.step()
+    scheduler.step()
+```
+
+## Add in 🤗 Accelerate
+
+To start using 🤗 Accelerate, first import and create an [`Accelerator`] instance:
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+```
+[`Accelerator`] is the main force behind utilizing all the possible options for distributed training!
+
+### Setting the right device
+
+The [`Accelerator`] class knows the right device to move any PyTorch object to at any time, so you should
+change the definition of `device` to come from [`Accelerator`]:
+
+```diff
+- device = 'cuda'
++ device = accelerator.device
+  model.to(device)
+```
+
+### Preparing your objects
+
+Next you need to pass all of the important objects related to training into [`~Accelerator.prepare`]. 🤗 Accelerate will
+make sure everything is setup in the current environment for you to start training:
+
+```
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+```
+These objects are returned in the same order they were sent in with. By default when using `device_placement=True`, all of the objects that can be sent to the right device will be.
+If you need to work with data that isn't passed to [~Accelerator.prepare] but should be on the active device, you should pass in the `device` you made earlier. 
+
+<Tip warning={true}>
+
+    Accelerate will only prepare objects that inherit from their respective PyTorch classes (such as `torch.optim.Optimizer`).
+
+</Tip>
+
+### Modifying the training loop
+
+Finally, three lines of code need to be changed in the training loop. 🤗 Accelerate's DataLoader classes will automatically handle the device placement by default,
+and [`~Accelerator.backward`] should be used for performing the backward pass:
+
+```diff
+-   inputs = inputs.to(device)
+-   targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+-   loss.backward()
++   accelerator.backward(loss)
+```
+
+With that, your training loop is now ready to use 🤗 Accelerate!
+
+## The finished code
+
+Below is the final version of the converted code: 
+
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    accelerator.backward(loss)
+    optimizer.step()
+    scheduler.step()
+```
+
diff --git a/docs/source/basic_tutorials/notebook.mdx b/docs/source/basic_tutorials/notebook.mdx
new file mode 100644
index 000000000..903a992d3
--- /dev/null
+++ b/docs/source/basic_tutorials/notebook.mdx
@@ -0,0 +1,429 @@
+<!--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.
+-->
+
+# Launching Multi-Node Training from a Jupyter Environment
+
+This tutorial teaches you how to fine tune a computer vision model with 🤗 Accelerate from a Jupyter Notebook on a distributed system.
+You will also learn how to setup a few requirements needed for ensuring your environment is configured properly, your data has been prepared properly, and finally how to launch training.
+
+<Tip>
+
+    This tutorial is also available as a Jupyter Notebook [here](https://github.com/huggingface/notebooks/blob/main/examples/accelerate_examples/simple_cv_example.ipynb)
+
+</Tip>
+
+## Configuring the Environment
+
+Before any training can be performed, a 🤗 Accelerate config file must exist in the system. Usually this can be done by running the following in a terminal and answering the prompts:
+
+```bash
+accelerate config
+```
+
+However, if general defaults are fine and you are *not* running on a TPU, 🤗Accelerate has a utility to quickly write your GPU configuration into a config file via [`utils.write_basic_config`].
+
+The following code will restart Jupyter after writing the configuration, as CUDA code was called to perform this. 
+
+<Tip warning={true}>
+
+    CUDA can't be initialized more than once on a multi-node system. It's fine to debug in the notebook and have calls to CUDA, but in order to finally train a full cleanup and restart will need to be performed.
+    
+</Tip>
+
+```python
+import os
+from accelerate.utils import write_basic_config
+
+write_basic_config()  # Write a config file
+os._exit(00)  # Restart the notebook
+```
+
+## Preparing the Dataset and Model
+
+Next you should prepare your dataset. As mentioned at earlier, great care should be taken when preparing the `DataLoaders` and model to make sure that **nothing** is put on *any* GPU. 
+
+If you do, it is recommended to put that specific code into a function and call that from within the notebook launcher interface, which will be shown later. 
+
+Make sure the dataset is downloaded based on the directions [here](https://github.com/huggingface/accelerate/tree/main/examples#simple-vision-example)
+
+```python
+import os, re, torch, PIL
+import numpy as np
+
+from torch.optim.lr_scheduler import OneCycleLR
+from torch.utils.data import DataLoader, Dataset
+from torchvision.transforms import Compose, RandomResizedCrop, Resize, ToTensor
+
+from accelerate import Accelerator
+from accelerate.utils import set_seed
+from timm import create_model
+```
+
+First you need to create a function to extract the class name based on a filename:
+
+```python
+import os
+
+data_dir = "../../images"
+fnames = os.listdir(data_dir)
+fname = fnames[0]
+print(fname)
+```
+
+```python out
+beagle_32.jpg
+```
+
+In the case here, the label is `beagle`. Using regex you can extract the label from the filename:
+
+```python
+import re
+
+
+def extract_label(fname):
+    stem = fname.split(os.path.sep)[-1]
+    return re.search(r"^(.*)_\d+\.jpg$", stem).groups()[0]
+```
+
+```python
+extract_label(fname)
+```
+
+And you can see it properly returned the right name for our file:
+
+```python out
+"beagle"
+```
+
+Next a `Dataset` class should be made to handle grabbing the image and the label:
+
+```python
+class PetsDataset(Dataset):
+    def __init__(self, file_names, image_transform=None, label_to_id=None):
+        self.file_names = file_names
+        self.image_transform = image_transform
+        self.label_to_id = label_to_id
+
+    def __len__(self):
+        return len(self.file_names)
+
+    def __getitem__(self, idx):
+        fname = self.file_names[idx]
+        raw_image = PIL.Image.open(fname)
+        image = raw_image.convert("RGB")
+        if self.image_transform is not None:
+            image = self.image_transform(image)
+        label = extract_label(fname)
+        if self.label_to_id is not None:
+            label = self.label_to_id[label]
+        return {"image": image, "label": label}
+```
+
+Now to build the dataset. Outside the training function you can find and declare all the filenames and labels and use them as references inside the 
+launched function:
+
+```python
+fnames = [os.path.join("../../images", fname) for fname in fnames if fname.endswith(".jpg")]
+```
+
+Next gather all the labels:
+
+```python
+all_labels = [extract_label(fname) for fname in fnames]
+id_to_label = list(set(all_labels))
+id_to_label.sort()
+label_to_id = {lbl: i for i, lbl in enumerate(id_to_label)}
+```
+
+Next, you should make a `get_dataloaders` function that will return your built dataloaders for you. As mentioned earlier, if data is automatically 
+sent to the GPU or a TPU device when building your `DataLoaders`, they must be built using this method. 
+
+```python
+def get_dataloaders(batch_size: int = 64):
+    "Builds a set of dataloaders with a batch_size"
+    random_perm = np.random.permutation(len(fnames))
+    cut = int(0.8 * len(fnames))
+    train_split = random_perm[:cut]
+    eval_split = random_perm[:cut]
+
+    # For training a simple RandomResizedCrop will be used
+    train_tfm = Compose([RandomResizedCrop((224, 224), scale=(0.5, 1.0)), ToTensor()])
+    train_dataset = PetsDataset([fnames[i] for i in train_split], image_transform=train_tfm, label_to_id=label_to_id)
+
+    # For evaluation a deterministic Resize will be used
+    eval_tfm = Compose([Resize((224, 224)), ToTensor()])
+    eval_dataset = PetsDataset([fnames[i] for i in eval_split], image_transform=eval_tfm, label_to_id=label_to_id)
+
+    # Instantiate dataloaders
+    train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size, num_workers=4)
+    eval_dataloader = DataLoader(eval_dataset, shuffle=False, batch_size=batch_size * 2, num_workers=4)
+    return train_dataloader, eval_dataloader
+```
+
+Finally, you should import the scheduler to be used later:
+
+```python
+from torch.optim.lr_scheduler import CosineAnnealingLR
+```
+
+## Writing the Training Function
+
+Now you can build the training loop. [`notebook_launcher`] works by passing in a function to call that will be ran across the distributed system.
+
+Here is a basic training loop for the animal classification problem:
+
+<Tip>
+
+    The code has been split up to allow for explainations on each section. A full version that can be copy and pasted will be available at the end
+
+</Tip>
+
+
+```python
+def training_loop(mixed_precision="fp16", seed: int = 42, batch_size: int = 64):
+    set_seed(seed)
+    accelerator = Accelerator(mixed_precision=mixed_precision)
+```
+
+First you should set the seed and create an [`Accelerator`] object as early in the training loop as possible.
+
+<Tip warning={true}>
+
+    If training on the TPU, your training loop should take in the model as a parameter and it should be instantiated 
+    outside of the training loop function. See the [TPU best practices](../concept_guides/training_tpu) 
+    to learn why
+
+</Tip>
+
+Next you should build your dataloaders and create your model:
+
+```python
+    train_dataloader, eval_dataloader = get_dataloaders(batch_size)
+    model = create_model("resnet50d", pretrained=True, num_classes=len(label_to_id))
+```
+
+<Tip>
+
+    You build the model here so that the seed also controls the new weight initialization
+
+</Tip>
+
+As you are performing transfer learning in this example, the encoder of the model starts out frozen so the head of the model can be 
+trained only initially:
+
+```python
+    for param in model.parameters():
+        param.requires_grad = False
+    for param in model.get_classifier().parameters():
+        param.requires_grad = True
+```
+
+Normalizing the batches of images will make training a little faster:
+
+```python
+    mean = torch.tensor(model.default_cfg["mean"])[None, :, None, None]
+    std = torch.tensor(model.default_cfg["std"])[None, :, None, None]
+```
+
+To make these constants available on the active device, you should set it to the Accelerator's device:
+
+```python
+    mean = mean.to(accelerator.device)
+    std = std.to(accelerator.device)
+```
+
+Next instantiate the rest of the PyTorch classes used for training:
+
+```python
+    optimizer = torch.optim.Adam(params=model.parameters(), lr=3e-2 / 25)
+    lr_scheduler = OneCycleLR(optimizer=optimizer, max_lr=3e-2, epochs=5, steps_per_epoch=len(train_dataloader))
+```
+
+Before passing everything to [`~Accelerator.prepare`].
+
+<Tip>
+
+    There is no specific order to remember, you just need to unpack the objects in the same order you gave them to the prepare method.
+
+</Tip>
+
+```python
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+```
+
+Now train the model:
+
+```python
+    for epoch in range(5):
+        model.train()
+        for batch in train_dataloader:
+            inputs = (batch["image"] - mean) / std
+            outputs = model(inputs)
+            loss = torch.nn.functional.cross_entropy(outputs, batch["label"])
+            accelerator.backward(loss)
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
+```
+
+The evaluation loop will look slightly different compared to the training loop. The number of elements passed as well as the overall 
+total accuracy of each batch will be added to two constants:
+
+```python
+        model.eval()
+        accurate = 0
+        num_elems = 0
+```
+
+Next you have the rest of your standard PyTorch loop:
+
+```python
+        for batch in eval_dataloader:
+            inputs = (batch["image"] - mean) / std
+            with torch.no_grad():
+                outputs = model(inputs)
+            predictions = outputs.argmax(dim=-1)
+```
+
+Before finally the last major difference. 
+
+When performing distributed evaluation, the predictions and labels need to be passed through 
+[`~Accelerator.gather`] so that all of the data is available on the current device and a properly calculated metric can be achieved:
+
+```python
+            accurate_preds = accelerator.gather(predictions) == accelerator.gather(batch["label"])
+            num_elems += accurate_preds.shape[0]
+            accurate += accurate_preds.long().sum()
+```
+
+Now you just need to calculate the actual metric for this problem, and you can print it on the main process using [`~Accelerator.print`]:
+
+```python
+        eval_metric = accurate.item() / num_elems
+        accelerator.print(f"epoch {epoch}: {100 * eval_metric:.2f}")
+```
+
+A full version of this training loop is available below:
+
+```python
+def training_loop(mixed_precision="fp16", seed: int = 42, batch_size: int = 64):
+    set_seed(seed)
+    # Initialize accelerator
+    accelerator = Accelerator(mixed_precision=mixed_precision)
+    # Build dataloaders
+    train_dataloader, eval_dataloader = get_dataloaders(batch_size)
+
+    # Instantiate the model (you build the model here so that the seed also controls new weight initaliziations)
+    model = create_model("resnet50d", pretrained=True, num_classes=len(label_to_id))
+
+    # Freeze the base model
+    for param in model.parameters():
+        param.requires_grad = False
+    for param in model.get_classifier().parameters():
+        param.requires_grad = True
+
+    # You can normalize the batches of images to be a bit faster
+    mean = torch.tensor(model.default_cfg["mean"])[None, :, None, None]
+    std = torch.tensor(model.default_cfg["std"])[None, :, None, None]
+
+    # To make this constant available on the active device, set it to the accelerator device
+    mean = mean.to(accelerator.device)
+    std = std.to(accelerator.device)
+
+    # Intantiate the optimizer
+    optimizer = torch.optim.Adam(params=model.parameters(), lr=3e-2 / 25)
+
+    # Instantiate the learning rate scheduler
+    lr_scheduler = OneCycleLR(optimizer=optimizer, max_lr=3e-2, epochs=5, steps_per_epoch=len(train_dataloader))
+
+    # Prepare everything
+    # There is no specific order to remember, you just need to unpack the objects in the same order you gave them to the
+    # prepare method.
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    # Now you train the model
+    for epoch in range(5):
+        model.train()
+        for batch in train_dataloader:
+            inputs = (batch["image"] - mean) / std
+            outputs = model(inputs)
+            loss = torch.nn.functional.cross_entropy(outputs, batch["label"])
+            accelerator.backward(loss)
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
+
+        model.eval()
+        accurate = 0
+        num_elems = 0
+        for batch in eval_dataloader:
+            inputs = (batch["image"] - mean) / std
+            with torch.no_grad():
+                outputs = model(inputs)
+            predictions = outputs.argmax(dim=-1)
+            accurate_preds = accelerator.gather(predictions) == accelerator.gather(batch["label"])
+            num_elems += accurate_preds.shape[0]
+            accurate += accurate_preds.long().sum()
+
+        eval_metric = accurate.item() / num_elems
+        # Use accelerator.print to print only on the main process.
+        accelerator.print(f"epoch {epoch}: {100 * eval_metric:.2f}")
+```
+
+## Using the notebook_launcher
+
+All that's left is to use the [`notebook_launcher`].
+
+You pass in the function, the arguments (as a tuple), and the number of processes to train on. (See the [documentation](../package_reference/launchers) for more information)
+
+```python
+from accelerate import notebook_launcher
+```
+
+```python
+args = ("fp16", 42, 64)
+notebook_launcher(training_loop, args, num_processes=2)
+```
+
+In the case of running on the TPU, it would look like so:
+
+```python
+model = create_model("resnet50d", pretrained=True, num_classes=len(label_to_id))
+
+args = (model, "fp16", 42, 64)
+notebook_launcher(training_loop, args, num_processes=8)
+```
+
+As it's running it will print the progress as well as state how many devices you ran on. This tutorial was ran with two GPUs:
+
+```python out
+Launching training on 2 GPUs.
+epoch 0: 88.12
+epoch 1: 91.73
+epoch 2: 92.58
+epoch 3: 93.90
+epoch 4: 94.71
+```
+
+And that's it!
+
+## Conclusion
+
+This notebook showed how to perform distributed training from inside of a Jupyter Notebook. Some key notes to remember:
+
+- Make sure to save any code that use CUDA (or CUDA imports) for the function passed to [`notebook_launcher`]
+- Set the `num_processes` to be the number of devices used for training (such as number of GPUs, CPUs, TPUs, etc)
+- If using the TPU, declare your model outside the training loop function
\ No newline at end of file
diff --git a/docs/source/basic_tutorials/overview.mdx b/docs/source/basic_tutorials/overview.mdx
new file mode 100644
index 000000000..59ff9cbae
--- /dev/null
+++ b/docs/source/basic_tutorials/overview.mdx
@@ -0,0 +1,21 @@
+<!--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.
+-->
+
+# Overview
+
+Welcome to the 🤗 Accelerate tutorials! These introductory guides will help catch you up to speed on working with 🤗 Accelerate.
+You'll learn how to modify your code to have it work with the API seamlessly, how to launch your script properly,
+and more!
+
+These tutorials assume some basic knowledge of Python and familiarity with the PyTorch framework.
+
+If you have any questions about 🤗 Accelerate, feel free to join and ask the community on our [forum](https://discuss.huggingface.co/c/accelerate/18).
\ No newline at end of file
diff --git a/docs/source/concept_guides/deferring_execution.mdx b/docs/source/concept_guides/deferring_execution.mdx
new file mode 100644
index 000000000..4297e2567
--- /dev/null
+++ b/docs/source/concept_guides/deferring_execution.mdx
@@ -0,0 +1,95 @@
+# Deferring Executions
+
+When you run your usual script, instructions are executed in order. Using 🤗 Accelerate to deploy your script on several
+GPUs at the same time introduces a complication: while each process executes all instructions in order, some may be
+faster than others.
+
+You might need to wait for all processes to have reached a certain point before executing a given instruction. For
+instance, you shouldn't save a model before being sure every process is done with training, and you wouldn't want to 
+continue training before all the model weights have been loaded in. To do this, just write the following line in your code:
+
+```
+accelerator.wait_for_everyone()
+```
+
+This instruction will block all the processes that arrive first until all the other processes have reached that
+point (if you run your script on just one GPU or CPU, this won't do anything).
+
+A few example cases for when to use this utility are listed below:
+
+<Tip>
+
+    Some of these are utilized with the [`~Accelerator.main_process_first`] context manager, which utilizes [`~Accelerator.wait_for_everyone`] to 
+    run a particular set of code on the main process beforehand before triggering and launching the other processes
+
+</Tip>
+
+## Downloading a Dataset 
+
+When downloading a dataset, you should download it first on the main process and then loading the cached dataset in afterwards
+
+<Tip>
+
+    `load_dataset` will perform a lock under the hood to stop multiple downloads from happening at once, but if you are downloading something 
+    not using this library you should use this method.
+    
+</Tip>
+
+```python
+with accelerator.main_process_first():
+    datasets = load_dataset("glue", "mrpc")
+```
+
+Under the hood this is the same as calling: 
+
+```python
+# First do something on the main process
+if accelerator.is_main_process:
+    datasets = load_dataset("glue", "mrpc")
+else:
+    accelerator.wait_for_everyone()
+
+# And then send it to the rest of them
+if not accelerator.is_main_process:
+    datasets = load_dataset("glue", "mrpc")
+else:
+    accelerator.wait_for_everyone()
+```
+
+## Saving the `state_dict`
+
+When saving the `state_dict` of the model, since you would normally save one file on just the main process
+you should specify that:
+
+```python
+if accelerator.is_main_process:
+    model = accelerator.unwrap_model(model)
+    torch.save(model.state_dict(), "weights.pth")
+```
+
+## Loading in the `state_dict`
+
+When loading in the `state_dict` to a model, optimizer, or scheduler, you should wait 
+for all workers to have the weights loaded in before moving on to training
+
+```python
+with accelerator.main_process_first():
+    state = torch.load("weights.pth")
+    model.load_state_dict(state)
+```
+
+## Applying a multi-worker CPU operation 
+
+Applying a `map()` operation on multiple workers, such as tokenizing should be done on the 
+main process first, and then propagated to each one. 
+
+```python
+datasets = load_dataset("glue", "mrpc")
+
+with accelerator.main_process_first():
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["idx", "sentence1", "sentence2"],
+    )
+```
\ No newline at end of file
diff --git a/docs/source/concept_guides/gradient_synchronization.mdx b/docs/source/concept_guides/gradient_synchronization.mdx
new file mode 100644
index 000000000..02b5adf0d
--- /dev/null
+++ b/docs/source/concept_guides/gradient_synchronization.mdx
@@ -0,0 +1,117 @@
+<!--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.
+-->
+
+# Gradient Synchronization
+
+PyTorch's distributed module operates by communicating back and forth between all of the GPUs in your system.
+This communication takes time, and ensuring all processes know the states of each other happens at particular triggerpoints
+when using the `ddp` module. 
+
+These triggerpoints are added to the PyTorch model, specifically their `forward()` and `backward()` methods. 
+This happens when the model is wrapped with `DistributedDataParallel`:
+```python
+import torch.nn as nn
+from torch.nn.parallel import DistributedDataParallel
+
+model = nn.Linear(10, 10)
+ddp_model = DistributedDataParallel(model)
+```
+In 🤗 Accelerate this conversion happens automatically when calling [`~Accelerator.prepare`] and passing in your model.
+
+```diff
++ from accelerate import Accelerator
++ accelerator = Accelerator()
+  import torch.nn as nn
+- from torch.nn.parallel import DistributedDataParallel
+
+  model = nn.Linear(10,10)
++ model = accelerator.prepare(model)
+```
+
+## The slowdown in gradient accumulation
+
+You now understand that PyTorch adds hooks to the `forward` and `backward` method of your PyTorch model when 
+training in a distributed setup. But how does this risk slowing down your code?
+
+In DDP (distributed data parallel), the specific order in which processes are performed and ran are expected
+at specific points and these must also occur at roughly the same time before moving on.
+
+The most direct example is when you update all of the parameters in a model through `.backward()`. All instances of the model
+need to have updated their gradients, collated, and updated again before moving onto the next batch of data. But when performing 
+gradient accumulation, you accumulate `n` losses and skip `.backward()` until `n` batches have been reached. This 
+can cause a significant slowdown since all the processes need to communicate with them more times than needed. How 
+can you avoid this overhead?
+
+## Solving the slowdown problem
+
+Since you are skipping these batches, their gradients do not need to be synchronized until the point where `.backward()` is actually called. 
+PyTorch cannot automagically tell when you need to do this, but they do provide a tool to help through the [`no_sync`](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel.no_sync) context manager
+that is added to your model after converting it to DDP.
+
+Under this context manager, PyTorch will skip synchronizing the gradients when `.backward()` is called, and the first call to `.backward()` outside this 
+context manager will trigger the synchronization. See an example below:
+```python
+ddp_model, dataloader = accelerator.prepare(model, dataloader)
+
+for index, batch in enumerate(dataloader):
+    inputs, targets = batch
+    # Trigger gradient synchronization on the last batch
+    if index != (len(dataloader) - 1):
+        with ddp_model.no_sync():
+            # Gradients only accumulate
+            outputs = ddp_model(inputs)
+            loss = loss_func(outputs)
+            accelerator.backward(loss)
+    else:
+        # Gradients finally sync
+        outputs = ddp_model(inputs)
+        loss = loss_func(outputs)
+        accelerator.backward(loss)
+```
+
+In 🤗 Accelerate to make this an API that can be called no matter the training device (though it may not do anything if you are not in a distributed system!),
+`ddp_model.no_sync` gets replaced with [`~Accelerator.no_sync`] and operates the same way:
+
+```diff
+  ddp_model, dataloader = accelerator.prepare(model, dataloader)
+
+  for index, batch in enumerate(dataloader):
+      inputs, targets = batch
+      # Trigger gradient synchronization on the last batch
+      if index != (len(dataloader)-1):
+-         with ddp_model.no_sync():
++         with accelerator.no_sync(model):
+              # Gradients only accumulate
+              outputs = ddp_model(inputs)
+              loss = loss_func(outputs, targets)
+              accelerator.backward(loss)
+      else:
+          # Gradients finally sync
+          outputs = ddp_model(inputs)
+          loss = loss_func(outputs)
+          accelerator.backward(loss)
+```
+
+As you may expect, the [`~Accelerator.accumulate`] function wraps around this conditional check by keeping track of the current batch number, leaving you with the final
+gradient accumulation API:
+
+```python
+ddp_model, dataloader = accelerator.prepare(model, dataloader)
+
+for batch in dataloader:
+    with accelerator.accumulate(model):
+        optimizer.zero_grad()
+        inputs, targets = batch
+        outputs = model(inputs)
+        loss = loss_function(outputs, targets)
+        accelerator.backward(loss)
+```
\ No newline at end of file
diff --git a/docs/source/concept_guides/training_tpu.mdx b/docs/source/concept_guides/training_tpu.mdx
new file mode 100644
index 000000000..32736ce31
--- /dev/null
+++ b/docs/source/concept_guides/training_tpu.mdx
@@ -0,0 +1,164 @@
+<!--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.
+-->
+
+# Training on TPUs with 🤗 Accelerate
+
+Training on TPUs can be slightly different than training on multi-gpu, even with 🤗 Accelerate. This guide aims to show you 
+where you should be careful and why, as well as the best practices in general.
+
+## Training in a Notebook
+
+The main carepoint when training on TPUs comes from the [`notebook_launcher`]. As mentioned in the [notebook tutorial](../usage_guides/notebook), you need to 
+restructure your training code into a function that can get passed to the [`notebook_launcher`] function and be careful about not declaring any tensors on the GPU.
+
+While on a TPU that last part is not as important, a critical part to understand is that when you launch code from a notebook you do so through a process called **forking**. 
+When launching from the command-line, you perform **spawning**, where a python process is not currently running and you *spawn* a new process in. Since your Jupyter notebook is already 
+utilizing a python process, you need to *fork* a new process from it to launch your code. 
+
+Where this becomes important is in regards to declaring your model. On forked TPU processes, it is recommended that you instantiate your model *once* and pass this into your 
+training function. This is different than training on GPUs where you create `n` models that have their gradients synced and back-propagated at certain moments. Instead one 
+model instance is shared between all the nodes and it is passed back and forth. This is important especially when training on low-resource TPUs such as those provided in Kaggle kernels or
+on Google Colaboratory. 
+
+Below is an example of a training function passed to the [`notebook_launcher`] if training on CPUs or GPUs:
+
+<Tip>
+
+    This code snippet is based off the one from the `simple_nlp_example` notebook found [here](https://github.com/huggingface/notebooks/blob/main/examples/accelerate/simple_nlp_example.ipynb) with slight 
+    modifications for the sake of simplicity
+
+</Tip>
+
+```python
+def training_function():
+    # Initialize accelerator
+    accelerator = Accelerator()
+    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)
+    train_dataloader, eval_dataloader = create_dataloaders(
+        train_batch_size=hyperparameters["train_batch_size"], eval_batch_size=hyperparameters["eval_batch_size"]
+    )
+
+    # Instantiate optimizer
+    optimizer = AdamW(params=model.parameters(), lr=hyperparameters["learning_rate"])
+
+    # Prepare everything
+    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
+    # prepare method.
+    model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader
+    )
+
+    num_epochs = hyperparameters["num_epochs"]
+    # Now we train the model
+    for epoch in range(num_epochs):
+        model.train()
+        for step, batch in enumerate(train_dataloader):
+            outputs = model(**batch)
+            loss = outputs.loss
+            accelerator.backward(loss)
+
+            optimizer.step()
+            optimizer.zero_grad()
+```
+
+```python
+from accelerate import notebook_launcher
+
+notebook_launcher(training_function)
+```
+
+<Tip>
+
+    The `notebook_launcher` will default to 8 processes if 🤗 Accelerate has been configured for a TPU
+
+</Tip>
+
+If you use this example and declare the model *inside* the training loop, then on a low-resource system you will potentially see an error 
+like:
+
+```
+ProcessExitedException: process 0 terminated with signal SIGSEGV
+```
+
+This error is *extremely* cryptic but the basic explaination is you ran out of system RAM. You can avoid this entirely by reconfiguring the training function to 
+accept a single `model` argument, and declare it in an outside cell:
+
+```python
+# In another Jupyter cell
+model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)
+```
+
+```diff
++ def training_function(model):
+      # Initialize accelerator
+      accelerator = Accelerator()
+-     model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)
+      train_dataloader, eval_dataloader = create_dataloaders(
+          train_batch_size=hyperparameters["train_batch_size"], eval_batch_size=hyperparameters["eval_batch_size"]
+      )
+  ...
+```
+
+And finally calling the training function with:
+
+```diff
+  from accelerate import notebook_launcher
+- notebook_launcher(training_function)
++ notebook_launcher(training_function, (model,))
+```
+
+<Tip>
+
+    The above workaround is only needed when launching a TPU instance from a Jupyter Notebook on a low-resource server such as Google Colaboratory or Kaggle. If 
+    using a script or launching on a much beefier server declaring the model beforehand is not needed.
+
+</Tip>
+
+## Mixed Precision and Global Variables 
+
+As mentioned in the [mixed precision tutorial](../usage_guides/mixed_precision), 🤗 Accelerate supports fp16 and bf16, both of which can be used on TPUs.
+That being said, ideally `bf16` should be utilized as it is extremely efficient to use.
+
+There are two "layers" when using `bf16` and 🤗 Accelerate on TPUs, at the base level and at the operation level. 
+
+At the base level, this is enabled when passing `mixed_precision="bf16"` to `Accelerator`, such as:
+```python
+accelerator = Accelerator(mixed_precision="bf16")
+```
+By default this will cast `torch.float` and `torch.double` to `bfloat16` on TPUs. 
+The specific configuration being set is an environmental variable of `XLA_USE_BF16` is set to `1`.
+
+There is a futher configuration you can perform which is setting the `XLA_DOWNCAST_BF16` environmental variable. If set to `1`, then 
+`torch.float` is `bfloat16` and `torch.double` is `float32`.
+
+This is performed in the `Accelerator` object when passing `downcast_bf16=True`:
+```python
+accelerator = Accelerator(mixed_precision="bf16", downcast_bf16=True)
+```
+
+Using downcasting instead of bf16 everywhere is good for when you are trying to calculate metrics, log values, and more where raw bf16 tensors would be unusable. 
+
+## Training Times on TPUs
+
+As you launch your script, you may notice that training seems exceptionally slow at first. This is because TPUs
+first run through a few batches of data to see how much memory to allocate before finally utilizing this configured 
+memory allocation extremely efficiently. 
+
+If you notice that your evaluation code to calculate the metrics of your model takes longer due to a larger batch size being used, 
+it is recommended to keep the batch size the same as the training data if it is too slow. Otherwise the memory will reallocate to this 
+new batch size after the first few iterations. 
+
+<Tip>
+
+    Just because the memory is allocated does not mean it will be used or that the batch size will increase when going back to your training dataloader.
+
+</Tip>
\ No newline at end of file
diff --git a/docs/source/index.mdx b/docs/source/index.mdx
index 16593d554..3f0f58802 100644
--- a/docs/source/index.mdx
+++ b/docs/source/index.mdx
@@ -1,4 +1,4 @@
-<!--Copyright 2021 The HuggingFace Team. All rights reserved.
+<!--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
@@ -12,121 +12,60 @@ specific language governing permissions and limitations under the License.
 
 # Accelerate
 
-Run your *raw* PyTorch training script on any kind of device.
-
-## Features
-
-- 🤗 Accelerate provides an easy API to make your scripts run with mixed precision and in any kind of distributed
-  setting (multi-GPUs, TPUs etc.) while still letting you write your own training loop. The same code can then run
-  seamlessly on your local machine for debugging or your training environment.
-
-- 🤗 Accelerate also provides a CLI tool that allows you to quickly configure and test your training environment and
-  then launch the scripts.
-
-
-## Easy to integrate
-
-A traditional training loop in PyTorch looks like this:
-
-```python
-my_model.to(device)
-
-for batch in my_training_dataloader:
-    my_optimizer.zero_grad()
-    inputs, targets = batch
-    inputs = inputs.to(device)
-    targets = targets.to(device)
-    outputs = my_model(inputs)
-    loss = my_loss_function(outputs, targets)
-    loss.backward()
-    my_optimizer.step()
-```
-
-Changing it to work with accelerate is really easy and only adds a few lines of code:
+🤗 Accelerate is a library that enables the same PyTorch code to be run across any distributed configuration by adding just four lines of code! In short, training and inference at scale made simple, efficient and adaptable.
 
 ```diff
 + from accelerate import Accelerator
-
 + accelerator = Accelerator()
-  # Use the device given by the *accelerator* object.
-+ device = accelerator.device
-  my_model.to(device)
-  # Pass every important object (model, optimizer, dataloader) to *accelerator.prepare*
-+ my_model, my_optimizer, my_training_dataloader = accelerator.prepare(
-+     my_model, my_optimizer, my_training_dataloader
+
++ model, optimizer, training_dataloader, scheduler = accelerator.prepare(
++     model, optimizer, training_dataloader, scheduler
 + )
 
-  for batch in my_training_dataloader:
-      my_optimizer.zero_grad()
+  for batch in training_dataloader:
+      optimizer.zero_grad()
       inputs, targets = batch
       inputs = inputs.to(device)
       targets = targets.to(device)
-      outputs = my_model(inputs)
-      loss = my_loss_function(outputs, targets)
-      # Just a small change for the backward instruction
--     loss.backward()
+      outputs = model(inputs)
+      loss = loss_function(outputs, targets)
 +     accelerator.backward(loss)
-      my_optimizer.step()
+      optimizer.step()
+      scheduler.step()
 ```
 
-and with this, your script can now run in a distributed environment (multi-GPU, TPU).
+Built on `torch_xla` and `torch.distributed`, 🤗 Accelerate takes care of the heavy lifting, so you don't have to write any custom code to adapt to these platforms.
+Convert existing codebases to utilize [DeepSpeed](usage_guides/deepspeed), perform [fully sharded data parallelism](usage_guides/fsdp), and have automatic support for mixed-precision training! 
 
-You can even simplify your script a bit by letting 🤗 Accelerate handle the device placement for you (which is safer,
-especially for TPU training):
+<Tip> 
 
-```diff
-+ from accelerate import Accelerator
+  To get a better idea of this process, make sure to check out the [Tutorials](basic_tutorials/overview)! 
 
-+ accelerator = Accelerator()
-- my_model.to(device)
-  # Pass every important object (model, optimizer, dataloader) to *accelerator.prepare*
-+ my_model, my_optimizer, my_training_dataloader = accelerate.prepare(
-+     my_model, my_optimizer, my_training_dataloader
-+ )
+</Tip>
 
-  for batch in my_training_dataloader:
-      my_optimizer.zero_grad()
-      inputs, targets = batch
--     inputs = inputs.to(device)
--     targets = targets.to(device)
-      outputs = my_model(inputs)
-      loss = my_loss_function(outputs, targets)
-      # Just a small change for the backward instruction
--     loss.backward()
-+     accelerator.backward(loss)
-      my_optimizer.step()
-```
-
-## Script launcher
-
-No need to remember how to use `torch.distributed.launch` or to write a specific launcher for TPU training! 🤗
-Accelerate comes with a CLI tool that will make your life easier when launching distributed scripts.
-
-On your machine(s) just run:
 
+This code can then be launched on any system through Accelerate's CLI interface:
 ```bash
-accelerate config
+accelerate launch {my_script.py}
 ```
 
-and answer the questions asked. This will generate a config file that will be used automatically to properly set the
-default options when doing
-
-```bash
-accelerate launch my_script.py --args_to_my_script
-```
-
-For instance, here is how you would run the NLP example (from the root of the repo):
-
-```bash
-accelerate launch examples/nlp_example.py
-```
-
-## Supported integrations
-
-- CPU only
-- single GPU
-- multi-GPU on one node (machine)
-- multi-GPU on several nodes (machines)
-- TPU
-- FP16 with native AMP (apex on the roadmap)
-- DeepSpeed (experimental support)
+<div class="mt-10">
+  <div class="w-full flex flex-col space-y-4 md:space-y-0 md:grid md:grid-cols-2 md:gap-y-4 md:gap-x-5">
+    <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="/docs/accelerate/basic_tutorials/overview"
+      ><div class="w-full text-center bg-gradient-to-br from-blue-400 to-blue-500 rounded-lg py-1.5 font-semibold mb-5 text-white text-lg leading-relaxed">Tutorials</div>
+      <p class="text-gray-700">Learn the basics and become familiar with using 🤗 Accelerate. Start here if you are using 🤗 Accelerate for the first time!</p>
+    </a>
+    <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="/docs/accelerate/utility_guides/gradient_accumulation"
+      ><div class="w-full text-center bg-gradient-to-br from-indigo-400 to-indigo-500 rounded-lg py-1.5 font-semibold mb-5 text-white text-lg leading-relaxed">How-to guides</div>
+      <p class="text-gray-700">Practical guides to help you achieve a specific goal. Take a look at these guides to learn how to use 🤗 Accelerate to solve real-world problems.</p>
+    </a>
+    <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="/docs/accelerate/concept_guides/gradient_synchronization"
+      ><div class="w-full text-center bg-gradient-to-br from-pink-400 to-pink-500 rounded-lg py-1.5 font-semibold mb-5 text-white text-lg leading-relaxed">Conceptual guides</div>
+      <p class="text-gray-700">High-level explanations for building a better understanding of important topics such as avoiding subtle nuances and pitfalls in distributed training and DeepSpeed.</p>
+   </a>
+    <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="/docs/accelerate/package_reference/accelerator"
+      ><div class="w-full text-center bg-gradient-to-br from-purple-400 to-purple-500 rounded-lg py-1.5 font-semibold mb-5 text-white text-lg leading-relaxed">Reference</div>
+      <p class="text-gray-700">Technical descriptions of how 🤗 Accelerate classes and methods work.</p>
+    </a>
+  </div>
+</div>
\ No newline at end of file
diff --git a/docs/source/installation.mdx b/docs/source/installation.mdx
deleted file mode 100644
index 26c2698aa..000000000
--- a/docs/source/installation.mdx
+++ /dev/null
@@ -1,96 +0,0 @@
-<!---
-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.
--->
-
-# Installation
-
-🤗 Accelerate is tested on Python 3.6+, and PyTorch 1.6.0+.
-
-You should install 🤗 Accelerate in a [virtual environment](https://docs.python.org/3/library/venv.html). If you're
-unfamiliar with Python virtual environments, check out the [user guide](https://packaging.python.org/guides/installing-using-pip-and-virtual-environments/). Create a virtual environment with the version of Python you're going
-to use and activate it.
-
-Now, if you want to use 🤗 Accelerate, you can install it with pip.
-
-## Installation with pip
-
-First you need to install PyTorch. Please refer to the
-[PyTorch installation page](https://pytorch.org/get-started/locally/#start-locally) regarding the specific install command for your platform.
-
-When PyTorch has been installed, 🤗 Accelerate can be installed using pip as follows:
-
-```bash
-pip install accelerate
-```
-
-Alternatively, for CPU-support only, you can install 🤗 Accelerate and PyTorch in one line with:
-
-```bash
-pip install accelerate[torch]
-```
-
-To check 🤗 Accelerate is properly installed, run the following command:
-
-```bash
-python -c "TODO write"
-```
-
-## Installing from source
-
-Here is how to quickly install `accelerate` from source:
-
-```bash
-pip install git+https://github.com/huggingface/accelerate
-```
-
-Note that this will install not the latest released version, but the bleeding edge `main` version, which you may want to use in case a bug has been fixed since the last official release and a new release hasn't  been yet rolled out.
-
-While we strive to keep `main` operational at all times, if you notice some issues, they usually get fixed within a few hours or a day and you're more than welcome to help us detect any problems by opening an [Issue](https://github.com/huggingface/accelerate/issues) and this way, things will get fixed even sooner.
-
-Again, you can run:
-
-```bash
-python -c "TODO write"
-```
-
-to check 🤗 Accelerate is properly installed.
-
-## Editable install
-
-If you want to constantly use the bleeding edge `main` version of the source code, or if you want to contribute to the library and need to test the changes in the code you're making, you will need an editable install. This is done by cloning the repository and installing with the following commands:
-
-``` bash
-git clone https://github.com/huggingface/accelerate.git
-cd accelerate
-pip install -e .
-```
-
-This command performs a magical link between the folder you cloned the repository to and your python library paths, and it'll look inside this folder in addition to the normal library-wide paths. So if normally your python packages get installed into:
-```
-~/anaconda3/envs/main/lib/python3.7/site-packages/
-```
-now this editable install will reside where you clone the folder to, e.g. `~/accelerate/` and python will search it too.
-
-Do note that you have to keep that `accelerate` folder around and not delete it to continue using the 🤗 Accelerate library.
-
-Now, let's get to the real benefit of this installation approach. Say, you saw some new feature just has been committed into `main`. If you have already performed all the steps above, to update your accelerate repo to include all the latest commits, all you need to do is to `cd` into that cloned repository folder and update the clone to the latest version:
-
-```bash
-cd ~/accelerate/
-git pull
-```
-
-There is nothing else to do. Your python environment will find the bleeding edge version of 🤗 Accelerate on the next run.
-
diff --git a/docs/source/launcher.mdx b/docs/source/launcher.mdx
deleted file mode 100644
index fb672d5ad..000000000
--- a/docs/source/launcher.mdx
+++ /dev/null
@@ -1,28 +0,0 @@
-<!--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.
--->
-
-# Notebook Launcher
-
-Launch your training function inside a notebook. Currently supports launching a training with TPUs on [Google
-Colab](https://colab.research.google.com/) and [Kaggle kernels](https://www.kaggle.com/code), as well as training on
-several GPUs (if the machine on which you are running your notebook has them).
-
-An example can be found in [this notebook](https://github.com/huggingface/notebooks/blob/master/examples/accelerate/simple_nlp_example.ipynb).
-
-<Tip warning={true}>
-
-Your `Accelerator` object should only be defined inside the training function. This is because the
-initialization should be done inside the launcher only.
-
-</Tip>
-
-[[autodoc]] notebook_launcher
diff --git a/docs/source/package_reference/accelerator.mdx b/docs/source/package_reference/accelerator.mdx
new file mode 100644
index 000000000..fb20f1a66
--- /dev/null
+++ b/docs/source/package_reference/accelerator.mdx
@@ -0,0 +1,163 @@
+<!--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.
+-->
+
+# Accelerator
+
+The [`Accelerator`] is the main class provided by 🤗 Accelerate. 
+It serves at the main entrypoint for the API. 
+
+## Quick adaptation of your code
+
+To quickly adapt your script to work on any kind of setup with 🤗 Accelerate just:
+
+1. Initialize an [`Accelerator`] object (that we will call `accelerator` throughout this page) as early as possible in your script.
+2. Pass your dataloader(s), model(s), optimizer(s), and scheduler(s) to the [`~Accelerator.prepare`] method.
+3. Remove all the `.cuda()` or `.to(device)` from your code and let the `accelerator` handle the device placement for you. 
+
+<Tip>
+
+    Step three is optional, but considered a best practice.
+
+</Tip>
+
+4. Replace `loss.backward()` in your code with `accelerator.backward(loss)`
+5. Gather your predictions and labels before storing them or using them for metric computation using [`~Accelerator.gather`]
+
+<Tip warning={true}>
+
+    Step five is mandatory when using distributed evaluation
+    
+</Tip>
+
+In most cases this is all that is needed. The next section lists a few more advanced use cases and nice features
+you should search for and replace by the corresponding methods of your `accelerator`:
+
+## Advanced recommendations
+
+### Printing
+
+`print` statements should be replaced by [`~Accelerator.print`] to be printed once per process
+
+```diff
+- print("My thing I want to print!")
++ accelerator.print("My thing I want to print!")
+```
+
+### Executing processes
+
+#### Once on a single server
+
+For statements that should be executed once per server, use [`~Accelerator.is_local_main_process`]:
+
+```python
+if accelerator.is_local_main_process:
+    do_thing_once_per_server()
+```
+
+A function can be wrapped using the [`~Accelerator.on_local_main_process`] function to achieve the same 
+behavior on a function's execution:
+
+```python
+@accelerator.on_local_main_process
+def do_my_thing():
+    "Something done once per server"
+    do_thing_once_per_server()
+```
+
+#### Only ever once across all servers
+
+For statements that should only ever be executed once, use [`~Accelerator.is_main_process`]:
+
+```python
+if accelerator.is_main_process:
+    do_thing_once()
+```
+
+A function can be wrapped using the [`~Accelerator.on_main_process`] function to achieve the same 
+behavior on a function's execution:
+
+```python
+@accelerator.on_main_process
+def do_my_thing():
+    "Something done once per server"
+    do_thing_once()
+```
+
+#### On specific processes
+
+If a function should be ran on a specific overall or local process index, there are similar decorators 
+to achieve this:
+
+```python
+@accelerator.on_local_process(local_process_idx=0)
+def do_my_thing():
+    "Something done on process index 0 on each server"
+    do_thing_on_index_zero_on_each_server()
+```
+
+```python
+@accelerator.on_process(process_index=0)
+def do_my_thing():
+    "Something done on process index 0"
+    do_thing_on_index_zero()
+```
+
+### Synchronicity control
+
+Use [`~Accelerator.wait_for_everyone`] to make sure all processes join that point before continuing. (Useful before a model save for instance)
+
+### Saving and loading
+
+Use [`~Accelerator.unwrap_model`] before saving to remove all special model wrappers added during the distributed process. 
+
+```python
+model = MyModel()
+model = accelerator.prepare(model)
+# Unwrap
+model = accelerator.unwrap_model(model)
+```
+
+Use [`~Accelerator.save`] instead of `torch.save`:
+
+```diff
+  state_dict = model.state_dict()
+- torch.save(state_dict, "my_state.pkl")
++ accelerator.save(state_dict, "my_state.pkl")
+```
+
+### Operations
+
+Use [`~Accelerator.clip_grad_norm_`] instead of ``torch.nn.utils.clip_grad_norm_`` and [`~Accelerator.clip_grad_value_`] instead of ``torch.nn.utils.clip_grad_value``
+
+### Gradient Accumulation
+
+To perform gradient accumulation use [`~Accelerator.accumulate`] and specify a gradient_accumulation_steps. 
+This will also automatically ensure the gradients are synced or unsynced when on 
+multi-device training, check if the step should actually be performed, and auto-scale the loss:
+
+```diff
+- accelerator = Accelerator()
++ accelerator = Accelerator(gradient_accumulation_steps=2)
+
+  for (input, label) in training_dataloader:
++     with accelerator.accumulate(model):
+          predictions = model(input)
+          loss = loss_function(predictions, labels)
+          accelerator.backward(loss)
+          optimizer.step()
+          scheduler.step()
+          optimizer.zero_grad()
+```
+
+## Overall API documentation:
+
+[[autodoc]] Accelerator
\ No newline at end of file
diff --git a/docs/source/package_reference/big_modeling.mdx b/docs/source/package_reference/big_modeling.mdx
new file mode 100644
index 000000000..e54ac0807
--- /dev/null
+++ b/docs/source/package_reference/big_modeling.mdx
@@ -0,0 +1,41 @@
+<!--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.
+-->
+
+# Working with large models
+
+## Dispatching and Offloading Models
+
+[[autodoc]] big_modeling.init_empty_weights
+[[autodoc]] big_modeling.cpu_offload
+[[autodoc]] big_modeling.disk_offload
+[[autodoc]] big_modeling.dispatch_model
+[[autodoc]] big_modeling.load_checkpoint_and_dispatch
+
+## Model Hooks
+
+### Hook Classes
+
+[[autodoc]] hooks.ModelHook
+[[autodoc]] hooks.AlignDevicesHook
+[[autodoc]] hooks.SequentialHook
+
+### Adding Hooks
+
+[[autodoc]] hooks.add_hook_to_module
+[[autodoc]] hooks.attach_execution_device_hook
+[[autodoc]] hooks.attach_align_device_hook
+[[autodoc]] hooks.attach_align_device_hook_on_blocks
+
+### Removing Hooks
+
+[[autodoc]] hooks.remove_hook_from_module
+[[autodoc]] hooks.remove_hook_from_submodules
\ No newline at end of file
diff --git a/docs/source/package_reference/cli.mdx b/docs/source/package_reference/cli.mdx
new file mode 100644
index 000000000..092b72850
--- /dev/null
+++ b/docs/source/package_reference/cli.mdx
@@ -0,0 +1,152 @@
+<!--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.
+-->
+
+# The Command Line 
+
+Below is a list of all the available commands 🤗 Accelerate with their parameters
+
+## accelerate config
+
+**Command**:
+
+`accelerate config` or `accelerate-config`
+
+Launches a series of prompts to create and save a `default_config.yml` configuration file for your training system. Should 
+always be ran first on your machine.
+
+**Usage**: 
+
+```bash
+accelerate config [arguments]
+```
+
+**Optional Arguments**:
+* `--config_file CONFIG_FILE` (`str`) -- The path to use to store the config file. Will default to a file named default_config.yaml in the cache location, which is the content
+                        of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have such an environment variable, your cache directory
+                        (`~/.cache` or the content of `XDG_CACHE_HOME`) suffixed with `huggingface`.
+* `-h`, `--help` (`bool`) -- Show a help message and exit
+
+## accelerate env
+
+**Command**:
+
+`accelerate env` or `accelerate-env`
+
+Lists the contents of the passed 🤗 Accelerate configuration file. Should always be used when opening an issue on the [GitHub repository](https://github.com/huggingface/accelerate).
+
+**Usage**:
+
+```bash
+accelerate env [arguments]
+```
+
+**Optional Arguments**:
+* `--config_file CONFIG_FILE` (`str`) -- The path to use to store the config file. Will default to a file named default_config.yaml in the cache location, which is the content
+                        of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have such an environment variable, your cache directory
+                        (`~/.cache` or the content of `XDG_CACHE_HOME`) suffixed with `huggingface`.
+* `-h`, `--help` (`bool`) -- Show a help message and exit
+
+## accelerate launch
+
+**Command**:
+
+`accelerate launch` or `accelerate-launch`
+
+Launches a specified script on a distributed system with the right parameters.
+
+**Usage**: 
+
+```bash
+accelerate launch [arguments] {training_script} --{training_script-argument-1} --{training_script-argument-2} ...
+```
+
+**Positional Arguments**:
+
+- `{training_script}` -- The full path to the script to be launched in parallel
+- `--{training_script-argument-1}` -- Arguments of the training script
+
+**Optional Arguments**:
+
+* `-h`, `--help` (`bool`) -- Show a help message and exit
+* `--config_file CONFIG_FILE` (`str`)-- The config file to use for the default values in the launching script.
+* `--cpu` (`bool`) -- Whether or not to force the training on the CPU.
+* `--mixed_precision {no,fp16,bf16}` (`str`) -- Whether or not to use mixed precision training. Choose between FP16 and BF16 (bfloat16) training. BF16 training is only supported on
+                    Nvidia Ampere GPUs and PyTorch 1.10 or later.
+* `--multi_gpu` (`bool`, defaults to `False`) -- Whether or not this should launch a distributed GPU training.
+* `-m`, `--module` (`bool`) -- Change each process to interpret the launch script as a Python module, executing with the same behavior as 'python -m'.
+* `--no_python` (`bool`) -- Skip prepending the training script with 'python' - just execute it directly. Useful when the script is not a Python script.
+
+The rest of these arguments are configured through `accelerate config` and are read in from the specified `--config_file` (or default configuration) for their 
+values. They can also be passed in manually.
+
+**Machine Configuration Arguments**:
+
+The following arguments are useful for customization of worker machines
+* `--machine_rank MACHINE_RANK` (`int`) -- The rank of the machine on which this script is launched.
+* `--num_machines NUM_MACHINES` (`int`) -- The total number of machines used in this training.
+* `--num_processes NUM_PROCESSES` (`int`) -- The total number of processes to be launched in parallel.
+* `--main_process_ip MAIN_PROCESS_IP` (`str`) -- The IP address of the machine of rank 0.
+* `--main_process_port MAIN_PROCESS_PORT` (`int`) -- The port to use to communicate with the machine of rank 0.
+* `--num_cpu_threads_per_process NUM_CPU_THREADS_PER_PROCESS` (`int`) -- The number of CPU threads per process. Can be tuned for optimal performance.
+
+
+**DeepSpeed Arguments**:
+
+The following arguments are only useful when `use_deepspeed` is passed: 
+* `--use_deepspeed` (`bool`) -- Whether to use deepspeed.
+* `--deepspeed_config_file DEEPSPEED_CONFIG_FILE` (`str`) -- DeepSpeed config file.
+* `--zero_stage ZERO_STAGE` (`str`) -- DeepSpeed's ZeRO optimization stage
+* `--offload_optimizer_device OFFLOAD_OPTIMIZER_DEVICE` (`str`) -- Decides where (none|cpu|nvme) to offload optimizer states
+* `--offload_param_device OFFLOAD_PARAM_DEVICE` (`str`) -- Decides where (none|cpu|nvme) to offload parameters
+* `--gradient_accumulation_steps GRADIENT_ACCUMULATION_STEPS` (`int`) -- Number of gradient_accumulation_steps used in your training script
+* `--gradient_clipping GRADIENT_CLIPPING` (`float`) -- gradient clipping value used in your training script
+The following arguments are related to using ZeRO Stage-3
+* `--zero3_init_flag ZERO3_INIT_FLAG` (`bool`) -- Decides Whether (true|false) to enable `deepspeed.zero.Init` for constructing massive models
+* `--zero3_save_16bit_model ZERO3_SAVE_16BIT_MODEL` (`bool`) -- Decides Whether (true|false) to save 16-bit model weights when using ZeRO Stage-3
+
+**Fully Sharded Data Parallelism Arguments**:
+
+The following arguments are only useful when `use_fdsp` is passed:
+* `--use_fsdp` (`bool`) -- Whether to use fsdp.
+* `--offload_params OFFLOAD_PARAMS` (`bool`) -- Decides Whether (true|false) to offload parameters and gradients to CPU.
+* `--min_num_params MIN_NUM_PARAMS` (`int`) -- FSDP's minimum number of parameters for Default Auto Wrapping.
+* `--sharding_strategy SHARDING_STRATEGY` (`str`) -- FSDP's Sharding Strategy.
+
+**TPU Arguments**:
+
+The following arguments are only useful when `tpu` is passed:
+* `--tpu` (`bool`) - Whether or not this should launch a TPU training.
+* `--main_training_function MAIN_TRAINING_FUNCTION` (`str`) -- The name of the main function to be executed in your script.
+
+**AWS SageMaker Arguments**:
+
+The following arguments are only useful when training in SageMaker
+* `--aws_access_key_id AWS_ACCESS_KEY_ID` (`str`) -- The AWS_ACCESS_KEY_ID used to launch the Amazon SageMaker training job
+* `--aws_secret_access_key AWS_SECRET_ACCESS_KEY` (`str`) -- The AWS_SECRET_ACCESS_KEY used to launch the Amazon SageMaker training job
+
+## accelerate test
+
+`accelerate test` or `accelerate-test`
+
+Runs `accelerate/test_utils/test_script.py` to verify that 🤗 Accelerate has been properly configured on your system and runs. 
+
+**Usage**: 
+
+```bash
+accelerate test [arguments]
+```
+
+**Optional Arguments**:
+* `--config_file CONFIG_FILE` (`str`) -- The path to use to store the config file. Will default to a file named default_config.yaml in the cache location, which is the content
+                        of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have such an environment variable, your cache directory
+                        (`~/.cache` or the content of `XDG_CACHE_HOME`) suffixed with `huggingface`.
+* `-h`, `--help` (`bool`) -- Show a help message and exit
diff --git a/docs/source/package_reference/deepspeed.mdx b/docs/source/package_reference/deepspeed.mdx
new file mode 100644
index 000000000..fee886f4e
--- /dev/null
+++ b/docs/source/package_reference/deepspeed.mdx
@@ -0,0 +1,25 @@
+<!--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.
+-->
+
+# Utilities for DeepSpeed
+
+[[autodoc]] utils.DeepSpeedPlugin
+
+[[autodoc]] utils.DummyOptim
+
+[[autodoc]] utils.DummyScheduler
+
+[[autodoc]] utils.DeepSpeedEngineWrapper
+
+[[autodoc]] utils.DeepSpeedOptimizerWrapper
+
+[[autodoc]] utils.DeepSpeedSchedulerWrapper
diff --git a/docs/source/kwargs.mdx b/docs/source/package_reference/kwargs.mdx
similarity index 100%
rename from docs/source/kwargs.mdx
rename to docs/source/package_reference/kwargs.mdx
diff --git a/docs/source/package_reference/launchers.mdx b/docs/source/package_reference/launchers.mdx
new file mode 100644
index 000000000..6f37f0af0
--- /dev/null
+++ b/docs/source/package_reference/launchers.mdx
@@ -0,0 +1,19 @@
+<!--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.
+-->
+
+# Launchers
+
+Functions for launching training on distributed processes.
+
+
+[[autodoc]] accelerate.notebook_launcher
+[[autodoc]] accelerate.debug_launcher
\ No newline at end of file
diff --git a/docs/source/package_reference/logging.mdx b/docs/source/package_reference/logging.mdx
new file mode 100644
index 000000000..675af41ee
--- /dev/null
+++ b/docs/source/package_reference/logging.mdx
@@ -0,0 +1,24 @@
+<!--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.
+-->
+
+# Logging with Accelerate
+
+Accelerate has its own logging utility to handle logging while in a distributed system.
+To utilize this replace cases of `logging` with `accelerate.logging`:
+```diff
+- import logging
++ from accelerate.logging import get_logger
+- logger = logging.getLogger(__name__)
++ logger = get_logger(__name__)
+```
+
+[[autodoc]] logging.get_logger
\ No newline at end of file
diff --git a/docs/source/package_reference/state.mdx b/docs/source/package_reference/state.mdx
new file mode 100644
index 000000000..f1f5ef974
--- /dev/null
+++ b/docs/source/package_reference/state.mdx
@@ -0,0 +1,23 @@
+<!--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.
+-->
+
+# Stateful Classes
+
+Below are variations of a [singleton class](https://en.wikipedia.org/wiki/Singleton_pattern) in the sense that all
+instances share the same state, which is initialized on the first instantiation.
+
+These classes are immutable and store information about certain configurations or 
+states.
+
+[[autodoc]] state.AcceleratorState
+
+[[autodoc]] state.GradientState
\ No newline at end of file
diff --git a/docs/source/internal.mdx b/docs/source/package_reference/torch_wrappers.mdx
similarity index 64%
rename from docs/source/internal.mdx
rename to docs/source/package_reference/torch_wrappers.mdx
index 4ff1d6ff9..4ac8ae572 100644
--- a/docs/source/internal.mdx
+++ b/docs/source/package_reference/torch_wrappers.mdx
@@ -10,44 +10,24 @@ an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express o
 specific language governing permissions and limitations under the License.
 -->
 
-# Internals
+# Wrapper classes for torch Dataloaders, Optimizers, and Schedulers
 
-## Gradient Accumulation states
+The internal classes Accelerate uses to prepare objects for distributed training
+when calling [`~Accelerator.prepare`].
 
-[[autodoc]] state.GradientState
-
-## Optimizer
-
-[[autodoc]] optimizer.AcceleratedOptimizer
-
-## DataLoader
-
-The main work on your PyTorch `DataLoader` is done by the following function:
+## Datasets and DataLoaders
 
 [[autodoc]] data_loader.prepare_data_loader
 
-### DataLoaderShard
-
-[[autodoc]] data_loader.DataLoaderShard
-
-### BatchSamplerShard
-
 [[autodoc]] data_loader.BatchSamplerShard
-
-### IterableDatasetShard
-
 [[autodoc]] data_loader.IterableDatasetShard
+[[autodoc]] data_loader.DataLoaderShard
+[[autodoc]] data_loader.DataLoaderDispatcher
 
-## Scheduler
-
-[[autodoc]] scheduler.AcceleratedScheduler
-
-## Distributed Config
-
-### AcceleratorState
+## Optimizers 
 
-[[autodoc]] state.AcceleratorState
+[[autodoc]] optimizer.AcceleratedOptimizer
 
-## Tracking
+## Schedulers 
 
-[[autodoc]] tracking.GeneralTracker
+[[autodoc]] scheduler.AcceleratedScheduler
\ No newline at end of file
diff --git a/docs/source/package_reference/tracking.mdx b/docs/source/package_reference/tracking.mdx
new file mode 100644
index 000000000..5e7a97f24
--- /dev/null
+++ b/docs/source/package_reference/tracking.mdx
@@ -0,0 +1,26 @@
+<!--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.
+-->
+
+# Experiment Tracking
+
+## The Base Tracker Class
+
+[[autodoc]] tracking.GeneralTracker
+
+## Integrated Trackers
+
+[[autodoc]] tracking.TensorBoardTracker
+    - __init__
+[[autodoc]] tracking.WandBTracker
+    - __init__
+[[autodoc]] tracking.CometMLTracker
+    - __init__
diff --git a/docs/source/utilities.mdx b/docs/source/package_reference/utilities.mdx
similarity index 97%
rename from docs/source/utilities.mdx
rename to docs/source/package_reference/utilities.mdx
index bfcaeb60e..0c64953fc 100644
--- a/docs/source/utilities.mdx
+++ b/docs/source/package_reference/utilities.mdx
@@ -44,8 +44,6 @@ These include data operations that mimic the same `torch` ops but can be used on
 
 These functionalities check the state of the current working environment including information about the operating system itself, what it can support, and if particular dependencies are installed. 
 
-[[autodoc]] utils.get_max_memory
-
 [[autodoc]] utils.is_bf16_available
 
 [[autodoc]] utils.is_torch_version
@@ -58,6 +56,12 @@ These functionalities check the state of the current working environment includi
 
 When setting up 🤗 Accelerate for the first time, rather than running `accelerate config` [~utils.write_basic_config] can be used as an alternative for quick configuration.
 
+## Memory
+
+[[autodoc]] utils.get_max_memory
+
+[[autodoc]] utils.find_executable_batch_size
+
 ## Modeling
 
 These utilities relate to interacting with PyTorch models
diff --git a/docs/source/quicktour.mdx b/docs/source/quicktour.mdx
index 4b144ba05..812b20226 100644
--- a/docs/source/quicktour.mdx
+++ b/docs/source/quicktour.mdx
@@ -18,7 +18,7 @@ Let's have a look at the 🤗 Accelerate main features and traps to avoid.
 
 To use 🤗 Accelerate in your own script, you have to change four things:
 
-1. Import the [`Accelerator`] main class instantiate one in an `accelerator` object:
+1. Import the [`Accelerator`] main class and instantiate one in an `accelerator` object:
 
 ```python
 from accelerate import Accelerator
@@ -28,7 +28,7 @@ accelerator = Accelerator()
 
 This should happen as early as possible in your training script as it will initialize everything necessary for
 distributed training. You don't need to indicate the kind of environment you are in (just one machine with a GPU, one
-match with several GPUs, several machines with multiple GPUs or a TPU), the library will detect this automatically.
+machines with several GPUs, several machines with multiple GPUs or a TPU), the library will detect this automatically.
 
 2. Remove the call `.to(device)` or `.cuda()` for your model and input data. The `accelerator` object
 will handle this for you and place all those objects on the right device for you. If you know what you're doing, you
@@ -40,8 +40,8 @@ To fully deactivate the automatic device placement, pass along `device_placement
 
 <Tip warning={true}>
 
-If you place your objects manually on the proper device, be careful to create your optimizer after putting your
-model on `accelerator.device` or your training will fail on TPU.
+    If you place your objects manually on the proper device, be careful to create your optimizer after putting your
+    model on `accelerator.device` or your training will fail on TPU.
 
 </Tip>
 
@@ -61,9 +61,9 @@ use `shuffle=True` or any kind of random sampler).
 
 <Tip>
 
-The actual batch size for your training will be the number of devices used multiplied by the batch size you set in
-your script: for instance training on 4 GPUs with a batch size of 16 set when creating the training dataloader will
-train at an actual batch size of 64.
+    The actual batch size for your training will be the number of devices used multiplied by the batch size you set in
+    your script: for instance training on 4 GPUs with a batch size of 16 set when creating the training dataloader will
+    train at an actual batch size of 64.
 
 </Tip>
 
@@ -76,16 +76,16 @@ training loop.
 
 <Tip warning={true}>
 
-You should only pass the learning rate scheduler to [`~Accelerator.prepare`] when the scheduler needs to be stepped
-at each optimizer step.
+    You should only pass the learning rate scheduler to [`~Accelerator.prepare`] when the scheduler needs to be stepped
+    at each optimizer step.
 
 </Tip>
 
 <Tip warning={true}>
 
-Your training dataloader may change length when going through this method: if you run on X GPUs, it will have its
-length divided by X (since your actual batch size will be multiplied by X), unless you set
-`split_batches=True`.
+    Your training dataloader may change length when going through this method: if you run on X GPUs, it will have its
+    length divided by X (since your actual batch size will be multiplied by X), unless you set
+    `split_batches=True`.
 
 </Tip>
 
@@ -133,27 +133,32 @@ for inputs, targets in validation_dataloader:
 
 <Tip warning={true}>
 
-As for the training dataloader, passing your validation dataloader through
-[`~Accelerator.prepare`] may change it: if you run on X GPUs, it will have its length divided by X
-(since your actual batch size will be multiplied by X), unless you set `split_batches=True`.
+    Similar to the training dataloader, passing your validation dataloader through
+    [`~Accelerator.prepare`] may change it: if you run on X GPUs, it will have its length divided by X
+    (since your actual batch size will be multiplied by X), unless you set `split_batches=True`.
+
+</Tip>
 
 Any instruction using your training dataloader length (for instance if you need the number of total training steps
-to create a learning rate scheduler) should go after the call to [`~Accelerator.prepare`].
+to create a learning rate scheduler) should go after the call to [`~Accelerator.prepare`]. 
 
-As some data at the end of the dataset may be duplicated so the batch can divide equally to all workers, metrics should be 
-calculated through the [`~Accelerator.gather_for_metrics`] method to automatically remove the duplicated data.
+Some data at the end of the dataset may be duplicated so the batch can be divided equally among all workers. As a result, metrics
+should be calculated through the [`~Accelerator.gather_for_metrics`] method to automatically remove the duplicated data while gathering.
 
-If for some reason you don't wish to have this automatically done, [`~Accelerator.gather`] can be used instead to gather 
-the data across all processes and this can manually be done instead.
+<Tip>
+
+    If for some reason you don't wish to have this automatically done, [`~Accelerator.gather`] can be used instead to gather 
+    the data across all processes and this can manually be done instead.
 
 </Tip>
 
+
 <Tip warning={true}>
 
-The [`~Accelerator.gather`] and [`~Accelerator.gather_for_metrics`] methods require the tensors to be all the same size on each process. If
-you have tensors of different sizes on each process (for instance when dynamically padding to the maximum length in
-a batch), you should use the [`~Accelerator.pad_across_processes`] method to pad you tensor to the
-biggest size across processes.
+    The [`~Accelerator.gather`] and [`~Accelerator.gather_for_metrics`] methods require the tensors to be all the same size on each process. If
+    you have tensors of different sizes on each process (for instance when dynamically padding to the maximum length in
+    a batch), you should use the [`~Accelerator.pad_across_processes`] method to pad you tensor to the
+    biggest size across processes.
 
 </Tip>
 
@@ -163,7 +168,7 @@ You can use the regular commands to launch your distributed training (like `torc
 PyTorch), they are fully compatible with 🤗 Accelerate. The only caveat here is that 🤗 Accelerate uses the environment
 to determine all useful information, so `torch.distributed.launch` should be used with the flag `--use_env`.
 
-🤗 Accelerate also provides a CLI tool that unifies all launcher, so you only have to remember one command. To use it,
+🤗 Accelerate also provides a CLI tool that unifies all launchers, so you only have to remember one command. To use it,
 just run:
 
 ```bash
@@ -207,7 +212,10 @@ If you stored the config file in a non-default location, you can indicate it to
 accelerate launch --config_file path_to_config.yaml path_to_script.py --args_for_the_script
 ```
 
-You can also override any of the arguments determined by your config file, see TODO: insert ref here.
+You can also override any of the arguments determined by your config file. 
+To see the complete list of parameters that you can pass in, run `accelerate launch -h`. 
+
+Check out the [Launch tutorial](basic_tutorials/launch) for more information about launching your scripts. 
 
 
 ## Launching training from a notebook
@@ -227,11 +235,14 @@ notebook_launcher(training_function)
 
 <Tip warning={true}>
 
-Your `Accelerator` object should only be defined inside the training function. This is because the
-initialization should be done inside the launcher only.
+    Your [`Accelerator`] object should only be defined inside the training function. This is because the
+    initialization should be done inside the launcher only.
 
 </Tip>
 
+Check out the [Notebook Launcher tutorial](basic_tutorials/notebook) for more information about training on TPUs. 
+
+
 ## Training on TPU
 
 If you want to launch your script on TPUs, there are a few caveats you should be aware of. Behind the scenes, the TPUs
@@ -242,7 +253,7 @@ optimizations takes some time.
 The good news is that this compilation will be cached so the second step and all the following will be much faster. The
 bad news is that it only applies if all of your steps do exactly the same operations, which implies:
 
-- having all tensors of the same length in all your lengths
+- having all tensors of the same length in all your batches
 - having static code (i.e., not a for loop of length that could change from step to step)
 
 Having any of the things above change between two steps will trigger a new compilation which will, once again, take a
@@ -271,6 +282,8 @@ passed your model to [`~Accelerator.prepare`]) will break the tying. You will ne
 after. You can find an example of this in the [run_clm_no_trainer](https://github.com/huggingface/transformers/blob/master/examples/pytorch/language-modeling/run_clm.py) script in
 the Transformers repository.
 
+Check out the [TPU tutorial](concept_guides/training_tpu) for more information about training on TPUs. 
+
 
 ## Other caveats
 
@@ -357,18 +370,21 @@ Note that since all the model parameters are references to tensors, this will lo
 ## Saving/loading entire states
 
 When training your model, you may want to save the current state of the model, optimizer, random generators, and potentially LR schedulers to be restored in the _same script_.
-You can use `accelerator.save_state` and `accelerator.load_state` respectively to do so, just by simply passing in a save location. 
-If you have registered any other stateful items to be stored through `accelerator.register_for_checkpointing` they will also be saved and/or loaded.
+You can use [`~Accelerator.save_state`] and [`~Accelerator.load_state`] respectively to do so, just by simply passing in a save location. 
+If you have registered any other stateful items to be stored through [`~Accelerator.register_for_checkpointing`] they will also be saved and/or loaded.
+
 <Tip>
-    Every object passed to `register_for_checkpointing` must have a `load_state_dict` and `save_dict` function to be stored
+
+    Every object passed to [`~Accelerator.register_for_checkpointing`] must have a `load_state_dict` and `state_dict` function to be stored
+
 </Tip>
 
 
 ### Gradient clipping
 
 If you are using gradient clipping in your script, you should replace the calls to
-`torch.nn.utils.clip_grad_norm_` or `torch.nn.utils.clip_grad_value_` with `accelerator.clip_grad_norm_`
-and `accelerator.clip_grad_value_` respectively.
+`torch.nn.utils.clip_grad_norm_` or `torch.nn.utils.clip_grad_value_` with [`~Accelerator.clip_grad_norm_`]
+and [`~Accelerator.clip_grad_value_`] respectively.
 
 
 ### Mixed Precision training
@@ -397,6 +413,26 @@ if not accelerator.optimizer_step_was_skipped:
     lr_scheduler.step()
 ```
 
+### Gradient Accumulation 
+
+To perform gradient accumulation use [`~Accelerator.accumulate`] and specify a `gradient_accumulation_steps`. 
+This will also automatically ensure the gradients are synced or unsynced when on multi-device training, check if the step should
+actually be performed, and auto-scale the loss:
+
+```python
+accelerator = Accelerator(gradient_accumulation_steps=2)
+model, optimizer, training_dataloader = accelerator.prepare(model, optimizer, training_dataloader)
+
+for input, label in training_dataloader:
+    with accelerator.accumulate(model):
+        predictions = model(input)
+        loss = loss_function(predictions, label)
+        accelerator.backward(loss)
+        optimizer.step()
+        scheduler.step()
+        optimizer.zero_grad()
+```
+
 ### DeepSpeed
 
 DeepSpeed support is experimental, so the underlying API will evolve in the near future and may have some slight
@@ -405,7 +441,7 @@ will be added in a next version.
 
 <Tip warning={true}>
 
-The [`notebook_launcher`] does not support the DeepSpeed integration yet.
+    The [`notebook_launcher`] does not support the DeepSpeed integration yet.
 
 </Tip>
 
@@ -415,7 +451,7 @@ Internally, the library works by first analyzing the environment in which the sc
 kind of distributed setup is used, how many different processes there are and which one the current script is in. All
 that information is stored in the [`~AcceleratorState`].
 
-This class is initialized the first time you instantiate a [`Accelerator`] as well as performing any
+This class is initialized the first time you instantiate an [`~Accelerator`] as well as performing any
 specific initialization your distributed setup needs. Its state is then uniquely shared through all instances of
 [`~state.AcceleratorState`].
 
@@ -448,18 +484,18 @@ setting the same seed in the main random number generator in all processes.
 
 <Tip warning={true}>
 
-Synchronization of the main torch (or CUDA or XLA) random number generator will affect any other potential random
-artifacts you could have in your dataset (like random data augmentation) in the sense that all processes will get
-the same random numbers from the torch random modules (so will apply the same random data augmentation if it's
-controlled by torch).
+    Synchronization of the main torch (or CUDA or XLA) random number generator will affect any other potential random
+    artifacts you could have in your dataset (like random data augmentation) in the sense that all processes will get
+    the same random numbers from the torch random modules (so will apply the same random data augmentation if it's
+    controlled by torch).
 
 </Tip>
 
 <Tip>
 
-The randomization part of your custom sampler, batch sampler or iterable dataset should be done using a local
-`torch.Generator` object (in PyTorch >= 1.6), see the traditional `RandomSampler`, as an example.
+    The randomization part of your custom sampler, batch sampler or iterable dataset should be done using a local
+    `torch.Generator` object (in PyTorch >= 1.6), see the traditional `RandomSampler`, as an example.
 
 </Tip>
 
-For more details about the internals, see the [Internals page](internal).
+For more details about the internals, see the [Internals page](package_reference/torch_wrappers).
diff --git a/docs/source/big_modeling.mdx b/docs/source/usage_guides/big_modeling.mdx
similarity index 87%
rename from docs/source/big_modeling.mdx
rename to docs/source/usage_guides/big_modeling.mdx
index 006c8bc4e..243b354d0 100644
--- a/docs/source/big_modeling.mdx
+++ b/docs/source/usage_guides/big_modeling.mdx
@@ -31,7 +31,7 @@ While this works very well for regularly sized models, this workflow has some cl
 
 <Tip warning={true}>
 
-This API is quite new and still in its experimental stage. While we strive to provide a stable API, it's possible some small parts of the public API will change in the future.
+    This API is quite new and still in its experimental stage. While we strive to provide a stable API, it's possible some small parts of the public API will change in the future.
 
 </Tip>
 
@@ -57,7 +57,7 @@ initializes an empty model with a bit more than 100B parameters. Behind the scen
 
 <Tip warning={true}>
 
-You can't move a model initialized like this on CPU or another device directly, since it doesn't have any data. It's also very likely that a forward pass with that empty model will fail, as not all operations are supported on the meta device.
+    You can't move a model initialized like this on CPU or another device directly, since it doesn't have any data. It's also very likely that a forward pass with that empty model will fail, as not all operations are supported on the meta device.
 
 </Tip>
 
@@ -199,7 +199,7 @@ This way, you model can run for inference even if it doesn't fit on one of the G
 
 <Tip warning={true}>
 
-This only supports inference of your model, not training. Most of the computation happens behind `torch.no_grad()` context managers to avoid spending some GPU memory with intermediate activations.
+    This only supports inference of your model, not training. Most of the computation happens behind `torch.no_grad()` context managers to avoid spending some GPU memory with intermediate activations.
 
 </Tip>
 
@@ -209,7 +209,7 @@ You can let 🤗 Accelerate handle the device map computation by setting `device
 
 <Tip>
 
-You can derive all sizes of the model (and thus compute a `device_map`) on a model that is on the meta device.
+    You can derive all sizes of the model (and thus compute a `device_map`) on a model that is on the meta device.
 
 </Tip>
 
@@ -222,7 +222,7 @@ When you have more GPU memory available than the model size, here the difference
 
 <Tip>
 
-The options `"auto"` and `"balanced"` produce the same results for now, but the behavior of `"auto"` might change in the future if we find a strategy that makes more sense, while `"balanced"` will stay stable.
+    The options `"auto"` and `"balanced"` produce the same results for now, but the behavior of `"auto"` might change in the future if we find a strategy that makes more sense, while `"balanced"` will stay stable.
 
 </Tip>
 
@@ -238,9 +238,9 @@ device_map = infer_auto_device_map(my_model, max_memory={0: "10GiB", 1: "10GiB",
 
 <Tip warning={true}>
 
-When a first allocation happens in PyTorch, it loads CUDA kernels which take about 1-2GB of memory depending on the GPU. Therefore you always have less usable memory than the actual size of the GPU. To see how much memory is actually used do `torch.ones(1).cuda()` and look at the memory usage.
+    When a first allocation happens in PyTorch, it loads CUDA kernels which take about 1-2GB of memory depending on the GPU. Therefore you always have less usable memory than the actual size of the GPU. To see how much memory is actually used do `torch.ones(1).cuda()` and look at the memory usage.
 
-Therefore when you create memory maps with `max_memory` make sure to adjust the avaialble memory accordingly to avoid out-of-memory errors.
+    Therefore when you create memory maps with `max_memory` make sure to adjust the avaialble memory accordingly to avoid out-of-memory errors.
 
 </Tip>
 
@@ -271,7 +271,7 @@ device_map = {"block1": 0, "block2.linear1": 1, "block2.linear2": 1}
 
 <Tip>
 
-To be the most efficient, make sure your device map puts the parameters on the GPUs in a sequential manner (e.g. don't put one of the first weights on GPU 0, then weights on GPU 1 and the last weight back to GPU 0) to avoid making many transfers of data between the GPUs.
+    To be the most efficient, make sure your device map puts the parameters on the GPUs in a sequential manner (e.g. don't put one of the first weights on GPU 0, then weights on GPU 1 and the last weight back to GPU 0) to avoid making many transfers of data between the GPUs.
 
 </Tip>
 
@@ -285,20 +285,4 @@ We are aware of the current limitations in the API:
 - [`load_checkpoint_and_dispatch`] and [`load_checkpoint_in_model`] do not perform any check on the correctness of your state dict compared to your model at the moment (this will be fixed in a future version), so you may get some weird errors if trying to load a checkpoint with mismatched or missing keys.
 - The model parallelism used when your model is split on several GPUs is naive and not optimized, meaning that only one GPU works at a given time and the other sits idle.
 - When weights are offloaded on the CPU/hard drive, there is no pre-fetching (yet, we will work on this for future versions) which means the weights are put on the GPU when they are needed and not before.
-- Hard-drive offloading might be very slow if the hardware you run on does not have fast communication between disk and CPU (like NVMes).
-
-## API doc
-
-[[autodoc]] cpu_offload
-
-[[autodoc]] disk_offload
-
-[[autodoc]] dispatch_model
-
-[[autodoc]] infer_auto_device_map
-
-[[autodoc]] init_empty_weights
-
-[[autodoc]] load_checkpoint_and_dispatch
-
-[[autodoc]] load_checkpoint_in_model
+- Hard-drive offloading might be very slow if the hardware you run on does not have fast communication between disk and CPU (like NVMes).
\ No newline at end of file
diff --git a/docs/source/checkpoint.mdx b/docs/source/usage_guides/checkpoint.mdx
similarity index 100%
rename from docs/source/checkpoint.mdx
rename to docs/source/usage_guides/checkpoint.mdx
diff --git a/docs/source/deepspeed.mdx b/docs/source/usage_guides/deepspeed.mdx
similarity index 98%
rename from docs/source/deepspeed.mdx
rename to docs/source/usage_guides/deepspeed.mdx
index c0385c352..8a4238732 100644
--- a/docs/source/deepspeed.mdx
+++ b/docs/source/usage_guides/deepspeed.mdx
@@ -472,25 +472,11 @@ model, eval_dataloader = accelerator.prepare(model, eval_dataloader)
 
 1. Current integration doesn’t support Pipeline Parallelism of DeepSpeed.
 2. Current integration doesn’t support `mpu`, limiting the tensor parallelism which is supported in Megatron-LM. 
-3. Current integration doesn’t support multiple models for a given `accelerator` object. 
+3. Current integration doesn’t support multiple models. 
 
+## DeepSpeed Resources
 
-## Internals
-
-[[autodoc]] utils.DeepSpeedPlugin
-
-[[autodoc]] utils.DummyOptim
-
-[[autodoc]] utils.DummyScheduler
-
-[[autodoc]] utils.DeepSpeedEngineWrapper
-
-[[autodoc]] utils.DeepSpeedOptimizerWrapper
-
-[[autodoc]] utils.DeepSpeedSchedulerWrapper
-
-
-## Main DeepSpeed Resources
+The documentation for the internals related to deepspeed can be found [here](../package_reference/deepspeed).
 
 - [Project's github](https://github.com/microsoft/deepspeed)
 - [Usage docs](https://www.deepspeed.ai/getting-started/)
diff --git a/docs/source/fsdp.mdx b/docs/source/usage_guides/fsdp.mdx
similarity index 73%
rename from docs/source/fsdp.mdx
rename to docs/source/usage_guides/fsdp.mdx
index 32f982786..a561e4ecd 100644
--- a/docs/source/fsdp.mdx
+++ b/docs/source/usage_guides/fsdp.mdx
@@ -39,10 +39,14 @@ For instance, here is how you would run the NLP example (from the root of the re
 compute_environment: LOCAL_MACHINE
 deepspeed_config: {}
 distributed_type: FSDP
+downcast_bf16: 'no'
 fsdp_config:
-  min_num_params: 2000
-  offload_params: false
-  sharding_strategy: 1
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_backward_prefetch_policy: BACKWARD_PRE
+  fsdp_offload_params: false
+  fsdp_sharding_strategy: 1
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_transformer_layer_cls_to_wrap: GPT2Block
 machine_rank: 0
 main_process_ip: null
 main_process_port: null
@@ -60,9 +64,13 @@ accelerate launch examples/nlp_example.py
 Currently, `Accelerate` supports the following config through the CLI:
 
 ```bash
-`Sharding Strategy`: [1] FULL_SHARD, [2] SHARD_GRAD_OP
-`Min Num Params`: FSDP\'s minimum number of parameters for Default Auto Wrapping.
-`Offload Params`: Decides Whether to offload parameters and gradients to CPU.
+`Sharding Strategy`: [1] FULL_SHARD (shards optimizer states, gradients and parameters), [2] SHARD_GRAD_OP (shards optimizer states and gradients), [3] NO_SHARD
+`Offload Params`: Decides Whether to offload parameters and gradients to CPU
+`Auto Wrap Policy`: [1] TRANSFORMER_BASED_WRAP, [2] SIZE_BASED_WRAP, [3] NO_WRAP 
+`Transformer Layer Class to Wrap`: When using `TRANSFORMER_BASED_WRAP`, user specifies transformer layer class name (case-sensitive) to wrap ,e.g, `BertLayer`, `GPTJBlock`, `T5Block`...
+`Min Num Params`: minimum number of parameters when using `SIZE_BASED_WRAP`
+`Backward Prefetch`: [1] BACKWARD_PRE, [2] BACKWARD_POST, [3] NO_PREFETCH
+`State Dict Type`: [1] FULL_STATE_DICT, [2] LOCAL_STATE_DICT, [3] SHARDED_STATE_DICT  
 ```
 
 ## A few caveats to be aware of
@@ -76,19 +84,18 @@ Currently, `Accelerate` supports the following config through the CLI:
 However, below is the recommended way to prepare model and optimizer while using FSDP:
 
 ```diff
-model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
+  model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
 + model = accelerator.prepare(model)
 
-optimizer = torch.optim.AdamW(params=model.parameters(), lr=lr)
+  optimizer = torch.optim.AdamW(params=model.parameters(), lr=lr)
 
-- model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(model,
--        optimizer, train_dataloader, eval_dataloader, lr_scheduler
+- model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+-        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
 -    )
 
 + optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
 +         optimizer, train_dataloader, eval_dataloader, lr_scheduler
-+        )
-
++    )
 ```
 
 - In case of a single model, if you have created the optimizer with multiple parameter groups and called prepare with them together,
@@ -110,11 +117,9 @@ optimizer = torch.optim.AdamW(params=model.parameters(), lr=lr)
   ```
 
 
-- In case of multiple models, it is necessary to prepare the models before creating optimizers or else it will throw an error.
-- Mixed precision is currently not supported with FSDP.
+- In case of multiple models, it is necessary to prepare the models before creating optimizers or else it will throw an error. 
+Then pass the optimizers to the prepare call in the same order as corresponding models else `accelerator.save_state()` and `accelerator.load_state()` will result in wrong/unexpected behaviour.
+- This feature is incompatible with `--predict_with_generate` in the `run_translation.py` script of 🤗 `Transformers` library.
 
-For more control, users can leverage the `FullyShardedDataParallelPlugin` wherein they can specify `auto_wrap_policy`, `backward_prefetch` and `ignored_modules`.
-After creating an instance of this class, users can pass it to the Accelerator class instantiation.
+For more control, users can leverage the `FullyShardedDataParallelPlugin`. After creating an instance of this class, users can pass it to the Accelerator class instantiation.
 For more information on these options, please refer to the PyTorch [FullyShardedDataParallel](https://github.com/pytorch/pytorch/blob/0df2e863fbd5993a7b9e652910792bd21a516ff3/torch/distributed/fsdp/fully_sharded_data_parallel.py#L236) code.
-
-[[autodoc]] utils.FullyShardedDataParallelPlugin
diff --git a/docs/source/gradient_accumulation.mdx b/docs/source/usage_guides/gradient_accumulation.mdx
similarity index 86%
rename from docs/source/gradient_accumulation.mdx
rename to docs/source/usage_guides/gradient_accumulation.mdx
index f6ab857d8..f9f5dc62a 100644
--- a/docs/source/gradient_accumulation.mdx
+++ b/docs/source/usage_guides/gradient_accumulation.mdx
@@ -71,13 +71,15 @@ First the code shown earlier will be converted to utilize 🤗 Accelerate withou
 ```
 
 <Tip warning={true}>
-In its current state, this code is not going to perform gradient accumulation efficiently due to a process called gradient synchronization.
+
+  In its current state, this code is not going to perform gradient accumulation efficiently due to a process called gradient synchronization. Read more about that in the [Concepts tutorial](concept_guides/gradient_synchronization)!
+
 </Tip>
 
 ## Letting 🤗 Accelerate handle gradient accumulation
 
 All that is left now is to let 🤗 Accelerate handle the gradient accumulation for us. To do so you should pass in a `gradient_accumulation_steps` parameter to [`Accelerator`], dictating the number 
-of steps to perform before each call to `step()` and how to automatically adjust the loss during the call to [`Accelerator.backward`]:
+of steps to perform before each call to `step()` and how to automatically adjust the loss during the call to [`~Accelerator.backward`]:
 
 ```diff
   from accelerate import Accelerator
@@ -85,19 +87,19 @@ of steps to perform before each call to `step()` and how to automatically adjust
 + accelerator = Accelerator(gradient_accumulation_steps=2)
 ```
 
-From here you can use the [`Accelerator.accumulate`] context manager from inside your training loop to automatically perform the gradient accumulation for you!
-You just wrap it around the entire training part of your code: 
+From here you can use the [`~Accelerator.accumulate`] context manager from inside your training loop to automatically perform the gradient accumulation for you!
+You just wrap it around the entire training part of our code: 
 
 ```diff
 - for index, batch in enumerate(training_dataloader):
 + for batch in training_dataloader:
-+    with accelerator.accumulate(model):
++     with accelerator.accumulate(model):
           optimizer.zero_grad()
           inputs, targets = batch
           outputs = model(inputs)
 ```
 
-and you can remove all the special checks for the step number and the loss adjustment:
+You can remove all the special checks for the step number and the loss adjustment:
 
 ```diff
 - loss = loss / gradient_accumulation_steps
diff --git a/docs/source/memory.mdx b/docs/source/usage_guides/memory.mdx
similarity index 95%
rename from docs/source/memory.mdx
rename to docs/source/usage_guides/memory.mdx
index e578e43d7..9c5674977 100644
--- a/docs/source/memory.mdx
+++ b/docs/source/usage_guides/memory.mdx
@@ -48,4 +48,4 @@ def training_function(args):
 +   inner_training_loop()
 ```
 
-[[autodoc]] utils.find_executable_batch_size
\ No newline at end of file
+To find out more, check the documentation [here](package_reference/utilities#accelerate.find_executable_batch_size)
\ No newline at end of file
diff --git a/docs/source/sagemaker.mdx b/docs/source/usage_guides/sagemaker.mdx
similarity index 90%
rename from docs/source/sagemaker.mdx
rename to docs/source/usage_guides/sagemaker.mdx
index 563f0c742..6195783e9 100644
--- a/docs/source/sagemaker.mdx
+++ b/docs/source/usage_guides/sagemaker.mdx
@@ -54,7 +54,7 @@ accelerate config
 
 <Tip>
 
-🤗 Accelerate is not saving any of your credentials.
+    🤗 Accelerate is not saving any of your credentials.
 
 </Tip>
 
@@ -72,8 +72,8 @@ directory. After training, artifacts in this directory are uploaded to S3:
 
 <Tip warning={true}>
 
-SageMaker doesn’t support argparse actions. If you want to use, for example, boolean hyperparameters, you need to
-specify type as bool in your script and provide an explicit True or False value for this hyperparameter. [[REF]](https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/using_pytorch.html#prepare-a-pytorch-training-script).
+    SageMaker doesn’t support argparse actions. If you want to use, for example, boolean hyperparameters, you need to
+    specify type as bool in your script and provide an explicit True or False value for this hyperparameter. [[REF]](https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/using_pytorch.html#prepare-a-pytorch-training-script).
 
 </Tip>
 
@@ -92,7 +92,7 @@ arguments needed by your training script as named arguments.
 
 <Tip>
 
-If you run one of the example scripts, don't forget to add `accelerator.save('/opt/ml/model')` to it.
+    If you run one of the example scripts, don't forget to add `accelerator.save('/opt/ml/model')` to it.
 
 </Tip>
 
diff --git a/docs/source/tracking.mdx b/docs/source/usage_guides/tracking.mdx
similarity index 85%
rename from docs/source/tracking.mdx
rename to docs/source/usage_guides/tracking.mdx
index 25e44b3c9..2fb68c0f2 100644
--- a/docs/source/tracking.mdx
+++ b/docs/source/usage_guides/tracking.mdx
@@ -19,12 +19,9 @@ There are a large number of experiment tracking API's available, however getting
 
 Currently `Accelerate` supports three trackers out-of-the-box:
 
-
-[[autodoc]] tracking.TensorBoardTracker
-
-[[autodoc]] tracking.WandBTracker
-
-[[autodoc]] tracking.CometMLTracker
+- TensorBoard
+- WandB
+- CometML
 
 To use any of them, pass in the selected type(s) to the `log_with` parameter in [`Accelerate`]:
 ```python
@@ -107,6 +104,9 @@ import wandb
 
 
 class MyCustomTracker(GeneralTracker):
+    name = "wandb"
+    requires_logging_directory = False
+
     def __init__(self, run_name: str):
         self.run_name = run_name
         wandb.init(self.run_name)
@@ -137,27 +137,27 @@ accelerator = Accelerator(log_with=[tracker, "all"])
 
 If a library has an API that does not follow a strict `.log` with an overall dictionary such as Neptune.AI, logging can be done manually under an `if accelerator.is_main_process` statement:
 ```diff
-from accelerate import Accelerator
+  from accelerate import Accelerator
 + import neptune.new as neptune
 
-accelerator = Accelerator()
+  accelerator = Accelerator()
 + run = neptune.init(...)
 
-my_model, my_optimizer, my_training_dataloader = accelerate.prepare(my_model, my_optimizer, my_training_dataloader)
-device = accelerator.device
-my_model.to(device)
-
-for iteration in config["num_iterations"]:
-    for batch in my_training_dataloader:
-        my_optimizer.zero_grad()
-        inputs, targets = batch
-        inputs = inputs.to(device)
-        targets = targets.to(device)
-        outputs = my_model(inputs)
-        loss = my_loss_function(outputs, targets)
-        total_loss += loss
-        accelerator.backward(loss)
-        my_optimizer.step()
-+       if accelerator.is_main_process:
-+           run["logs/training/batch/loss"].log(loss)
+  my_model, my_optimizer, my_training_dataloader = accelerate.prepare(my_model, my_optimizer, my_training_dataloader)
+  device = accelerator.device
+  my_model.to(device)
+
+  for iteration in config["num_iterations"]:
+      for batch in my_training_dataloader:
+          my_optimizer.zero_grad()
+          inputs, targets = batch
+          inputs = inputs.to(device)
+          targets = targets.to(device)
+          outputs = my_model(inputs)
+          loss = my_loss_function(outputs, targets)
+          total_loss += loss
+          accelerator.backward(loss)
+          my_optimizer.step()
++         if accelerator.is_main_process:
++             run["logs/training/batch/loss"].log(loss)
 ```
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 977cd5d24..41fc1e3d4 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -140,10 +140,20 @@ class Accelerator:
             A list of `KwargHandler` to customize how the objects related to distributed training or mixed precision
             are created. See [kwargs](kwargs) for more information.
 
-    Attributes
+    **Attributes:**
 
         - **device** (`torch.device`) -- The device to use.
+        - **distributed_type** ([`~utils.DistributedType`]) -- The distributed training configuration.
+        - **local_process_index** (`int`) -- The process index on the current machine.
+        - **mixed_precision** (`str`) -- The configured mixed precision mode.
+        - **num_processes** (`int`) -- The total number of processes used for training.
+        - **optimizer_step_was_skipped** (`bool`) -- Whether or not the optimizer update was skipped (because of
+          gradient overflow in mixed precision), in which
+        case the learning rate should not be changed.
+        - **process_index** (`int`) -- The overall index of the current process among all processes.
         - **state** ([`~state.AcceleratorState`]) -- The distributed setup state.
+        - **sync_gradients** (`bool`) -- Whether the gradients are currently being synced across all processes.
+        - **use_distributed** (`bool`) -- Whether the current configuration is for distributed training.
     """
 
     def __init__(
@@ -308,6 +318,9 @@ def __init__(
 
     @property
     def use_distributed(self):
+        """
+        Whether the Accelerator is configured for distributed training
+        """
         return self.distributed_type != DistributedType.NO and self.num_processes > 1
 
     @property
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index 27a627ae9..a4f6bac2a 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -47,20 +47,18 @@ def parse_choice_from_env(key, default="no"):
 # Inspired by Alex Martelli's 'Borg'.
 class AcceleratorState:
     """
-    This is a variation of a [singleton class](https://en.wikipedia.org/wiki/Singleton_pattern) in the sense that all
-    instance of `AcceleratorState` share the same state, which is initialized on the first instantiation.
+    Singleton class that has information about the current training environment.
 
-    Attributes:
+    **Attributes:**
 
         - **device** (`torch.device`) -- The device to use.
-        - **sync_gradients** (`bool`) -- Whether to sync the gradients or not
-        - **distributed_type** (`~accelerate.state.DistributedType`) -- The type of distributed environment currently
+        - **distributed_type** ([`~accelerate.state.DistributedType`]) -- The type of distributed environment currently
           in use.
-        - **num_processes** (`int`) -- The number of processes currently launched in parallel.
-        - **process_index** (`int`) -- The index of the current process.
         - **local_process_index** (`int`) -- The index of the current process on the current server.
         - **mixed_precision** (`str`) -- Whether or not the current script will use mixed precision. If you are using
           mixed precision, define if you want to use FP16 or BF16 (bfloat16) as the floating point.
+        - **num_processes** (`int`) -- The number of processes currently launched in parallel.
+        - **process_index** (`int`) -- The index of the current process.
     """
 
     _shared_state = {}
@@ -250,13 +248,10 @@ def _check_initialized(self, mixed_precision=None, cpu=None):
 
 class GradientState:
     """
-    This is a variation of a [singleton class](https://en.wikipedia.org/wiki/Singleton_pattern) in the sense that all
-    instance of `GradientState` share the same state, which is initialized on the first instantiation.
+    Singleton class that has information related to gradient synchronization for gradient accumulation
 
-    This specific state revolves around whether gradients should be synced and if we have reached the end of a prepared
-    dataloader Attributes:
+    **Attributes:**
 
-        - **sync_gradients** (`bool`) -- Whether the gradients should be synced
         - **end_of_dataloader** (`bool`) -- Whether we have reached the end the current dataloader
         - **remainder** (`int`) -- The number of extra samples that were added from padding the dataloader
     """

missing the use of the context manager below
@@ -0,0 +1,109 @@
+<!--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.
+-->
+
+# Performing gradient accumulation with 🤗 Accelerate
+
+Gradient accumulation is a technique where you can train on bigger batch sizes than 
+your machine would normally be able to fit into memory. This is done by accumulating gradients over
+several batches, and only performing the step through the optimizer after a certain number of batches have been performed.
+
+While technically standard gradient accumulation code would work fine in a distributed setup, it is not the most efficient
+method for doing so and you may experience considerable slowdowns!
+
+In this tutorial you will see how to quickly setup gradient accumulation and perform it with the utilties provided in 🤗 Accelerate,
+which can total to adding just one line of code!
+
+This example will use a very simplistic PyTorch training loop that performs gradient accumulation every two batches:
+
+```python
+device = 'cuda'
+model.to(device)
+
+gradient_accumulation_steps = 2
+
+for index, batch in enumerate(training_dataloader):
+    optimizer.zero_grad()
+    inputs, targets = batch
+    inputs = inputs.to(device)
+    targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    loss = loss / gradient_accumulation_steps
+    loss.backward()
+    if (index+1) % gradient_accumulation_steps == 0:
+        optimizer.step()
+        scheduler.step()
+```
+
+## Converting it to 🤗 Accelerate
+
+First the code shown earlier will be converted to utilize 🤗 Accelerate without the special gradient accumulation helper:
+
+```python
+from accelerate import Accelerator
+accelerator = Accelerator()
+
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+
+for index, batch in enumerate(training_dataloader):
+    optimizer.zero_grad()
+    inputs, targets = batch
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    loss = loss / gradient_accumulation_steps
+    accelerator.backward(loss)
+    if (index+1) % gradient_accumulation_steps == 0:
+        optimizer.step()
+        scheduler.step()
+```
+
+<Tip warning={true}>
+In its current state, this code is not going to perform gradient accumulation efficiently due to a process called gradient synchronization. Read more about that in the Concepts tutorial!
+</Tip>
+
+## Letting 🤗 Accelerate handle gradient accumulation
+
+All that is left now is to let 🤗 Accelerate handle the gradient accumulation for us. To do so you should pass in a `gradient_accumulation_steps` parameter to `Accelerator`, dictating the number 
+of steps to perform before each call to `step()` and `backward()`:
+
+```diff
+from accelerate import Accelerator
+-accelerator = Accelerator()
++accelerator = Accelerator(gradient_accumulation_steps=2)
+```
+
+From here you can use the [`Accelerator.accumulate`] context manager from inside your training loop to automatically perform the gradient accumulation for you!
Thanks 😄 *Should* be set now, appreciate your thorough review on it! 🙇 
@@ -0,0 +1,109 @@
+<!--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.
+-->
+
+# Performing gradient accumulation with 🤗 Accelerate
+
+Gradient accumulation is a technique where you can train on bigger batch sizes than 
+your machine would normally be able to fit into memory. This is done by accumulating gradients over
+several batches, and only performing the step through the optimizer after a certain number of batches have been performed.
+
+While technically standard gradient accumulation code would work fine in a distributed setup, it is not the most efficient
+method for doing so and you may experience considerable slowdowns!
+
+In this tutorial you will see how to quickly setup gradient accumulation and perform it with the utilties provided in 🤗 Accelerate,
+which can total to adding just one line of code!
+
+This example will use a very simplistic PyTorch training loop that performs gradient accumulation every two batches:
+
+```python
+device = 'cuda'
+model.to(device)
+
+gradient_accumulation_steps = 2
+
+for index, batch in enumerate(training_dataloader):
+    optimizer.zero_grad()
+    inputs, targets = batch
+    inputs = inputs.to(device)
+    targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    loss = loss / gradient_accumulation_steps
+    loss.backward()
+    if (index+1) % gradient_accumulation_steps == 0:
+        optimizer.step()
+        scheduler.step()
+```
+
+## Converting it to 🤗 Accelerate
+
+First the code shown earlier will be converted to utilize 🤗 Accelerate without the special gradient accumulation helper:
+
+```python
+from accelerate import Accelerator
+accelerator = Accelerator()
+
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+
+for index, batch in enumerate(training_dataloader):
+    optimizer.zero_grad()
+    inputs, targets = batch
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    loss = loss / gradient_accumulation_steps
+    accelerator.backward(loss)
+    if (index+1) % gradient_accumulation_steps == 0:
+        optimizer.step()
+        scheduler.step()
+```
+
+<Tip warning={true}>
+In its current state, this code is not going to perform gradient accumulation efficiently due to a process called gradient synchronization. Read more about that in the Concepts tutorial!
+</Tip>
+
+## Letting 🤗 Accelerate handle gradient accumulation
+
+All that is left now is to let 🤗 Accelerate handle the gradient accumulation for us. To do so you should pass in a `gradient_accumulation_steps` parameter to `Accelerator`, dictating the number 
+of steps to perform before each call to `step()` and `backward()`:
+
+```diff
+from accelerate import Accelerator
+-accelerator = Accelerator()
++accelerator = Accelerator(gradient_accumulation_steps=2)
+```
+
+From here you can use the [`Accelerator.accumulate`] context manager from inside your training loop to automatically perform the gradient accumulation for you!
looking forward to this feature!
@@ -0,0 +1,109 @@
+<!--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.
+-->
+
+# Performing gradient accumulation with 🤗 Accelerate
+
+Gradient accumulation is a technique where you can train on bigger batch sizes than 
+your machine would normally be able to fit into memory. This is done by accumulating gradients over
+several batches, and only performing the step through the optimizer after a certain number of batches have been performed.
+
+While technically standard gradient accumulation code would work fine in a distributed setup, it is not the most efficient
+method for doing so and you may experience considerable slowdowns!
+
+In this tutorial you will see how to quickly setup gradient accumulation and perform it with the utilties provided in 🤗 Accelerate,
+which can total to adding just one line of code!
+
+This example will use a very simplistic PyTorch training loop that performs gradient accumulation every two batches:
+
+```python
+device = 'cuda'
+model.to(device)
+
+gradient_accumulation_steps = 2
+
+for index, batch in enumerate(training_dataloader):
+    optimizer.zero_grad()
+    inputs, targets = batch
+    inputs = inputs.to(device)
+    targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    loss = loss / gradient_accumulation_steps
+    loss.backward()
+    if (index+1) % gradient_accumulation_steps == 0:
+        optimizer.step()
+        scheduler.step()
+```
+
+## Converting it to 🤗 Accelerate
+
+First the code shown earlier will be converted to utilize 🤗 Accelerate without the special gradient accumulation helper:
+
+```python
+from accelerate import Accelerator
+accelerator = Accelerator()
+
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+
+for index, batch in enumerate(training_dataloader):
+    optimizer.zero_grad()
+    inputs, targets = batch
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    loss = loss / gradient_accumulation_steps
+    accelerator.backward(loss)
+    if (index+1) % gradient_accumulation_steps == 0:
+        optimizer.step()
+        scheduler.step()
+```
+
+<Tip warning={true}>
+In its current state, this code is not going to perform gradient accumulation efficiently due to a process called gradient synchronization. Read more about that in the Concepts tutorial!
+</Tip>
+
+## Letting 🤗 Accelerate handle gradient accumulation
+
+All that is left now is to let 🤗 Accelerate handle the gradient accumulation for us. To do so you should pass in a `gradient_accumulation_steps` parameter to `Accelerator`, dictating the number 
+of steps to perform before each call to `step()` and `backward()`:
+
+```diff
+from accelerate import Accelerator
+-accelerator = Accelerator()
++accelerator = Accelerator(gradient_accumulation_steps=2)
+```
+
+From here you can use the [`Accelerator.accumulate`] context manager from inside your training loop to automatically perform the gradient accumulation for you!
```suggestion
🤗 Accelerate is a library that enables the same PyTorch code to be run across any distributed configuration by adding just four lines of code! In short, training and inference at scale made simple, efficient and adaptable.
```
@@ -12,121 +12,56 @@ specific language governing permissions and limitations under the License.
 
 # Accelerate
 
-Run your *raw* PyTorch training script on any kind of device.
+🤗 Accelerate is a library that allows the same PyTorch code be ran across any distributed configuration by adding just four lines of code!
```suggestion
platforms separately so you don't have to. In minutes, existing codebases can be converted to utilize [DeepSpeed](https://huggingface.co/docs/accelerate/deepspeed), perform [fully sharded data parallelism](https://huggingface.co/docs/accelerate/fsdp), and have automatic support for mixed-precision training! 
```
@@ -12,121 +12,56 @@ specific language governing permissions and limitations under the License.
 
 # Accelerate
 
-Run your *raw* PyTorch training script on any kind of device.
+🤗 Accelerate is a library that allows the same PyTorch code be ran across any distributed configuration by adding just four lines of code!
 
-## Features
+Built on `torch_xla` and `torch.distributed`, Accelerate wraps around the tough and problematic code changes that need to occur on each of these
+platforms separately so you don't have to. In minutes existing codebases can be converted to utilize [DeepSpeed](https://github.com/microsoft/DeepSpeed), perform [fully sharded data parallelism](https://huggingface.co/docs/accelerate/fsdp), and have automatic support for mixed-precision training! 
```suggestion
<Tip> To get a better idea of this process, make sure to check out the [Tutorials](https://huggingface.co/docs/accelerate/basic_tutorials/overview) ! </Tip>
```
@@ -12,121 +12,56 @@ specific language governing permissions and limitations under the License.
 
 # Accelerate
 
-Run your *raw* PyTorch training script on any kind of device.
+🤗 Accelerate is a library that allows the same PyTorch code be ran across any distributed configuration by adding just four lines of code!
 
-## Features
+Built on `torch_xla` and `torch.distributed`, Accelerate wraps around the tough and problematic code changes that need to occur on each of these
+platforms separately so you don't have to. In minutes existing codebases can be converted to utilize [DeepSpeed](https://github.com/microsoft/DeepSpeed), perform [fully sharded data parallelism](https://huggingface.co/docs/accelerate/fsdp), and have automatic support for mixed-precision training! 
 
-- 🤗 Accelerate provides an easy API to make your scripts run with mixed precision and in any kind of distributed
-  setting (multi-GPUs, TPUs etc.) while still letting you write your own training loop. The same code can then run
-  seamlessly on your local machine for debugging or your training environment.
-
-- 🤗 Accelerate also provides a CLI tool that allows you to quickly configure and test your training environment and
-  then launch the scripts.
-
-
-## Easy to integrate
-
-A traditional training loop in PyTorch looks like this:
-
-```python
-my_model.to(device)
-
-for batch in my_training_dataloader:
-    my_optimizer.zero_grad()
-    inputs, targets = batch
-    inputs = inputs.to(device)
-    targets = targets.to(device)
-    outputs = my_model(inputs)
-    loss = my_loss_function(outputs, targets)
-    loss.backward()
-    my_optimizer.step()
-```
-
-Changing it to work with accelerate is really easy and only adds a few lines of code:
+For an idea of just how easy it can be, here's a brief example of the changes needed on a basic PyTorch script:
+<Tip>To get a better idea of this process, make sure to check out the [Tutorials](https://huggingface.co/docs/accelerate/basic_tutorials/overview)!</Tip>
I think for even greater impact, you can move the code snippet up here so users can immediately see the four lines of code and how easy it is to launch.
@@ -12,121 +12,56 @@ specific language governing permissions and limitations under the License.
 
 # Accelerate
 
-Run your *raw* PyTorch training script on any kind of device.
+🤗 Accelerate is a library that enables the same PyTorch code to be run across any distributed configuration by adding just four lines of code! In short, training and inference at scale made simple, efficient and adaptable.
Maybe to avoid sounding too negative, this is a bit better?

```suggestion
Built on `torch_xla` and `torch.distributed`, 🤗 Accelerate takes care of the heavy lifting, so you don't have to write any custom code to adapt to these platforms.
```
@@ -12,121 +12,56 @@ specific language governing permissions and limitations under the License.
 
 # Accelerate
 
-Run your *raw* PyTorch training script on any kind of device.
+🤗 Accelerate is a library that enables the same PyTorch code to be run across any distributed configuration by adding just four lines of code! In short, training and inference at scale made simple, efficient and adaptable.
 
-## Features
+Built on `torch_xla` and `torch.distributed`, Accelerate wraps around the tough and problematic code changes that need to occur on each of these
```suggestion
platforms separately so you don't have to. Convert existing codebases to utilize [DeepSpeed](https://huggingface.co/docs/accelerate/deepspeed), perform [fully sharded data parallelism](https://huggingface.co/docs/accelerate/fsdp), and have automatic support for mixed-precision training! 
```
@@ -12,121 +12,56 @@ specific language governing permissions and limitations under the License.
 
 # Accelerate
 
-Run your *raw* PyTorch training script on any kind of device.
+🤗 Accelerate is a library that enables the same PyTorch code to be run across any distributed configuration by adding just four lines of code! In short, training and inference at scale made simple, efficient and adaptable.
 
-## Features
+Built on `torch_xla` and `torch.distributed`, Accelerate wraps around the tough and problematic code changes that need to occur on each of these
+platforms separately so you don't have to. In minutes, existing codebases can be converted to utilize [DeepSpeed](https://huggingface.co/docs/accelerate/deepspeed), perform [fully sharded data parallelism](https://huggingface.co/docs/accelerate/fsdp), and have automatic support for mixed-precision training! 
I think the Markdown link isn't rendered because there needs to be a newline after `<Tip>`.

```suggestion
<Tip> 

To get a better idea of this process, make sure to check out the [Tutorials](https://huggingface.co/docs/accelerate/basic_tutorials/overview)! 

</Tip>
```
@@ -12,121 +12,56 @@ specific language governing permissions and limitations under the License.
 
 # Accelerate
 
-Run your *raw* PyTorch training script on any kind of device.
+🤗 Accelerate is a library that enables the same PyTorch code to be run across any distributed configuration by adding just four lines of code! In short, training and inference at scale made simple, efficient and adaptable.
 
-## Features
+Built on `torch_xla` and `torch.distributed`, Accelerate wraps around the tough and problematic code changes that need to occur on each of these
+platforms separately so you don't have to. In minutes, existing codebases can be converted to utilize [DeepSpeed](https://huggingface.co/docs/accelerate/deepspeed), perform [fully sharded data parallelism](https://huggingface.co/docs/accelerate/fsdp), and have automatic support for mixed-precision training! 
 
-- 🤗 Accelerate provides an easy API to make your scripts run with mixed precision and in any kind of distributed
-  setting (multi-GPUs, TPUs etc.) while still letting you write your own training loop. The same code can then run
-  seamlessly on your local machine for debugging or your training environment.
-
-- 🤗 Accelerate also provides a CLI tool that allows you to quickly configure and test your training environment and
-  then launch the scripts.
-
-
-## Easy to integrate
-
-A traditional training loop in PyTorch looks like this:
-
-```python
-my_model.to(device)
-
-for batch in my_training_dataloader:
-    my_optimizer.zero_grad()
-    inputs, targets = batch
-    inputs = inputs.to(device)
-    targets = targets.to(device)
-    outputs = my_model(inputs)
-    loss = my_loss_function(outputs, targets)
-    loss.backward()
-    my_optimizer.step()
-```
-
-Changing it to work with accelerate is really easy and only adds a few lines of code:
+For an idea of just how easy it can be, here's a brief example of the changes needed on a basic PyTorch script:
+<Tip> To get a better idea of this process, make sure to check out the [Tutorials](https://huggingface.co/docs/accelerate/basic_tutorials/overview) ! </Tip>
The spacing is slightly off in this snippet for the lines removed and added.

```suggestion
-   inputs = inputs.to(device)
```
@@ -12,121 +12,56 @@ specific language governing permissions and limitations under the License.
 
 # Accelerate
 
-Run your *raw* PyTorch training script on any kind of device.
+🤗 Accelerate is a library that enables the same PyTorch code to be run across any distributed configuration by adding just four lines of code! In short, training and inference at scale made simple, efficient and adaptable.
 
-## Features
+Built on `torch_xla` and `torch.distributed`, Accelerate wraps around the tough and problematic code changes that need to occur on each of these
+platforms separately so you don't have to. In minutes, existing codebases can be converted to utilize [DeepSpeed](https://huggingface.co/docs/accelerate/deepspeed), perform [fully sharded data parallelism](https://huggingface.co/docs/accelerate/fsdp), and have automatic support for mixed-precision training! 
 
-- 🤗 Accelerate provides an easy API to make your scripts run with mixed precision and in any kind of distributed
-  setting (multi-GPUs, TPUs etc.) while still letting you write your own training loop. The same code can then run
-  seamlessly on your local machine for debugging or your training environment.
-
-- 🤗 Accelerate also provides a CLI tool that allows you to quickly configure and test your training environment and
-  then launch the scripts.
-
-
-## Easy to integrate
-
-A traditional training loop in PyTorch looks like this:
-
-```python
-my_model.to(device)
-
-for batch in my_training_dataloader:
-    my_optimizer.zero_grad()
-    inputs, targets = batch
-    inputs = inputs.to(device)
-    targets = targets.to(device)
-    outputs = my_model(inputs)
-    loss = my_loss_function(outputs, targets)
-    loss.backward()
-    my_optimizer.step()
-```
-
-Changing it to work with accelerate is really easy and only adds a few lines of code:
+For an idea of just how easy it can be, here's a brief example of the changes needed on a basic PyTorch script:
+<Tip> To get a better idea of this process, make sure to check out the [Tutorials](https://huggingface.co/docs/accelerate/basic_tutorials/overview) ! </Tip>
 
 ```diff
-+ from accelerate import Accelerator
-
-+ accelerator = Accelerator()
-  # Use the device given by the *accelerator* object.
-+ device = accelerator.device
-  my_model.to(device)
-  # Pass every important object (model, optimizer, dataloader) to *accelerator.prepare*
-+ my_model, my_optimizer, my_training_dataloader = accelerator.prepare(
-+     my_model, my_optimizer, my_training_dataloader
-+ )
-
-  for batch in my_training_dataloader:
-      my_optimizer.zero_grad()
-      inputs, targets = batch
-      inputs = inputs.to(device)
-      targets = targets.to(device)
-      outputs = my_model(inputs)
-      loss = my_loss_function(outputs, targets)
-      # Just a small change for the backward instruction
--     loss.backward()
-+     accelerator.backward(loss)
-      my_optimizer.step()
-```
++from accelerate import Accelerator
++accelerator = Accelerator()
 
-and with this, your script can now run in a distributed environment (multi-GPU, TPU).
-
-You can even simplify your script a bit by letting 🤗 Accelerate handle the device placement for you (which is safer,
-especially for TPU training):
-
-```diff
-+ from accelerate import Accelerator
++model, optimizer, training_dataloader, scheduler = accelerator.prepare(
++    model, optimizer, training_dataloader, scheduler
++)
 
-+ accelerator = Accelerator()
-- my_model.to(device)
-  # Pass every important object (model, optimizer, dataloader) to *accelerator.prepare*
-+ my_model, my_optimizer, my_training_dataloader = accelerate.prepare(
-+     my_model, my_optimizer, my_training_dataloader
-+ )
-
-  for batch in my_training_dataloader:
-      my_optimizer.zero_grad()
-      inputs, targets = batch
--     inputs = inputs.to(device)
--     targets = targets.to(device)
-      outputs = my_model(inputs)
-      loss = my_loss_function(outputs, targets)
-      # Just a small change for the backward instruction
--     loss.backward()
-+     accelerator.backward(loss)
-      my_optimizer.step()
-```
-
-## Script launcher
-
-No need to remember how to use `torch.distributed.launch` or to write a specific launcher for TPU training! 🤗
-Accelerate comes with a CLI tool that will make your life easier when launching distributed scripts.
-
-On your machine(s) just run:
-
-```bash
-accelerate config
-```
-
-and answer the questions asked. This will generate a config file that will be used automatically to properly set the
-default options when doing
-
-```bash
-accelerate launch my_script.py --args_to_my_script
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+-    inputs = inputs.to(device)
```suggestion
New features are added every day that haven't been released yet. To try them out yourself, install
```
@@ -0,0 +1,91 @@
+<!--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.
+-->
+
+# Installation and Configuration
+
+Before you start, you will need to setup your environment, install the appropriate packages, and configure 🤗 Accelerate. 🤗 Accelerate is tested on **Python 3.7+**.
+
+## Installing 🤗 Accelerate
+
+🤗 Accelerate is available on pypi and conda, as well as on GitHub. Details to install from each are below:
+
+### pip 
+
+To install 🤗 Accelerate from pypi, perform:
+
+```bash
+pip install accelerate
+```
+
+### conda
+
+🤗 Accelerate can also be installed with conda with:
+
+```bash
+conda install -c conda-forge accelerate
+```
+
+### Source
+
+As 🤗 Accelerate is rapidly under development, new features are added every day that haven't been released yet. To try them out yourself install
```suggestion
from the GitHub repository:
```
@@ -0,0 +1,91 @@
+<!--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.
+-->
+
+# Installation and Configuration
+
+Before you start, you will need to setup your environment, install the appropriate packages, and configure 🤗 Accelerate. 🤗 Accelerate is tested on **Python 3.7+**.
+
+## Installing 🤗 Accelerate
+
+🤗 Accelerate is available on pypi and conda, as well as on GitHub. Details to install from each are below:
+
+### pip 
+
+To install 🤗 Accelerate from pypi, perform:
+
+```bash
+pip install accelerate
+```
+
+### conda
+
+🤗 Accelerate can also be installed with conda with:
+
+```bash
+conda install -c conda-forge accelerate
+```
+
+### Source
+
+As 🤗 Accelerate is rapidly under development, new features are added every day that haven't been released yet. To try them out yourself install
+from the Github Repository:
```suggestion
After installing, you need to configure 🤗 Accelerate for how the current system is setup for training. 
```
@@ -0,0 +1,91 @@
+<!--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.
+-->
+
+# Installation and Configuration
+
+Before you start, you will need to setup your environment, install the appropriate packages, and configure 🤗 Accelerate. 🤗 Accelerate is tested on **Python 3.7+**.
+
+## Installing 🤗 Accelerate
+
+🤗 Accelerate is available on pypi and conda, as well as on GitHub. Details to install from each are below:
+
+### pip 
+
+To install 🤗 Accelerate from pypi, perform:
+
+```bash
+pip install accelerate
+```
+
+### conda
+
+🤗 Accelerate can also be installed with conda with:
+
+```bash
+conda install -c conda-forge accelerate
+```
+
+### Source
+
+As 🤗 Accelerate is rapidly under development, new features are added every day that haven't been released yet. To try them out yourself install
+from the Github Repository:
+
+```bash
+git clone https://github.com/huggingface/accelerate
+cd accelerate
+pip install -e .
+```
+
+## Configuring 🤗 Accelerate
+
+After installing, you need to tell Accelerate how the current system is setup for training. 
```suggestion
This tutorial will detail how to easily convert existing PyTorch code to use 🤗 Accelerate!
```
@@ -0,0 +1,119 @@
+<!--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.
+-->
+
+# Migrating your code to 🤗 Accelerate
+
+This tutorial will detail how to easily convert existing PyTorch code into utilizing 🤗 Accelerate!
```suggestion
You'll see that by just changing a few lines of code, 🤗 Accelerate can perform its magic and get you on 
```
@@ -0,0 +1,119 @@
+<!--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.
+-->
+
+# Migrating your code to 🤗 Accelerate
+
+This tutorial will detail how to easily convert existing PyTorch code into utilizing 🤗 Accelerate!
+You'll be able to see that by just changing a few lines of code, 🤗 Accelerate can perform its magic and get you on 
```suggestion
To begin, write out a very basic PyTorch training loop. 
```
@@ -0,0 +1,119 @@
+<!--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.
+-->
+
+# Migrating your code to 🤗 Accelerate
+
+This tutorial will detail how to easily convert existing PyTorch code into utilizing 🤗 Accelerate!
+You'll be able to see that by just changing a few lines of code, 🤗 Accelerate can perform its magic and get you on 
+your way towards running your code on distributed systems with ease!
+
+## The base training loop
+
+To begin, first write out a very basic PyTorch training loop. 
```suggestion
We are under the presumption that `training_dataloader`, `model`, `optimizer`, `scheduler`, and `loss_function` have been defined beforehand.
```
@@ -0,0 +1,119 @@
+<!--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.
+-->
+
+# Migrating your code to 🤗 Accelerate
+
+This tutorial will detail how to easily convert existing PyTorch code into utilizing 🤗 Accelerate!
+You'll be able to see that by just changing a few lines of code, 🤗 Accelerate can perform its magic and get you on 
+your way towards running your code on distributed systems with ease!
+
+## The base training loop
+
+To begin, first write out a very basic PyTorch training loop. 
+
+<Tip>
+We are under the presumption that `training_dataloader`, `model`, `optimizer`, `scheduler`, and `loss_function` have been defined beforehand
```suggestion
[`Accelerator`] is the main force behind utilizing all the possible options for distributed training!
```
@@ -0,0 +1,119 @@
+<!--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.
+-->
+
+# Migrating your code to 🤗 Accelerate
+
+This tutorial will detail how to easily convert existing PyTorch code into utilizing 🤗 Accelerate!
+You'll be able to see that by just changing a few lines of code, 🤗 Accelerate can perform its magic and get you on 
+your way towards running your code on distributed systems with ease!
+
+## The base training loop
+
+To begin, first write out a very basic PyTorch training loop. 
+
+<Tip>
+We are under the presumption that `training_dataloader`, `model`, `optimizer`, `scheduler`, and `loss_function` have been defined beforehand
+</Tip>
+
+```python
+device = "cuda"
+model.to(device)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    inputs = inputs.to(device)
+    targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    loss.backward()
+    optimizer.step()
+    scheduler.step()
+```
+
+## Add in 🤗 Accelerate
+
+To start using 🤗 Accelerate, first import and create an [`Accelerator`] instance:
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+```
+This class will be the main force helping you with utilizing everything possible for distributed training!
```suggestion
make sure everything is setup in the current environment for you to start training:
```
@@ -0,0 +1,119 @@
+<!--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.
+-->
+
+# Migrating your code to 🤗 Accelerate
+
+This tutorial will detail how to easily convert existing PyTorch code into utilizing 🤗 Accelerate!
+You'll be able to see that by just changing a few lines of code, 🤗 Accelerate can perform its magic and get you on 
+your way towards running your code on distributed systems with ease!
+
+## The base training loop
+
+To begin, first write out a very basic PyTorch training loop. 
+
+<Tip>
+We are under the presumption that `training_dataloader`, `model`, `optimizer`, `scheduler`, and `loss_function` have been defined beforehand
+</Tip>
+
+```python
+device = "cuda"
+model.to(device)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    inputs = inputs.to(device)
+    targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    loss.backward()
+    optimizer.step()
+    scheduler.step()
+```
+
+## Add in 🤗 Accelerate
+
+To start using 🤗 Accelerate, first import and create an [`Accelerator`] instance:
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+```
+This class will be the main force helping you with utilizing everything possible for distributed training!
+
+### Setting the right device
+
+The [`Accelerator`] class knows the right device to move any PyTorch object to at any time, so you should
+change the definition of `device` to come from [`Accelerator`]:
+
+```diff
+-device = 'cuda'
++device = accelerator.device
+model.to(device)
+```
+
+### Preparing your objects
+
+Next you need to pass all of the important objects related to training into [`Accelerator.prepare`]. 🤗 Accelerate will
+then make sure that everything is setup for the environment you are currently in to start your training:
```suggestion
These objects are returned in the same order they were sent in with. By default, `device_placement=True`, so your model is automatically moved to the proper device.
```
@@ -0,0 +1,119 @@
+<!--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.
+-->
+
+# Migrating your code to 🤗 Accelerate
+
+This tutorial will detail how to easily convert existing PyTorch code into utilizing 🤗 Accelerate!
+You'll be able to see that by just changing a few lines of code, 🤗 Accelerate can perform its magic and get you on 
+your way towards running your code on distributed systems with ease!
+
+## The base training loop
+
+To begin, first write out a very basic PyTorch training loop. 
+
+<Tip>
+We are under the presumption that `training_dataloader`, `model`, `optimizer`, `scheduler`, and `loss_function` have been defined beforehand
+</Tip>
+
+```python
+device = "cuda"
+model.to(device)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    inputs = inputs.to(device)
+    targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    loss.backward()
+    optimizer.step()
+    scheduler.step()
+```
+
+## Add in 🤗 Accelerate
+
+To start using 🤗 Accelerate, first import and create an [`Accelerator`] instance:
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+```
+This class will be the main force helping you with utilizing everything possible for distributed training!
+
+### Setting the right device
+
+The [`Accelerator`] class knows the right device to move any PyTorch object to at any time, so you should
+change the definition of `device` to come from [`Accelerator`]:
+
+```diff
+-device = 'cuda'
++device = accelerator.device
+model.to(device)
+```
+
+### Preparing your objects
+
+Next you need to pass all of the important objects related to training into [`Accelerator.prepare`]. 🤗 Accelerate will
+then make sure that everything is setup for the environment you are currently in to start your training:
+
+```
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+```
+These objects are returned in the same order they were sent in with. If `device_placement=True` is passed into the [`Accelerator`] object (the default), 
I think you also need to add a newline before and after the `<Tip>` here so the code format is rendered.

```suggestion
Accelerate will only prepare objects that inherit from their respective PyTorch classes (such as `torch.optim.Optimizer`).
```
@@ -0,0 +1,119 @@
+<!--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.
+-->
+
+# Migrating your code to 🤗 Accelerate
+
+This tutorial will detail how to easily convert existing PyTorch code into utilizing 🤗 Accelerate!
+You'll be able to see that by just changing a few lines of code, 🤗 Accelerate can perform its magic and get you on 
+your way towards running your code on distributed systems with ease!
+
+## The base training loop
+
+To begin, first write out a very basic PyTorch training loop. 
+
+<Tip>
+We are under the presumption that `training_dataloader`, `model`, `optimizer`, `scheduler`, and `loss_function` have been defined beforehand
+</Tip>
+
+```python
+device = "cuda"
+model.to(device)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    inputs = inputs.to(device)
+    targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    loss.backward()
+    optimizer.step()
+    scheduler.step()
+```
+
+## Add in 🤗 Accelerate
+
+To start using 🤗 Accelerate, first import and create an [`Accelerator`] instance:
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+```
+This class will be the main force helping you with utilizing everything possible for distributed training!
+
+### Setting the right device
+
+The [`Accelerator`] class knows the right device to move any PyTorch object to at any time, so you should
+change the definition of `device` to come from [`Accelerator`]:
+
+```diff
+-device = 'cuda'
++device = accelerator.device
+model.to(device)
+```
+
+### Preparing your objects
+
+Next you need to pass all of the important objects related to training into [`Accelerator.prepare`]. 🤗 Accelerate will
+then make sure that everything is setup for the environment you are currently in to start your training:
+
+```
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+```
+These objects are returned in the same order they were sent in with. If `device_placement=True` is passed into the [`Accelerator`] object (the default), 
+performing this action will also automatically move your model to the proper device as well.
+
+<Tip warning={true}>
+Accelerate will only prepare objects that inherit from their respective PyTorch classes (such as `torch.optim.Optimizer`)
```suggestion
Finally, three lines of code need to be changed in the training loop. 🤗 Accelerate's DataLoader classes will automatically handle the device placement by default,
```
@@ -0,0 +1,119 @@
+<!--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.
+-->
+
+# Migrating your code to 🤗 Accelerate
+
+This tutorial will detail how to easily convert existing PyTorch code into utilizing 🤗 Accelerate!
+You'll be able to see that by just changing a few lines of code, 🤗 Accelerate can perform its magic and get you on 
+your way towards running your code on distributed systems with ease!
+
+## The base training loop
+
+To begin, first write out a very basic PyTorch training loop. 
+
+<Tip>
+We are under the presumption that `training_dataloader`, `model`, `optimizer`, `scheduler`, and `loss_function` have been defined beforehand
+</Tip>
+
+```python
+device = "cuda"
+model.to(device)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    inputs = inputs.to(device)
+    targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    loss.backward()
+    optimizer.step()
+    scheduler.step()
+```
+
+## Add in 🤗 Accelerate
+
+To start using 🤗 Accelerate, first import and create an [`Accelerator`] instance:
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+```
+This class will be the main force helping you with utilizing everything possible for distributed training!
+
+### Setting the right device
+
+The [`Accelerator`] class knows the right device to move any PyTorch object to at any time, so you should
+change the definition of `device` to come from [`Accelerator`]:
+
+```diff
+-device = 'cuda'
++device = accelerator.device
+model.to(device)
+```
+
+### Preparing your objects
+
+Next you need to pass all of the important objects related to training into [`Accelerator.prepare`]. 🤗 Accelerate will
+then make sure that everything is setup for the environment you are currently in to start your training:
+
+```
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+```
+These objects are returned in the same order they were sent in with. If `device_placement=True` is passed into the [`Accelerator`] object (the default), 
+performing this action will also automatically move your model to the proper device as well.
+
+<Tip warning={true}>
+Accelerate will only prepare objects that inherit from their respective PyTorch classes (such as `torch.optim.Optimizer`)
+</Tip>
+
+### Modifying the training loop
+
+Finally, three lines of code need changing in our training loop. 🤗 Accelerate's DataLoader classes will automatically handle the device placement by default,
Love that you added a complete version of all the snippets so users have the option of just copying the code once!
@@ -0,0 +1,119 @@
+<!--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.
+-->
+
+# Migrating your code to 🤗 Accelerate
+
+This tutorial will detail how to easily convert existing PyTorch code into utilizing 🤗 Accelerate!
+You'll be able to see that by just changing a few lines of code, 🤗 Accelerate can perform its magic and get you on 
+your way towards running your code on distributed systems with ease!
+
+## The base training loop
+
+To begin, first write out a very basic PyTorch training loop. 
+
+<Tip>
+We are under the presumption that `training_dataloader`, `model`, `optimizer`, `scheduler`, and `loss_function` have been defined beforehand
+</Tip>
+
+```python
+device = "cuda"
+model.to(device)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    inputs = inputs.to(device)
+    targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    loss.backward()
+    optimizer.step()
+    scheduler.step()
+```
+
+## Add in 🤗 Accelerate
+
+To start using 🤗 Accelerate, first import and create an [`Accelerator`] instance:
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+```
+This class will be the main force helping you with utilizing everything possible for distributed training!
+
+### Setting the right device
+
+The [`Accelerator`] class knows the right device to move any PyTorch object to at any time, so you should
+change the definition of `device` to come from [`Accelerator`]:
+
+```diff
+-device = 'cuda'
++device = accelerator.device
+model.to(device)
+```
+
+### Preparing your objects
+
+Next you need to pass all of the important objects related to training into [`Accelerator.prepare`]. 🤗 Accelerate will
+then make sure that everything is setup for the environment you are currently in to start your training:
+
+```
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+```
+These objects are returned in the same order they were sent in with. If `device_placement=True` is passed into the [`Accelerator`] object (the default), 
+performing this action will also automatically move your model to the proper device as well.
+
+<Tip warning={true}>
+Accelerate will only prepare objects that inherit from their respective PyTorch classes (such as `torch.optim.Optimizer`)
+</Tip>
+
+### Modifying the training loop
+
+Finally, three lines of code need changing in our training loop. 🤗 Accelerate's DataLoader classes will automatically handle the device placement by default,
+and [`Accelerator.backward`] should be used for performing the backward pass:
+
+```diff
+-   inputs = inputs.to(device)
+-   targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+-   loss.backward()
++   accelerator.backward(loss)
+```
+
+With that, your training loop is now ready to use 🤗 Accelerate!
+
+## The finished code
```suggestion
In the previous tutorial, you were introduced to how to modify your current training script to use 🤗 Accelerate.
```
@@ -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.
+-->
+
+# Launching your 🤗 Accelerate scripts
+
+In the previous tutorial you were introduced with how to modify your current training script to use 🤗 Accelerate.
The right navbar doesn't display code format so the ticks will just show up :/

```suggestion
## Using accelerate launch
```
@@ -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.
+-->
+
+# Launching your 🤗 Accelerate scripts
+
+In the previous tutorial you were introduced with how to modify your current training script to use 🤗 Accelerate.
+The final version of that code is shown below:
+
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    accelerator.backward(loss)
+    optimizer.step()
+    scheduler.step()
+```
+
+But how do you run this code and have it utilize the special hardware available to it?
+
+First you should rewrite the above code into a function, and make it callable as a script. For example:
+
+```diff
+  from accelerate import Accelerator
+  
++ def main():
+
+    accelerator = Accelerator()
+
+    model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+        model, optimizer, training_dataloader, scheduler
+    )
+
+    for batch in training_dataloader:
+        optimizer.zero_grad()
+        inputs, targets = batch
+        outputs = model(inputs)
+        loss = loss_function(outputs, targets)
+        accelerator.backward(loss)
+        optimizer.step()
+        scheduler.step()
+
++ if __name__ == "__main__":
++   main()
+```
+
+Next you need to launch it with `accelerate launch`. 
+
+## Using `accelerate launch`
Need some newlines here too :)
@@ -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.
+-->
+
+# Launching your 🤗 Accelerate scripts
+
+In the previous tutorial you were introduced with how to modify your current training script to use 🤗 Accelerate.
+The final version of that code is shown below:
+
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    accelerator.backward(loss)
+    optimizer.step()
+    scheduler.step()
+```
+
+But how do you run this code and have it utilize the special hardware available to it?
+
+First you should rewrite the above code into a function, and make it callable as a script. For example:
+
+```diff
+  from accelerate import Accelerator
+  
++ def main():
+
+    accelerator = Accelerator()
+
+    model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+        model, optimizer, training_dataloader, scheduler
+    )
+
+    for batch in training_dataloader:
+        optimizer.zero_grad()
+        inputs, targets = batch
+        outputs = model(inputs)
+        loss = loss_function(outputs, targets)
+        accelerator.backward(loss)
+        optimizer.step()
+        scheduler.step()
+
++ if __name__ == "__main__":
++   main()
+```
+
+Next you need to launch it with `accelerate launch`. 
+
+## Using `accelerate launch`
+
+🤗 Accelerate has a special CLI command to help you launch your code in your system through `accelerate launch`.
+This command wraps around all of the different commands needed to launch your script on various platforms, without you
+having to remember what each of them are.
+
+<Tip>
```suggestion
Since this runs the various torch spawn methods, all of the expected environment variables can be modified here as well.
```
@@ -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.
+-->
+
+# Launching your 🤗 Accelerate scripts
+
+In the previous tutorial you were introduced with how to modify your current training script to use 🤗 Accelerate.
+The final version of that code is shown below:
+
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    accelerator.backward(loss)
+    optimizer.step()
+    scheduler.step()
+```
+
+But how do you run this code and have it utilize the special hardware available to it?
+
+First you should rewrite the above code into a function, and make it callable as a script. For example:
+
+```diff
+  from accelerate import Accelerator
+  
++ def main():
+
+    accelerator = Accelerator()
+
+    model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+        model, optimizer, training_dataloader, scheduler
+    )
+
+    for batch in training_dataloader:
+        optimizer.zero_grad()
+        inputs, targets = batch
+        outputs = model(inputs)
+        loss = loss_function(outputs, targets)
+        accelerator.backward(loss)
+        optimizer.step()
+        scheduler.step()
+
++ if __name__ == "__main__":
++   main()
+```
+
+Next you need to launch it with `accelerate launch`. 
+
+## Using `accelerate launch`
+
+🤗 Accelerate has a special CLI command to help you launch your code in your system through `accelerate launch`.
+This command wraps around all of the different commands needed to launch your script on various platforms, without you
+having to remember what each of them are.
+
+<Tip>
+If you are familiar with launching scripts in PyTorch yourself such as with `torchrun`, you can still do this. It is not required to use `accelerate launch`
+</Tip>
+
+You can launch your script quickly by using:
+```bash
+accelerate launch {script_name.py} --arg1 --arg2 ...
+```
+
+Just put `accelerate launch` at the start of your command, and pass in additional arguments and parameters to your script afterwards like normal!
+
+Since this runs the various different torch spawn methods, all of the environmental variables that would be expected for them can be modified here as well.
```suggestion
For example, here is how to use `accelerate launch` with a single GPU:
```
@@ -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.
+-->
+
+# Launching your 🤗 Accelerate scripts
+
+In the previous tutorial you were introduced with how to modify your current training script to use 🤗 Accelerate.
+The final version of that code is shown below:
+
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    accelerator.backward(loss)
+    optimizer.step()
+    scheduler.step()
+```
+
+But how do you run this code and have it utilize the special hardware available to it?
+
+First you should rewrite the above code into a function, and make it callable as a script. For example:
+
+```diff
+  from accelerate import Accelerator
+  
++ def main():
+
+    accelerator = Accelerator()
+
+    model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+        model, optimizer, training_dataloader, scheduler
+    )
+
+    for batch in training_dataloader:
+        optimizer.zero_grad()
+        inputs, targets = batch
+        outputs = model(inputs)
+        loss = loss_function(outputs, targets)
+        accelerator.backward(loss)
+        optimizer.step()
+        scheduler.step()
+
++ if __name__ == "__main__":
++   main()
+```
+
+Next you need to launch it with `accelerate launch`. 
+
+## Using `accelerate launch`
+
+🤗 Accelerate has a special CLI command to help you launch your code in your system through `accelerate launch`.
+This command wraps around all of the different commands needed to launch your script on various platforms, without you
+having to remember what each of them are.
+
+<Tip>
+If you are familiar with launching scripts in PyTorch yourself such as with `torchrun`, you can still do this. It is not required to use `accelerate launch`
+</Tip>
+
+You can launch your script quickly by using:
+```bash
+accelerate launch {script_name.py} --arg1 --arg2 ...
+```
+
+Just put `accelerate launch` at the start of your command, and pass in additional arguments and parameters to your script afterwards like normal!
+
+Since this runs the various different torch spawn methods, all of the environmental variables that would be expected for them can be modified here as well.
+For example, here is how to use `accelerate launch` but only use a single GPU:
```suggestion
In this case, 🤗 Accelerate will make some hyperparameter decisions for you, such as if a GPU is available, it will use all of them by default, and no mixed precision will be used. 
```
@@ -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.
+-->
+
+# Launching your 🤗 Accelerate scripts
+
+In the previous tutorial you were introduced with how to modify your current training script to use 🤗 Accelerate.
+The final version of that code is shown below:
+
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    accelerator.backward(loss)
+    optimizer.step()
+    scheduler.step()
+```
+
+But how do you run this code and have it utilize the special hardware available to it?
+
+First you should rewrite the above code into a function, and make it callable as a script. For example:
+
+```diff
+  from accelerate import Accelerator
+  
++ def main():
+
+    accelerator = Accelerator()
+
+    model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+        model, optimizer, training_dataloader, scheduler
+    )
+
+    for batch in training_dataloader:
+        optimizer.zero_grad()
+        inputs, targets = batch
+        outputs = model(inputs)
+        loss = loss_function(outputs, targets)
+        accelerator.backward(loss)
+        optimizer.step()
+        scheduler.step()
+
++ if __name__ == "__main__":
++   main()
+```
+
+Next you need to launch it with `accelerate launch`. 
+
+## Using `accelerate launch`
+
+🤗 Accelerate has a special CLI command to help you launch your code in your system through `accelerate launch`.
+This command wraps around all of the different commands needed to launch your script on various platforms, without you
+having to remember what each of them are.
+
+<Tip>
+If you are familiar with launching scripts in PyTorch yourself such as with `torchrun`, you can still do this. It is not required to use `accelerate launch`
+</Tip>
+
+You can launch your script quickly by using:
+```bash
+accelerate launch {script_name.py} --arg1 --arg2 ...
+```
+
+Just put `accelerate launch` at the start of your command, and pass in additional arguments and parameters to your script afterwards like normal!
+
+Since this runs the various different torch spawn methods, all of the environmental variables that would be expected for them can be modified here as well.
+For example, here is how to use `accelerate launch` but only use a single GPU:
+```bash
+CUDA_VISIBLE_DEVICES="0" accelerate launch {script_name.py} --arg1 --arg2 ...
+```
+
+You can also use `accelerate launch` without performing `accelerate config` first, but you may need to manually pass in the right configuration parameters.
+Accelerate can take some liberties towards a few hyperparameters, such as if a GPU is available it will use all of them by default, and no mixed precision will be used. 
I would maybe clarify that you'll get a warning if you don't pass in the needed parameters. This is kind of unclear until now.

```suggestion
would also launch that same script on two GPUs using mixed precision while avoiding all of the warnings: 
```
@@ -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.
+-->
+
+# Launching your 🤗 Accelerate scripts
+
+In the previous tutorial you were introduced with how to modify your current training script to use 🤗 Accelerate.
+The final version of that code is shown below:
+
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    accelerator.backward(loss)
+    optimizer.step()
+    scheduler.step()
+```
+
+But how do you run this code and have it utilize the special hardware available to it?
+
+First you should rewrite the above code into a function, and make it callable as a script. For example:
+
+```diff
+  from accelerate import Accelerator
+  
++ def main():
+
+    accelerator = Accelerator()
+
+    model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+        model, optimizer, training_dataloader, scheduler
+    )
+
+    for batch in training_dataloader:
+        optimizer.zero_grad()
+        inputs, targets = batch
+        outputs = model(inputs)
+        loss = loss_function(outputs, targets)
+        accelerator.backward(loss)
+        optimizer.step()
+        scheduler.step()
+
++ if __name__ == "__main__":
++   main()
+```
+
+Next you need to launch it with `accelerate launch`. 
+
+## Using `accelerate launch`
+
+🤗 Accelerate has a special CLI command to help you launch your code in your system through `accelerate launch`.
+This command wraps around all of the different commands needed to launch your script on various platforms, without you
+having to remember what each of them are.
+
+<Tip>
+If you are familiar with launching scripts in PyTorch yourself such as with `torchrun`, you can still do this. It is not required to use `accelerate launch`
+</Tip>
+
+You can launch your script quickly by using:
+```bash
+accelerate launch {script_name.py} --arg1 --arg2 ...
+```
+
+Just put `accelerate launch` at the start of your command, and pass in additional arguments and parameters to your script afterwards like normal!
+
+Since this runs the various different torch spawn methods, all of the environmental variables that would be expected for them can be modified here as well.
+For example, here is how to use `accelerate launch` but only use a single GPU:
+```bash
+CUDA_VISIBLE_DEVICES="0" accelerate launch {script_name.py} --arg1 --arg2 ...
+```
+
+You can also use `accelerate launch` without performing `accelerate config` first, but you may need to manually pass in the right configuration parameters.
+Accelerate can take some liberties towards a few hyperparameters, such as if a GPU is available it will use all of them by default, and no mixed precision will be used. 
+Here is how you would use all GPUs and train on no mixed precision:
+```bash
+accelerate launch --multi_gpu {script_name.py} --arg1 --arg2 ...
+```
+To get more specific you should pass in the needed parameters yourself. For instance, here is how you 
+would also launch that same script on two GPUs using mixed precision, but also avoiding all of the warnings: 
```suggestion
For a complete list of parameters you can pass in, run:
```
@@ -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.
+-->
+
+# Launching your 🤗 Accelerate scripts
+
+In the previous tutorial you were introduced with how to modify your current training script to use 🤗 Accelerate.
+The final version of that code is shown below:
+
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    accelerator.backward(loss)
+    optimizer.step()
+    scheduler.step()
+```
+
+But how do you run this code and have it utilize the special hardware available to it?
+
+First you should rewrite the above code into a function, and make it callable as a script. For example:
+
+```diff
+  from accelerate import Accelerator
+  
++ def main():
+
+    accelerator = Accelerator()
+
+    model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+        model, optimizer, training_dataloader, scheduler
+    )
+
+    for batch in training_dataloader:
+        optimizer.zero_grad()
+        inputs, targets = batch
+        outputs = model(inputs)
+        loss = loss_function(outputs, targets)
+        accelerator.backward(loss)
+        optimizer.step()
+        scheduler.step()
+
++ if __name__ == "__main__":
++   main()
+```
+
+Next you need to launch it with `accelerate launch`. 
+
+## Using `accelerate launch`
+
+🤗 Accelerate has a special CLI command to help you launch your code in your system through `accelerate launch`.
+This command wraps around all of the different commands needed to launch your script on various platforms, without you
+having to remember what each of them are.
+
+<Tip>
+If you are familiar with launching scripts in PyTorch yourself such as with `torchrun`, you can still do this. It is not required to use `accelerate launch`
+</Tip>
+
+You can launch your script quickly by using:
+```bash
+accelerate launch {script_name.py} --arg1 --arg2 ...
+```
+
+Just put `accelerate launch` at the start of your command, and pass in additional arguments and parameters to your script afterwards like normal!
+
+Since this runs the various different torch spawn methods, all of the environmental variables that would be expected for them can be modified here as well.
+For example, here is how to use `accelerate launch` but only use a single GPU:
+```bash
+CUDA_VISIBLE_DEVICES="0" accelerate launch {script_name.py} --arg1 --arg2 ...
+```
+
+You can also use `accelerate launch` without performing `accelerate config` first, but you may need to manually pass in the right configuration parameters.
+Accelerate can take some liberties towards a few hyperparameters, such as if a GPU is available it will use all of them by default, and no mixed precision will be used. 
+Here is how you would use all GPUs and train on no mixed precision:
+```bash
+accelerate launch --multi_gpu {script_name.py} --arg1 --arg2 ...
+```
+To get more specific you should pass in the needed parameters yourself. For instance, here is how you 
+would also launch that same script on two GPUs using mixed precision, but also avoiding all of the warnings: 
+```bash
+accelerate launch --multi_gpu --mixed_precision=fp16 --num_processes=2 {script_name.py} --arg1 --arg2 ...
+```
+
+For a full list of all parameters you can pass in, run:
```suggestion
- If this does not exist either, the folder `~/.cache/huggingface/accelerate`.
```
@@ -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.
+-->
+
+# Launching your 🤗 Accelerate scripts
+
+In the previous tutorial you were introduced with how to modify your current training script to use 🤗 Accelerate.
+The final version of that code is shown below:
+
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    accelerator.backward(loss)
+    optimizer.step()
+    scheduler.step()
+```
+
+But how do you run this code and have it utilize the special hardware available to it?
+
+First you should rewrite the above code into a function, and make it callable as a script. For example:
+
+```diff
+  from accelerate import Accelerator
+  
++ def main():
+
+    accelerator = Accelerator()
+
+    model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+        model, optimizer, training_dataloader, scheduler
+    )
+
+    for batch in training_dataloader:
+        optimizer.zero_grad()
+        inputs, targets = batch
+        outputs = model(inputs)
+        loss = loss_function(outputs, targets)
+        accelerator.backward(loss)
+        optimizer.step()
+        scheduler.step()
+
++ if __name__ == "__main__":
++   main()
+```
+
+Next you need to launch it with `accelerate launch`. 
+
+## Using `accelerate launch`
+
+🤗 Accelerate has a special CLI command to help you launch your code in your system through `accelerate launch`.
+This command wraps around all of the different commands needed to launch your script on various platforms, without you
+having to remember what each of them are.
+
+<Tip>
+If you are familiar with launching scripts in PyTorch yourself such as with `torchrun`, you can still do this. It is not required to use `accelerate launch`
+</Tip>
+
+You can launch your script quickly by using:
+```bash
+accelerate launch {script_name.py} --arg1 --arg2 ...
+```
+
+Just put `accelerate launch` at the start of your command, and pass in additional arguments and parameters to your script afterwards like normal!
+
+Since this runs the various different torch spawn methods, all of the environmental variables that would be expected for them can be modified here as well.
+For example, here is how to use `accelerate launch` but only use a single GPU:
+```bash
+CUDA_VISIBLE_DEVICES="0" accelerate launch {script_name.py} --arg1 --arg2 ...
+```
+
+You can also use `accelerate launch` without performing `accelerate config` first, but you may need to manually pass in the right configuration parameters.
+Accelerate can take some liberties towards a few hyperparameters, such as if a GPU is available it will use all of them by default, and no mixed precision will be used. 
+Here is how you would use all GPUs and train on no mixed precision:
+```bash
+accelerate launch --multi_gpu {script_name.py} --arg1 --arg2 ...
+```
+To get more specific you should pass in the needed parameters yourself. For instance, here is how you 
+would also launch that same script on two GPUs using mixed precision, but also avoiding all of the warnings: 
+```bash
+accelerate launch --multi_gpu --mixed_precision=fp16 --num_processes=2 {script_name.py} --arg1 --arg2 ...
+```
+
+For a full list of all parameters you can pass in, run:
+```bash
+accelerate launch -h
+```
+
+<Tip>
+Even if you are not using 🤗 Accelerate in your code, you can still use the launcher for starting your scripts!
+</Tip>
+
+## Custom Configurations
+
+As briefly mentioned earlier, `accelerate launch` should be mostly used through combining set configurations 
+made with the `accelerate config` command. These configs are saved to a `default_config.yaml` file in your cache folder for 🤗 Accelerate. 
+This cache folder is located at (with decreasing order of priority):
+
+- The content of your environment variable `HF_HOME` suffixed with `accelerate`.
+- If it does not exist, the content of your environment variable `XDG_CACHE_HOME` suffixed with
+  `huggingface/accelerate`.
+- If this does not exist either, the folder `~/.cache/huggingface/accelerate`
I would also move the snippet at the end of this page under here since that's what it's showing here right?
@@ -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.
+-->
+
+# Launching your 🤗 Accelerate scripts
+
+In the previous tutorial you were introduced with how to modify your current training script to use 🤗 Accelerate.
+The final version of that code is shown below:
+
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    accelerator.backward(loss)
+    optimizer.step()
+    scheduler.step()
+```
+
+But how do you run this code and have it utilize the special hardware available to it?
+
+First you should rewrite the above code into a function, and make it callable as a script. For example:
+
+```diff
+  from accelerate import Accelerator
+  
++ def main():
+
+    accelerator = Accelerator()
+
+    model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+        model, optimizer, training_dataloader, scheduler
+    )
+
+    for batch in training_dataloader:
+        optimizer.zero_grad()
+        inputs, targets = batch
+        outputs = model(inputs)
+        loss = loss_function(outputs, targets)
+        accelerator.backward(loss)
+        optimizer.step()
+        scheduler.step()
+
++ if __name__ == "__main__":
++   main()
+```
+
+Next you need to launch it with `accelerate launch`. 
+
+## Using `accelerate launch`
+
+🤗 Accelerate has a special CLI command to help you launch your code in your system through `accelerate launch`.
+This command wraps around all of the different commands needed to launch your script on various platforms, without you
+having to remember what each of them are.
+
+<Tip>
+If you are familiar with launching scripts in PyTorch yourself such as with `torchrun`, you can still do this. It is not required to use `accelerate launch`
+</Tip>
+
+You can launch your script quickly by using:
+```bash
+accelerate launch {script_name.py} --arg1 --arg2 ...
+```
+
+Just put `accelerate launch` at the start of your command, and pass in additional arguments and parameters to your script afterwards like normal!
+
+Since this runs the various different torch spawn methods, all of the environmental variables that would be expected for them can be modified here as well.
+For example, here is how to use `accelerate launch` but only use a single GPU:
+```bash
+CUDA_VISIBLE_DEVICES="0" accelerate launch {script_name.py} --arg1 --arg2 ...
+```
+
+You can also use `accelerate launch` without performing `accelerate config` first, but you may need to manually pass in the right configuration parameters.
+Accelerate can take some liberties towards a few hyperparameters, such as if a GPU is available it will use all of them by default, and no mixed precision will be used. 
+Here is how you would use all GPUs and train on no mixed precision:
+```bash
+accelerate launch --multi_gpu {script_name.py} --arg1 --arg2 ...
+```
+To get more specific you should pass in the needed parameters yourself. For instance, here is how you 
+would also launch that same script on two GPUs using mixed precision, but also avoiding all of the warnings: 
+```bash
+accelerate launch --multi_gpu --mixed_precision=fp16 --num_processes=2 {script_name.py} --arg1 --arg2 ...
+```
+
+For a full list of all parameters you can pass in, run:
+```bash
+accelerate launch -h
+```
+
+<Tip>
+Even if you are not using 🤗 Accelerate in your code, you can still use the launcher for starting your scripts!
+</Tip>
+
+## Custom Configurations
+
+As briefly mentioned earlier, `accelerate launch` should be mostly used through combining set configurations 
+made with the `accelerate config` command. These configs are saved to a `default_config.yaml` file in your cache folder for 🤗 Accelerate. 
+This cache folder is located at (with decreasing order of priority):
+
+- The content of your environment variable `HF_HOME` suffixed with `accelerate`.
+- If it does not exist, the content of your environment variable `XDG_CACHE_HOME` suffixed with
+  `huggingface/accelerate`.
+- If this does not exist either, the folder `~/.cache/huggingface/accelerate`
+
+To have multiple configurations, the flag `--config_file` can be passed to the `accelerate launch` command paired 
I think you can create a separate Troubleshoot page for all these caveats. I think a user just getting started doesn't necessarily need to know all these smaller details. Likewise, I think the Internal mechanism section is probably better in the Concepts and Fundamentals section since the user doesn't need to know exactly how it works under the hood. 
@@ -271,6 +276,8 @@ passed your model to [`~Accelerator.prepare`]) will break the tying. You will ne
 after. You can find an example of this in the [run_clm_no_trainer](https://github.com/huggingface/transformers/blob/master/examples/pytorch/language-modeling/run_clm.py) script in
 the Transformers repository.
 
+Check out the [TPU tutorial](https://huggingface.co/docs/accelerate/concept_guides/training_tpu) for more information about training on TPUs. 
+
 
 ## Other caveats
Adding a link to it would make it easier for users to navigate there.
@@ -49,29 +49,29 @@ for index, batch in enumerate(training_dataloader):
 First the code shown earlier will be converted to utilize 🤗 Accelerate without the special gradient accumulation helper:
 
 ```diff
-+ from accelerate import Accelerator
-+ accelerator = Accelerator()
-
-+ model, optimizer, training_dataloader, scheduler = accelerator.prepare(
-+     model, optimizer, training_dataloader, scheduler
-+ )
-
-  for index, batch in enumerate(training_dataloader):
-      optimizer.zero_grad()
-      inputs, targets = batch
--     inputs = inputs.to(device)
--     targets = targets.to(device)
-      outputs = model(inputs)
-      loss = loss_function(outputs, targets)
-      loss = loss / gradient_accumulation_steps
-+     accelerator.backward(loss)
-      if (index+1) % gradient_accumulation_steps == 0:
-          optimizer.step()
-          scheduler.step()
++from accelerate import Accelerator
++accelerator = Accelerator()
+
++model, optimizer, training_dataloader, scheduler = accelerator.prepare(
++    model, optimizer, training_dataloader, scheduler
++)
+
+for index, batch in enumerate(training_dataloader):
+    optimizer.zero_grad()
+    inputs, targets = batch
+-    inputs = inputs.to(device)
+-    targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    loss = loss / gradient_accumulation_steps
++    accelerator.backward(loss)
+    if (index+1) % gradient_accumulation_steps == 0:
+        optimizer.step()
+        scheduler.step()
 ```
 
 <Tip warning={true}>
-In its current state, this code is not going to perform gradient accumulation efficiently due to a process called gradient synchronization.
+In its current state, this code is not going to perform gradient accumulation efficiently due to a process called gradient synchronization. Read more about that in the Concepts tutorial!
```suggestion
You can remove all the special checks for the step number and the loss adjustment:
```
@@ -80,35 +80,39 @@ All that is left now is to let 🤗 Accelerate handle the gradient accumulation
 of steps to perform before each call to `step()` and how to automatically adjust the loss during the call to [`Accelerator.backward`]:
 
 ```diff
-  from accelerate import Accelerator
-- accelerator = Accelerator()
-+ accelerator = Accelerator(gradient_accumulation_steps=2)
+from accelerate import Accelerator
+-accelerator = Accelerator()
++accelerator = Accelerator(gradient_accumulation_steps=2)
 ```
 
 From here you can use the [`Accelerator.accumulate`] context manager from inside your training loop to automatically perform the gradient accumulation for you!
-You just wrap it around the entire training part of your code: 
+You just wrap it around the entire training part of our code: 
 
 ```diff
-- for index, batch in enumerate(training_dataloader):
-+ for batch in training_dataloader:
-+    with accelerator.accumulate(model):
-          optimizer.zero_grad()
-          inputs, targets = batch
-          outputs = model(inputs)
+-for index, batch in enumerate(training_dataloader):
++for batch in training_dataloader:
++   with accelerator.accumulate(model):
+        optimizer.zero_grad()
+        inputs, targets = batch
+        outputs = model(inputs)
 ```
 
 and you can remove all the special checks for the step number and the loss adjustment:
Add some newlines so `DataLoaders` is properly rendered, and I'd also add a link to the Gradient Sync concept guide.
@@ -80,35 +80,39 @@ All that is left now is to let 🤗 Accelerate handle the gradient accumulation
 of steps to perform before each call to `step()` and how to automatically adjust the loss during the call to [`Accelerator.backward`]:
 
 ```diff
-  from accelerate import Accelerator
-- accelerator = Accelerator()
-+ accelerator = Accelerator(gradient_accumulation_steps=2)
+from accelerate import Accelerator
+-accelerator = Accelerator()
++accelerator = Accelerator(gradient_accumulation_steps=2)
 ```
 
 From here you can use the [`Accelerator.accumulate`] context manager from inside your training loop to automatically perform the gradient accumulation for you!
-You just wrap it around the entire training part of your code: 
+You just wrap it around the entire training part of our code: 
 
 ```diff
-- for index, batch in enumerate(training_dataloader):
-+ for batch in training_dataloader:
-+    with accelerator.accumulate(model):
-          optimizer.zero_grad()
-          inputs, targets = batch
-          outputs = model(inputs)
+-for index, batch in enumerate(training_dataloader):
++for batch in training_dataloader:
++   with accelerator.accumulate(model):
+        optimizer.zero_grad()
+        inputs, targets = batch
+        outputs = model(inputs)
 ```
 
 and you can remove all the special checks for the step number and the loss adjustment:
 
 ```diff
-- loss = loss / gradient_accumulation_steps
-  accelerator.backward(loss)
-- if (index+1) % gradient_accumulation_steps == 0:
-  optimizer.step()
-  scheduler.step()
+-loss = loss / gradient_accumulation_steps
+accelerator.backward(loss)
+-if (index+1) % gradient_accumulation_steps == 0:
+optimizer.step()
+scheduler.step()
 ```
 
 As you can see the [`Accelerator`] is able to keep track of the batch number you are on and it will automatically know whether to step through the prepared optimizer and how to adjust the loss. 
 
+<Tip warning={true}>
I think we can probably remove this header.
@@ -0,0 +1,111 @@
+<!--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.
+-->
+
+# Gradient Synchronization
+
+## Background into PyTorch distributed
```suggestion
training in a distributed setup. But how does this risk slowing down your code?
```
@@ -0,0 +1,111 @@
+<!--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.
+-->
+
+# Gradient Synchronization
+
+## Background into PyTorch distributed
+
+PyTorch's distributed module operates by communicating back and forth between all of the GPUs in your system.
+This communication takes time, and ensuring all processes know the states of each other happens at particular triggerpoints
+when using the `ddp` module. 
+
+These triggerpoints are added to the PyTorch model, specifically their `forward()` and `backward()` methods. 
+This happens when the model is wrapped with `DistributedDataParallel`:
+```python
+import torch.nn as nn
+from torch.nn.parallel import DistributedDataParallel
+
+model = nn.Linear(10, 10)
+ddp_model = DistributedDataParallel(model)
+```
+In 🤗 Accelerate this conversion happens automatically when calling [`Accelerator.prepare`] and passing in your model.
+
+```diff
++from accelerate import Accelerator
++accelerator = Accelerator()
+import torch.nn as nn
+-from torch.nn.parallel import DistributedDataParallel
+
+model = nn.Linear(10,10)
++model = accelerator.prepare(model)
+```
+
+## The slowdown in gradient accumulation
+
+You now understand that PyTorch adds hooks to the `forward` and `backward` method of your PyTorch model when 
+training in a distributed setup. But how does this automatically cause a risk of slowing down your code?
```suggestion
The most direct example is when you update all of the parameters in a model through `.backward()`. All instances of the model
```
@@ -0,0 +1,111 @@
+<!--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.
+-->
+
+# Gradient Synchronization
+
+## Background into PyTorch distributed
+
+PyTorch's distributed module operates by communicating back and forth between all of the GPUs in your system.
+This communication takes time, and ensuring all processes know the states of each other happens at particular triggerpoints
+when using the `ddp` module. 
+
+These triggerpoints are added to the PyTorch model, specifically their `forward()` and `backward()` methods. 
+This happens when the model is wrapped with `DistributedDataParallel`:
+```python
+import torch.nn as nn
+from torch.nn.parallel import DistributedDataParallel
+
+model = nn.Linear(10, 10)
+ddp_model = DistributedDataParallel(model)
+```
+In 🤗 Accelerate this conversion happens automatically when calling [`Accelerator.prepare`] and passing in your model.
+
+```diff
++from accelerate import Accelerator
++accelerator = Accelerator()
+import torch.nn as nn
+-from torch.nn.parallel import DistributedDataParallel
+
+model = nn.Linear(10,10)
++model = accelerator.prepare(model)
+```
+
+## The slowdown in gradient accumulation
+
+You now understand that PyTorch adds hooks to the `forward` and `backward` method of your PyTorch model when 
+training in a distributed setup. But how does this automatically cause a risk of slowing down your code?
+
+In DDP (distributed data parallel), the specific order in which processes are performed and ran are expected
+at specific points and these must also occur at roughly the same time before moving on.
+
+The most direct example is when you update all of the parameters in a model through `.backward()`, all instances of the model
This was a little hard to follow, is this what you meant?

```suggestion
need to have updated their gradients, collated, and updated again before moving onto the next batch of data. But when performing 
```
@@ -0,0 +1,111 @@
+<!--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.
+-->
+
+# Gradient Synchronization
+
+## Background into PyTorch distributed
+
+PyTorch's distributed module operates by communicating back and forth between all of the GPUs in your system.
+This communication takes time, and ensuring all processes know the states of each other happens at particular triggerpoints
+when using the `ddp` module. 
+
+These triggerpoints are added to the PyTorch model, specifically their `forward()` and `backward()` methods. 
+This happens when the model is wrapped with `DistributedDataParallel`:
+```python
+import torch.nn as nn
+from torch.nn.parallel import DistributedDataParallel
+
+model = nn.Linear(10, 10)
+ddp_model = DistributedDataParallel(model)
+```
+In 🤗 Accelerate this conversion happens automatically when calling [`Accelerator.prepare`] and passing in your model.
+
+```diff
++from accelerate import Accelerator
++accelerator = Accelerator()
+import torch.nn as nn
+-from torch.nn.parallel import DistributedDataParallel
+
+model = nn.Linear(10,10)
++model = accelerator.prepare(model)
+```
+
+## The slowdown in gradient accumulation
+
+You now understand that PyTorch adds hooks to the `forward` and `backward` method of your PyTorch model when 
+training in a distributed setup. But how does this automatically cause a risk of slowing down your code?
+
+In DDP (distributed data parallel), the specific order in which processes are performed and ran are expected
+at specific points and these must also occur at roughly the same time before moving on.
+
+The most direct example is when you update all of the parameters in a model through `.backward()`, all instances of the model
+need to have updated their gradients collated and updated before moving onto the next batch of data. But when performing 
```suggestion
gradient accumulation, you accumulate `n` losses and skip `.backward()` until `n` batches have been reached. This 
```
@@ -0,0 +1,111 @@
+<!--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.
+-->
+
+# Gradient Synchronization
+
+## Background into PyTorch distributed
+
+PyTorch's distributed module operates by communicating back and forth between all of the GPUs in your system.
+This communication takes time, and ensuring all processes know the states of each other happens at particular triggerpoints
+when using the `ddp` module. 
+
+These triggerpoints are added to the PyTorch model, specifically their `forward()` and `backward()` methods. 
+This happens when the model is wrapped with `DistributedDataParallel`:
+```python
+import torch.nn as nn
+from torch.nn.parallel import DistributedDataParallel
+
+model = nn.Linear(10, 10)
+ddp_model = DistributedDataParallel(model)
+```
+In 🤗 Accelerate this conversion happens automatically when calling [`Accelerator.prepare`] and passing in your model.
+
+```diff
++from accelerate import Accelerator
++accelerator = Accelerator()
+import torch.nn as nn
+-from torch.nn.parallel import DistributedDataParallel
+
+model = nn.Linear(10,10)
++model = accelerator.prepare(model)
+```
+
+## The slowdown in gradient accumulation
+
+You now understand that PyTorch adds hooks to the `forward` and `backward` method of your PyTorch model when 
+training in a distributed setup. But how does this automatically cause a risk of slowing down your code?
+
+In DDP (distributed data parallel), the specific order in which processes are performed and ran are expected
+at specific points and these must also occur at roughly the same time before moving on.
+
+The most direct example is when you update all of the parameters in a model through `.backward()`, all instances of the model
+need to have updated their gradients collated and updated before moving onto the next batch of data. But when performing 
+gradient accumulation, you accumulate `n` losses and skipping `.backward()` until `n` batches has been reached. This 
Is `them` referring to the batches of data?

```suggestion
can cause a significant slowdown since all the processes need to communicate with them more times than needed. How 
```
@@ -0,0 +1,111 @@
+<!--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.
+-->
+
+# Gradient Synchronization
+
+## Background into PyTorch distributed
+
+PyTorch's distributed module operates by communicating back and forth between all of the GPUs in your system.
+This communication takes time, and ensuring all processes know the states of each other happens at particular triggerpoints
+when using the `ddp` module. 
+
+These triggerpoints are added to the PyTorch model, specifically their `forward()` and `backward()` methods. 
+This happens when the model is wrapped with `DistributedDataParallel`:
+```python
+import torch.nn as nn
+from torch.nn.parallel import DistributedDataParallel
+
+model = nn.Linear(10, 10)
+ddp_model = DistributedDataParallel(model)
+```
+In 🤗 Accelerate this conversion happens automatically when calling [`Accelerator.prepare`] and passing in your model.
+
+```diff
++from accelerate import Accelerator
++accelerator = Accelerator()
+import torch.nn as nn
+-from torch.nn.parallel import DistributedDataParallel
+
+model = nn.Linear(10,10)
++model = accelerator.prepare(model)
+```
+
+## The slowdown in gradient accumulation
+
+You now understand that PyTorch adds hooks to the `forward` and `backward` method of your PyTorch model when 
+training in a distributed setup. But how does this automatically cause a risk of slowing down your code?
+
+In DDP (distributed data parallel), the specific order in which processes are performed and ran are expected
+at specific points and these must also occur at roughly the same time before moving on.
+
+The most direct example is when you update all of the parameters in a model through `.backward()`, all instances of the model
+need to have updated their gradients collated and updated before moving onto the next batch of data. But when performing 
+gradient accumulation, you accumulate `n` losses and skipping `.backward()` until `n` batches has been reached. This 
+can cause a significant slowdown since all the processes need to communicate with them more times than needed. But how 
Add newlines here as well :)
@@ -0,0 +1,27 @@
+# Deferring Executions
+
+When you run your usual script, instructions are executed in order. Using 🤗 Accelerate to deploy your script on several
+GPUs at the same time introduces a complication: while each process executes all instructions in order, some may be
+faster than others.
+
+You might need to wait for all processes to have reached a certain point before executing a given instruction. For
+instance, you shouldn't save a model before being sure every process is done with training, and you wouldn't want to 
+continue training before all the model weights have been loaded in. To do this, just write the following line in your code:
+
+```
+accelerator.wait_for_everyone()
+```
+
+This instruction will block all the processes that arrive first until all the other processes have reached that
+point (if you run your script on just one GPU or CPU, this won't do anything).
+
+A few example cases for when to use this utility are listed below:
+
+<Tip>
Add newlines here too
@@ -0,0 +1,146 @@
+<!--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.
+-->
+
+# Training on TPUs with 🤗 Accelerate
+
+Training on TPUs can be slightly different than training on multi-gpu, even with 🤗 Accelerate. This guide aims to show you 
+where you should be careful and why, as well as the best practices in general.
+
+## Training in a Notebook
+
+The main carepoint when training on TPUs comes from the `notebook_launcher`. As mentioned in the [notebook tutorial](https://huggingface.co/docs/accelerate/basic_tutorials/notebook), you need to 
+restructure your training code into a function that can get passed to the `notebook_launcher` function and be careful about not declaring any tensors on the GPU.
+
+While on a TPU that last part is not as important, a critical part to understand is that when you launch code from a notebook you do so through a process called **forking**. 
+When launching from the command-line, you perform **spawning**, where a python process is not currently running and you *spawn* a new process in. Since your Jupyter notebook is already 
+utilizing a python process, you need to *fork* a new process from it to launch your code. 
+
+Where this becomes important is in regards to declaring your model. On forked TPU processes, it is recommended that you instantiate your model *once* and pass this into your 
+training function. This is different than training on GPUs where you create `n` models that have their gradients synced and back-propagated at certain moments. Instead one 
+model instance is shared between all the nodes and it is passed back and forth. This is important especially when training on low-resource TPUs such as those provided in Kaggle kernels or
+on Google Colaboratory. 
+
+Below is an example of a training function passed to the `notebook_launcher` if training on CPUs or GPUs:
+
+<Tip>
At the time the concepts tutorial didn't exist 😃  will definitely add it in!
@@ -49,29 +49,29 @@ for index, batch in enumerate(training_dataloader):
 First the code shown earlier will be converted to utilize 🤗 Accelerate without the special gradient accumulation helper:
 
 ```diff
-+ from accelerate import Accelerator
-+ accelerator = Accelerator()
-
-+ model, optimizer, training_dataloader, scheduler = accelerator.prepare(
-+     model, optimizer, training_dataloader, scheduler
-+ )
-
-  for index, batch in enumerate(training_dataloader):
-      optimizer.zero_grad()
-      inputs, targets = batch
--     inputs = inputs.to(device)
--     targets = targets.to(device)
-      outputs = model(inputs)
-      loss = loss_function(outputs, targets)
-      loss = loss / gradient_accumulation_steps
-+     accelerator.backward(loss)
-      if (index+1) % gradient_accumulation_steps == 0:
-          optimizer.step()
-          scheduler.step()
++from accelerate import Accelerator
++accelerator = Accelerator()
+
++model, optimizer, training_dataloader, scheduler = accelerator.prepare(
++    model, optimizer, training_dataloader, scheduler
++)
+
+for index, batch in enumerate(training_dataloader):
+    optimizer.zero_grad()
+    inputs, targets = batch
+-    inputs = inputs.to(device)
+-    targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    loss = loss / gradient_accumulation_steps
++    accelerator.backward(loss)
+    if (index+1) % gradient_accumulation_steps == 0:
+        optimizer.step()
+        scheduler.step()
 ```
 
 <Tip warning={true}>
-In its current state, this code is not going to perform gradient accumulation efficiently due to a process called gradient synchronization.
+In its current state, this code is not going to perform gradient accumulation efficiently due to a process called gradient synchronization. Read more about that in the Concepts tutorial!
This actually is a mistake because this warning is no longer needed! But I added a link to the gradient sync tutorial int he earlier mention of it
@@ -80,35 +80,39 @@ All that is left now is to let 🤗 Accelerate handle the gradient accumulation
 of steps to perform before each call to `step()` and how to automatically adjust the loss during the call to [`Accelerator.backward`]:
 
 ```diff
-  from accelerate import Accelerator
-- accelerator = Accelerator()
-+ accelerator = Accelerator(gradient_accumulation_steps=2)
+from accelerate import Accelerator
+-accelerator = Accelerator()
++accelerator = Accelerator(gradient_accumulation_steps=2)
 ```
 
 From here you can use the [`Accelerator.accumulate`] context manager from inside your training loop to automatically perform the gradient accumulation for you!
-You just wrap it around the entire training part of your code: 
+You just wrap it around the entire training part of our code: 
 
 ```diff
-- for index, batch in enumerate(training_dataloader):
-+ for batch in training_dataloader:
-+    with accelerator.accumulate(model):
-          optimizer.zero_grad()
-          inputs, targets = batch
-          outputs = model(inputs)
+-for index, batch in enumerate(training_dataloader):
++for batch in training_dataloader:
++   with accelerator.accumulate(model):
+        optimizer.zero_grad()
+        inputs, targets = batch
+        outputs = model(inputs)
 ```
 
 and you can remove all the special checks for the step number and the loss adjustment:
 
 ```diff
-- loss = loss / gradient_accumulation_steps
-  accelerator.backward(loss)
-- if (index+1) % gradient_accumulation_steps == 0:
-  optimizer.step()
-  scheduler.step()
+-loss = loss / gradient_accumulation_steps
+accelerator.backward(loss)
+-if (index+1) % gradient_accumulation_steps == 0:
+optimizer.step()
+scheduler.step()
 ```
 
 As you can see the [`Accelerator`] is able to keep track of the batch number you are on and it will automatically know whether to step through the prepared optimizer and how to adjust the loss. 
 
+<Tip warning={true}>
```suggestion
Convert existing codebases to utilize [DeepSpeed](https://huggingface.co/docs/accelerate/deepspeed), perform [fully sharded data parallelism](https://huggingface.co/docs/accelerate/fsdp), and have automatic support for mixed-precision training! 
```
@@ -12,121 +12,60 @@ specific language governing permissions and limitations under the License.
 
 # Accelerate
 
-Run your *raw* PyTorch training script on any kind of device.
-
-## Features
-
-- 🤗 Accelerate provides an easy API to make your scripts run with mixed precision and in any kind of distributed
-  setting (multi-GPUs, TPUs etc.) while still letting you write your own training loop. The same code can then run
-  seamlessly on your local machine for debugging or your training environment.
-
-- 🤗 Accelerate also provides a CLI tool that allows you to quickly configure and test your training environment and
-  then launch the scripts.
-
-
-## Easy to integrate
-
-A traditional training loop in PyTorch looks like this:
-
-```python
-my_model.to(device)
-
-for batch in my_training_dataloader:
-    my_optimizer.zero_grad()
-    inputs, targets = batch
-    inputs = inputs.to(device)
-    targets = targets.to(device)
-    outputs = my_model(inputs)
-    loss = my_loss_function(outputs, targets)
-    loss.backward()
-    my_optimizer.step()
-```
-
-Changing it to work with accelerate is really easy and only adds a few lines of code:
+🤗 Accelerate is a library that enables the same PyTorch code to be run across any distributed configuration by adding just four lines of code! In short, training and inference at scale made simple, efficient and adaptable.
 
 ```diff
 + from accelerate import Accelerator
-
 + accelerator = Accelerator()
-  # Use the device given by the *accelerator* object.
-+ device = accelerator.device
-  my_model.to(device)
-  # Pass every important object (model, optimizer, dataloader) to *accelerator.prepare*
-+ my_model, my_optimizer, my_training_dataloader = accelerator.prepare(
-+     my_model, my_optimizer, my_training_dataloader
-+ )
-
-  for batch in my_training_dataloader:
-      my_optimizer.zero_grad()
-      inputs, targets = batch
-      inputs = inputs.to(device)
-      targets = targets.to(device)
-      outputs = my_model(inputs)
-      loss = my_loss_function(outputs, targets)
-      # Just a small change for the backward instruction
--     loss.backward()
-+     accelerator.backward(loss)
-      my_optimizer.step()
-```
 
-and with this, your script can now run in a distributed environment (multi-GPU, TPU).
-
-You can even simplify your script a bit by letting 🤗 Accelerate handle the device placement for you (which is safer,
-especially for TPU training):
-
-```diff
-+ from accelerate import Accelerator
-
-+ accelerator = Accelerator()
-- my_model.to(device)
-  # Pass every important object (model, optimizer, dataloader) to *accelerator.prepare*
-+ my_model, my_optimizer, my_training_dataloader = accelerate.prepare(
-+     my_model, my_optimizer, my_training_dataloader
++ model, optimizer, training_dataloader, scheduler = accelerator.prepare(
++     model, optimizer, training_dataloader, scheduler
 + )
 
-  for batch in my_training_dataloader:
-      my_optimizer.zero_grad()
-      inputs, targets = batch
--     inputs = inputs.to(device)
--     targets = targets.to(device)
-      outputs = my_model(inputs)
-      loss = my_loss_function(outputs, targets)
-      # Just a small change for the backward instruction
--     loss.backward()
-+     accelerator.backward(loss)
-      my_optimizer.step()
+  for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+-   inputs = inputs.to(device)
+-   targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
++   accelerator.backward(loss)
+    optimizer.step()
+    scheduler.step()
 ```
 
-## Script launcher
+Built on `torch_xla` and `torch.distributed`, 🤗 Accelerate takes care of the heavy lifting, so you don't have to write any custom code to adapt to these platforms.
+platforms separately so you don't have to. Convert existing codebases to utilize [DeepSpeed](https://huggingface.co/docs/accelerate/deepspeed), perform [fully sharded data parallelism](https://huggingface.co/docs/accelerate/fsdp), and have automatic support for mixed-precision training! 
```suggestion
Some data at the end of the dataset may be duplicated so the batch can be divided equally among all workers. As a result, metrics
```
@@ -137,17 +137,22 @@ As for the training dataloader, passing your validation dataloader through
 [`~Accelerator.prepare`] may change it: if you run on X GPUs, it will have its length divided by X
 (since your actual batch size will be multiplied by X), unless you set `split_batches=True`.
 
+</Tip>
+
 Any instruction using your training dataloader length (for instance if you need the number of total training steps
-to create a learning rate scheduler) should go after the call to [`~Accelerator.prepare`].
+to create a learning rate scheduler) should go after the call to [`~Accelerator.prepare`]. 
 
-As some data at the end of the dataset may be duplicated so the batch can divide equally to all workers, metrics should be 
-calculated through the [`~Accelerator.gather_for_metrics`] method to automatically remove the duplicated data.
+Some data at the end of the dataset may be duplicated so the batch can divide equally to all workers. As a result metrics
Should be un-commented before merging.
@@ -1,6 +1,6 @@
 name: Run Tests
 
-on: [pull_request]
+# on: [pull_request]
Installation should stay in Getting Started.
@@ -1,38 +1,67 @@
-- sections: 
+- sections:
   - local: index
     title: 🤗 Accelerate
   - local: quicktour
-    title: Quick tour
-  - local: installation
+    title: Quicktour of 🤗 Accelerate
+  title: Getting started
+- sections:
+  - local: basic_tutorials/overview
+    title: Overview
+  - local: basic_tutorials/install
```suggestion
    title: Quicktour
```
I'm pretty sure the user won't expect a quicktour of another lib here ;-)
@@ -1,38 +1,67 @@
-- sections: 
+- sections:
   - local: index
     title: 🤗 Accelerate
   - local: quicktour
-    title: Quick tour
-  - local: installation
+    title: Quicktour of 🤗 Accelerate
That's an editable install, a source install is just installing from git and should be detailed here, before you speak about editable install for developping.
@@ -0,0 +1,91 @@
+<!--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.
+-->
+
+# Installation and Configuration
+
+Before you start, you will need to setup your environment, install the appropriate packages, and configure 🤗 Accelerate. 🤗 Accelerate is tested on **Python 3.7+**.
+
+## Installing 🤗 Accelerate
+
+🤗 Accelerate is available on pypi and conda, as well as on GitHub. Details to install from each are below:
+
+### pip 
+
+To install 🤗 Accelerate from pypi, perform:
+
+```bash
+pip install accelerate
+```
+
+### conda
+
+🤗 Accelerate can also be installed with conda with:
+
+```bash
+conda install -c conda-forge accelerate
+```
+
+### Source
+
+New features are added every day that haven't been released yet. To try them out yourself, install
+from the GitHub repository:
+
+```bash
+git clone https://github.com/huggingface/accelerate
+cd accelerate
+pip install -e .
+```
You should insist here that running `accelerate config` beforehand is recommended.
@@ -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.
+-->
+
+# Launching your 🤗 Accelerate scripts
+
+In the previous tutorial, you were introduced to how to modify your current training script to use 🤗 Accelerate.
+The final version of that code is shown below:
+
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    accelerator.backward(loss)
+    optimizer.step()
+    scheduler.step()
+```
+
+But how do you run this code and have it utilize the special hardware available to it?
+
+First you should rewrite the above code into a function, and make it callable as a script. For example:
+
+```diff
+  from accelerate import Accelerator
+  
++ def main():
+
+    accelerator = Accelerator()
+
+    model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+        model, optimizer, training_dataloader, scheduler
+    )
+
+    for batch in training_dataloader:
+        optimizer.zero_grad()
+        inputs, targets = batch
+        outputs = model(inputs)
+        loss = loss_function(outputs, targets)
+        accelerator.backward(loss)
+        optimizer.step()
+        scheduler.step()
+
++ if __name__ == "__main__":
++   main()
+```
+
+Next you need to launch it with `accelerate launch`. 
You should give a code example of this.
@@ -0,0 +1,29 @@
+# Deferring Executions
+
+When you run your usual script, instructions are executed in order. Using 🤗 Accelerate to deploy your script on several
+GPUs at the same time introduces a complication: while each process executes all instructions in order, some may be
+faster than others.
+
+You might need to wait for all processes to have reached a certain point before executing a given instruction. For
+instance, you shouldn't save a model before being sure every process is done with training, and you wouldn't want to 
+continue training before all the model weights have been loaded in. To do this, just write the following line in your code:
+
+```
+accelerator.wait_for_everyone()
+```
+
+This instruction will block all the processes that arrive first until all the other processes have reached that
+point (if you run your script on just one GPU or CPU, this won't do anything).
+
+A few example cases for when to use this utility are listed below:
+
+<Tip>
+
+Some of these are utilized with the [`Accelerator.main_process_first`] context manager, which utilizes [`Accelerator.wait_for_everyone`] to 
+run a particular set of code on the main process beforehand before triggering and launching the other processes
+
+</Tip>
+
+- Downloading a dataset on the main process first and then loading the cached dataset in afterwards
Use 4 lines of indentation.
@@ -12,121 +12,60 @@ specific language governing permissions and limitations under the License.
 
 # Accelerate
 
-Run your *raw* PyTorch training script on any kind of device.
-
-## Features
-
-- 🤗 Accelerate provides an easy API to make your scripts run with mixed precision and in any kind of distributed
-  setting (multi-GPUs, TPUs etc.) while still letting you write your own training loop. The same code can then run
-  seamlessly on your local machine for debugging or your training environment.
-
-- 🤗 Accelerate also provides a CLI tool that allows you to quickly configure and test your training environment and
-  then launch the scripts.
-
-
-## Easy to integrate
-
-A traditional training loop in PyTorch looks like this:
-
-```python
-my_model.to(device)
-
-for batch in my_training_dataloader:
-    my_optimizer.zero_grad()
-    inputs, targets = batch
-    inputs = inputs.to(device)
-    targets = targets.to(device)
-    outputs = my_model(inputs)
-    loss = my_loss_function(outputs, targets)
-    loss.backward()
-    my_optimizer.step()
-```
-
-Changing it to work with accelerate is really easy and only adds a few lines of code:
+🤗 Accelerate is a library that enables the same PyTorch code to be run across any distributed configuration by adding just four lines of code! In short, training and inference at scale made simple, efficient and adaptable.
 
 ```diff
 + from accelerate import Accelerator
-
 + accelerator = Accelerator()
-  # Use the device given by the *accelerator* object.
-+ device = accelerator.device
-  my_model.to(device)
-  # Pass every important object (model, optimizer, dataloader) to *accelerator.prepare*
-+ my_model, my_optimizer, my_training_dataloader = accelerator.prepare(
-+     my_model, my_optimizer, my_training_dataloader
-+ )
-
-  for batch in my_training_dataloader:
-      my_optimizer.zero_grad()
-      inputs, targets = batch
-      inputs = inputs.to(device)
-      targets = targets.to(device)
-      outputs = my_model(inputs)
-      loss = my_loss_function(outputs, targets)
-      # Just a small change for the backward instruction
--     loss.backward()
-+     accelerator.backward(loss)
-      my_optimizer.step()
-```
 
-and with this, your script can now run in a distributed environment (multi-GPU, TPU).
-
-You can even simplify your script a bit by letting 🤗 Accelerate handle the device placement for you (which is safer,
-especially for TPU training):
-
-```diff
-+ from accelerate import Accelerator
-
-+ accelerator = Accelerator()
-- my_model.to(device)
-  # Pass every important object (model, optimizer, dataloader) to *accelerator.prepare*
-+ my_model, my_optimizer, my_training_dataloader = accelerate.prepare(
-+     my_model, my_optimizer, my_training_dataloader
++ model, optimizer, training_dataloader, scheduler = accelerator.prepare(
++     model, optimizer, training_dataloader, scheduler
 + )
 
-  for batch in my_training_dataloader:
-      my_optimizer.zero_grad()
-      inputs, targets = batch
--     inputs = inputs.to(device)
--     targets = targets.to(device)
-      outputs = my_model(inputs)
-      loss = my_loss_function(outputs, targets)
-      # Just a small change for the backward instruction
--     loss.backward()
-+     accelerator.backward(loss)
-      my_optimizer.step()
+  for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+-   inputs = inputs.to(device)
+-   targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
++   accelerator.backward(loss)
+    optimizer.step()
+    scheduler.step()
I would show first the sample where we don't change the device as letting Accelerate handle device placement changes more lines of code.
@@ -12,121 +12,60 @@ specific language governing permissions and limitations under the License.
 
 # Accelerate
 
-Run your *raw* PyTorch training script on any kind of device.
-
-## Features
-
-- 🤗 Accelerate provides an easy API to make your scripts run with mixed precision and in any kind of distributed
-  setting (multi-GPUs, TPUs etc.) while still letting you write your own training loop. The same code can then run
-  seamlessly on your local machine for debugging or your training environment.
-
-- 🤗 Accelerate also provides a CLI tool that allows you to quickly configure and test your training environment and
-  then launch the scripts.
-
-
-## Easy to integrate
-
-A traditional training loop in PyTorch looks like this:
-
-```python
-my_model.to(device)
-
-for batch in my_training_dataloader:
-    my_optimizer.zero_grad()
-    inputs, targets = batch
-    inputs = inputs.to(device)
-    targets = targets.to(device)
-    outputs = my_model(inputs)
-    loss = my_loss_function(outputs, targets)
-    loss.backward()
-    my_optimizer.step()
-```
-
-Changing it to work with accelerate is really easy and only adds a few lines of code:
+🤗 Accelerate is a library that enables the same PyTorch code to be run across any distributed configuration by adding just four lines of code! In short, training and inference at scale made simple, efficient and adaptable.
 
 ```diff
 + from accelerate import Accelerator
-
 + accelerator = Accelerator()
-  # Use the device given by the *accelerator* object.
-+ device = accelerator.device
-  my_model.to(device)
-  # Pass every important object (model, optimizer, dataloader) to *accelerator.prepare*
-+ my_model, my_optimizer, my_training_dataloader = accelerator.prepare(
-+     my_model, my_optimizer, my_training_dataloader
-+ )
-
-  for batch in my_training_dataloader:
-      my_optimizer.zero_grad()
-      inputs, targets = batch
-      inputs = inputs.to(device)
-      targets = targets.to(device)
-      outputs = my_model(inputs)
-      loss = my_loss_function(outputs, targets)
-      # Just a small change for the backward instruction
--     loss.backward()
-+     accelerator.backward(loss)
-      my_optimizer.step()
-```
 
-and with this, your script can now run in a distributed environment (multi-GPU, TPU).
-
-You can even simplify your script a bit by letting 🤗 Accelerate handle the device placement for you (which is safer,
-especially for TPU training):
-
-```diff
-+ from accelerate import Accelerator
-
-+ accelerator = Accelerator()
-- my_model.to(device)
-  # Pass every important object (model, optimizer, dataloader) to *accelerator.prepare*
-+ my_model, my_optimizer, my_training_dataloader = accelerate.prepare(
-+     my_model, my_optimizer, my_training_dataloader
++ model, optimizer, training_dataloader, scheduler = accelerator.prepare(
++     model, optimizer, training_dataloader, scheduler
 + )
 
-  for batch in my_training_dataloader:
-      my_optimizer.zero_grad()
-      inputs, targets = batch
--     inputs = inputs.to(device)
--     targets = targets.to(device)
-      outputs = my_model(inputs)
-      loss = my_loss_function(outputs, targets)
-      # Just a small change for the backward instruction
--     loss.backward()
-+     accelerator.backward(loss)
-      my_optimizer.step()
+  for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+-   inputs = inputs.to(device)
The old page had way more information, which we should keep.
@@ -0,0 +1,17 @@
+<!--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.
+-->
+
+# Accelerator
+
+The main [`Accelerator`] class:
````suggestion
```python
````
Why remove this? We lose all formatting after...
@@ -24,8 +24,8 @@ which can total to adding just one new line of code!
 
 This example will use a very simplistic PyTorch training loop that performs gradient accumulation every two batches:
 
-```python
-device = "cuda"
+```
Please use 4 spaces to indent
@@ -86,18 +88,18 @@ of steps to perform before each call to `step()` and how to automatically adjust
 ```
 
 From here you can use the [`Accelerator.accumulate`] context manager from inside your training loop to automatically perform the gradient accumulation for you!
-You just wrap it around the entire training part of your code: 
+You just wrap it around the entire training part of our code: 
 
 ```diff
 - for index, batch in enumerate(training_dataloader):
 + for batch in training_dataloader:
 +    with accelerator.accumulate(model):
-          optimizer.zero_grad()
-          inputs, targets = batch
-          outputs = model(inputs)
+       optimizer.zero_grad()
Some content seems to be lost here as there is no usage guide.
@@ -0,0 +1,19 @@
+<!--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.
+-->
+
+# Launchers
+
+Functions for launching training on distributed processes.
```suggestion
In this case, 🤗 Accelerate will make some hyperparameter decisions for you, e.g., if GPUs are available, it will use all of them by default without the mixed precision.
```
@@ -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.
+-->
+
+# Launching your 🤗 Accelerate scripts
+
+In the previous tutorial, you were introduced to how to modify your current training script to use 🤗 Accelerate.
+The final version of that code is shown below:
+
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    accelerator.backward(loss)
+    optimizer.step()
+    scheduler.step()
+```
+
+But how do you run this code and have it utilize the special hardware available to it?
+
+First you should rewrite the above code into a function, and make it callable as a script. For example:
+
+```diff
+  from accelerate import Accelerator
+  
++ def main():
+
+    accelerator = Accelerator()
+
+    model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+        model, optimizer, training_dataloader, scheduler
+    )
+
+    for batch in training_dataloader:
+        optimizer.zero_grad()
+        inputs, targets = batch
+        outputs = model(inputs)
+        loss = loss_function(outputs, targets)
+        accelerator.backward(loss)
+        optimizer.step()
+        scheduler.step()
+
++ if __name__ == "__main__":
++   main()
+```
+
+Next you need to launch it with `accelerate launch`. 
+
+## Using accelerate launch
+
+🤗 Accelerate has a special CLI command to help you launch your code in your system through `accelerate launch`.
+This command wraps around all of the different commands needed to launch your script on various platforms, without you
+having to remember what each of them are.
+
+<Tip>
+
+If you are familiar with launching scripts in PyTorch yourself such as with `torchrun`, you can still do this. It is not required to use `accelerate launch`
+
+</Tip>
+
+You can launch your script quickly by using:
+```bash
+accelerate launch {script_name.py} --arg1 --arg2 ...
+```
+
+Just put `accelerate launch` at the start of your command, and pass in additional arguments and parameters to your script afterwards like normal!
+
+Since this runs the various torch spawn methods, all of the expected environment variables can be modified here as well.
+For example, here is how to use `accelerate launch` with a single GPU:
+```bash
+CUDA_VISIBLE_DEVICES="0" accelerate launch {script_name.py} --arg1 --arg2 ...
+```
+
+You can also use `accelerate launch` without performing `accelerate config` first, but you may need to manually pass in the right configuration parameters.
+In this case, 🤗 Accelerate will make some hyperparameter decisions for you, such as if a GPU is available, it will use all of them by default, and no mixed precision will be used. 
```suggestion
Here is how you would use all GPUs and train with mixed precision disabled:
```
@@ -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.
+-->
+
+# Launching your 🤗 Accelerate scripts
+
+In the previous tutorial, you were introduced to how to modify your current training script to use 🤗 Accelerate.
+The final version of that code is shown below:
+
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    accelerator.backward(loss)
+    optimizer.step()
+    scheduler.step()
+```
+
+But how do you run this code and have it utilize the special hardware available to it?
+
+First you should rewrite the above code into a function, and make it callable as a script. For example:
+
+```diff
+  from accelerate import Accelerator
+  
++ def main():
+
+    accelerator = Accelerator()
+
+    model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+        model, optimizer, training_dataloader, scheduler
+    )
+
+    for batch in training_dataloader:
+        optimizer.zero_grad()
+        inputs, targets = batch
+        outputs = model(inputs)
+        loss = loss_function(outputs, targets)
+        accelerator.backward(loss)
+        optimizer.step()
+        scheduler.step()
+
++ if __name__ == "__main__":
++   main()
+```
+
+Next you need to launch it with `accelerate launch`. 
+
+## Using accelerate launch
+
+🤗 Accelerate has a special CLI command to help you launch your code in your system through `accelerate launch`.
+This command wraps around all of the different commands needed to launch your script on various platforms, without you
+having to remember what each of them are.
+
+<Tip>
+
+If you are familiar with launching scripts in PyTorch yourself such as with `torchrun`, you can still do this. It is not required to use `accelerate launch`
+
+</Tip>
+
+You can launch your script quickly by using:
+```bash
+accelerate launch {script_name.py} --arg1 --arg2 ...
+```
+
+Just put `accelerate launch` at the start of your command, and pass in additional arguments and parameters to your script afterwards like normal!
+
+Since this runs the various torch spawn methods, all of the expected environment variables can be modified here as well.
+For example, here is how to use `accelerate launch` with a single GPU:
+```bash
+CUDA_VISIBLE_DEVICES="0" accelerate launch {script_name.py} --arg1 --arg2 ...
+```
+
+You can also use `accelerate launch` without performing `accelerate config` first, but you may need to manually pass in the right configuration parameters.
+In this case, 🤗 Accelerate will make some hyperparameter decisions for you, such as if a GPU is available, it will use all of them by default, and no mixed precision will be used. 
+Here is how you would use all GPUs and train on no mixed precision:
maybe specify that this is one machine with two GPUs?
@@ -0,0 +1,91 @@
+<!--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.
+-->
+
+# Installation and Configuration
+
+Before you start, you will need to setup your environment, install the appropriate packages, and configure 🤗 Accelerate. 🤗 Accelerate is tested on **Python 3.7+**.
+
+## Installing 🤗 Accelerate
+
+🤗 Accelerate is available on pypi and conda, as well as on GitHub. Details to install from each are below:
+
+### pip 
+
+To install 🤗 Accelerate from pypi, perform:
+
+```bash
+pip install accelerate
+```
+
+### conda
+
+🤗 Accelerate can also be installed with conda with:
+
+```bash
+conda install -c conda-forge accelerate
+```
+
+### Source
+
+New features are added every day that haven't been released yet. To try them out yourself, install
+from the GitHub repository:
+
+```bash
+git clone https://github.com/huggingface/accelerate
+cd accelerate
+pip install -e .
+```
+
+## Configuring 🤗 Accelerate
+
+After installing, you need to configure 🤗 Accelerate for how the current system is setup for training. 
+To do so run the following and answer the questions prompted to you:
+
+```bash
+accelerate config
+```
+
+To write a barebones configuration that doesn't include options such as DeepSpeed configuration or running on TPUs, you can quickly run:
+
+```bash
+python -c "from accelerate.utils import write_basic_config; write_basic_config(mixed_precision='fp16')"
+```
+🤗 Accelerate will automatically utilize the maximum number of GPUs available and set the mixed precision mode.
+
+To check that your configuration looks fine, run:
+
+```bash
+accelerate env
+```
+
+An example output is shown below:
```suggestion
This command wraps around all of the different commands needed to launch your script on various platforms, without you having to remember what each of them are.
```
@@ -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.
+-->
+
+# Launching your 🤗 Accelerate scripts
+
+In the previous tutorial, you were introduced to how to modify your current training script to use 🤗 Accelerate.
+The final version of that code is shown below:
+
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    accelerator.backward(loss)
+    optimizer.step()
+    scheduler.step()
+```
+
+But how do you run this code and have it utilize the special hardware available to it?
+
+First you should rewrite the above code into a function, and make it callable as a script. For example:
+
+```diff
+  from accelerate import Accelerator
+  
++ def main():
+
+    accelerator = Accelerator()
+
+    model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+        model, optimizer, training_dataloader, scheduler
+    )
+
+    for batch in training_dataloader:
+        optimizer.zero_grad()
+        inputs, targets = batch
+        outputs = model(inputs)
+        loss = loss_function(outputs, targets)
+        accelerator.backward(loss)
+        optimizer.step()
+        scheduler.step()
+
++ if __name__ == "__main__":
++   main()
+```
+
+Next you need to launch it with `accelerate launch`. 
+
+## Using accelerate launch
+
+🤗 Accelerate has a special CLI command to help you launch your code in your system through `accelerate launch`.
+This command wraps around all of the different commands needed to launch your script on various platforms, without you
+having to remember what each of them are.
```suggestion
If you are familiar with launching scripts in PyTorch yourself such as with `torchrun`, you can still do this. It is not required to use `accelerate launch`.
```
@@ -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.
+-->
+
+# Launching your 🤗 Accelerate scripts
+
+In the previous tutorial, you were introduced to how to modify your current training script to use 🤗 Accelerate.
+The final version of that code is shown below:
+
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    accelerator.backward(loss)
+    optimizer.step()
+    scheduler.step()
+```
+
+But how do you run this code and have it utilize the special hardware available to it?
+
+First you should rewrite the above code into a function, and make it callable as a script. For example:
+
+```diff
+  from accelerate import Accelerator
+  
++ def main():
+
+    accelerator = Accelerator()
+
+    model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+        model, optimizer, training_dataloader, scheduler
+    )
+
+    for batch in training_dataloader:
+        optimizer.zero_grad()
+        inputs, targets = batch
+        outputs = model(inputs)
+        loss = loss_function(outputs, targets)
+        accelerator.backward(loss)
+        optimizer.step()
+        scheduler.step()
+
++ if __name__ == "__main__":
++   main()
+```
+
+Next you need to launch it with `accelerate launch`. 
+
+## Using accelerate launch
+
+🤗 Accelerate has a special CLI command to help you launch your code in your system through `accelerate launch`.
+This command wraps around all of the different commands needed to launch your script on various platforms, without you
+having to remember what each of them are.
+
+<Tip>
+
+If you are familiar with launching scripts in PyTorch yourself such as with `torchrun`, you can still do this. It is not required to use `accelerate launch`
why curly brackets here? i think also something like that would be clear enough:
```suggestion
accelerate launch your_script.py --arg1 --arg2 ...
```
@@ -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.
+-->
+
+# Launching your 🤗 Accelerate scripts
+
+In the previous tutorial, you were introduced to how to modify your current training script to use 🤗 Accelerate.
+The final version of that code is shown below:
+
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    accelerator.backward(loss)
+    optimizer.step()
+    scheduler.step()
+```
+
+But how do you run this code and have it utilize the special hardware available to it?
+
+First you should rewrite the above code into a function, and make it callable as a script. For example:
+
+```diff
+  from accelerate import Accelerator
+  
++ def main():
+
+    accelerator = Accelerator()
+
+    model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+        model, optimizer, training_dataloader, scheduler
+    )
+
+    for batch in training_dataloader:
+        optimizer.zero_grad()
+        inputs, targets = batch
+        outputs = model(inputs)
+        loss = loss_function(outputs, targets)
+        accelerator.backward(loss)
+        optimizer.step()
+        scheduler.step()
+
++ if __name__ == "__main__":
++   main()
+```
+
+Next you need to launch it with `accelerate launch`. 
+
+## Using accelerate launch
+
+🤗 Accelerate has a special CLI command to help you launch your code in your system through `accelerate launch`.
+This command wraps around all of the different commands needed to launch your script on various platforms, without you
+having to remember what each of them are.
+
+<Tip>
+
+If you are familiar with launching scripts in PyTorch yourself such as with `torchrun`, you can still do this. It is not required to use `accelerate launch`
+
+</Tip>
+
+You can launch your script quickly by using:
+```bash
+accelerate launch {script_name.py} --arg1 --arg2 ...
what's the benefit of that? what happens if you start a random script with accelerate?
@@ -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.
+-->
+
+# Launching your 🤗 Accelerate scripts
+
+In the previous tutorial, you were introduced to how to modify your current training script to use 🤗 Accelerate.
+The final version of that code is shown below:
+
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    accelerator.backward(loss)
+    optimizer.step()
+    scheduler.step()
+```
+
+But how do you run this code and have it utilize the special hardware available to it?
+
+First you should rewrite the above code into a function, and make it callable as a script. For example:
+
+```diff
+  from accelerate import Accelerator
+  
++ def main():
+
+    accelerator = Accelerator()
+
+    model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+        model, optimizer, training_dataloader, scheduler
+    )
+
+    for batch in training_dataloader:
+        optimizer.zero_grad()
+        inputs, targets = batch
+        outputs = model(inputs)
+        loss = loss_function(outputs, targets)
+        accelerator.backward(loss)
+        optimizer.step()
+        scheduler.step()
+
++ if __name__ == "__main__":
++   main()
+```
+
+Next you need to launch it with `accelerate launch`. 
+
+## Using accelerate launch
+
+🤗 Accelerate has a special CLI command to help you launch your code in your system through `accelerate launch`.
+This command wraps around all of the different commands needed to launch your script on various platforms, without you
+having to remember what each of them are.
+
+<Tip>
+
+If you are familiar with launching scripts in PyTorch yourself such as with `torchrun`, you can still do this. It is not required to use `accelerate launch`
+
+</Tip>
+
+You can launch your script quickly by using:
+```bash
+accelerate launch {script_name.py} --arg1 --arg2 ...
+```
+
+Just put `accelerate launch` at the start of your command, and pass in additional arguments and parameters to your script afterwards like normal!
+
+Since this runs the various torch spawn methods, all of the expected environment variables can be modified here as well.
+For example, here is how to use `accelerate launch` with a single GPU:
+```bash
+CUDA_VISIBLE_DEVICES="0" accelerate launch {script_name.py} --arg1 --arg2 ...
+```
+
+You can also use `accelerate launch` without performing `accelerate config` first, but you may need to manually pass in the right configuration parameters.
+In this case, 🤗 Accelerate will make some hyperparameter decisions for you, such as if a GPU is available, it will use all of them by default, and no mixed precision will be used. 
+Here is how you would use all GPUs and train on no mixed precision:
+```bash
+accelerate launch --multi_gpu {script_name.py} --arg1 --arg2 ...
+```
+To get more specific you should pass in the needed parameters yourself. For instance, here is how you 
+would also launch that same script on two GPUs using mixed precision while avoiding all of the warnings: 
+```bash
+accelerate launch --multi_gpu --mixed_precision=fp16 --num_processes=2 {script_name.py} --arg1 --arg2 ...
+```
+
+For a complete list of parameters you can pass in, run:
+```bash
+accelerate launch -h
+```
+
+<Tip>
+
+Even if you are not using 🤗 Accelerate in your code, you can still use the launcher for starting your scripts!
regarding path_to_script.py: maybe just be consistent with the above name used.
@@ -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.
+-->
+
+# Launching your 🤗 Accelerate scripts
+
+In the previous tutorial, you were introduced to how to modify your current training script to use 🤗 Accelerate.
+The final version of that code is shown below:
+
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    accelerator.backward(loss)
+    optimizer.step()
+    scheduler.step()
+```
+
+But how do you run this code and have it utilize the special hardware available to it?
+
+First you should rewrite the above code into a function, and make it callable as a script. For example:
+
+```diff
+  from accelerate import Accelerator
+  
++ def main():
+
+    accelerator = Accelerator()
+
+    model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+        model, optimizer, training_dataloader, scheduler
+    )
+
+    for batch in training_dataloader:
+        optimizer.zero_grad()
+        inputs, targets = batch
+        outputs = model(inputs)
+        loss = loss_function(outputs, targets)
+        accelerator.backward(loss)
+        optimizer.step()
+        scheduler.step()
+
++ if __name__ == "__main__":
++   main()
+```
+
+Next you need to launch it with `accelerate launch`. 
+
+## Using accelerate launch
+
+🤗 Accelerate has a special CLI command to help you launch your code in your system through `accelerate launch`.
+This command wraps around all of the different commands needed to launch your script on various platforms, without you
+having to remember what each of them are.
+
+<Tip>
+
+If you are familiar with launching scripts in PyTorch yourself such as with `torchrun`, you can still do this. It is not required to use `accelerate launch`
+
+</Tip>
+
+You can launch your script quickly by using:
+```bash
+accelerate launch {script_name.py} --arg1 --arg2 ...
+```
+
+Just put `accelerate launch` at the start of your command, and pass in additional arguments and parameters to your script afterwards like normal!
+
+Since this runs the various torch spawn methods, all of the expected environment variables can be modified here as well.
+For example, here is how to use `accelerate launch` with a single GPU:
+```bash
+CUDA_VISIBLE_DEVICES="0" accelerate launch {script_name.py} --arg1 --arg2 ...
+```
+
+You can also use `accelerate launch` without performing `accelerate config` first, but you may need to manually pass in the right configuration parameters.
+In this case, 🤗 Accelerate will make some hyperparameter decisions for you, such as if a GPU is available, it will use all of them by default, and no mixed precision will be used. 
+Here is how you would use all GPUs and train on no mixed precision:
+```bash
+accelerate launch --multi_gpu {script_name.py} --arg1 --arg2 ...
+```
+To get more specific you should pass in the needed parameters yourself. For instance, here is how you 
+would also launch that same script on two GPUs using mixed precision while avoiding all of the warnings: 
+```bash
+accelerate launch --multi_gpu --mixed_precision=fp16 --num_processes=2 {script_name.py} --arg1 --arg2 ...
+```
+
+For a complete list of parameters you can pass in, run:
+```bash
+accelerate launch -h
+```
+
+<Tip>
+
+Even if you are not using 🤗 Accelerate in your code, you can still use the launcher for starting your scripts!
+
+</Tip>
+
+## Custom Configurations
+
+As briefly mentioned earlier, `accelerate launch` should be mostly used through combining set configurations 
+made with the `accelerate config` command. These configs are saved to a `default_config.yaml` file in your cache folder for 🤗 Accelerate. 
+This cache folder is located at (with decreasing order of priority):
+
+- The content of your environment variable `HF_HOME` suffixed with `accelerate`.
+- If it does not exist, the content of your environment variable `XDG_CACHE_HOME` suffixed with
+  `huggingface/accelerate`.
+- If this does not exist either, the folder `~/.cache/huggingface/accelerate`.
+
+To have multiple configurations, the flag `--config_file` can be passed to the `accelerate launch` command paired 
+with the location of the custom yaml. 
+
+An example yaml may look something like the following:
+```yaml
+compute_environment: LOCAL_MACHINE
+deepspeed_config: {}
+distributed_type: MULTI_GPU
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 1
+use_cpu: false
+```
+
+An example of launching a script with a yaml located in a custom location would look something like so:
+```bash
+accelerate launch --config_file path/to/config/my_config_file.yaml path_to_script.py --args_for_the_script
`prepare` makes the "Setting the right device" a bit obsolete no? maybe make a note when to use `prepare` or `model.to`
@@ -0,0 +1,119 @@
+<!--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.
+-->
+
+# Migrating your code to 🤗 Accelerate
+
+This tutorial will detail how to easily convert existing PyTorch code to use 🤗 Accelerate!
+You'll see that by just changing a few lines of code, 🤗 Accelerate can perform its magic and get you on 
+your way towards running your code on distributed systems with ease!
+
+## The base training loop
+
+To begin, write out a very basic PyTorch training loop. 
+
+<Tip>
+We are under the presumption that `training_dataloader`, `model`, `optimizer`, `scheduler`, and `loss_function` have been defined beforehand.
+</Tip>
+
+```python
+device = "cuda"
+model.to(device)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    inputs = inputs.to(device)
+    targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    loss.backward()
+    optimizer.step()
+    scheduler.step()
+```
+
+## Add in 🤗 Accelerate
+
+To start using 🤗 Accelerate, first import and create an [`Accelerator`] instance:
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+```
+[`Accelerator`] is the main force behind utilizing all the possible options for distributed training!
+
+### Setting the right device
+
+The [`Accelerator`] class knows the right device to move any PyTorch object to at any time, so you should
+change the definition of `device` to come from [`Accelerator`]:
+
+```diff
+- device = 'cuda'
++ device = accelerator.device
+  model.to(device)
+```
+
+### Preparing your objects
+
+Next you need to pass all of the important objects related to training into [`Accelerator.prepare`]. 🤗 Accelerate will
+make sure everything is setup in the current environment for you to start training:
+
+```
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+```
+These objects are returned in the same order they were sent in with. By default, `device_placement=True`, so your model is automatically moved to the proper device.
+performing this action will also automatically move your model to the proper device as well.
agree, a few examples would be useful here. also a graphic could be useful.
@@ -0,0 +1,29 @@
+# Deferring Executions
+
+When you run your usual script, instructions are executed in order. Using 🤗 Accelerate to deploy your script on several
+GPUs at the same time introduces a complication: while each process executes all instructions in order, some may be
+faster than others.
+
+You might need to wait for all processes to have reached a certain point before executing a given instruction. For
+instance, you shouldn't save a model before being sure every process is done with training, and you wouldn't want to 
+continue training before all the model weights have been loaded in. To do this, just write the following line in your code:
+
+```
+accelerator.wait_for_everyone()
+```
+
+This instruction will block all the processes that arrive first until all the other processes have reached that
+point (if you run your script on just one GPU or CPU, this won't do anything).
+
+A few example cases for when to use this utility are listed below:
+
+<Tip>
+
+Some of these are utilized with the [`Accelerator.main_process_first`] context manager, which utilizes [`Accelerator.wait_for_everyone`] to 
+run a particular set of code on the main process beforehand before triggering and launching the other processes
+
+</Tip>
+
+- Downloading a dataset on the main process first and then loading the cached dataset in afterwards
shouldn't you use the target to get a loss?
```suggestion
            ddp_model(input, target).backward()
```
@@ -0,0 +1,109 @@
+<!--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.
+-->
+
+# Gradient Synchronization
+
+PyTorch's distributed module operates by communicating back and forth between all of the GPUs in your system.
+This communication takes time, and ensuring all processes know the states of each other happens at particular triggerpoints
+when using the `ddp` module. 
+
+These triggerpoints are added to the PyTorch model, specifically their `forward()` and `backward()` methods. 
+This happens when the model is wrapped with `DistributedDataParallel`:
+```python
+import torch.nn as nn
+from torch.nn.parallel import DistributedDataParallel
+
+model = nn.Linear(10, 10)
+ddp_model = DistributedDataParallel(model)
+```
+In 🤗 Accelerate this conversion happens automatically when calling [`Accelerator.prepare`] and passing in your model.
+
+```diff
++ from accelerate import Accelerator
++ accelerator = Accelerator()
+  import torch.nn as nn
+- from torch.nn.parallel import DistributedDataParallel
+
+  model = nn.Linear(10,10)
++ model = accelerator.prepare(model)
+```
+
+## The slowdown in gradient accumulation
+
+You now understand that PyTorch adds hooks to the `forward` and `backward` method of your PyTorch model when 
+training in a distributed setup. But how does this risk slowing down your code?
+
+In DDP (distributed data parallel), the specific order in which processes are performed and ran are expected
+at specific points and these must also occur at roughly the same time before moving on.
+
+The most direct example is when you update all of the parameters in a model through `.backward()`. All instances of the model
+need to have updated their gradients, collated, and updated again before moving onto the next batch of data. But when performing 
+gradient accumulation, you accumulate `n` losses and skip `.backward()` until `n` batches have been reached. This 
+can cause a significant slowdown since all the processes need to communicate with them more times than needed. How 
+can you avoid this overhead?
+
+## Solving the slowdown problem
+
+Since you are skipping these batches, their gradients do not need to be synchronized until the point where `.backward()` is actually called. 
+PyTorch cannot automagically tell when you need to do this, but they do provide a tool to help through the [`no_sync`](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel.no_sync) context manager
+that is added to your model after converting it to DDP.
+
+Under this context manager, PyTorch will skip synchronizing the gradients when `.backward()` is called, and the first call to `.backward()` outside this 
+context manager will trigger the synchronization. See an example below:
+```python
+ddp_model, dataloader = accelerator.prepare(model, dataloader)
+
+for index, batch in enumerate(dataloader):
+    input, target = batch
+    # Trigger gradient synchronization on the last batch
+    if index != (len(dataloader) - 1):
+        with ddp_model.no_sync():
+            # Gradients only accumulate
+            ddp_model(input).backward()
same as above
```suggestion
        ddp_model(input, target).backward()
```
@@ -0,0 +1,109 @@
+<!--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.
+-->
+
+# Gradient Synchronization
+
+PyTorch's distributed module operates by communicating back and forth between all of the GPUs in your system.
+This communication takes time, and ensuring all processes know the states of each other happens at particular triggerpoints
+when using the `ddp` module. 
+
+These triggerpoints are added to the PyTorch model, specifically their `forward()` and `backward()` methods. 
+This happens when the model is wrapped with `DistributedDataParallel`:
+```python
+import torch.nn as nn
+from torch.nn.parallel import DistributedDataParallel
+
+model = nn.Linear(10, 10)
+ddp_model = DistributedDataParallel(model)
+```
+In 🤗 Accelerate this conversion happens automatically when calling [`Accelerator.prepare`] and passing in your model.
+
+```diff
++ from accelerate import Accelerator
++ accelerator = Accelerator()
+  import torch.nn as nn
+- from torch.nn.parallel import DistributedDataParallel
+
+  model = nn.Linear(10,10)
++ model = accelerator.prepare(model)
+```
+
+## The slowdown in gradient accumulation
+
+You now understand that PyTorch adds hooks to the `forward` and `backward` method of your PyTorch model when 
+training in a distributed setup. But how does this risk slowing down your code?
+
+In DDP (distributed data parallel), the specific order in which processes are performed and ran are expected
+at specific points and these must also occur at roughly the same time before moving on.
+
+The most direct example is when you update all of the parameters in a model through `.backward()`. All instances of the model
+need to have updated their gradients, collated, and updated again before moving onto the next batch of data. But when performing 
+gradient accumulation, you accumulate `n` losses and skip `.backward()` until `n` batches have been reached. This 
+can cause a significant slowdown since all the processes need to communicate with them more times than needed. How 
+can you avoid this overhead?
+
+## Solving the slowdown problem
+
+Since you are skipping these batches, their gradients do not need to be synchronized until the point where `.backward()` is actually called. 
+PyTorch cannot automagically tell when you need to do this, but they do provide a tool to help through the [`no_sync`](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel.no_sync) context manager
+that is added to your model after converting it to DDP.
+
+Under this context manager, PyTorch will skip synchronizing the gradients when `.backward()` is called, and the first call to `.backward()` outside this 
+context manager will trigger the synchronization. See an example below:
+```python
+ddp_model, dataloader = accelerator.prepare(model, dataloader)
+
+for index, batch in enumerate(dataloader):
+    input, target = batch
+    # Trigger gradient synchronization on the last batch
+    if index != (len(dataloader) - 1):
+        with ddp_model.no_sync():
+            # Gradients only accumulate
+            ddp_model(input).backward()
+    else:
+        # Gradients finally sync
+        ddp_model(input).backward()
same as above
@@ -0,0 +1,109 @@
+<!--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.
+-->
+
+# Gradient Synchronization
+
+PyTorch's distributed module operates by communicating back and forth between all of the GPUs in your system.
+This communication takes time, and ensuring all processes know the states of each other happens at particular triggerpoints
+when using the `ddp` module. 
+
+These triggerpoints are added to the PyTorch model, specifically their `forward()` and `backward()` methods. 
+This happens when the model is wrapped with `DistributedDataParallel`:
+```python
+import torch.nn as nn
+from torch.nn.parallel import DistributedDataParallel
+
+model = nn.Linear(10, 10)
+ddp_model = DistributedDataParallel(model)
+```
+In 🤗 Accelerate this conversion happens automatically when calling [`Accelerator.prepare`] and passing in your model.
+
+```diff
++ from accelerate import Accelerator
++ accelerator = Accelerator()
+  import torch.nn as nn
+- from torch.nn.parallel import DistributedDataParallel
+
+  model = nn.Linear(10,10)
++ model = accelerator.prepare(model)
+```
+
+## The slowdown in gradient accumulation
+
+You now understand that PyTorch adds hooks to the `forward` and `backward` method of your PyTorch model when 
+training in a distributed setup. But how does this risk slowing down your code?
+
+In DDP (distributed data parallel), the specific order in which processes are performed and ran are expected
+at specific points and these must also occur at roughly the same time before moving on.
+
+The most direct example is when you update all of the parameters in a model through `.backward()`. All instances of the model
+need to have updated their gradients, collated, and updated again before moving onto the next batch of data. But when performing 
+gradient accumulation, you accumulate `n` losses and skip `.backward()` until `n` batches have been reached. This 
+can cause a significant slowdown since all the processes need to communicate with them more times than needed. How 
+can you avoid this overhead?
+
+## Solving the slowdown problem
+
+Since you are skipping these batches, their gradients do not need to be synchronized until the point where `.backward()` is actually called. 
+PyTorch cannot automagically tell when you need to do this, but they do provide a tool to help through the [`no_sync`](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel.no_sync) context manager
+that is added to your model after converting it to DDP.
+
+Under this context manager, PyTorch will skip synchronizing the gradients when `.backward()` is called, and the first call to `.backward()` outside this 
+context manager will trigger the synchronization. See an example below:
+```python
+ddp_model, dataloader = accelerator.prepare(model, dataloader)
+
+for index, batch in enumerate(dataloader):
+    input, target = batch
+    # Trigger gradient synchronization on the last batch
+    if index != (len(dataloader) - 1):
+        with ddp_model.no_sync():
+            # Gradients only accumulate
+            ddp_model(input).backward()
+    else:
+        # Gradients finally sync
+        ddp_model(input).backward()
+```
+
+In 🤗 Accelerate to make this an API that can be called no matter the training device (though it may not do anything if you are not in a distributed system!),
+`ddp_model.no_sync` gets replaced with [`Accelerator.no_sync`] and operates the same way:
+
+```diff
+  ddp_model, dataloader = accelerator.prepare(model, dataloader)
+
+  for index, batch in enumerate(dataloader):
+      input, target = batch
+      # Trigger gradient synchronization on the last batch
+      if index != (len(dataloader)-1):
+-         with ddp_model.no_sync():
++         with accelerator.no_sync(model):
+              # Gradients only accumulate
+              ddp_model(input).backward()
If you want to make it a bit more concise here, you could just show the right way here and just say: if you initialize the model inside the function this might happen <show error message and explanation>.
@@ -0,0 +1,154 @@
+<!--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.
+-->
+
+# Training on TPUs with 🤗 Accelerate
+
+Training on TPUs can be slightly different than training on multi-gpu, even with 🤗 Accelerate. This guide aims to show you 
+where you should be careful and why, as well as the best practices in general.
+
+## Training in a Notebook
+
+The main carepoint when training on TPUs comes from the `notebook_launcher`. As mentioned in the [notebook tutorial](https://huggingface.co/docs/accelerate/basic_tutorials/notebook), you need to 
+restructure your training code into a function that can get passed to the `notebook_launcher` function and be careful about not declaring any tensors on the GPU.
+
+While on a TPU that last part is not as important, a critical part to understand is that when you launch code from a notebook you do so through a process called **forking**. 
+When launching from the command-line, you perform **spawning**, where a python process is not currently running and you *spawn* a new process in. Since your Jupyter notebook is already 
+utilizing a python process, you need to *fork* a new process from it to launch your code. 
+
+Where this becomes important is in regards to declaring your model. On forked TPU processes, it is recommended that you instantiate your model *once* and pass this into your 
+training function. This is different than training on GPUs where you create `n` models that have their gradients synced and back-propagated at certain moments. Instead one 
+model instance is shared between all the nodes and it is passed back and forth. This is important especially when training on low-resource TPUs such as those provided in Kaggle kernels or
+on Google Colaboratory. 
+
+Below is an example of a training function passed to the `notebook_launcher` if training on CPUs or GPUs:
+
+<Tip>
+
+This code snippet is based off the one from the `simple_nlp_example` notebook found [here](https://github.com/huggingface/notebooks/blob/main/examples/accelerate/simple_nlp_example.ipynb) with slight 
+modifications for the sake of simplicity
+
+</Tip>
+
+```python
+def training_function():
+    # Initialize accelerator
+    accelerator = Accelerator()
+    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)
+    train_dataloader, eval_dataloader = create_dataloaders(
+        train_batch_size=hyperparameters["train_batch_size"], eval_batch_size=hyperparameters["eval_batch_size"]
+    )
+
+    # Instantiate optimizer
+    optimizer = AdamW(params=model.parameters(), lr=hyperparameters["learning_rate"])
+
+    # Prepare everything
+    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
+    # prepare method.
+    model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader
+    )
+
+    num_epochs = hyperparameters["num_epochs"]
+    # Now we train the model
+    for epoch in range(num_epochs):
+        model.train()
+        for step, batch in enumerate(train_dataloader):
+            outputs = model(**batch)
+            loss = outputs.loss
+            accelerator.backward(loss)
+
+            optimizer.step()
+            optimizer.zero_grad()
+```
+
+```python
+from accelerate import notebook_launcher
+
+notebook_launcher(training_function)
+```
+
+<Tip>
+
+The `notebook_launcher` will default to 8 processes if 🤗 Accelerate has been configured for a TPU
+
+</Tip>
+
+If you use this example and declare the model *inside* the training loop, then on a low-resource system you will potentially see an error 
+like:
+
+```
+ProcessExitedException: process 0 terminated with signal SIGSEGV
+```
+
+This error is *extremely* cryptic but the basic explaination is you ran out of system RAM. You can avoid this entirely by reconfiguring the training function to 
+accept a single `model` argument, and declare it in an outside cell:
+
+```python
+# In another Jupyter cell
+model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)
+```
+
+```diff
++ def training_function(model):
+      # Initialize accelerator
+      accelerator = Accelerator()
+-     model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)
+      train_dataloader, eval_dataloader = create_dataloaders(
+          train_batch_size=hyperparameters["train_batch_size"], eval_batch_size=hyperparameters["eval_batch_size"]
+      )
+  ...
+```
+And finally calling the training function with:
+
+```diff
+  from accelerate import notebook_launcher
+- notebook_launcher(training_function)
++ notebook_launcher(training_function, (model,))
+```
+
can you maybe explain a bit what the difference to the first setting is and what the impact is and when you should use which?
@@ -0,0 +1,154 @@
+<!--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.
+-->
+
+# Training on TPUs with 🤗 Accelerate
+
+Training on TPUs can be slightly different than training on multi-gpu, even with 🤗 Accelerate. This guide aims to show you 
+where you should be careful and why, as well as the best practices in general.
+
+## Training in a Notebook
+
+The main carepoint when training on TPUs comes from the `notebook_launcher`. As mentioned in the [notebook tutorial](https://huggingface.co/docs/accelerate/basic_tutorials/notebook), you need to 
+restructure your training code into a function that can get passed to the `notebook_launcher` function and be careful about not declaring any tensors on the GPU.
+
+While on a TPU that last part is not as important, a critical part to understand is that when you launch code from a notebook you do so through a process called **forking**. 
+When launching from the command-line, you perform **spawning**, where a python process is not currently running and you *spawn* a new process in. Since your Jupyter notebook is already 
+utilizing a python process, you need to *fork* a new process from it to launch your code. 
+
+Where this becomes important is in regards to declaring your model. On forked TPU processes, it is recommended that you instantiate your model *once* and pass this into your 
+training function. This is different than training on GPUs where you create `n` models that have their gradients synced and back-propagated at certain moments. Instead one 
+model instance is shared between all the nodes and it is passed back and forth. This is important especially when training on low-resource TPUs such as those provided in Kaggle kernels or
+on Google Colaboratory. 
+
+Below is an example of a training function passed to the `notebook_launcher` if training on CPUs or GPUs:
+
+<Tip>
+
+This code snippet is based off the one from the `simple_nlp_example` notebook found [here](https://github.com/huggingface/notebooks/blob/main/examples/accelerate/simple_nlp_example.ipynb) with slight 
+modifications for the sake of simplicity
+
+</Tip>
+
+```python
+def training_function():
+    # Initialize accelerator
+    accelerator = Accelerator()
+    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)
+    train_dataloader, eval_dataloader = create_dataloaders(
+        train_batch_size=hyperparameters["train_batch_size"], eval_batch_size=hyperparameters["eval_batch_size"]
+    )
+
+    # Instantiate optimizer
+    optimizer = AdamW(params=model.parameters(), lr=hyperparameters["learning_rate"])
+
+    # Prepare everything
+    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
+    # prepare method.
+    model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader
+    )
+
+    num_epochs = hyperparameters["num_epochs"]
+    # Now we train the model
+    for epoch in range(num_epochs):
+        model.train()
+        for step, batch in enumerate(train_dataloader):
+            outputs = model(**batch)
+            loss = outputs.loss
+            accelerator.backward(loss)
+
+            optimizer.step()
+            optimizer.zero_grad()
+```
+
+```python
+from accelerate import notebook_launcher
+
+notebook_launcher(training_function)
+```
+
+<Tip>
+
+The `notebook_launcher` will default to 8 processes if 🤗 Accelerate has been configured for a TPU
+
+</Tip>
+
+If you use this example and declare the model *inside* the training loop, then on a low-resource system you will potentially see an error 
+like:
+
+```
+ProcessExitedException: process 0 terminated with signal SIGSEGV
+```
+
+This error is *extremely* cryptic but the basic explaination is you ran out of system RAM. You can avoid this entirely by reconfiguring the training function to 
+accept a single `model` argument, and declare it in an outside cell:
+
+```python
+# In another Jupyter cell
+model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)
+```
+
+```diff
++ def training_function(model):
+      # Initialize accelerator
+      accelerator = Accelerator()
+-     model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)
+      train_dataloader, eval_dataloader = create_dataloaders(
+          train_batch_size=hyperparameters["train_batch_size"], eval_batch_size=hyperparameters["eval_batch_size"]
+      )
+  ...
+```
+And finally calling the training function with:
+
+```diff
+  from accelerate import notebook_launcher
+- notebook_launcher(training_function)
++ notebook_launcher(training_function, (model,))
+```
+
+## Mixed Precision and Global Variables 
+
+As mentioned in the [mixed precision tutorial](https://huggingface.co/docs/accelerate/usage_guides/mixed_precision), 🤗 Accelerate supports fp16 and bf16, both of which can be used on TPUs.
+That being said, ideally `bf16` should be utilized as it is extremely efficient to use.
+
+There are two "layers" when using `bf16` and 🤗 Accelerate on TPUs, at the base level and at the operation level. 
+
+At the base level, this is enabled when passing `mixed_precision="bf16"` to `Accelerator`, such as:
+```python
+accelerator = Accelerator(mixed_precision="bf16")
+```
+By default this will cast `torch.float` and `torch.double` to `bfloat16` on TPUs. 
+The specific configuration being set is an environmental variable of `XLA_USE_BF16` is set to `1`.
+
+There is a futher configuration you can perform which is setting the `XLA_DOWNCAST_BF16` environmental variable. If set to `1`, then 
+`torch.float` is `bfloat16` and `torch.double` is `float32`.
It will remove the boilerplate needed when calling `torchrun` XYZ. I can expand this further to point it out, by translating the earlier bash with accelerate launch to show the equivalence. 
@@ -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.
+-->
+
+# Launching your 🤗 Accelerate scripts
+
+In the previous tutorial, you were introduced to how to modify your current training script to use 🤗 Accelerate.
+The final version of that code is shown below:
+
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    accelerator.backward(loss)
+    optimizer.step()
+    scheduler.step()
+```
+
+But how do you run this code and have it utilize the special hardware available to it?
+
+First you should rewrite the above code into a function, and make it callable as a script. For example:
+
+```diff
+  from accelerate import Accelerator
+  
++ def main():
+
+    accelerator = Accelerator()
+
+    model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+        model, optimizer, training_dataloader, scheduler
+    )
+
+    for batch in training_dataloader:
+        optimizer.zero_grad()
+        inputs, targets = batch
+        outputs = model(inputs)
+        loss = loss_function(outputs, targets)
+        accelerator.backward(loss)
+        optimizer.step()
+        scheduler.step()
+
++ if __name__ == "__main__":
++   main()
+```
+
+Next you need to launch it with `accelerate launch`. 
+
+## Using accelerate launch
+
+🤗 Accelerate has a special CLI command to help you launch your code in your system through `accelerate launch`.
+This command wraps around all of the different commands needed to launch your script on various platforms, without you
+having to remember what each of them are.
+
+<Tip>
+
+If you are familiar with launching scripts in PyTorch yourself such as with `torchrun`, you can still do this. It is not required to use `accelerate launch`
+
+</Tip>
+
+You can launch your script quickly by using:
+```bash
+accelerate launch {script_name.py} --arg1 --arg2 ...
+```
+
+Just put `accelerate launch` at the start of your command, and pass in additional arguments and parameters to your script afterwards like normal!
+
+Since this runs the various torch spawn methods, all of the expected environment variables can be modified here as well.
+For example, here is how to use `accelerate launch` with a single GPU:
+```bash
+CUDA_VISIBLE_DEVICES="0" accelerate launch {script_name.py} --arg1 --arg2 ...
+```
+
+You can also use `accelerate launch` without performing `accelerate config` first, but you may need to manually pass in the right configuration parameters.
+In this case, 🤗 Accelerate will make some hyperparameter decisions for you, such as if a GPU is available, it will use all of them by default, and no mixed precision will be used. 
+Here is how you would use all GPUs and train on no mixed precision:
+```bash
+accelerate launch --multi_gpu {script_name.py} --arg1 --arg2 ...
+```
+To get more specific you should pass in the needed parameters yourself. For instance, here is how you 
+would also launch that same script on two GPUs using mixed precision while avoiding all of the warnings: 
+```bash
+accelerate launch --multi_gpu --mixed_precision=fp16 --num_processes=2 {script_name.py} --arg1 --arg2 ...
+```
+
+For a complete list of parameters you can pass in, run:
+```bash
+accelerate launch -h
+```
+
+<Tip>
+
+Even if you are not using 🤗 Accelerate in your code, you can still use the launcher for starting your scripts!
I skipped over calculating the loss in this case, to simply show calling `.backward()`. but I'll include a loss calculation to make it a bit more clear
@@ -0,0 +1,109 @@
+<!--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.
+-->
+
+# Gradient Synchronization
+
+PyTorch's distributed module operates by communicating back and forth between all of the GPUs in your system.
+This communication takes time, and ensuring all processes know the states of each other happens at particular triggerpoints
+when using the `ddp` module. 
+
+These triggerpoints are added to the PyTorch model, specifically their `forward()` and `backward()` methods. 
+This happens when the model is wrapped with `DistributedDataParallel`:
+```python
+import torch.nn as nn
+from torch.nn.parallel import DistributedDataParallel
+
+model = nn.Linear(10, 10)
+ddp_model = DistributedDataParallel(model)
+```
+In 🤗 Accelerate this conversion happens automatically when calling [`Accelerator.prepare`] and passing in your model.
+
+```diff
++ from accelerate import Accelerator
++ accelerator = Accelerator()
+  import torch.nn as nn
+- from torch.nn.parallel import DistributedDataParallel
+
+  model = nn.Linear(10,10)
++ model = accelerator.prepare(model)
+```
+
+## The slowdown in gradient accumulation
+
+You now understand that PyTorch adds hooks to the `forward` and `backward` method of your PyTorch model when 
+training in a distributed setup. But how does this risk slowing down your code?
+
+In DDP (distributed data parallel), the specific order in which processes are performed and ran are expected
+at specific points and these must also occur at roughly the same time before moving on.
+
+The most direct example is when you update all of the parameters in a model through `.backward()`. All instances of the model
+need to have updated their gradients, collated, and updated again before moving onto the next batch of data. But when performing 
+gradient accumulation, you accumulate `n` losses and skip `.backward()` until `n` batches have been reached. This 
+can cause a significant slowdown since all the processes need to communicate with them more times than needed. How 
+can you avoid this overhead?
+
+## Solving the slowdown problem
+
+Since you are skipping these batches, their gradients do not need to be synchronized until the point where `.backward()` is actually called. 
+PyTorch cannot automagically tell when you need to do this, but they do provide a tool to help through the [`no_sync`](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel.no_sync) context manager
+that is added to your model after converting it to DDP.
+
+Under this context manager, PyTorch will skip synchronizing the gradients when `.backward()` is called, and the first call to `.backward()` outside this 
+context manager will trigger the synchronization. See an example below:
+```python
+ddp_model, dataloader = accelerator.prepare(model, dataloader)
+
+for index, batch in enumerate(dataloader):
+    input, target = batch
+    # Trigger gradient synchronization on the last batch
+    if index != (len(dataloader) - 1):
+        with ddp_model.no_sync():
+            # Gradients only accumulate
+            ddp_model(input).backward()
I made it less concise since one example shows running a regular script vs this one is for the tpu launcher in a notebook, but will think on this a bit more
@@ -0,0 +1,154 @@
+<!--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.
+-->
+
+# Training on TPUs with 🤗 Accelerate
+
+Training on TPUs can be slightly different than training on multi-gpu, even with 🤗 Accelerate. This guide aims to show you 
+where you should be careful and why, as well as the best practices in general.
+
+## Training in a Notebook
+
+The main carepoint when training on TPUs comes from the `notebook_launcher`. As mentioned in the [notebook tutorial](https://huggingface.co/docs/accelerate/basic_tutorials/notebook), you need to 
+restructure your training code into a function that can get passed to the `notebook_launcher` function and be careful about not declaring any tensors on the GPU.
+
+While on a TPU that last part is not as important, a critical part to understand is that when you launch code from a notebook you do so through a process called **forking**. 
+When launching from the command-line, you perform **spawning**, where a python process is not currently running and you *spawn* a new process in. Since your Jupyter notebook is already 
+utilizing a python process, you need to *fork* a new process from it to launch your code. 
+
+Where this becomes important is in regards to declaring your model. On forked TPU processes, it is recommended that you instantiate your model *once* and pass this into your 
+training function. This is different than training on GPUs where you create `n` models that have their gradients synced and back-propagated at certain moments. Instead one 
+model instance is shared between all the nodes and it is passed back and forth. This is important especially when training on low-resource TPUs such as those provided in Kaggle kernels or
+on Google Colaboratory. 
+
+Below is an example of a training function passed to the `notebook_launcher` if training on CPUs or GPUs:
+
+<Tip>
+
+This code snippet is based off the one from the `simple_nlp_example` notebook found [here](https://github.com/huggingface/notebooks/blob/main/examples/accelerate/simple_nlp_example.ipynb) with slight 
+modifications for the sake of simplicity
+
+</Tip>
+
+```python
+def training_function():
+    # Initialize accelerator
+    accelerator = Accelerator()
+    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)
+    train_dataloader, eval_dataloader = create_dataloaders(
+        train_batch_size=hyperparameters["train_batch_size"], eval_batch_size=hyperparameters["eval_batch_size"]
+    )
+
+    # Instantiate optimizer
+    optimizer = AdamW(params=model.parameters(), lr=hyperparameters["learning_rate"])
+
+    # Prepare everything
+    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
+    # prepare method.
+    model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader
+    )
+
+    num_epochs = hyperparameters["num_epochs"]
+    # Now we train the model
+    for epoch in range(num_epochs):
+        model.train()
+        for step, batch in enumerate(train_dataloader):
+            outputs = model(**batch)
+            loss = outputs.loss
+            accelerator.backward(loss)
+
+            optimizer.step()
+            optimizer.zero_grad()
+```
+
+```python
+from accelerate import notebook_launcher
+
+notebook_launcher(training_function)
+```
+
+<Tip>
+
+The `notebook_launcher` will default to 8 processes if 🤗 Accelerate has been configured for a TPU
+
+</Tip>
+
+If you use this example and declare the model *inside* the training loop, then on a low-resource system you will potentially see an error 
+like:
+
+```
+ProcessExitedException: process 0 terminated with signal SIGSEGV
+```
+
+This error is *extremely* cryptic but the basic explaination is you ran out of system RAM. You can avoid this entirely by reconfiguring the training function to 
+accept a single `model` argument, and declare it in an outside cell:
+
+```python
+# In another Jupyter cell
+model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)
+```
+
+```diff
++ def training_function(model):
+      # Initialize accelerator
+      accelerator = Accelerator()
+-     model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)
+      train_dataloader, eval_dataloader = create_dataloaders(
+          train_batch_size=hyperparameters["train_batch_size"], eval_batch_size=hyperparameters["eval_batch_size"]
+      )
+  ...
+```
+And finally calling the training function with:
+
+```diff
+  from accelerate import notebook_launcher
+- notebook_launcher(training_function)
++ notebook_launcher(training_function, (model,))
+```
+
Thanks for pointing this out! Will make sure to have a `notebook.mdx` as part of the `basic_tutorials` before merging
@@ -0,0 +1,19 @@
+<!--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.
+-->
+
+# Launchers
+
+Functions for launching training on distributed processes.
Accident 😓 
@@ -24,8 +24,8 @@ which can total to adding just one new line of code!
 
 This example will use a very simplistic PyTorch training loop that performs gradient accumulation every two batches:
 
-```python
-device = "cuda"
+```
With curly brackets the reader knows right away that it's meant to represent templating, though I didn't propagate this well enough through the docs so will make sure to do so
@@ -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.
+-->
+
+# Launching your 🤗 Accelerate scripts
+
+In the previous tutorial, you were introduced to how to modify your current training script to use 🤗 Accelerate.
+The final version of that code is shown below:
+
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    accelerator.backward(loss)
+    optimizer.step()
+    scheduler.step()
+```
+
+But how do you run this code and have it utilize the special hardware available to it?
+
+First you should rewrite the above code into a function, and make it callable as a script. For example:
+
+```diff
+  from accelerate import Accelerator
+  
++ def main():
+
+    accelerator = Accelerator()
+
+    model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+        model, optimizer, training_dataloader, scheduler
+    )
+
+    for batch in training_dataloader:
+        optimizer.zero_grad()
+        inputs, targets = batch
+        outputs = model(inputs)
+        loss = loss_function(outputs, targets)
+        accelerator.backward(loss)
+        optimizer.step()
+        scheduler.step()
+
++ if __name__ == "__main__":
++   main()
+```
+
+Next you need to launch it with `accelerate launch`. 
+
+## Using accelerate launch
+
+🤗 Accelerate has a special CLI command to help you launch your code in your system through `accelerate launch`.
+This command wraps around all of the different commands needed to launch your script on various platforms, without you
+having to remember what each of them are.
+
+<Tip>
+
+If you are familiar with launching scripts in PyTorch yourself such as with `torchrun`, you can still do this. It is not required to use `accelerate launch`
+
+</Tip>
+
+You can launch your script quickly by using:
+```bash
+accelerate launch {script_name.py} --arg1 --arg2 ...
Move the file back to the root maybe? It would be more consistent.
@@ -1,38 +1,68 @@
-- sections: 
+- sections:
   - local: index
     title: 🤗 Accelerate
-  - local: quicktour
-    title: Quick tour
-  - local: installation
+  - local: basic_tutorials/install
_The documentation is not available anymore as the PR was closed or merged._
@pacman100 the difference between the two is something like the following:
- How-To guides are how do you do things with **Accelerate**, e.g. how to perform gradient accumulation, how to perform logging, how to perform tracking (think our example scripts!)
- Concepts and Fundamentals is given this thing accelerate does, here's an understanding of some of the harder-to-grasp magic it does underneath that you should know more about as you become more familiar with the library. E.g. Gradient synchronization is how we do gradient accumulation, so that's one example where we took an applied concept taught in the higher-level tutorial and took it to a lower level.

- Re: example scripts, it'll be in this PR eventually, if not a follow up. Still thinking on a good design for integration. Right now this is a very tl;dr version to make sure everyone is aligned (which we are!) before getting too deep, and now I'll move towards migrating the full docs in (avoiding your reference issues, etc, etc)
https://api.github.com/repos/huggingface/accelerate/pulls/495
diff --git a/examples/requirements.txt b/examples/requirements.txt
index 65b6b04c3..912986bf1 100644
--- a/examples/requirements.txt
+++ b/examples/requirements.txt
@@ -1 +1,3 @@
-accelerate # used to be installed in Amazon SageMaker environment
\ No newline at end of file
+accelerate # used to be installed in Amazon SageMaker environment
+evaluate
+datasets==2.3.2
\ No newline at end of file
diff --git a/src/accelerate/commands/config/config_args.py b/src/accelerate/commands/config/config_args.py
index fcb527aeb..fe3d05bf3 100644
--- a/src/accelerate/commands/config/config_args.py
+++ b/src/accelerate/commands/config/config_args.py
@@ -23,6 +23,7 @@
 import yaml
 
 from ...utils import ComputeEnvironment, DistributedType, SageMakerDistributedType
+from ...utils.constants import SAGEMAKER_PYTHON_VERSION, SAGEMAKER_PYTORCH_VERSION, SAGEMAKER_TRANSFORMERS_VERSION
 
 
 hf_cache_home = os.path.expanduser(
@@ -123,7 +124,10 @@ def __post_init__(self):
         if isinstance(self.compute_environment, str):
             self.compute_environment = ComputeEnvironment(self.compute_environment)
         if isinstance(self.distributed_type, str):
-            self.distributed_type = DistributedType(self.distributed_type)
+            if self.compute_environment == ComputeEnvironment.AMAZON_SAGEMAKER:
+                self.distributed_type = SageMakerDistributedType(self.distributed_type)
+            else:
+                self.distributed_type = DistributedType(self.distributed_type)
 
 
 @dataclass
@@ -156,5 +160,6 @@ class SageMakerConfig(BaseConfig):
     region: str = "us-east-1"
     num_machines: int = 1
     base_job_name: str = f"accelerate-sagemaker-{num_machines}"
-    pytorch_version: str = "1.6"
-    transformers_version: str = "4.4"
+    pytorch_version: str = SAGEMAKER_PYTORCH_VERSION
+    transformers_version: str = SAGEMAKER_TRANSFORMERS_VERSION
+    py_version: str = SAGEMAKER_PYTHON_VERSION
diff --git a/src/accelerate/commands/config/sagemaker.py b/src/accelerate/commands/config/sagemaker.py
index f54659548..c0b17bc88 100644
--- a/src/accelerate/commands/config/sagemaker.py
+++ b/src/accelerate/commands/config/sagemaker.py
@@ -16,6 +16,7 @@
 import json
 import os
 
+from ...utils.constants import SAGEMAKER_PARALLEL_EC2_INSTANCES
 from ...utils.dataclasses import ComputeEnvironment, SageMakerDistributedType
 from ...utils.imports import is_boto3_available
 from .config_args import SageMakerConfig
@@ -120,23 +121,31 @@ def get_sagemaker_input():
         _create_iam_role_for_sagemaker(iam_role_name)
 
     distributed_type = _ask_field(
-        "Which type of machine are you using? ([0] No distributed training, [1] data parallelism, [2] model parallelism): ",
+        "Which type of machine are you using? ([0] No distributed training, [1] data parallelism): ",
         _convert_sagemaker_distributed_mode,
-        error_message="Please enter 0, 1 or 2",
+        error_message="Please enter 0 or 1",
     )
 
-    # using the best two instances for single-gpu training or multi-gpu -> can turn into question to make it more diverse
-    ec2_instance_type = "ml.p3.2xlarge" if distributed_type == SageMakerDistributedType.NO else "ml.p3dn.24xlarge"
+    ec2_instance_query = "Which EC2 instance type you want to use for your training "
+    if distributed_type != SageMakerDistributedType.NO:
+        ec2_instance_query += "("
+        for i, instance_type in enumerate(SAGEMAKER_PARALLEL_EC2_INSTANCES):
+            ec2_instance_query += f"[{i}] {instance_type}, "
+        ec2_instance_query = ec2_instance_query[:-2] + ")? [0]: "
+        ec2_instance_type = _ask_field(ec2_instance_query, lambda x: SAGEMAKER_PARALLEL_EC2_INSTANCES[int(x)])
+    else:
+        ec2_instance_query += "? [ml.p3.2xlarge]:"
+        ec2_instance_type = _ask_field(ec2_instance_query, lambda x: str(x).lower(), default="ml.p3.2xlarge")
+
     num_machines = 1
     if (
         distributed_type == SageMakerDistributedType.DATA_PARALLEL
         or distributed_type == SageMakerDistributedType.MODEL_PARALLEL
     ):
-        raise NotImplementedError("Model or Data Parallelism is not implemented yet. We are working on it")
         num_machines = _ask_field(
-            "How many machines do you want use? [2]: ",
+            "How many machines do you want use? [1]: ",
             lambda x: int(x),
-            default=2,
+            default=1,
         )
 
     mixed_precision = _ask_field(
@@ -148,6 +157,7 @@ def get_sagemaker_input():
     return SageMakerConfig(
         compute_environment=ComputeEnvironment.AMAZON_SAGEMAKER,
         distributed_type=distributed_type,
+        use_cpu=False,
         ec2_instance_type=ec2_instance_type,
         profile=aws_profile,
         region=aws_region,
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 2bc6fbe4d..421bfe064 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -35,6 +35,7 @@
     is_deepspeed_available,
     is_sagemaker_available,
 )
+from accelerate.utils.dataclasses import SageMakerDistributedType
 
 
 def launch_command_parser(subparsers=None):
@@ -486,9 +487,15 @@ def sagemaker_launcher(sagemaker_config: SageMakerConfig, args):
         mixed_precision = "fp16"
 
     # Environment variables to be set for use during training job
-    environment = {"MIXED_PRECISION": str(mixed_precision)}
+    environment = {
+        "USE_SAGEMAKER": "true",
+        "MIXED_PRECISION": str(mixed_precision),
+        "SAGEMAKER_DISTRIBUTED_TYPE": sagemaker_config.distributed_type.value,
+    }
     # configure distribution set up
     distribution = None  # TODO: not yet implemented
+    if sagemaker_config.distributed_type == SageMakerDistributedType.DATA_PARALLEL:
+        distribution = {"smdistributed": {"dataparallel": {"enabled": True}}}
 
     # configure session
     print("Creating Estimator")
@@ -496,9 +503,9 @@ def sagemaker_launcher(sagemaker_config: SageMakerConfig, args):
         entry_point=entry_point,
         source_dir=source_dir,
         role=sagemaker_config.iam_role_name,
-        transformers_version="4.4",
-        pytorch_version="1.6",
-        py_version="py36",
+        transformers_version=sagemaker_config.transformers_version,
+        pytorch_version=sagemaker_config.pytorch_version,
+        py_version=sagemaker_config.py_version,
         base_job_name=sagemaker_config.base_job_name,
         instance_count=sagemaker_config.num_machines,
         instance_type=sagemaker_config.ec2_instance_type,
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index 330830eb5..4482341f3 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -18,6 +18,7 @@
 import torch
 
 from .utils import DistributedType, is_ccl_available, is_deepspeed_available, is_tpu_available
+from .utils.dataclasses import SageMakerDistributedType
 
 
 if is_tpu_available(check_device=False):
@@ -89,7 +90,25 @@ def __init__(
                     "Please make sure to properly initialize your accelerator via `accelerator = Accelerator()` "
                     "before using any functionality from the `accelerate` library."
                 )
-            if is_tpu_available() and not cpu:
+            if (
+                os.environ.get("USE_SAGEMAKER", "false") == "true"
+                and os.environ.get("SAGEMAKER_DISTRIBUTED_TYPE") != SageMakerDistributedType.NO
+                and not cpu
+            ):
+                if os.environ.get("SAGEMAKER_DISTRIBUTED_TYPE") == SageMakerDistributedType.DATA_PARALLEL:
+                    self.distributed_type = DistributedType.MULTI_GPU
+                    import smdistributed.dataparallel.torch.torch_smddp  # noqa
+
+                    if not torch.distributed.is_initialized():
+                        torch.distributed.init_process_group(backend="smddp")
+                    self.backend = "smddp"
+                    self.num_processes = torch.distributed.get_world_size()
+                    self.process_index = torch.distributed.get_rank()
+                    self.local_process_index = int(os.environ.get("LOCAL_RANK", -1))
+                    self.device = torch.device("cuda", self.local_process_index)
+                    torch.cuda.set_device(self.device)
+                    self.mixed_precision = mixed_precision
+            elif is_tpu_available() and not cpu:
                 self.distributed_type = DistributedType.TPU
                 self.num_processes = xm.xrt_world_size()
                 self.process_index = xm.get_ordinal()
diff --git a/src/accelerate/utils/constants.py b/src/accelerate/utils/constants.py
index af8a6bf35..173cfa790 100644
--- a/src/accelerate/utils/constants.py
+++ b/src/accelerate/utils/constants.py
@@ -20,5 +20,9 @@
 RNG_STATE_NAME = "random_states"
 OPTIMIZER_NAME = "optimizer"
 SCHEDULER_NAME = "scheduler"
+SAGEMAKER_PYTORCH_VERSION = "1.10.2"
+SAGEMAKER_PYTHON_VERSION = "py38"
+SAGEMAKER_TRANSFORMERS_VERSION = "4.17.0"
+SAGEMAKER_PARALLEL_EC2_INSTANCES = ["ml.p3.16xlarge", "ml.p3dn.24xlarge", "ml.p4dn.24xlarge"]
 
 STR_OPERATION_TO_FUNC = {">": op.gt, ">=": op.ge, "==": op.eq, "!=": op.ne, "<=": op.le, "<": op.lt}

Would remove the choice and error message for now as I don't think we will add this any time soon :-)
@@ -89,7 +90,27 @@ def __init__(
                     "Please make sure to properly initialize your accelerator via `accelerator = Accelerator()` "
                     "before using any functionality from the `accelerate` library."
                 )
-            if is_tpu_available() and not cpu:
+            if (
+                os.environ.get("USE_SAGEMAKER", "false") == "true"
+                and os.environ.get("SAGEMAKER_DISTRIBUTED_TYPE") != SageMakerDistributedType.NO
+                and not cpu
+            ):
+                if os.environ.get("SAGEMAKER_DISTRIBUTED_TYPE") == SageMakerDistributedType.DATA_PARALLEL:
+                    self.distributed_type = DistributedType.MULTI_GPU
+                    import smdistributed.dataparallel.torch.torch_smddp  # noqa
+
+                    if not torch.distributed.is_initialized():
+                        torch.distributed.init_process_group(backend="smddp")
+                    self.backend = "smddp"
+                    self.num_processes = torch.distributed.get_world_size()
+                    self.process_index = torch.distributed.get_rank()
+                    self.local_process_index = int(os.environ.get("LOCAL_RANK", -1))
+                    self.device = torch.device("cuda", self.local_process_index)
+                    torch.cuda.set_device(self.device)
+                    self.mixed_precision = mixed_precision
+                elif os.environ.get("SAGEMAKER_DISTRIBUTED_TYPE") == SageMakerDistributedType.MODEL_PARALLEL:
+                    raise NotImplementedError("Model Parallelism is not implemented yet. We are working on it")
_The documentation is not available anymore as the PR was closed or merged._
I will test this with my use case as soon as it gets merged into `main`
> LGTM! How should we maintain the `SAGEMAKER_PYTHON_VERSION, SAGEMAKER_PYTORCH_VERSION, SAGEMAKER_TRANSFORMERS_VERSION` versions once we release new DLCs?

Hello, as it is not a frequent change, it can be done manually for time being whenever new releases of DLC happen. Any suggestions/best practices for automating it? 
Hello @pacman100, what is the sagemaker SDK estimator supposed to be used with accelerate?
https://api.github.com/repos/huggingface/accelerate/pulls/494
diff --git a/docs/source/internal.mdx b/docs/source/internal.mdx
index 8da23bee8..4ff1d6ff9 100644
--- a/docs/source/internal.mdx
+++ b/docs/source/internal.mdx
@@ -26,7 +26,7 @@ The main work on your PyTorch `DataLoader` is done by the following function:
 
 [[autodoc]] data_loader.prepare_data_loader
 
-### BatchSamplerShard
+### DataLoaderShard
 
 [[autodoc]] data_loader.DataLoaderShard
 
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index 524034eb1..f58752e53 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -38,6 +38,22 @@
 if is_tpu_available(check_device=False):
     import torch_xla.distributed.parallel_loader as xpl
 
+    class MpDeviceLoaderWrapper(xpl.MpDeviceLoader):
+        """
+        Wrapper for the xpl.MpDeviceLoader class that knows the total batch size.
+
+        **Available attributes:**
+
+        - **total_batch_size** (`int`) -- Total batch size of the dataloader across all processes.
+            Equal to the original batch size when `split_batches=True`; otherwise the original batch size * the total
+            number of processes
+        """
+
+        @property
+        def total_batch_size(self):
+            return self._loader.total_batch_size
+
+
 logger = get_logger(__name__)
 
 # kwargs of the DataLoader in min version 1.4.0.
@@ -289,6 +305,12 @@ class DataLoaderShard(DataLoader):
             A random number generator to keep synchronized across processes.
         kwargs:
             All other keyword arguments to pass to the regular `DataLoader` initialization.
+
+    **Available attributes:**
+
+        - **total_batch_size** (`int`) -- Total batch size of the dataloader across all processes.
+            Equal to the original batch size when `split_batches=True`; otherwise the original batch size * the total
+            number of processes
     """
 
     def __init__(self, dataset, device=None, rng_types=None, generator=None, **kwargs):
@@ -321,6 +343,14 @@ def __iter__(self):
                 yield current_batch
                 break
 
+    @property
+    def total_batch_size(self):
+        return (
+            self.batch_sampler.batch_size
+            if self.batch_sampler.split_batches
+            else (self.batch_sampler.batch_size * self.batch_sampler.num_processes)
+        )
+
 
 class DataLoaderDispatcher(DataLoader):
     """
@@ -334,6 +364,12 @@ class DataLoaderDispatcher(DataLoader):
             the same as the initial `dataloader` if this option is set to `True`, the batch size of the initial
             `dataloader` multiplied by `num_processes` otherwise. Setting this option to `True` requires that the batch
             size of the `dataloader` is a round multiple of `batch_size`.
+
+    **Available attributes:**
+
+        - **total_batch_size** (`int`) -- Total batch size of the dataloader across all processes.
+            Equal to the original batch size when `split_batches=True`; otherwise the original batch size * the total
+            number of processes
     """
 
     def __init__(self, dataset, split_batches: bool = False, **kwargs):
@@ -432,6 +468,12 @@ def __len__(self):
         else:
             return math.ceil(whole_length / self.state.num_processes)
 
+    @property
+    def total_batch_size(self):
+        return (
+            self.dataset.batch_size if self.split_batches else (self.dataset.batch_size * self.dataset.num_processes)
+        )
+
 
 def prepare_data_loader(
     dataloader: DataLoader,
@@ -577,7 +619,10 @@ def prepare_data_loader(
 
     if dispatch_batches:
         dataloader = DataLoaderDispatcher(
-            new_dataset, split_batches=split_batches, batch_sampler=new_batch_sampler, **kwargs
+            new_dataset,
+            split_batches=split_batches,
+            batch_sampler=new_batch_sampler,
+            **kwargs,
         )
     else:
         dataloader = DataLoaderShard(
@@ -590,5 +635,5 @@ def prepare_data_loader(
         )
 
     if state.distributed_type == DistributedType.TPU:
-        return xpl.MpDeviceLoader(dataloader, device)
+        return MpDeviceLoaderWrapper(dataloader, device)
     return dataloader

Here we should probably define a `property for `total_batch_size` that goes read it in the `loader`. It's not something that should be settable anyway.
@@ -38,6 +38,12 @@
 if is_tpu_available(check_device=False):
     import torch_xla.distributed.parallel_loader as xpl
 
+    class MpDeviceLoaderWrapper(xpl.MpDeviceLoader):
+        def __init__(self, loader, device, **kwargs):
+            super().__init__(loader, device, **kwargs)
+            self.total_batch_size = loader.total_batch_size
Same comments here except I think `total_batch_size` can be computed in a property since we have access to the `batch_size` and `split_batches`.
@@ -336,7 +343,7 @@ class DataLoaderDispatcher(DataLoader):
             size of the `dataloader` is a round multiple of `batch_size`.
     """
 
-    def __init__(self, dataset, split_batches: bool = False, **kwargs):
+    def __init__(self, dataset, split_batches: bool = False, total_batch_size=None, **kwargs):
```suggestion
```
Can be removed entirely ;-)
@@ -38,6 +38,15 @@
 if is_tpu_available(check_device=False):
     import torch_xla.distributed.parallel_loader as xpl
 
+    class MpDeviceLoaderWrapper(xpl.MpDeviceLoader):
+        def __init__(self, loader, device, **kwargs):
+            super().__init__(loader, device, **kwargs)
Let's add a docstring to the class here stating briefly what it is and the new property, mimicking how [Accelerator](https://github.com/huggingface/accelerate/blob/docs-revamp/src/accelerate/accelerator.py#L140-L152) has it, since this will be useful for folks and it should show in the documentation as it's different from the base torch dataloader.

Similarly add this property doc to the other two classes as well.
@@ -38,6 +38,12 @@
 if is_tpu_available(check_device=False):
     import torch_xla.distributed.parallel_loader as xpl
 
+    class MpDeviceLoaderWrapper(xpl.MpDeviceLoader):
+        @property
+        def total_batch_size(self):
+            return self._loader.total_batch_size
```suggestion
        """
        Wrapper for the xpl.MpDeviceLoader class that knows the total batch size. 
        **Available attributes:**
        - **total_batch_size** (`int`) -- Total batch size of the dataloader across all processes. Equal to the original batch size when `split_batches=True` otherwise the original batch size * the total number of processes
        """
```

Since property documentation doesn't get automatically documented when we do autodoc, it's better to include it in the main class docstring instead. 
@@ -39,8 +39,18 @@
     import torch_xla.distributed.parallel_loader as xpl
 
     class MpDeviceLoaderWrapper(xpl.MpDeviceLoader):
+        """
+        Wrapper for the xpl.MpDeviceLoader class. This class is used to add `total_batch_size` property to the
+        xpl.MpDeviceLoader class.
+        """
+
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/493
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 7e99c76ef..7f04c719a 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -977,7 +977,7 @@ def save_state(self, output_dir: str):
         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]
+        weights = [self.get_state_dict(m, unwrap=False) for m in self._models]
         save_location = save_accelerator_state(
             output_dir, weights, self._optimizers, self._schedulers, self.state.process_index, self.scaler
         )
@@ -1056,7 +1056,7 @@ def _get_devices(self, *args):
                         break
         return (model_device, optimizer_device)
 
-    def get_state_dict(self, model):
+    def get_state_dict(self, model, unwrap=True):
         is_zero_3 = False
         if self.distributed_type == DistributedType.DEEPSPEED:
             is_zero_3 = self.deepspeed_config["zero_optimization"]["stage"] == 3
@@ -1072,7 +1072,8 @@ def get_state_dict(self, model):
                     "To save the full checkpoint, run `model.save_checkpoint(save_dir)` and use `zero_to_fp32.py` to recover weights."
                 )
         else:
-            model = self.unwrap_model(model)
+            if unwrap:
+                model = self.unwrap_model(model)
             state_dict = model.state_dict()
 
         if state_dict is not None:

_The documentation is not available anymore as the PR was closed or merged._
Since this went unnoticed, might be good to add in a deepspeed test to the checkpointing test available here: https://github.com/huggingface/accelerate/blob/main/tests/test_state_checkpointing.py

If you think you can manage that feel free to add it to this PR, worst case I can add it in as a follow up
Thanks! I think I can add a test tomorrow.
Hi, although this solves the model weights issue, but now it fails when loading optimizer states.

Because deepspeed `optimizer.state_dict()` returns a dictionary.
https://github.com/microsoft/DeepSpeed/blob/9305916d6b858b20fca3ebe03940ceaaa13cdfbd/deepspeed/runtime/zero/stage_1_and_2.py#L1998

But deepspeed `optimizer.load_state_dict()`  expects a list of dictionary.
https://github.com/microsoft/DeepSpeed/blob/9305916d6b858b20fca3ebe03940ceaaa13cdfbd/deepspeed/runtime/zero/stage_1_and_2.py#L2135-L2167
because deepspeed `load_checkpoint` actually gets multiple state_dict (one for each partition or something) https://github.com/microsoft/DeepSpeed/blob/9305916d6b858b20fca3ebe03940ceaaa13cdfbd/deepspeed/runtime/engine.py#L2650

So I think I will just use this
```
if accelerator.distributed_type == DistributedType.DEEPSPEED:
    model.load_checkpoint(path)
    model.save_checkpoint(path)
else:
    accelerator.load_state(path)
    accelerator.save_state(path)
```
Since the test I tried to add fails when loading optimizer, I won't push it here. Please feel free to merge or close this PR.
Let's merge this first fix (thanks a lot!). Cc @pacman100 for the second test.
https://api.github.com/repos/huggingface/accelerate/pulls/489
diff --git a/examples/by_feature/fsdp_with_peak_mem_tracking.py b/examples/by_feature/fsdp_with_peak_mem_tracking.py
index 8e3b9e130..09199154a 100644
--- a/examples/by_feature/fsdp_with_peak_mem_tracking.py
+++ b/examples/by_feature/fsdp_with_peak_mem_tracking.py
@@ -110,10 +110,9 @@ def training_function(config, args):
     batch_size = int(config["batch_size"])
 
     # We need to initialize the trackers we use, and also store our configuration
-    if args.with_tracking:
-        if accelerator.is_main_process:
-            experiment_config = vars(args)
-            accelerator.init_trackers("fsdp_glue_no_trainer", experiment_config)
+    if args.with_tracking and accelerator.is_main_process:
+        experiment_config = vars(args)
+        accelerator.init_trackers("fsdp_glue_no_trainer", experiment_config)
 
     tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
     datasets = load_dataset("glue", "mrpc")
diff --git a/examples/by_feature/tracking.py b/examples/by_feature/tracking.py
index 78d4af422..00c8665dd 100644
--- a/examples/by_feature/tracking.py
+++ b/examples/by_feature/tracking.py
@@ -162,10 +162,9 @@ def training_function(config, args):
 
     # New Code #
     # We need to initalize the trackers we use. Overall configurations can also be stored
-    if args.with_tracking:
-        if accelerator.is_main_process:
-            run = os.path.split(__file__)[-1].split(".")[0]
-            accelerator.init_trackers(run, config)
+    if args.with_tracking and accelerator.is_main_process:
+        run = os.path.split(__file__)[-1].split(".")[0]
+        accelerator.init_trackers(run, config)
 
     # Now we train the model
     for epoch in range(num_epochs):
diff --git a/examples/complete_cv_example.py b/examples/complete_cv_example.py
index 8f893b7c4..17cc83170 100644
--- a/examples/complete_cv_example.py
+++ b/examples/complete_cv_example.py
@@ -103,12 +103,11 @@ def training_function(config, args):
         checkpointing_steps = None
 
     # We need to initialize the trackers we use, and also store our configuration
-    if args.with_tracking:
-        if accelerator.is_main_process:
-            run = os.path.split(__file__)[-1].split(".")[0]
-            if args.logging_dir:
-                run = os.path.join(args.logging_dir, run)
-            accelerator.init_trackers(run, config)
+    if args.with_tracking and accelerator.is_main_process:
+        run = os.path.split(__file__)[-1].split(".")[0]
+        if args.logging_dir:
+            run = os.path.join(args.logging_dir, run)
+        accelerator.init_trackers(run, config)
 
     # Grab all the image filenames
     file_names = [os.path.join(args.data_dir, fname) for fname in os.listdir(args.data_dir) if fname.endswith(".jpg")]
diff --git a/examples/complete_nlp_example.py b/examples/complete_nlp_example.py
index 572bc9a23..dc0cf43ed 100644
--- a/examples/complete_nlp_example.py
+++ b/examples/complete_nlp_example.py
@@ -75,12 +75,11 @@ def training_function(config, args):
     batch_size = int(config["batch_size"])
 
     # We need to initialize the trackers we use, and also store our configuration
-    if args.with_tracking:
-        if accelerator.is_main_process:
-            run = os.path.split(__file__)[-1].split(".")[0]
-            if args.logging_dir:
-                run = os.path.join(args.logging_dir, run)
-            accelerator.init_trackers(run, config)
+    if args.with_tracking and accelerator.is_main_process:
+        run = os.path.split(__file__)[-1].split(".")[0]
+        if args.logging_dir:
+            run = os.path.join(args.logging_dir, run)
+        accelerator.init_trackers(run, config)
 
     tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
     datasets = load_dataset("glue", "mrpc")
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 35afd6c5d..5f32a83d1 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -19,6 +19,7 @@
 import sys
 import warnings
 from contextlib import contextmanager
+from functools import wraps
 from typing import List, Optional, Union
 
 import torch
@@ -356,23 +357,59 @@ def mixed_precision(self):
             mixed_precision = self.state.mixed_precision
         return mixed_precision
 
-    @contextmanager
-    def local_main_process_first(self):
+    def on_main_process(func):
+        """
+        A decorator that will run the decorated function on the main process only.
         """
-        Lets the local main process go inside a with block.
 
-        The other processes will enter the with block after the main process exits.
+        @wraps(func)
+        def wrapper(self, *args, **kwargs):
+            if self.is_main_process or not self.use_distributed:
+                return func(self, *args, **kwargs)
+
+        return wrapper
+
+    def on_local_main_process(func):
+        """
+        A decorator that will run the decorated function on the local main process only.
         """
-        yield from self._goes_first(self.is_local_main_process)
 
-    @contextmanager
-    def main_process_first(self):
+        @wraps(func)
+        def wrapper(self, *args, **kwargs):
+            if self.is_local_main_process or not self.use_distributed:
+                return func(self, *args, **kwargs)
+
+        return wrapper
+
+    def on_process(process_idx):
+        """
+        A decorator that will run the decorated function on a given process index only.
         """
-        Lets the main process go first inside a with block.
 
-        The other processes will enter the with block after the main process exits.
+        def decorator(func):
+            @wraps(func)
+            def wrapper(self, *args, **kwargs):
+                if self.process_idx == process_idx or not self.use_distributed:
+                    return func(self, *args, **kwargs)
+
+            return wrapper
+
+        return decorator
+
+    def on_local_process(local_process_idx):
         """
-        yield from self._goes_first(self.is_main_process)
+        Run func on certain local process only
+        """
+
+        def decorator(func):
+            @wraps(func)
+            def wrapper(self, *args, **kwargs):
+                if self.local_process_idx == local_process_idx or not self.use_distributed:
+                    return func(self, *args, **kwargs)
+
+            return wrapper
+
+        return decorator
 
     def _goes_first(self, is_main):
         if not is_main:
@@ -383,6 +420,24 @@ def _goes_first(self, is_main):
         if is_main:
             self.wait_for_everyone()
 
+    @contextmanager
+    def main_process_first(self):
+        """
+        Lets the main process go first inside a with block.
+
+        The other processes will enter the with block after the main process exits.
+        """
+        yield from self._goes_first(self.is_main_process)
+
+    @contextmanager
+    def local_main_process_first(self):
+        """
+        Lets the local main process go inside a with block.
+
+        The other processes will enter the with block after the main process exits.
+        """
+        yield from self._goes_first(self.is_local_main_process)
+
     @contextmanager
     def no_sync(self, model):
         """
@@ -991,6 +1046,7 @@ def init_trackers(self, project_name: str, config: Optional[dict] = None, init_k
             for tracker in self.trackers:
                 tracker.store_init_configuration(config)
 
+    @on_main_process
     def log(self, values: dict, step: Optional[int] = None, log_kwargs: Optional[dict] = {}):
         """
         Logs `values` to all stored trackers in `self.trackers`.
@@ -1007,17 +1063,16 @@ def log(self, values: dict, step: Optional[int] = None, log_kwargs: Optional[dic
                 {"wandb": {"tags": ["tag_a", "tag_b"]}}
                 ```
         """
-        if self.is_main_process:
-            for tracker in self.trackers:
-                tracker.log(values, step=step, **log_kwargs.get(tracker.name, {}))
+        for tracker in self.trackers:
+            tracker.log(values, step=step, **log_kwargs.get(tracker.name, {}))
 
+    @on_main_process
     def end_training(self):
         """
         Runs any special end training behaviors, such as stopping trackers
         """
-        if self.is_main_process:
-            for tracker in self.trackers:
-                tracker.finish()
+        for tracker in self.trackers:
+            tracker.finish()
 
     def save(self, obj, f):
         """

```suggestion
        A decorator that will run the decorated function on the main process only.
```
@@ -356,23 +357,59 @@ def mixed_precision(self):
             mixed_precision = self.state.mixed_precision
         return mixed_precision
 
-    @contextmanager
-    def local_main_process_first(self):
+    def on_main_process(func):
+        """
+        Run func on main process only
```suggestion
        A decorator that will run the decorated function on the local main process only.
```
@@ -356,23 +357,59 @@ def mixed_precision(self):
             mixed_precision = self.state.mixed_precision
         return mixed_precision
 
-    @contextmanager
-    def local_main_process_first(self):
+    def on_main_process(func):
+        """
+        Run func on main process only
         """
-        Lets the local main process go inside a with block.
 
-        The other processes will enter the with block after the main process exits.
+        @wraps(func)
+        def wrapper(self, *args, **kwargs):
+            if self.is_main_process or not self.use_distributed:
+                return func(self, *args, **kwargs)
+
+        return wrapper
+
+    def on_local_main_process(func):
+        """
+        Run func on local main process only
```suggestion
        A decorator that will run the decorated function on a given process index only.
```
@@ -356,23 +357,59 @@ def mixed_precision(self):
             mixed_precision = self.state.mixed_precision
         return mixed_precision
 
-    @contextmanager
-    def local_main_process_first(self):
+    def on_main_process(func):
+        """
+        Run func on main process only
         """
-        Lets the local main process go inside a with block.
 
-        The other processes will enter the with block after the main process exits.
+        @wraps(func)
+        def wrapper(self, *args, **kwargs):
+            if self.is_main_process or not self.use_distributed:
+                return func(self, *args, **kwargs)
+
+        return wrapper
+
+    def on_local_main_process(func):
+        """
+        Run func on local main process only
         """
-        yield from self._goes_first(self.is_local_main_process)
 
-    @contextmanager
-    def main_process_first(self):
+        @wraps(func)
+        def wrapper(self, *args, **kwargs):
+            if self.is_local_main_process or not self.use_distributed:
+                return func(self, *args, **kwargs)
+
+        return wrapper
+
+    def on_process(process_idx):
+        """
+        Run func on certain process only
```suggestion
    def on_process(process_idx, local=False):
```
Maybe we could group this one and the text in one decorator since it's one that takes arguments? 
@@ -356,23 +357,59 @@ def mixed_precision(self):
             mixed_precision = self.state.mixed_precision
         return mixed_precision
 
-    @contextmanager
-    def local_main_process_first(self):
+    def on_main_process(func):
+        """
+        Run func on main process only
         """
-        Lets the local main process go inside a with block.
 
-        The other processes will enter the with block after the main process exits.
+        @wraps(func)
+        def wrapper(self, *args, **kwargs):
+            if self.is_main_process or not self.use_distributed:
+                return func(self, *args, **kwargs)
+
+        return wrapper
+
+    def on_local_main_process(func):
+        """
+        Run func on local main process only
         """
-        yield from self._goes_first(self.is_local_main_process)
 
-    @contextmanager
-    def main_process_first(self):
+        @wraps(func)
+        def wrapper(self, *args, **kwargs):
+            if self.is_local_main_process or not self.use_distributed:
+                return func(self, *args, **kwargs)
+
+        return wrapper
+
+    def on_process(process_idx):
_The documentation is not available anymore as the PR was closed or merged._
Hi @sgugger , sorry for the delay, hope this is better now
@ZhiyuanChen could you run "make style; make quality" to solve the Quality Check issue? 😃 

Thanks!
> @ZhiyuanChen could you run "make style; make quality" to solve the Quality Check issue? 😃
> 
> Thanks!

Sorry, I thought I have fixed it but somehow didn't pushed... 
> Thanks a lot for iterating! Almost ready to merge, I just left some comments on the docstrings and a suggestion to group the decorators `on_process` and `on_local_process` together. Let me know your thoughts!

Thank you very much tor your comments and suggestion. 

For the docstrings, i have accepted rll suggestions. 

For the group... I think its rather important to ensure api in a similar organisation, and grouping in this way break the consistency with `is_local_main_process` and `on_local_main_process`. 
Though we could use alias/partial/impl to make the implementation more compact
Let's roll with your choice and we'll see what users think then. We can always add aliases in the future :-)

Thanks again for your contribution!
> Let's roll with your choice and we'll see what users think then. We can always add aliases in the future :-)

Haha, sure thing~

> Thanks again for your contribution!

No worries~

https://api.github.com/repos/huggingface/accelerate/pulls/488
diff --git a/examples/by_feature/fsdp_with_peak_mem_tracking.py b/examples/by_feature/fsdp_with_peak_mem_tracking.py
index 196aa0b06..eedbf3392 100644
--- a/examples/by_feature/fsdp_with_peak_mem_tracking.py
+++ b/examples/by_feature/fsdp_with_peak_mem_tracking.py
@@ -280,7 +280,7 @@ def collate_fn(examples):
                 predictions, references = accelerator.gather(
                     (predictions, batch["labels"])
                 )  # If we are in a multiprocess environment, the last batch has duplicates
-                if accelerator.num_processes > 1:
+                if accelerator.use_distributed:
                     if step == len(eval_dataloader) - 1:
                         predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
                         references = references[: len(eval_dataloader.dataset) - samples_seen]
diff --git a/examples/by_feature/multi_process_metrics.py b/examples/by_feature/multi_process_metrics.py
index c9445ee81..741e62b8c 100644
--- a/examples/by_feature/multi_process_metrics.py
+++ b/examples/by_feature/multi_process_metrics.py
@@ -178,7 +178,7 @@ def training_function(config, args):
             predictions, references = accelerator.gather((predictions, batch["labels"]))
             # New Code #
             # First we check if it's a distributed system
-            if accelerator.num_processes > 1:
+            if accelerator.use_distributed:
                 # Then see if we're on the last batch of our eval dataloader
                 if step == len(eval_dataloader) - 1:
                     # Last batch needs to be truncated on distributed systems as it contains additional samples
diff --git a/examples/complete_cv_example.py b/examples/complete_cv_example.py
index bd2ed3b29..6e414bc9c 100644
--- a/examples/complete_cv_example.py
+++ b/examples/complete_cv_example.py
@@ -242,7 +242,7 @@ def training_function(config, args):
                 outputs = model(inputs)
             predictions = outputs.argmax(dim=-1)
             predictions, references = accelerator.gather((predictions, batch["label"]))
-            if accelerator.num_processes > 1:
+            if accelerator.use_distributed:
                 if step == len(eval_dataloader) - 1:
                     predictions = predictions[: len(eval_dataloader) - samples_seen]
                     references = references[: len(eval_dataloader) - samples_seen]
diff --git a/examples/complete_nlp_example.py b/examples/complete_nlp_example.py
index 0433ed9a1..6e6aec0ff 100644
--- a/examples/complete_nlp_example.py
+++ b/examples/complete_nlp_example.py
@@ -222,7 +222,7 @@ def collate_fn(examples):
             predictions, references = accelerator.gather(
                 (predictions, batch["labels"])
             )  # If we are in a multiprocess environment, the last batch has duplicates
-            if accelerator.num_processes > 1:
+            if accelerator.use_distributed:
                 if step == len(eval_dataloader) - 1:
                     predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
                     references = references[: len(eval_dataloader.dataset) - samples_seen]
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 99afddeb8..0a2e5bb35 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -273,6 +273,10 @@ def __init__(
         if self.rng_types is None:
             self.rng_types = ["torch"] if is_torch_version("<=", "1.5.1") else ["generator"]
 
+    @property
+    def use_distributed(self):
+        return self.distributed_type != DistributedType.NO and self.num_processes > 1
+
     @property
     def distributed_type(self):
         return self.state.distributed_type
@@ -361,7 +365,7 @@ def no_sync(self, model):
                 PyTorch Module that was prepared with `Accelerator.prepare`
         """
         context = contextlib.nullcontext
-        if self.num_processes > 1:
+        if self.use_distributed:
             context = getattr(model, "no_sync", context)
 
         with context():

Maybe also check for the number of processes to be `> 1` here, since you are replacing all those tests?
@@ -273,6 +273,10 @@ def __init__(
         if self.rng_types is None:
             self.rng_types = ["torch"] if is_torch_version("<=", "1.5.1") else ["generator"]
 
+    @property
+    def use_distributed(self):
+        return self.distributed_type != DistributedType.NO
_The documentation is not available anymore as the PR was closed or merged._
I noticed the rest modification proposed in the original issues have been resolved, so just add this last property

> Thanks for your PR! I left a small comment.

Thank you for the comments, I have updated the pr
Thanks!
https://api.github.com/repos/huggingface/accelerate/pulls/487
diff --git a/docs/source/accelerator.mdx b/docs/source/accelerator.mdx
index 6e7a93a3d..5202aeb8c 100644
--- a/docs/source/accelerator.mdx
+++ b/docs/source/accelerator.mdx
@@ -38,4 +38,27 @@ should search for and replace by the corresponding methods of your `accelerator`
 - Use [`~Accelerator.clip_grad_norm_`] instead of `torch.nn.utils.clip_grad_norm_` and
   [`~Accelerator.clip_grad_value_`] instead of `torch.nn.utils.clip_grad_value_`.
 
+To perform gradient accumulation use [`~Accelerator.accumulate`] and specify a `gradient_accumulation_steps`. 
+This will also automatically ensure the gradients are synced or unsynced when on multi-device training, check if the step should
+actually be performed, and auto-scale the loss:
+
+```python
+accelerator = Accelerator(gradient_accumulation_steps=2)
+
+for (input, label) in enumerate(training_dataloader):
+    with accelerator.accumulate(model):
+        predictions = model(input)
+        loss = loss_function(predictions, labels)
+        accelerator.backward(loss)
+        optimizer.step()
+        scheduler.step()
+        optimizer.zero_grad()
+```
+
+<Tip warning={true}>
+
+Using this with `dispatch_batches=True`  (which is the default for iterable datasets) is currently not supported.
+
+</Tip>
+
 [[autodoc]] Accelerator
diff --git a/docs/source/internal.mdx b/docs/source/internal.mdx
index e1eec36be..8da23bee8 100644
--- a/docs/source/internal.mdx
+++ b/docs/source/internal.mdx
@@ -12,6 +12,10 @@ specific language governing permissions and limitations under the License.
 
 # Internals
 
+## Gradient Accumulation states
+
+[[autodoc]] state.GradientState
+
 ## Optimizer
 
 [[autodoc]] optimizer.AcceleratedOptimizer
diff --git a/examples/by_feature/gradient_accumulation.py b/examples/by_feature/gradient_accumulation.py
index b69978023..d3b253bce 100644
--- a/examples/by_feature/gradient_accumulation.py
+++ b/examples/by_feature/gradient_accumulation.py
@@ -102,19 +102,21 @@ def collate_fn(examples):
 
 
 def training_function(config, args):
-    # Initialize accelerator
-    accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
-    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
-    lr = config["lr"]
-    num_epochs = int(config["num_epochs"])
-    seed = int(config["seed"])
-    batch_size = int(config["batch_size"])
     # New Code #
     gradient_accumulation_steps = int(args.gradient_accumulation_steps)
+    # Initialize accelerator
+    accelerator = Accelerator(
+        cpu=args.cpu, mixed_precision=args.mixed_precision, gradient_accumulation_steps=gradient_accumulation_steps
+    )
     if accelerator.distributed_type == DistributedType.TPU and gradient_accumulation_steps > 1:
         raise NotImplementedError(
             "Gradient accumulation on TPUs is currently not supported. Pass `gradient_accumulation_steps=1`"
         )
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
 
     metric = evaluate.load("glue", "mrpc")
 
@@ -152,20 +154,11 @@ def training_function(config, args):
             # We could avoid this line since we set the accelerator with `device_placement=True`.
             batch.to(accelerator.device)
             # New code #
-            # We use the new `no_sync` context manager to prevent gradient averaging
-            # until we want to at the proper step if we happen to be in a distributed setup
-            # otherwise it does nothing
+            # We use the new `accumulate` context manager to perform gradient accumulation
             # We also currently do not support TPUs nor advise it as bugs were found on the XLA side when running our tests.
-            if step % gradient_accumulation_steps != 0:
-                # Accumulate gradients locally
-                with accelerator.no_sync(model):
-                    output = model(**batch)
-                    loss = output.loss / gradient_accumulation_steps
-                    accelerator.backward(loss)
-            else:
-                # Sync gradients and step
+            with accelerator.accumulate(model):
                 output = model(**batch)
-                loss = output.loss / gradient_accumulation_steps
+                loss = output.loss
                 accelerator.backward(loss)
                 optimizer.step()
                 lr_scheduler.step()
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 1b3fa27e9..1148262af 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -28,7 +28,7 @@
 from .logging import get_logger
 from .optimizer import AcceleratedOptimizer
 from .scheduler import AcceleratedScheduler
-from .state import AcceleratorState
+from .state import AcceleratorState, GradientState
 from .tracking import LOGGER_TYPE_TO_CLASS, GeneralTracker, filter_trackers
 from .utils import (
     DeepSpeedPlugin,
@@ -92,6 +92,9 @@ class Accelerator:
             default to the value in the environment variable `MIXED_PRECISION`, which will use the default value in the
             accelerate config of the current system or the flag passed with the `accelerate.launch` command. 'fp16'
             requires pytorch 1.6 or higher. 'bf16' requires pytorch 1.10 or higher.
+        gradient_accumulation_steps (`int`, *optional*, default to 1):
+            The number of steps that should pass before gradients are accumulated. A number > 1 should be combined with
+            `Accelerator.accumulate`.
         cpu (`bool`, *optional*):
             Whether or not to force the script to execute on CPU. Will ignore GPU available if set to `True` and force
             the execution on one process only.
@@ -146,6 +149,7 @@ def __init__(
         split_batches: bool = False,
         fp16: bool = None,
         mixed_precision: Union[PrecisionType, str] = None,
+        gradient_accumulation_steps: int = 1,
         cpu: bool = False,
         deepspeed_plugin: DeepSpeedPlugin = None,
         fsdp_plugin: FullyShardedDataParallelPlugin = None,
@@ -231,6 +235,17 @@ def __init__(
             **kwargs,
         )
 
+        if gradient_accumulation_steps > 1:
+            if self.state.distributed_type == DistributedType.TPU:
+                raise NotImplementedError(
+                    "Gradient accumulation on TPU is not supported. Pass in `gradient_accumulation_steps=1`"
+                )
+            if dispatch_batches:
+                raise NotImplementedError(
+                    "Gradient accumulation with dispatched dataloaders is not supported. Pass in `gradient_accumulation_steps=1` or `dispatch_batches=False`"
+                )
+
+        self.gradient_accumulation_steps = gradient_accumulation_steps
         self.device_placement = device_placement
         self.split_batches = split_batches
         self.dispatch_batches = dispatch_batches
@@ -262,6 +277,10 @@ def __init__(
                 kwargs = self.scaler_handler.to_kwargs() if self.scaler_handler is not None else {}
                 self.scaler = torch.cuda.amp.GradScaler(**kwargs)
 
+        # Start of internal step tracking
+        self.step = 0
+        self.gradient_state = GradientState()
+
         # Internal references to the training objects
         self._optimizers = []
         self._models = []
@@ -367,6 +386,37 @@ def no_sync(self, model):
         with context():
             yield
 
+    def _do_sync(self):
+        "Sets the right `sync_gradients` context and either resets or increases `self.step`"
+        if self.gradient_state.end_of_dataloader:
+            self.step = 0
+            self.gradient_state._set_sync_gradients(True)
+        else:
+            self.step += 1
+        self.gradient_state._set_sync_gradients((self.step % self.gradient_accumulation_steps) == 0)
+
+    @property
+    def sync_gradients(self):
+        return self.gradient_state.sync_gradients
+
+    @contextmanager
+    def accumulate(self, model):
+        """
+        A context manager that will lightly wrap around and perform gradient accumulation automatically
+
+        Args:
+            model (`torch.nn.Module`):
+                PyTorch Module that was prepared with `Accelerator.prepare`
+        """
+        self._do_sync()
+        if self.sync_gradients:
+            context = contextlib.nullcontext
+        else:
+            context = self.no_sync
+
+        with context(model):
+            yield
+
     def print(self, *args, **kwargs):
         """
         Use in replacement of `print()` to only print once per server.
@@ -734,6 +784,7 @@ def backward(self, loss, **kwargs):
         """
         Use `accelerator.backward(loss)` in lieu of `loss.backward()`.
         """
+        loss /= self.gradient_accumulation_steps
         if self.distributed_type == DistributedType.DEEPSPEED:
             self.deepspeed_engine_wrapped.backward(loss, **kwargs)
         elif self.scaler is not None:
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index 806fb7ddc..524034eb1 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -18,7 +18,8 @@
 import torch
 from torch.utils.data import BatchSampler, DataLoader, IterableDataset
 
-from .state import AcceleratorState, DistributedType, is_tpu_available
+from .logging import get_logger
+from .state import AcceleratorState, DistributedType, GradientState, is_tpu_available
 from .utils import (
     RNGType,
     broadcast,
@@ -37,6 +38,7 @@
 if is_tpu_available(check_device=False):
     import torch_xla.distributed.parallel_loader as xpl
 
+logger = get_logger(__name__)
 
 # kwargs of the DataLoader in min version 1.4.0.
 _PYTORCH_DATALOADER_KWARGS = {
@@ -294,31 +296,44 @@ def __init__(self, dataset, device=None, rng_types=None, generator=None, **kwarg
         self.device = device
         self.rng_types = rng_types
         self.generator = generator
+        self.gradient_state = GradientState()
 
     def __iter__(self):
         if self.rng_types is not None:
             synchronize_rng_states(self.rng_types, self.generator)
-        for batch in super().__iter__():
-            yield batch if self.device is None else send_to_device(batch, self.device)
+        self.gradient_state._set_end_of_dataloader(False)
+        dataloader_iter = super().__iter__()
+        # We iterate one batch ahead to check when we are at the end
+        try:
+            current_batch = next(dataloader_iter)
+        except StopIteration:
+            yield
+        while True:
+            try:
+                # But we still move it to the device so it is done before `StopIteration` is reached
+                if self.device is not None:
+                    current_batch = send_to_device(current_batch, self.device)
+                next_batch = next(dataloader_iter)
+                yield current_batch
+                current_batch = next_batch
+            except StopIteration:
+                self.gradient_state._set_end_of_dataloader(True)
+                yield current_batch
+                break
 
 
 class DataLoaderDispatcher(DataLoader):
     """
+    Args:
     Subclass of a PyTorch `DataLoader` that will iterate and preprocess on process 0 only, then dispatch on each
     process their part of the batch.
-
-    Args:
         split_batches (`bool`, *optional*, defaults to `False`):
             Whether the resulting `DataLoader` should split the batches of the original data loader across devices or
             yield full batches (in which case it will yield batches starting at the `process_index`-th and advancing of
-            `num_processes` batches at each iteration).
-
-            Another way to see this is that the observed batch size will be the same as the initial `dataloader` if
-            this option is set to `True`, the batch size of the initial `dataloader` multiplied by `num_processes`
-            otherwise.
-
-            Setting this option to `True` requires that the batch size of the `dataloader` is a round multiple of
-            `batch_size`.
+            `num_processes` batches at each iteration). Another way to see this is that the observed batch size will be
+            the same as the initial `dataloader` if this option is set to `True`, the batch size of the initial
+            `dataloader` multiplied by `num_processes` otherwise. Setting this option to `True` requires that the batch
+            size of the `dataloader` is a round multiple of `batch_size`.
     """
 
     def __init__(self, dataset, split_batches: bool = False, **kwargs):
@@ -338,6 +353,9 @@ def __init__(self, dataset, split_batches: bool = False, **kwargs):
         if shuffle:
             torch.utils.data.graph_settings.apply_shuffle_settings(dataset, shuffle=shuffle)
 
+        self.gradient_state = GradientState()
+        self.state = AcceleratorState()
+
     def __iter__(self):
         state = AcceleratorState()
         if state.process_index == 0:
@@ -408,12 +426,11 @@ def __iter__(self):
             yield slice_tensors(batch, data_slice)
 
     def __len__(self):
-        state = AcceleratorState()
         whole_length = super().__len__()
         if self.drop_last:
-            return whole_length // state.num_processes
+            return whole_length // self.state.num_processes
         else:
-            return math.ceil(whole_length / state.num_processes)
+            return math.ceil(whole_length / self.state.num_processes)
 
 
 def prepare_data_loader(
diff --git a/src/accelerate/optimizer.py b/src/accelerate/optimizer.py
index 046d5ce5b..c5ed66cd4 100644
--- a/src/accelerate/optimizer.py
+++ b/src/accelerate/optimizer.py
@@ -17,7 +17,7 @@
 
 import torch
 
-from .state import AcceleratorState
+from .state import AcceleratorState, GradientState
 from .utils import DistributedType, honor_type, is_torch_version, is_tpu_available
 
 
@@ -53,6 +53,7 @@ def __init__(self, optimizer, device_placement=True, scaler=None):
         self.optimizer = optimizer
         self.scaler = scaler
         self.accelerator_state = AcceleratorState()
+        self.gradient_state = GradientState()
         self.device_placement = device_placement
         self._is_overflow = False
 
@@ -101,37 +102,39 @@ def state_dict(self):
         return self.optimizer.state_dict()
 
     def zero_grad(self, set_to_none=None):
-        if is_torch_version("<", "1.7.0"):
-            if set_to_none is not None:
-                raise ValueError(
-                    "`set_to_none` for Optimizer.zero_grad` was introduced in PyTorch 1.7.0 and can't be used for "
-                    f"earlier versions (found version {torch.__version__})."
-                )
-            self.optimizer.zero_grad()
-        else:
-            accept_arg = "set_to_none" in inspect.signature(self.optimizer.zero_grad).parameters
-            if accept_arg:
-                if set_to_none is None:
-                    set_to_none = False
-                self.optimizer.zero_grad(set_to_none=set_to_none)
-            else:
+        if self.gradient_state.sync_gradients:
+            if is_torch_version("<", "1.7.0"):
                 if set_to_none is not None:
-                    raise ValueError("`set_to_none` for Optimizer.zero_grad` is not supported by this optimizer.")
+                    raise ValueError(
+                        "`set_to_none` for Optimizer.zero_grad` was introduced in PyTorch 1.7.0 and can't be used for "
+                        f"earlier versions (found version {torch.__version__})."
+                    )
                 self.optimizer.zero_grad()
+            else:
+                accept_arg = "set_to_none" in inspect.signature(self.optimizer.zero_grad).parameters
+                if accept_arg:
+                    if set_to_none is None:
+                        set_to_none = False
+                    self.optimizer.zero_grad(set_to_none=set_to_none)
+                else:
+                    if set_to_none is not None:
+                        raise ValueError("`set_to_none` for Optimizer.zero_grad` is not supported by this optimizer.")
+                    self.optimizer.zero_grad()
 
     def step(self, closure=None):
-        if self.accelerator_state.distributed_type == DistributedType.TPU:
-            optimizer_args = {"closure": closure} if closure is not None else {}
-            xm.optimizer_step(self.optimizer, optimizer_args=optimizer_args)
-        elif self.scaler is not None:
-            scale_before = self.scaler.get_scale()
-            self.scaler.step(self.optimizer, closure)
-            self.scaler.update()
-            scale_after = self.scaler.get_scale()
-            # If we reduced the loss scale, it means the optimizer step was skipped because of gradient overflow.
-            self._is_overflow = scale_after < scale_before
-        else:
-            self.optimizer.step(closure)
+        if self.gradient_state.sync_gradients:
+            if self.accelerator_state.distributed_type == DistributedType.TPU:
+                optimizer_args = {"closure": closure} if closure is not None else {}
+                xm.optimizer_step(self.optimizer, optimizer_args=optimizer_args)
+            elif self.scaler is not None:
+                scale_before = self.scaler.get_scale()
+                self.scaler.step(self.optimizer, closure)
+                self.scaler.update()
+                scale_after = self.scaler.get_scale()
+                # If we reduced the loss scale, it means the optimizer step was skipped because of gradient overflow.
+                self._is_overflow = scale_after < scale_before
+            else:
+                self.optimizer.step(closure)
 
     def _switch_parameters(self, parameters_map):
         for param_group in self.optimizer.param_groups:
diff --git a/src/accelerate/scheduler.py b/src/accelerate/scheduler.py
index 1fac5f012..cfd6d719f 100644
--- a/src/accelerate/scheduler.py
+++ b/src/accelerate/scheduler.py
@@ -12,7 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .state import AcceleratorState
+# We ignore warnings about stepping the scheduler since we step it ourselves during gradient accumulation
+
+import warnings
+
+from .state import AcceleratorState, GradientState
+
+
+warnings.filterwarnings("ignore", category=UserWarning, module="torch.optim.lr_scheduler")
 
 
 class AcceleratedScheduler:
@@ -40,6 +47,7 @@ def __init__(self, scheduler, optimizers, step_with_optimizer: bool = True, spli
         self.scheduler = scheduler
         self.optimizers = optimizers if isinstance(optimizers, (list, tuple)) else [optimizers]
         self.split_batches = split_batches
+        self.gradient_state = GradientState()
         self.step_with_optimizer = step_with_optimizer
 
     def step(self, *args, **kwargs):
@@ -52,7 +60,6 @@ def step(self, *args, **kwargs):
         for opt in self.optimizers:
             if opt.step_was_skipped:
                 return
-
         if self.split_batches:
             # Split batches -> the training dataloader batch size is not changed so one step per training step
             self.scheduler.step(*args, **kwargs)
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index 63527124c..330830eb5 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -52,6 +52,7 @@ class AcceleratorState:
     Attributes:
 
         - **device** (`torch.device`) -- The device to use.
+        - **sync_gradients** (`bool`) -- Whether to sync the gradients or not
         - **distributed_type** (`~accelerate.state.DistributedType`) -- The type of distributed environment currently
           in use.
         - **num_processes** (`int`) -- The number of processes currently launched in parallel.
@@ -209,3 +210,36 @@ def _check_initialized(self, mixed_precision=None, cpu=None):
                 raise ValueError(err.format(flag="cpu=True"))
             if mixed_precision is not None and mixed_precision != self.mixed_precision:
                 raise ValueError(err.format(flag=f"mixed_precision='{mixed_precision}'"))
+
+
+class GradientState:
+    """
+    This is a variation of a [singleton class](https://en.wikipedia.org/wiki/Singleton_pattern) in the sense that all
+    instance of `GradientState` share the same state, which is initialized on the first instantiation.
+
+    This specific state revolves around whether gradients should be synced and if we have reached the end of a prepared
+    dataloader Attributes:
+
+        - **sync_gradients** (`bool`) -- Whether the gradients should be synced
+        - **end_of_dataloader** (`bool`) -- Whether we have reached the end the current dataloader
+    """
+
+    _shared_state = {}
+
+    def __init__(self):
+        self.__dict__ = self._shared_state
+        if not getattr(self, "initialized", False):
+            self.sync_gradients = True
+            self.end_of_dataloader = False
+        self.initialized = True
+
+    def __repr__(self):
+        return f"Sync Gradients: {self.sync_gradients}\n" f"At end of current dataloader: {self.end_of_dataloader}\n"
+
+    def _set_sync_gradients(self, sync_gradients):
+        "Private function that sets whether gradients should be synchronized. Users should not have to call this."
+        self.sync_gradients = sync_gradients
+
+    def _set_end_of_dataloader(self, end_of_dataloader):
+        "Private function that sets whether the end of the current dataloader has been reached. Users should not have to call this."
+        self.end_of_dataloader = end_of_dataloader
diff --git a/src/accelerate/test_utils/scripts/test_sync.py b/src/accelerate/test_utils/scripts/test_sync.py
index b58a8c092..0b0dc7abd 100644
--- a/src/accelerate/test_utils/scripts/test_sync.py
+++ b/src/accelerate/test_utils/scripts/test_sync.py
@@ -16,6 +16,8 @@
 
 import torch
 import torch.nn.functional as F
+from torch.optim import AdamW
+from torch.optim.lr_scheduler import LambdaLR
 from torch.utils.data import DataLoader
 
 from accelerate import Accelerator
@@ -23,30 +25,61 @@
 from accelerate.utils import DistributedType, set_seed
 
 
-def step_model(model, input, target, accelerator):
+def check_model_parameters(model_a, model_b, did_step):
+    for param, grad_param in zip(model_a.parameters(), model_b.parameters()):
+        if not param.requires_grad:
+            continue
+        if not did_step:
+            # Grads should not be in sync
+            assert (
+                torch.allclose(param.grad, grad_param.grad) is False
+            ), f"Gradients in sync when they should not be:\nmodel_a grad ({param.grad}) == model_b grad ({grad_param.grad})"
+        else:
+            # Grads should be in sync
+            assert (
+                torch.allclose(param.grad, grad_param.grad) is True
+            ), f"Gradients not in sync when they should be:\nmodel_a grad ({param.grad}) != model_b grad ({grad_param.grad})"
+
+
+def step_model(model, input, target, accelerator, do_backward=True):
     model.train()
     output = model(input)
     loss = F.mse_loss(output, target.to(output.device))
-    accelerator.backward(loss)
+    if not do_backward:
+        loss /= accelerator.gradient_accumulation_steps
+        loss.backward()
+    else:
+        accelerator.backward(loss)
 
 
-def get_training_setup(accelerator):
+def get_training_setup(accelerator, sched=False):
     "Returns everything needed to perform basic training"
     set_seed(42)
     model = RegressionModel()
-    model.to(accelerator.device)
-    dset = RegressionDataset()
+    ddp_model = deepcopy(model)
+    dset = RegressionDataset(length=80)
     dataloader = DataLoader(dset, batch_size=16)
+    model.to(accelerator.device)
+    if sched:
+        opt = AdamW(params=model.parameters(), lr=1e-3)
+        ddp_opt = AdamW(params=ddp_model.parameters(), lr=1e-3)
+        sched = LambdaLR(opt, lr_lambda=lambda epoch: epoch**0.65)
+        ddp_sched = LambdaLR(ddp_opt, lr_lambda=lambda epoch: epoch**0.65)
     # Make a copy of `model`
-    ddp_model, dataloader = accelerator.prepare(deepcopy(model), dataloader)
-    # Use a single batch for all of the tests
-    ddp_input, ddp_target = next(iter(dataloader)).values()
-    return model, ddp_model, ddp_input, ddp_target
+    if sched:
+        ddp_model, ddp_opt, ddp_sched, dataloader = accelerator.prepare(ddp_model, ddp_opt, ddp_sched, dataloader)
+    else:
+        ddp_model, dataloader = accelerator.prepare(ddp_model, dataloader)
+    if sched:
+        return (model, opt, sched, dataloader, ddp_model, ddp_opt, ddp_sched)
+    return model, ddp_model, dataloader
 
 
 def test_noop_sync(accelerator):
     # Test when on a single CPU or GPU that the context manager does nothing
-    model, ddp_model, ddp_input, ddp_target = get_training_setup(accelerator)
+    model, ddp_model, dataloader = get_training_setup(accelerator)
+    # Use a single batch
+    ddp_input, ddp_target = next(iter(dataloader)).values()
     for iteration in range(3):
         # Gather the distributed inputs and targs for the base model
         input, target = accelerator.gather((ddp_input, ddp_target))
@@ -63,6 +96,7 @@ def test_noop_sync(accelerator):
             step_model(ddp_model, ddp_input, ddp_target, accelerator)
 
         # Since `no_sync` is a noop, `ddp_model` and `model` grads should always be in sync
+        check_model_parameters(model, ddp_model, True)
         for param, ddp_param in zip(model.parameters(), ddp_model.parameters()):
             if not param.requires_grad:
                 continue
@@ -77,7 +111,9 @@ def test_noop_sync(accelerator):
 
 def test_distributed_sync(accelerator):
     # Test on distributed setup that context manager behaves properly
-    model, ddp_model, ddp_input, ddp_target = get_training_setup(accelerator)
+    model, ddp_model, dataloader = get_training_setup(accelerator)
+    # Use a single batch
+    ddp_input, ddp_target = next(iter(dataloader)).values()
     for iteration in range(3):
         # Gather the distributed inputs and targs for the base model
         input, target = accelerator.gather((ddp_input, ddp_target))
@@ -113,17 +149,92 @@ def test_distributed_sync(accelerator):
         ddp_input = ddp_input[torch.randperm(16)]
 
 
+def test_gradient_accumulation():
+    accelerator = Accelerator(gradient_accumulation_steps=2)
+    # Test that context manager behaves properly
+    model, ddp_model, dataloader = get_training_setup(accelerator)
+    for iteration, batch in enumerate(dataloader):
+        ddp_input, ddp_target = batch.values()
+        # Gather the distributed inputs and targs for the base model
+        input, target = accelerator.gather((ddp_input, ddp_target))
+        input, target = input.to(accelerator.device), target.to(accelerator.device)
+        # Perform our initial ground truth step in non "DDP"
+        step_model(model, input, target, accelerator, False)
+        # Do "gradient accumulation" (noop)
+        with accelerator.accumulate(ddp_model):
+            step_model(ddp_model, ddp_input, ddp_target, accelerator)
+
+        # DDP model and model should only be in sync when not (iteration % 2 == 0)
+        for param, ddp_param in zip(model.parameters(), ddp_model.parameters()):
+            if not param.requires_grad:
+                continue
+            if ((iteration + 1) % 2 == 0) or (iteration == len(dataloader) - 1):
+                # Grads should be in sync
+                assert (
+                    torch.allclose(param.grad, ddp_param.grad) is True
+                ), f"Gradients not in sync when they should be at iteration {iteration}:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})"
+            else:
+                # Grads should not be in sync
+                assert (
+                    torch.allclose(param.grad, ddp_param.grad) is False
+                ), f"Gradients in sync when they should not be at iteration {iteration}:\nModel grad ({param.grad}) == DDP grad ({ddp_param.grad})"
+
+        # Shuffle ddp_input on each iteration
+        torch.manual_seed(1337 + iteration)
+        ddp_input = ddp_input[torch.randperm(16)]
+
+
+def test_gradient_accumulation_with_opt_and_scheduler():
+    accelerator = Accelerator(gradient_accumulation_steps=2)
+    # Test that context manager behaves properly
+    model, opt, sched, dataloader, ddp_model, ddp_opt, ddp_sched = get_training_setup(accelerator, True)
+    for iteration, batch in enumerate(dataloader):
+        ddp_input, ddp_target = batch.values()
+        # Gather the distributed inputs and targs for the base model
+        input, target = accelerator.gather((ddp_input, ddp_target))
+        input, target = input.to(accelerator.device), target.to(accelerator.device)
+        # Perform our initial ground truth step in non "DDP"
+        model.train()
+        ddp_model.train()
+        step_model(model, input, target, accelerator, False)
+        opt.step()
+        for _ in range(accelerator.num_processes):
+            sched.step()
+        opt.zero_grad()
+        # Perform gradient accumulation under wrapper
+        with accelerator.accumulate(ddp_model):
+            step_model(ddp_model, ddp_input, ddp_target, accelerator)
+            ddp_opt.step()
+            ddp_sched.step()
+            ddp_opt.zero_grad()
+
+        # Learning rates should be the same
+        assert opt.param_groups[0]["lr"] == ddp_opt.param_groups[0]["lr"]
+        did_step = (((iteration + 1) % 2) == 0) or (iteration == (len(dataloader) - 1))
+        if accelerator.num_processes > 1:
+            check_model_parameters(model, ddp_model, did_step)
+        # Shuffle ddp_input on each iteration
+        torch.manual_seed(1337 + iteration)
+
+
 def main():
     accelerator = Accelerator()
     state = accelerator.state
     if state.distributed_type == DistributedType.NO:
         if state.local_process_index == 0:
-            print("**NOOP `no_sync` gradient accumulation**")
+            print("**Test NOOP `no_sync` context manager**")
         test_noop_sync(accelerator)
     if state.distributed_type in (DistributedType.MULTI_GPU, DistributedType.MULTI_CPU):
         if state.local_process_index == 0:
-            print("**Distributed `no_sync` gradient accumulation**")
+            print("**Test Distributed `no_sync` context manager**")
         test_distributed_sync(accelerator)
+    if state.distributed_type == DistributedType.MULTI_GPU:
+        if state.local_process_index == 0:
+            print("**Test `accumulate` gradient accumulation**")
+        test_gradient_accumulation()
+    if state.local_process_index == 0:
+        print("**Test `accumulate` gradient accumulation with optimizer and scheduler**")
+    test_gradient_accumulation_with_opt_and_scheduler()
 
 
 def _mp_fn(index):
diff --git a/tests/test_grad_sync.py b/tests/test_grad_sync.py
index c066bd41a..182d3ef95 100644
--- a/tests/test_grad_sync.py
+++ b/tests/test_grad_sync.py
@@ -36,20 +36,19 @@ def setUp(self):
         self.test_file_path = os.path.sep.join(mod_file.split(os.path.sep)[:-1] + ["scripts", "test_sync.py"])
 
     @require_cpu
-    def test_gradient_sync_single_cpu_noop(self):
-        debug_launcher(test_sync.main)
+    def test_gradient_sync_cpu_noop(self):
         debug_launcher(test_sync.main, num_processes=1)
 
     @require_cpu
-    def test_gradient_sync_multi_cpu(self):
+    def test_gradient_sync_cpu_multi(self):
         debug_launcher(test_sync.main)
 
     @require_single_gpu
-    def test_gradient_sync_single_gpu(self):
-        debug_launcher(test_sync.main, num_processes=1)
+    def test_gradient_sync_gpu(self):
+        test_sync.main()
 
     @require_multi_gpu
-    def test_gradient_sync_multi_gpu(self):
+    def test_gradient_sync_gpu_multi(self):
         print(f"Found {torch.cuda.device_count()} devices.")
         cmd = get_launch_prefix() + [f"--nproc_per_node={torch.cuda.device_count()}", self.test_file_path]
         with patch_environment(omp_num_threads=1):
diff --git a/tests/test_kwargs_handlers.py b/tests/test_kwargs_handlers.py
index 8b438d1c6..d8c893234 100644
--- a/tests/test_kwargs_handlers.py
+++ b/tests/test_kwargs_handlers.py
@@ -21,6 +21,7 @@
 import torch
 
 from accelerate import Accelerator, DistributedDataParallelKwargs, GradScalerKwargs
+from accelerate.state import AcceleratorState
 from accelerate.test_utils import execute_subprocess_async, require_cuda, require_multi_gpu
 from accelerate.utils import KwargsHandler
 
@@ -44,7 +45,8 @@ def test_kwargs_handler(self):
     def test_grad_scaler_kwargs(self):
         # If no defaults are changed, `to_kwargs` returns an empty dict.
         scaler_handler = GradScalerKwargs(init_scale=1024, growth_factor=2)
-        accelerator = Accelerator(fp16=True, kwargs_handlers=[scaler_handler])
+        AcceleratorState._reset_state()
+        accelerator = Accelerator(mixed_precision="fp16", kwargs_handlers=[scaler_handler])
         print(accelerator.use_fp16)
         scaler = accelerator.scaler
 

Note that the first part of the test is implied by the second ;-)
@@ -367,6 +380,39 @@ def no_sync(self, model):
         with context():
             yield
 
+    @property
+    def _do_sync(self) -> bool:
+        "Checks if gradients should be synchronized and the optimizers + schedulers should be stepped"
+        if self.state.end_of_dataloader:
+            self.step = 0
+            return True
+        self.step += 1
+        if (self.gradient_accumulation_steps == 1) or ((self.step % self.gradient_accumulation_steps) == 0):
This shouldn't be in the state, the state only should contain the environment info, not some things that are subject to change. So store it directly in self.
@@ -367,6 +380,39 @@ def no_sync(self, model):
         with context():
             yield
 
+    @property
+    def _do_sync(self) -> bool:
+        "Checks if gradients should be synchronized and the optimizers + schedulers should be stepped"
+        if self.state.end_of_dataloader:
Same here, this should be stored somewhere else.
@@ -367,6 +380,39 @@ def no_sync(self, model):
         with context():
             yield
 
+    @property
+    def _do_sync(self) -> bool:
+        "Checks if gradients should be synchronized and the optimizers + schedulers should be stepped"
+        if self.state.end_of_dataloader:
+            self.step = 0
+            return True
+        self.step += 1
+        if (self.gradient_accumulation_steps == 1) or ((self.step % self.gradient_accumulation_steps) == 0):
+            return True
+        return False
+
+    @contextmanager
+    def accumulate(self, model):
+        """
+        A context manager that will lightly wrap around and perform gradient accumulation automatically
+
+        Args:
+            model (`torch.nn.Module`):
+                PyTorch Module that was prepared with `Accelerator.prepare`
+            dataloader (`torch.utils.DataLoader`)
+                PyTorch DataLoader that was prepared with `Accelerator.prepare`
+        """
+
+        if self._do_sync:
+            context = contextlib.nullcontext
+            AcceleratorState._set_state("sync_gradients", True)
+        else:
+            context = self.no_sync
+            AcceleratorState._set_state("sync_gradients", False)
```suggestion
        gradient_accumulation_steps (`int`, *optional*, defaults to 1):
```
@@ -92,6 +92,9 @@ class Accelerator:
             default to the value in the environment variable `MIXED_PRECISION`, which will use the default value in the
             accelerate config of the current system or the flag passed with the `accelerate.launch` command. 'fp16'
             requires pytorch 1.6 or higher. 'bf16' requires pytorch 1.10 or higher.
+        gradient_accumulation_steps (`int`, *optional*):
```suggestion
            [`Accelerator.accumulate`].
```
Link to the doc.
@@ -92,6 +92,9 @@ class Accelerator:
             default to the value in the environment variable `MIXED_PRECISION`, which will use the default value in the
             accelerate config of the current system or the flag passed with the `accelerate.launch` command. 'fp16'
             requires pytorch 1.6 or higher. 'bf16' requires pytorch 1.10 or higher.
+        gradient_accumulation_steps (`int`, *optional*):
+            The number of steps that should pass before gradients are accumulated. Should be combined with
+            `Accelerator.accumulate`
I think the class should have a static method to set this.
@@ -367,6 +380,37 @@ def no_sync(self, model):
         with context():
             yield
 
+    def _do_sync(self) -> bool:
+        "Checks if gradients should be synchronized and the optimizers + schedulers should be stepped"
+        if GradientState().end_of_dataloader:
+            self.step = 0
+            return True
+        else:
+            self.step += 1
+        return (self.step % self.gradient_accumulation_steps) == 0
+
+    @contextmanager
+    def accumulate(self, model):
+        """
+        A context manager that will lightly wrap around and perform gradient accumulation automatically
+
+        Args:
+            model (`torch.nn.Module`):
+                PyTorch Module that was prepared with `Accelerator.prepare`
+            dataloader (`torch.utils.DataLoader`)
+                PyTorch DataLoader that was prepared with `Accelerator.prepare`
+        """
+
+        if self._do_sync():
+            context = contextlib.nullcontext
+            GradientState._set_state("sync_gradients", True)
Same here
@@ -367,6 +380,37 @@ def no_sync(self, model):
         with context():
             yield
 
+    def _do_sync(self) -> bool:
+        "Checks if gradients should be synchronized and the optimizers + schedulers should be stepped"
+        if GradientState().end_of_dataloader:
+            self.step = 0
+            return True
+        else:
+            self.step += 1
+        return (self.step % self.gradient_accumulation_steps) == 0
+
+    @contextmanager
+    def accumulate(self, model):
+        """
+        A context manager that will lightly wrap around and perform gradient accumulation automatically
+
+        Args:
+            model (`torch.nn.Module`):
+                PyTorch Module that was prepared with `Accelerator.prepare`
+            dataloader (`torch.utils.DataLoader`)
+                PyTorch DataLoader that was prepared with `Accelerator.prepare`
+        """
+
+        if self._do_sync():
+            context = contextlib.nullcontext
+            GradientState._set_state("sync_gradients", True)
+        else:
+            context = self.no_sync
+            GradientState._set_state("sync_gradients", False)
The dataloader may still have no length, so set to `False at the start of `__iter__` and to `True` at the end.
@@ -298,7 +298,8 @@ def __init__(self, dataset, device=None, rng_types=None, generator=None, **kwarg
     def __iter__(self):
         if self.rng_types is not None:
             synchronize_rng_states(self.rng_types, self.generator)
-        for batch in super().__iter__():
+        for i, batch in enumerate(super().__iter__()):
+            GradientState._set_state("end_of_dataloader", i == (len(self) - 1))
```suggestion
        GradientState._set_state("end_of_dataloader", True)
```
Like this.

@@ -406,6 +408,8 @@ def __iter__(self):
 
             data_slice = slice(state.process_index * batch_size, (state.process_index + 1) * batch_size)
             yield slice_tensors(batch, data_slice)
+            current_idx += 1
+            GradientState._set_state("end_of_dataloader", current_idx == len(self) - 1)
For the scheduler, I'm not entirely sure how to proceed: we step for each process to account for Distributed training, not sure if we should step all the time or just when synchronizing gradients (since the length set for the scheduler probably does not account for gradient accumulation steps).
@@ -52,18 +53,18 @@ def step(self, *args, **kwargs):
         for opt in self.optimizers:
             if opt.step_was_skipped:
                 return
-
-        if self.split_batches:
-            # Split batches -> the training dataloader batch size is not changed so one step per training step
-            self.scheduler.step(*args, **kwargs)
-        else:
-            # Otherwise the training dataloader batch size was multiplied by `num_processes`, so we need to do
-            # num_processes steps per training step
-            num_processes = AcceleratorState().num_processes
-            for _ in range(num_processes):
-                # Special case when using OneCycle and `drop_last` was not used
-                if getattr(self.scheduler, "total_steps", 0) <= self.scheduler.last_epoch:
-                    self.scheduler.step(*args, **kwargs)
+        if self.gradient_state.sync_gradients:
+            if self.split_batches:
+                # Split batches -> the training dataloader batch size is not changed so one step per training step
+                self.scheduler.step(*args, **kwargs)
+            else:
+                # Otherwise the training dataloader batch size was multiplied by `num_processes`, so we need to do
+                # num_processes steps per training step
+                num_processes = AcceleratorState().num_processes
+                for _ in range(num_processes):
+                    # Special case when using OneCycle and `drop_last` was not used
+                    if getattr(self.scheduler, "total_steps", 0) <= self.scheduler.last_epoch:
+                        self.scheduler.step(*args, **kwargs)
Add two static methods to quickly set `sync_gradients` and `end_of_dataloader` (also wondering why you can't jsut set them directly on an instance? It should automatically update everything no?)
@@ -209,3 +210,35 @@ def _check_initialized(self, mixed_precision=None, cpu=None):
                 raise ValueError(err.format(flag="cpu=True"))
             if mixed_precision is not None and mixed_precision != self.mixed_precision:
                 raise ValueError(err.format(flag=f"mixed_precision='{mixed_precision}'"))
+
+
+class GradientState:
+    """
+    This is a variation of a [singleton class](https://en.wikipedia.org/wiki/Singleton_pattern) in the sense that all
+    instance of `GradientState` share the same state, which is initialized on the first instantiation.
+
+    This specific state revolves around whether gradients should be synced and if we have reached the end of a prepared
+    dataloader Attributes:
+
+        - **sync_gradients** (`bool`) -- Whether the gradients should be synced
+        - **end_of_dataloader** (`bool`) -- Whether we have reached the end the current dataloader
+    """
+
+    _shared_state = {}
+
+    def __init__(self):
+        self.__dict__ = self._shared_state
+        if not getattr(self, "initialized", False):
+            self.sync_gradients = True
+            self.end_of_dataloader = False
+        self.initialized = True
+
+    @staticmethod
+    def _set_state(key, val):
+        "Sets `key` to `val` in GradientState"
+        if GradientState._shared_state != {} and key not in GradientState._shared_state.keys():
+            raise KeyError(f"{key} is not a valid key of `GradientState`, {GradientState._shared_state.keys()}")
+        GradientState._shared_state[key] = val
+
Resolved: we update every time. 
@@ -52,18 +53,18 @@ def step(self, *args, **kwargs):
         for opt in self.optimizers:
             if opt.step_was_skipped:
                 return
-
-        if self.split_batches:
-            # Split batches -> the training dataloader batch size is not changed so one step per training step
-            self.scheduler.step(*args, **kwargs)
-        else:
-            # Otherwise the training dataloader batch size was multiplied by `num_processes`, so we need to do
-            # num_processes steps per training step
-            num_processes = AcceleratorState().num_processes
-            for _ in range(num_processes):
-                # Special case when using OneCycle and `drop_last` was not used
-                if getattr(self.scheduler, "total_steps", 0) <= self.scheduler.last_epoch:
-                    self.scheduler.step(*args, **kwargs)
+        if self.gradient_state.sync_gradients:
+            if self.split_batches:
+                # Split batches -> the training dataloader batch size is not changed so one step per training step
+                self.scheduler.step(*args, **kwargs)
+            else:
+                # Otherwise the training dataloader batch size was multiplied by `num_processes`, so we need to do
+                # num_processes steps per training step
+                num_processes = AcceleratorState().num_processes
+                for _ in range(num_processes):
+                    # Special case when using OneCycle and `drop_last` was not used
+                    if getattr(self.scheduler, "total_steps", 0) <= self.scheduler.last_epoch:
+                        self.scheduler.step(*args, **kwargs)
This is unnecessary complex. Why are you not doing `self.gradient_state._set_end_of_dataloader(True)` at the end of `iter`?
@@ -294,12 +296,26 @@ def __init__(self, dataset, device=None, rng_types=None, generator=None, **kwarg
         self.device = device
         self.rng_types = rng_types
         self.generator = generator
+        self.gradient_state = GradientState()
 
     def __iter__(self):
         if self.rng_types is not None:
             synchronize_rng_states(self.rng_types, self.generator)
+        self.gradient_state._set_end_of_dataloader(False)
+        _current_batch = 0
+        _dataloader_length = len(self) - 1 if hasattr(self, "__len__") else False
         for batch in super().__iter__():
+            if _dataloader_length:
+                self.gradient_state._set_end_of_dataloader(_current_batch == _dataloader_length)
+                _current_batch += 1
+            else:
+                self.gradient_state._set_end_of_dataloader(False)
             yield batch if self.device is None else send_to_device(batch, self.device)
+        if not self.gradient_state.end_of_dataloader:
+            self.gradient_state._set_end_of_dataloader(True)
+            logger.warn(
+                "Warning! DataLoader had no length and finished iterating. Backwards pass and stepping must be manually performed one last time."
+            )
Not sure about this second test. If `step_with_optimizer=True` we then only step for sync gradients, which is not what we want.
@@ -40,10 +47,11 @@ def __init__(self, scheduler, optimizers, step_with_optimizer: bool = True, spli
         self.scheduler = scheduler
         self.optimizers = optimizers if isinstance(optimizers, (list, tuple)) else [optimizers]
         self.split_batches = split_batches
+        self.gradient_state = GradientState()
         self.step_with_optimizer = step_with_optimizer
 
     def step(self, *args, **kwargs):
-        if not self.step_with_optimizer:
+        if not self.step_with_optimizer or not self.gradient_state.sync_gradients:
The filter should be done in the scheduler, not the test scripts.
@@ -12,41 +12,62 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import warnings
 from copy import deepcopy
 
 import torch
 import torch.nn.functional as F
+from torch.optim import AdamW
+from torch.optim.lr_scheduler import LambdaLR
 from torch.utils.data import DataLoader
 
 from accelerate import Accelerator
 from accelerate.test_utils import RegressionDataset, RegressionModel
 from accelerate.utils import DistributedType, set_seed
 
 
-def step_model(model, input, target, accelerator):
+warnings.filterwarnings("ignore", category=UserWarning, module="torch.optim.lr_scheduler")
This check has to happen as the last batch is being pushed out, otherwise it doesn't occur until outside the iteration loop. 

Example in python of what doesn't happen:

```python
a = [0,1,2,3]
class A:
    def __init__(self,a):
        self.a,self.is_end = a,False
    def __iter__(self):
        for item in a: yield item
        self.is_end = True
aa = A(a)
for i,item in enumerate(aa):
    print(f'At end at {i}?: {aa.is_end}')
print(aa.is_end)
```
Result:
```
At end at 0?: False
At end at 1?: False
At end at 2?: False
At end at 3?: False
True
```
What we want to have happen:
```python
class A:
    def __init__(self,a):
        self.a,self.is_end = a,False
    def __iter__(self):
        for i,item in enumerate(a):
            if i == len(a)-1:
                self.is_end = True 
            yield item
        self.is_end = True
```
```
At end at 0?: False
At end at 1?: False
At end at 2?: False
At end at 3?: True
True
```
Since we cannot check the length sometimes, this extra step + warning needs to be done.
@@ -294,12 +296,26 @@ def __init__(self, dataset, device=None, rng_types=None, generator=None, **kwarg
         self.device = device
         self.rng_types = rng_types
         self.generator = generator
+        self.gradient_state = GradientState()
 
     def __iter__(self):
         if self.rng_types is not None:
             synchronize_rng_states(self.rng_types, self.generator)
+        self.gradient_state._set_end_of_dataloader(False)
+        _current_batch = 0
+        _dataloader_length = len(self) - 1 if hasattr(self, "__len__") else False
         for batch in super().__iter__():
+            if _dataloader_length:
+                self.gradient_state._set_end_of_dataloader(_current_batch == _dataloader_length)
+                _current_batch += 1
+            else:
+                self.gradient_state._set_end_of_dataloader(False)
             yield batch if self.device is None else send_to_device(batch, self.device)
+        if not self.gradient_state.end_of_dataloader:
+            self.gradient_state._set_end_of_dataloader(True)
+            logger.warn(
+                "Warning! DataLoader had no length and finished iterating. Backwards pass and stepping must be manually performed one last time."
+            )
I think there's a way to do this easier, will toy with it today. 
@@ -294,12 +296,26 @@ def __init__(self, dataset, device=None, rng_types=None, generator=None, **kwarg
         self.device = device
         self.rng_types = rng_types
         self.generator = generator
+        self.gradient_state = GradientState()
 
     def __iter__(self):
         if self.rng_types is not None:
             synchronize_rng_states(self.rng_types, self.generator)
+        self.gradient_state._set_end_of_dataloader(False)
+        _current_batch = 0
+        _dataloader_length = len(self) - 1 if hasattr(self, "__len__") else False
         for batch in super().__iter__():
+            if _dataloader_length:
+                self.gradient_state._set_end_of_dataloader(_current_batch == _dataloader_length)
+                _current_batch += 1
+            else:
+                self.gradient_state._set_end_of_dataloader(False)
             yield batch if self.device is None else send_to_device(batch, self.device)
+        if not self.gradient_state.end_of_dataloader:
+            self.gradient_state._set_end_of_dataloader(True)
+            logger.warn(
+                "Warning! DataLoader had no length and finished iterating. Backwards pass and stepping must be manually performed one last time."
+            )
Much simpler way for us to iterate over the dataloader, and given we only move to device right before we yield there shouldn't be too much trouble with multiple batches in memory at one time. If so let me know.
@@ -294,12 +296,27 @@ def __init__(self, dataset, device=None, rng_types=None, generator=None, **kwarg
         self.device = device
         self.rng_types = rng_types
         self.generator = generator
+        self.gradient_state = GradientState()
 
     def __iter__(self):
         if self.rng_types is not None:
             synchronize_rng_states(self.rng_types, self.generator)
-        for batch in super().__iter__():
-            yield batch if self.device is None else send_to_device(batch, self.device)
+        self.gradient_state._set_end_of_dataloader(False)
+        dataloader_iter = super().__iter__()
+        # We iterate one batch ahead to check when we are at the end
+        current_batch = next(dataloader_iter)
+        while True:
+            try:
+                # But we still move it to the device so it is done before `StopIteration` is reached
+                if self.device is not None:
+                    current_batch = send_to_device(current_batch, self.device)
+                next_batch = next(dataloader_iter)
+                yield current_batch
+                current_batch = next_batch
+            except StopIteration:
+                self.gradient_state._set_end_of_dataloader(True)
+                yield current_batch
+                break
We should have a try except here too in case of an empty dataloader.
@@ -294,12 +296,27 @@ def __init__(self, dataset, device=None, rng_types=None, generator=None, **kwarg
         self.device = device
         self.rng_types = rng_types
         self.generator = generator
+        self.gradient_state = GradientState()
 
     def __iter__(self):
         if self.rng_types is not None:
             synchronize_rng_states(self.rng_types, self.generator)
-        for batch in super().__iter__():
-            yield batch if self.device is None else send_to_device(batch, self.device)
+        self.gradient_state._set_end_of_dataloader(False)
+        dataloader_iter = super().__iter__()
+        # We iterate one batch ahead to check when we are at the end
+        current_batch = next(dataloader_iter)
Very clever!
@@ -294,12 +296,27 @@ def __init__(self, dataset, device=None, rng_types=None, generator=None, **kwarg
         self.device = device
         self.rng_types = rng_types
         self.generator = generator
+        self.gradient_state = GradientState()
 
     def __iter__(self):
         if self.rng_types is not None:
             synchronize_rng_states(self.rng_types, self.generator)
-        for batch in super().__iter__():
-            yield batch if self.device is None else send_to_device(batch, self.device)
+        self.gradient_state._set_end_of_dataloader(False)
+        dataloader_iter = super().__iter__()
+        # We iterate one batch ahead to check when we are at the end
+        current_batch = next(dataloader_iter)
+        while True:
+            try:
+                # But we still move it to the device so it is done before `StopIteration` is reached
+                if self.device is not None:
+                    current_batch = send_to_device(current_batch, self.device)
+                next_batch = next(dataloader_iter)
+                yield current_batch
+                current_batch = next_batch
+            except StopIteration:
+                self.gradient_state._set_end_of_dataloader(True)
+                yield current_batch
+                break
Small docstrings here would be great :-) (with a mention that it's private for a good reason and shouldn't be touched by the user directly!)
@@ -209,3 +210,34 @@ def _check_initialized(self, mixed_precision=None, cpu=None):
                 raise ValueError(err.format(flag="cpu=True"))
             if mixed_precision is not None and mixed_precision != self.mixed_precision:
                 raise ValueError(err.format(flag=f"mixed_precision='{mixed_precision}'"))
+
+
+class GradientState:
+    """
+    This is a variation of a [singleton class](https://en.wikipedia.org/wiki/Singleton_pattern) in the sense that all
+    instance of `GradientState` share the same state, which is initialized on the first instantiation.
+
+    This specific state revolves around whether gradients should be synced and if we have reached the end of a prepared
+    dataloader Attributes:
+
+        - **sync_gradients** (`bool`) -- Whether the gradients should be synced
+        - **end_of_dataloader** (`bool`) -- Whether we have reached the end the current dataloader
+    """
+
+    _shared_state = {}
+
+    def __init__(self):
+        self.__dict__ = self._shared_state
+        if not getattr(self, "initialized", False):
+            self.sync_gradients = True
+            self.end_of_dataloader = False
+        self.initialized = True
+
+    def __repr__(self):
+        return f"Sync Gradients: {self.sync_gradients}\n" f"At end of current dataloader: {self.end_of_dataloader}\n"
+
+    def _set_sync_gradients(self, sync_gradients):
+        self.sync_gradients = sync_gradients
+
+    def _set_end_of_dataloader(self, end_of_dataloader):
+        self.end_of_dataloader = end_of_dataloader
No we shouldn't raise here, just yield nothing. A training loop over an empty dataloader works, just doesn't do anything.
We can put a warning however, just in case the user didn't intend that.
@@ -304,7 +304,10 @@ def __iter__(self):
         self.gradient_state._set_end_of_dataloader(False)
         dataloader_iter = super().__iter__()
         # We iterate one batch ahead to check when we are at the end
-        current_batch = next(dataloader_iter)
+        try:
+            current_batch = next(dataloader_iter)
+        except StopIteration:
+            raise StopIteration("Tried iterating over an empty dataloader")
```suggestion
This will also automatically ensure the gradients are synced or unsynced when on multi-device training, check if the step should
```
@@ -38,4 +38,23 @@ should search for and replace by the corresponding methods of your `accelerator`
 - Use [`~Accelerator.clip_grad_norm_`] instead of `torch.nn.utils.clip_grad_norm_` and
   [`~Accelerator.clip_grad_value_`] instead of `torch.nn.utils.clip_grad_value_`.
 
+To perform gradient accumulation use [`~Accelerator.accumulate`] and specify a `gradient_accumulation_steps`. 
+This will also automatically ensure the gradients are synced or unsynced when on multi-node, check if the step should
Use the proper syntax for a warning here
```suggestion
<Tip warning={true}>

Using this with `dispatch_batches=True`  (which is the default for iterable datasets) is currently not supported.

</Tip>
```
@@ -38,4 +38,23 @@ should search for and replace by the corresponding methods of your `accelerator`
 - Use [`~Accelerator.clip_grad_norm_`] instead of `torch.nn.utils.clip_grad_norm_` and
   [`~Accelerator.clip_grad_value_`] instead of `torch.nn.utils.clip_grad_value_`.
 
+To perform gradient accumulation use [`~Accelerator.accumulate`] and specify a `gradient_accumulation_steps`. 
+This will also automatically ensure the gradients are synced or unsynced when on multi-node, check if the step should
+actually be performed, and auto-scale the loss:
+
+```python
+accelerator = Accelerator(gradient_accumulation_steps=2)
+
+for (input, label) in enumerate(training_dataloader):
+    with accelerator.accumulate(model):
+        predictions = model(input)
+        loss = loss_function(predictions, labels)
+        accelerator.backward(loss)
+        optimizer.step()
+        scheduler.step()
+        optimizer.zero_grad()
+```
+
+**Note**: Using this with `dispatch_batches=True` is currently not supported.
```suggestion
        gradient_accumulation_steps (`int`, *optional*, default to 1):
```
@@ -92,6 +92,9 @@ class Accelerator:
             default to the value in the environment variable `MIXED_PRECISION`, which will use the default value in the
             accelerate config of the current system or the flag passed with the `accelerate.launch` command. 'fp16'
             requires pytorch 1.6 or higher. 'bf16' requires pytorch 1.10 or higher.
+        gradient_accumulation_steps (`int`, *optional*):
```suggestion
            The number of steps that should pass before gradients are accumulated. A number > 1 should be combined with
```
@@ -92,6 +92,9 @@ class Accelerator:
             default to the value in the environment variable `MIXED_PRECISION`, which will use the default value in the
             accelerate config of the current system or the flag passed with the `accelerate.launch` command. 'fp16'
             requires pytorch 1.6 or higher. 'bf16' requires pytorch 1.10 or higher.
+        gradient_accumulation_steps (`int`, *optional*):
+            The number of steps that should pass before gradients are accumulated. Should be combined with
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/484
diff --git a/setup.py b/setup.py
index 01dc0439a..a24428c33 100644
--- a/setup.py
+++ b/setup.py
@@ -22,7 +22,7 @@
     "pytest",
     "pytest-xdist",
     "pytest-subtests",
-    "datasets",
+    "datasets<=2.2.2",
     "evaluate",
     "transformers",
     "scipy",

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/481
diff --git a/examples/by_feature/checkpointing.py b/examples/by_feature/checkpointing.py
index 3b8714911..6880f4626 100644
--- a/examples/by_feature/checkpointing.py
+++ b/examples/by_feature/checkpointing.py
@@ -136,7 +136,7 @@ def training_function(config, args):
 
     # If the batch size is too big we use gradient accumulation
     gradient_accumulation_steps = 1
-    if batch_size > MAX_GPU_BATCH_SIZE:
+    if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
         gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
         batch_size = MAX_GPU_BATCH_SIZE
 
diff --git a/examples/by_feature/cross_validation.py b/examples/by_feature/cross_validation.py
index f1f54d78d..e380d2d85 100644
--- a/examples/by_feature/cross_validation.py
+++ b/examples/by_feature/cross_validation.py
@@ -143,7 +143,7 @@ def training_function(config, args):
 
     # If the batch size is too big we use gradient accumulation
     gradient_accumulation_steps = 1
-    if batch_size > MAX_GPU_BATCH_SIZE:
+    if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
         gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
         batch_size = MAX_GPU_BATCH_SIZE
 
diff --git a/examples/by_feature/fsdp_with_peak_mem_tracking.py b/examples/by_feature/fsdp_with_peak_mem_tracking.py
index 4f0311785..2d1cf2f95 100644
--- a/examples/by_feature/fsdp_with_peak_mem_tracking.py
+++ b/examples/by_feature/fsdp_with_peak_mem_tracking.py
@@ -140,7 +140,7 @@ def tokenize_function(examples):
 
     # If the batch size is too big we use gradient accumulation
     gradient_accumulation_steps = 1
-    if batch_size > MAX_GPU_BATCH_SIZE:
+    if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
         gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
         batch_size = MAX_GPU_BATCH_SIZE
 
diff --git a/examples/by_feature/gradient_accumulation.py b/examples/by_feature/gradient_accumulation.py
index d0a044bcd..b69978023 100644
--- a/examples/by_feature/gradient_accumulation.py
+++ b/examples/by_feature/gradient_accumulation.py
@@ -111,6 +111,10 @@ def training_function(config, args):
     batch_size = int(config["batch_size"])
     # New Code #
     gradient_accumulation_steps = int(args.gradient_accumulation_steps)
+    if accelerator.distributed_type == DistributedType.TPU and gradient_accumulation_steps > 1:
+        raise NotImplementedError(
+            "Gradient accumulation on TPUs is currently not supported. Pass `gradient_accumulation_steps=1`"
+        )
 
     metric = evaluate.load("glue", "mrpc")
 
@@ -151,6 +155,7 @@ def training_function(config, args):
             # We use the new `no_sync` context manager to prevent gradient averaging
             # until we want to at the proper step if we happen to be in a distributed setup
             # otherwise it does nothing
+            # We also currently do not support TPUs nor advise it as bugs were found on the XLA side when running our tests.
             if step % gradient_accumulation_steps != 0:
                 # Accumulate gradients locally
                 with accelerator.no_sync(model):
diff --git a/examples/by_feature/memory.py b/examples/by_feature/memory.py
index 7aef8e1ff..65feed51c 100644
--- a/examples/by_feature/memory.py
+++ b/examples/by_feature/memory.py
@@ -119,7 +119,7 @@ def training_function(config, args):
 
     # If the batch size is too big we use gradient accumulation
     gradient_accumulation_steps = 1
-    if batch_size > MAX_GPU_BATCH_SIZE:
+    if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
         gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
         batch_size = MAX_GPU_BATCH_SIZE
 
diff --git a/examples/by_feature/multi_process_metrics.py b/examples/by_feature/multi_process_metrics.py
index 788c54f50..c9445ee81 100644
--- a/examples/by_feature/multi_process_metrics.py
+++ b/examples/by_feature/multi_process_metrics.py
@@ -121,7 +121,7 @@ def training_function(config, args):
 
     # If the batch size is too big we use gradient accumulation
     gradient_accumulation_steps = 1
-    if batch_size > MAX_GPU_BATCH_SIZE:
+    if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
         gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
         batch_size = MAX_GPU_BATCH_SIZE
 
diff --git a/examples/by_feature/tracking.py b/examples/by_feature/tracking.py
index 3b7b4c8d4..ef2b1388b 100644
--- a/examples/by_feature/tracking.py
+++ b/examples/by_feature/tracking.py
@@ -131,7 +131,7 @@ def training_function(config, args):
 
     # If the batch size is too big we use gradient accumulation
     gradient_accumulation_steps = 1
-    if batch_size > MAX_GPU_BATCH_SIZE:
+    if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
         gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
         batch_size = MAX_GPU_BATCH_SIZE
 
diff --git a/examples/complete_nlp_example.py b/examples/complete_nlp_example.py
index f13fdfd27..0433ed9a1 100644
--- a/examples/complete_nlp_example.py
+++ b/examples/complete_nlp_example.py
@@ -104,7 +104,7 @@ def tokenize_function(examples):
 
     # If the batch size is too big we use gradient accumulation
     gradient_accumulation_steps = 1
-    if batch_size > MAX_GPU_BATCH_SIZE:
+    if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
         gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
         batch_size = MAX_GPU_BATCH_SIZE
 
diff --git a/examples/nlp_example.py b/examples/nlp_example.py
index 00713e886..55dedae5e 100644
--- a/examples/nlp_example.py
+++ b/examples/nlp_example.py
@@ -105,7 +105,7 @@ def training_function(config, args):
 
     # If the batch size is too big we use gradient accumulation
     gradient_accumulation_steps = 1
-    if batch_size > MAX_GPU_BATCH_SIZE:
+    if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
         gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
         batch_size = MAX_GPU_BATCH_SIZE
 

No, you should raise an error at the beginning if on TPU with a gradient accumulation steps argument passed > 1.
@@ -110,7 +110,10 @@ def training_function(config, args):
     seed = int(config["seed"])
     batch_size = int(config["batch_size"])
     # New Code #
-    gradient_accumulation_steps = int(args.gradient_accumulation_steps)
+    if accelerator.distributed_type != DistributedType.TPU:
+        gradient_accumulation_steps = int(args.gradient_accumulation_steps)
+    else:
+        gradient_accumulation_steps = 1
Don't need to change here since `gradient_accumulation_steps` will always be 1.
@@ -196,7 +196,8 @@ def collate_fn(examples):
             if args.with_tracking:
                 total_loss += loss.detach().float()
             accelerator.backward(loss)
-            if step % gradient_accumulation_steps == 0:
+            # We currently do not support nor recommend gradient accumulation on the TPU as bugs were found on the XLA side when running our tests
+            if step % gradient_accumulation_steps == 0 or accelerator.distributed_type == DistributedType.TPU:
Same here.
@@ -146,7 +146,8 @@ def training_function(config, args):
             loss = outputs.loss
             loss = loss / gradient_accumulation_steps
             accelerator.backward(loss)
-            if step % gradient_accumulation_steps == 0:
+            # We currently do not support nor recommend gradient accumulation on the TPU as bugs were found on the XLA side when running our tests
+            if step % gradient_accumulation_steps == 0 or accelerator.distributed_type == DistributedType.TPU:
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/479
diff --git a/setup.py b/setup.py
index b48dc63ad..01dc0439a 100644
--- a/setup.py
+++ b/setup.py
@@ -22,7 +22,7 @@
     "pytest",
     "pytest-xdist",
     "pytest-subtests",
-    "datasets<=2.3.1",
+    "datasets",
     "evaluate",
     "transformers",
     "scipy",

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/478
diff --git a/setup.py b/setup.py
index 01dc0439a..b48dc63ad 100644
--- a/setup.py
+++ b/setup.py
@@ -22,7 +22,7 @@
     "pytest",
     "pytest-xdist",
     "pytest-subtests",
-    "datasets",
+    "datasets<=2.3.1",
     "evaluate",
     "transformers",
     "scipy",

_The documentation is not available anymore as the PR was closed or merged._
Looks like the PR with the fix will be merged soon though? But feel free to pin if you want in the meantime :-)
https://api.github.com/repos/huggingface/accelerate/pulls/477
diff --git a/examples/by_feature/fsdp_with_peak_mem_tracking.py b/examples/by_feature/fsdp_with_peak_mem_tracking.py
index 4f0311785..421a79d6b 100644
--- a/examples/by_feature/fsdp_with_peak_mem_tracking.py
+++ b/examples/by_feature/fsdp_with_peak_mem_tracking.py
@@ -112,11 +112,8 @@ def training_function(config, args):
     # We need to initialize the trackers we use, and also store our configuration
     if args.with_tracking:
         if accelerator.is_main_process:
-            run = os.path.split(__file__)[-1].split(".")[0]
-            if args.logging_dir:
-                run = os.path.join(args.logging_dir, run)
-                accelerator.print(run)
-            accelerator.init_trackers(run, config)
+            experiment_config = vars(args)
+            accelerator.init_trackers("fsdp_glue_no_trainer", experiment_config)
 
     tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
     datasets = load_dataset("glue", "mrpc")
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 1b3fa27e9..99afddeb8 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -188,8 +188,8 @@ def __init__(
             deepspeed_plugin.set_deepspeed_weakref()
 
         if os.environ.get("USE_FSDP", "false") == "true" or isinstance(fsdp_plugin, FullyShardedDataParallelPlugin):
-            if is_torch_version("<", "1.12.0.dev20220418+cu113"):
-                raise ValueError("FSDP requires PyTorch >= 1.12.0.dev20220418+cu113")
+            if is_torch_version("<", "1.12.0"):
+                raise ValueError("FSDP requires PyTorch >= 1.12.0")
 
         if fsdp_plugin is None:  # init from env variables
             fsdp_plugin = FullyShardedDataParallelPlugin() if os.environ.get("USE_FSDP", "false") == "true" else None
diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py
index 3ccf2012b..28fc3826d 100644
--- a/src/accelerate/utils/dataclasses.py
+++ b/src/accelerate/utils/dataclasses.py
@@ -442,7 +442,7 @@ class FullyShardedDataParallelPlugin:
 
     def __post_init__(self):
         from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload, ShardingStrategy
-        from torch.distributed.fsdp.wrap import default_auto_wrap_policy
+        from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy
 
         if self.sharding_strategy is None:
             self.sharding_strategy = ShardingStrategy(int(os.environ.get("FSDP_SHARDING_STRATEGY", 1)))
@@ -458,4 +458,6 @@ def __post_init__(self):
 
         if self.auto_wrap_policy is None:
             if self.min_num_params > 0:
-                self.auto_wrap_policy = functools.partial(default_auto_wrap_policy, min_num_params=self.min_num_params)
+                self.auto_wrap_policy = functools.partial(
+                    size_based_auto_wrap_policy, min_num_params=self.min_num_params
+                )

@pacman100 I'd recommend perhaps doing a version check to see what version of torch we're on to import one vs the other?
_The documentation is not available anymore as the PR was closed or merged._
> @pacman100 I'd recommend perhaps doing a version check to see what version of torch we're on to import one vs the other?

Done. As the FSDP support needs torch nightlies, just bumping up the version is viable.
https://api.github.com/repos/huggingface/accelerate/pulls/475
diff --git a/docs/source/accelerator.mdx b/docs/source/accelerator.mdx
index 990b75443..6e7a93a3d 100644
--- a/docs/source/accelerator.mdx
+++ b/docs/source/accelerator.mdx
@@ -13,7 +13,7 @@ specific language governing permissions and limitations under the License.
 # Accelerator
 
 The [`Accelerator`] is the main class provided by 🤗 Accelerate. It serves at the main entrypoint for
-the API. To quickly adapt your script to work on any kind of setup with 🤗 Accelerate juste:
+the API. To quickly adapt your script to work on any kind of setup with 🤗 Accelerate just:
 
 1. Initialize an [`Accelerator`] object (that we will call `accelerator` in the rest of this
    page) as early as possible in your script.
@@ -21,10 +21,10 @@ the API. To quickly adapt your script to work on any kind of setup with 🤗 Acc
 3. (Optional but best practice) Remove all the `.cuda()` or `.to(device)` in your code and let the
    `accelerator` handle device placement for you.
 4. Replace the `loss.backward()` in your code by `accelerator.backward(loss)`.
-5. (Optional, when using distributed evaluation) Gather your predictions and labelsbefore storing them or using them
-   for metric computation using [`~Accelerator.gather`].
+5. (Optional, when using distributed evaluation) Gather your predictions and labels before storing them or using
+   them for metric computation using [`~Accelerator.gather`].
 
-This is all what is needed in most cases. For more advanced case or a nicer experience here are the functions you
+This is all that is needed in most cases. For more advanced cases or a nicer experience here are the functions you
 should search for and replace by the corresponding methods of your `accelerator`:
 
 - `print` statements should be replaced by [`~Accelerator.print`] to be only printed once per
diff --git a/docs/source/big_modeling.mdx b/docs/source/big_modeling.mdx
index 4ca316870..a5cb3d4ed 100644
--- a/docs/source/big_modeling.mdx
+++ b/docs/source/big_modeling.mdx
@@ -27,7 +27,7 @@ In plain English, those steps are:
 2. Load the model weights (in a dictionary usually called a state dict) from the disk
 3. Load those weights inside the model
 
-While this works very well for regularly sized models, this workflow has some clear limitation when we deal with a huge model: in step 1, we load a full version of the model in RAM, and spend some time randomly initializing the weights (which will be discarded in step 3). In step 2, we load another full version of the model in RAM, with the pretrained weights. If you're loading a model with 6 billions parameters, this needs you will need 24GB of RAM for each copy of the model, so 48GB in total (half of it to load the model in FP16).
+While this works very well for regularly sized models, this workflow has some clear limitations when we deal with a huge model: in step 1, we load a full version of the model in RAM, and spend some time randomly initializing the weights (which will be discarded in step 3). In step 2, we load another full version of the model in RAM, with the pretrained weights. If you're loading a model with 6 billions parameters, this means you will need 24GB of RAM for each copy of the model, so 48GB in total (half of it to load the model in FP16).
 
 <Tip warning={true}>
 
@@ -37,7 +37,7 @@ This API is quite new and still in its experimental stage. While we strive to pr
 
 ## Instantiating an empty model
 
-The first tool Accelerate introduces to help with big models is a context manager [`init_empty_weights`] that helps you initialize a model without using any RAM, so that step 1 can be done on models of any size. Here is how it works:
+The first tool 🤗 Accelerate introduces to help with big models is a context manager [`init_empty_weights`] that helps you initialize a model without using any RAM, so that step 1 can be done on models of any size. Here is how it works:
 
 ```py
 from accelerate import init_empty_weights
@@ -65,7 +65,7 @@ You can't move a model initialized like this on CPU or another device directly,
 
 It's possible your model is so big that even a single copy won't fit in RAM. That doesn't mean it can't be loaded: if you have one or several GPUs, this is more memory available to store your model. In this case, it's better if your checkpoint is split in several smaller files that we call checkpoint shards.
 
-Accelerate will handle sharded checkpoints as long as you follow the following format: your checkpoint should be in a folder, with several files containing the partial state dicts, and there should be an index in the JSON format that contains a dictionary mapping parameter names to the file containing their weights. For instance we could have a folder containing:
+🤗 Accelerate will handle sharded checkpoints as long as you follow the following format: your checkpoint should be in a folder, with several files containing the partial state dicts, and there should be an index in the JSON format that contains a dictionary mapping parameter names to the file containing their weights. For instance we could have a folder containing:
 
 ```bash
 first_state_dict.bin
@@ -88,7 +88,7 @@ and `first_state_dict.bin` containing the weights for `"linear1.weight"` and `"l
 
 ## Loading weights
 
-The second tool Accelerate introduces is a function [`load_checkpoint_and_dispatch`], that will allow you to load a checkpoint inside your empty model. This supports full checkpoints (a single file containing the whole state dict) as well as sharded checkpoints. It will also automatically dispatch those weights across the devices you have available (GPUs, CPU RAM), so if you are loading a sharded checkpoint, the maximum RAM usage will be the size of the biggest shard.
+The second tool 🤗 Accelerate introduces is a function [`load_checkpoint_and_dispatch`], that will allow you to load a checkpoint inside your empty model. This supports full checkpoints (a single file containing the whole state dict) as well as sharded checkpoints. It will also automatically dispatch those weights across the devices you have available (GPUs, CPU RAM), so if you are loading a sharded checkpoint, the maximum RAM usage will be the size of the biggest shard.
 
 Here is how we can use this to load the [GPT-J-6B](https://huggingface.co/EleutherAI/gpt-j-6B) model. You clone the sharded version of this model with:
 
@@ -122,14 +122,14 @@ model = load_checkpoint_and_dispatch(
 )
 ```
 
-By passing `device_map="auto"`, we tell Accelerate to determine automatically where to put each layer of the model depending on the available resources:
+By passing `device_map="auto"`, we tell 🤗 Accelerate to determine automatically where to put each layer of the model depending on the available resources:
 - first we use the maximum space available on the GPU(s)
 - if we still need space, we store the remaining weights on the CPU
 - if there is not enough RAM, we store the remaining weights on the hard drive as memory-mapped tensors
 
 `no_split_module_classes=["GPTJBlock"]` indicates that the modules that are `GPTJBlock` should not be split on different devices. You should set here all blocks that include a residual connection of some kind.
 
-You can see the `device_map` that Accelerate picked by accessing the `hf_device_map` attribute of your model:
+You can see the `device_map` that 🤗 Accelerate picked by accessing the `hf_device_map` attribute of your model:
 
 ```py
 model.hf_device_map
@@ -190,7 +190,7 @@ output = model.generate(inputs["input_ids"])
 tokenizer.decode(output[0].tolist())
 ```
 
-Behind the scenes, Accelerate added hooks to the model, so that:
+Behind the scenes, 🤗 Accelerate added hooks to the model, so that:
 - at each layer, the inputs are put on the right device (so even if your model is spread across several GPUs, it works)
 - for the weights offloaded on the CPU, they are put on a GPU just before the forward pass, and cleaned up just after
 - for the weights offloaded on the hard drive, they are loaded in RAM then put on a GPU just before the forward pass, and cleaned up just after
@@ -207,7 +207,7 @@ This only supports inference of your model, not training. Most of the computatio
 
 We are aware of the current limitations in the API:
 
-- While this could theoretically work just one CPU with potential disk offload, you need at least one GPU to run this API. This will be fixed in further development.
+- While this could theoretically work on just one CPU with potential disk offload, you need at least one GPU to run this API. This will be fixed in further development.
 - [`infer_auto_device_map`] (or `device_map="auto"` in [`load_checkpoint_and_dispatch`]) tries to maximize GPU and CPU RAM it sees available when you execute it. While PyTorch is very good at managing GPU RAM efficiently (and giving it back when not needed), it's not entirely true with Python and CPU RAM. Therefore, an automatically computed device map might be too intense on the CPU. Move a few modules to the disk device if you get crashes due to lack of RAM.
 - [`infer_auto_device_map`] (or `device_map="auto"` in [`load_checkpoint_and_dispatch`]) attributes devices sequentially (to avoid moving things back and forth) so if your first layer is bigger than the size of the GPU you have, it will end up with everything on the CPU/Disk.
 - [`load_checkpoint_and_dispatch`] and [`load_checkpoint_in_model`] do not perform any check on the correctness of your state dict compared to your model at the moment (this will be fixed in a future version), so you may get some weird errors if trying to load a checkpoint with mismatched or missing keys.
diff --git a/docs/source/checkpoint.mdx b/docs/source/checkpoint.mdx
index 26ef21150..c818c9847 100644
--- a/docs/source/checkpoint.mdx
+++ b/docs/source/checkpoint.mdx
@@ -12,8 +12,8 @@ specific language governing permissions and limitations under the License.
 
 # Checkpointing
 
-When training a PyTorch model with Accelerate, you may often want to save and continue a state of training. Doing so requires
-saving and loading the model, optimizer, RNG generators, and the GradScaler. Inside Accelerate are two convience functions to achieve this quickly:
+When training a PyTorch model with 🤗 Accelerate, you may often want to save and continue a state of training. Doing so requires
+saving and loading the model, optimizer, RNG generators, and the GradScaler. Inside 🤗 Accelerate are two convience functions to achieve this quickly:
 - Use [`~Accelerator.save_state`] for saving everything mentioned above to a folder location
 - Use [`~Accelerator.load_state`] for loading everything stored from an earlier `save_state`
 
@@ -57,4 +57,4 @@ for epoch in range(num_epochs):
 
 # Restore previous state
 accelerate.load_state("my/save/path")
-```
\ No newline at end of file
+```
diff --git a/docs/source/fsdp.mdx b/docs/source/fsdp.mdx
index e2a0210c8..32f982786 100644
--- a/docs/source/fsdp.mdx
+++ b/docs/source/fsdp.mdx
@@ -18,7 +18,7 @@ To read more about it and the benefits, check out the [Fully Sharded Data Parall
 We have integrated the latest PyTorch's Fully Sharded Data Parallel (FSDP) training feature.
 All you need to do is enable it through the config.
 
-## How it works out the box
+## How it works out of the box
 
 On your machine(s) just run:
 
@@ -57,7 +57,7 @@ use_cpu: false
 accelerate launch examples/nlp_example.py
 ```
 
-Currently, `Accelerate` supports following config through the CLI:
+Currently, `Accelerate` supports the following config through the CLI:
 
 ```bash
 `Sharding Strategy`: [1] FULL_SHARD, [2] SHARD_GRAD_OP
@@ -65,11 +65,11 @@ Currently, `Accelerate` supports following config through the CLI:
 `Offload Params`: Decides Whether to offload parameters and gradients to CPU.
 ```
 
-## Few caveats to be aware of
+## A few caveats to be aware of
 
 - PyTorch FSDP auto wraps sub-modules, flattens the parameters and shards the parameters in place.
   Due to this, any optimizer created before model wrapping gets broken and occupies more memory.
-  Hence, it is highly recommended and efficient to prepare model before creating optimizer.
+  Hence, it is highly recommended and efficient to prepare the model before creating the optimizer.
   `Accelerate` will automatically wrap the model and create an optimizer for you in case of single model with a warning message.
   > FSDP Warning: When using FSDP, it is efficient and recommended to call prepare for the model before creating the optimizer
 
@@ -91,14 +91,14 @@ optimizer = torch.optim.AdamW(params=model.parameters(), lr=lr)
 
 ```
 
-- In case of a single model, if you have created optimizer with multiple parameter groups and called prepare with them together,
+- In case of a single model, if you have created the optimizer with multiple parameter groups and called prepare with them together,
   then the parameter groups will be lost and the following warning is displayed:
   > FSDP Warning: When using FSDP, several parameter groups will be conflated into
   > a single one due to nested module wrapping and parameter flattening.
   
-  This is because parameter groups created before wrapping will have no meaning post wrapping due parameter flattening of nested FSDP modules into 1D arrays (which can consume many layers).
-  For instance, below are the named parameters of FSDP model on GPU 0 (When using 2 GPUs. Around 55M (110M/2) params in 1D arrays as this will have the 1st shard of the parameters). 
-  Here, if one has applied no weight decay for [bias, LayerNorm.weight] named parameters of unwrapped BERT model, 
+  This is because parameter groups created before wrapping will have no meaning post wrapping due to parameter flattening of nested FSDP modules into 1D arrays (which can consume many layers).
+  For instance, below are the named parameters of an FSDP model on GPU 0 (When using 2 GPUs. Around 55M (110M/2) params in 1D arrays as this will have the 1st shard of the parameters). 
+  Here, if one has applied no weight decay for [bias, LayerNorm.weight] the named parameters of an unwrapped BERT model, 
   it can't be applied to the below FSDP wrapped model as there are no named parameters with either of those strings and 
   the parameters of those layers are concatenated with parameters of various other layers.
   ```
@@ -110,7 +110,7 @@ optimizer = torch.optim.AdamW(params=model.parameters(), lr=lr)
   ```
 
 
-- In case of multiple models, it is necessary to prepare the models before creating optimizers else it will throw an error.
+- In case of multiple models, it is necessary to prepare the models before creating optimizers or else it will throw an error.
 - Mixed precision is currently not supported with FSDP.
 
 For more control, users can leverage the `FullyShardedDataParallelPlugin` wherein they can specify `auto_wrap_policy`, `backward_prefetch` and `ignored_modules`.
diff --git a/docs/source/index.mdx b/docs/source/index.mdx
index 9d9000097..16593d554 100644
--- a/docs/source/index.mdx
+++ b/docs/source/index.mdx
@@ -12,16 +12,16 @@ specific language governing permissions and limitations under the License.
 
 # Accelerate
 
-Run your *raw* PyTorch training script on any kind of device
+Run your *raw* PyTorch training script on any kind of device.
 
 ## Features
 
-- 🤗 Accelerate provides an easy API to make your scripts run with mixed precision and on any kind of distributed
-  setting (multi-GPUs, TPUs etc.) while still letting you write your own training loop. The same code can then runs
+- 🤗 Accelerate provides an easy API to make your scripts run with mixed precision and in any kind of distributed
+  setting (multi-GPUs, TPUs etc.) while still letting you write your own training loop. The same code can then run
   seamlessly on your local machine for debugging or your training environment.
 
-- 🤗 Accelerate also provides a CLI tool that allows you to quickly configure and test your training environment then
-  launch the scripts.
+- 🤗 Accelerate also provides a CLI tool that allows you to quickly configure and test your training environment and
+  then launch the scripts.
 
 
 ## Easy to integrate
diff --git a/docs/source/installation.mdx b/docs/source/installation.mdx
index be8394af7..26c2698aa 100644
--- a/docs/source/installation.mdx
+++ b/docs/source/installation.mdx
@@ -57,7 +57,7 @@ pip install git+https://github.com/huggingface/accelerate
 
 Note that this will install not the latest released version, but the bleeding edge `main` version, which you may want to use in case a bug has been fixed since the last official release and a new release hasn't  been yet rolled out.
 
-While we strive to keep `main` operational at all times, if you notice some issues, they usually get fixed within a few hours or a day and and you're more than welcome to help us detect any problems by opening an [Issue](https://github.com/huggingface/accelerate/issues) and this way, things will get fixed even sooner.
+While we strive to keep `main` operational at all times, if you notice some issues, they usually get fixed within a few hours or a day and you're more than welcome to help us detect any problems by opening an [Issue](https://github.com/huggingface/accelerate/issues) and this way, things will get fixed even sooner.
 
 Again, you can run:
 
@@ -85,7 +85,7 @@ now this editable install will reside where you clone the folder to, e.g. `~/acc
 
 Do note that you have to keep that `accelerate` folder around and not delete it to continue using the 🤗 Accelerate library.
 
-Now, let's get to the real benefit of this installation approach. Say, you saw some new feature has been just committed into `main`. If you have already performed all the steps above, to update your accelerate repo to include all the latest commits, all you need to do is to `cd` into that cloned repository folder and update the clone to the latest version:
+Now, let's get to the real benefit of this installation approach. Say, you saw some new feature just has been committed into `main`. If you have already performed all the steps above, to update your accelerate repo to include all the latest commits, all you need to do is to `cd` into that cloned repository folder and update the clone to the latest version:
 
 ```bash
 cd ~/accelerate/
diff --git a/docs/source/quicktour.mdx b/docs/source/quicktour.mdx
index 1c62f05f5..2e9cbe203 100644
--- a/docs/source/quicktour.mdx
+++ b/docs/source/quicktour.mdx
@@ -12,7 +12,7 @@ specific language governing permissions and limitations under the License.
 
 # Quick tour
 
-Let's have a look at a look at 🤗 Accelerate main features and traps to avoid.
+Let's have a look at the 🤗 Accelerate main features and traps to avoid.
 
 ## Main use
 
@@ -54,7 +54,7 @@ model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
 )
 ```
 
-In particular, your training dataloader will be sharded accross all GPUs/TPU cores available so that each one sees a
+In particular, your training dataloader will be sharded across all GPUs/TPU cores available so that each one sees a
 different portion of the training dataset. Also, the random states of all processes will be synchronized at the
 beginning of each iteration through your dataloader, to make sure the data is shuffled the same way (if you decided to
 use `shuffle=True` or any kind of random sampler).
@@ -118,7 +118,7 @@ method:
 validation_dataloader = accelerator.prepare(validation_dataloader)
 ```
 
-Like for your training dataloader, it will mean that (should you run your script on multiple devices) each device will
+As for your training dataloader, it will mean that (should you run your script on multiple devices) each device will
 only see part of the evaluation data. This means you will need to group your predictions together. This is very easy to
 do with the [`~Accelerator.gather`] method.
 
@@ -134,8 +134,8 @@ for inputs, targets in validation_dataloader:
 
 <Tip warning={true}>
 
-Like for the training dataloader, passing your validation dataloader through
-[`~Accelerator.prepare`] may change its: if you run on X GPUs, it will have its length divided by X
+As for the training dataloader, passing your validation dataloader through
+[`~Accelerator.prepare`] may change it: if you run on X GPUs, it will have its length divided by X
 (since your actual batch size will be multiplied by X), unless you set `split_batches=True`.
 
 Any instruction using your training dataloader length (for instance if you need the number of total training steps
@@ -159,7 +159,7 @@ PyTorch), they are fully compatible with 🤗 Accelerate. The only caveat here i
 to determine all useful information, so `torch.distributed.launch` should be used with the flag `--use_env`.
 
 🤗 Accelerate also provides a CLI tool that unifies all launcher, so you only have to remember one command. To use it,
-just run
+just run:
 
 ```bash
 accelerate config
@@ -175,7 +175,7 @@ on your machine and reply to the questions asked. This will save a *default_conf
 
 You can also specify with the flag `--config_file` the location of the file you want to save.
 
-Once this is done, you can test everything is going well on your setup by running
+Once this is done, you can test everything is going well on your setup by running:
 
 ```bash
 accelerate test
@@ -235,14 +235,14 @@ step). This is why your first step of training will always be very long as build
 optimizations takes some time.
 
 The good news is that this compilation will be cached so the second step and all the following will be much faster. The
-bas news is that it only applies if all of your steps do exactly the same operations, which implies:
+bad news is that it only applies if all of your steps do exactly the same operations, which implies:
 
 - having all tensors of the same length in all your lengths
 - having static code (i.e., not a for loop of length that could change from step to step)
 
 Having any of the things above change between two steps will trigger a new compilation which will, once again, take a
 lot of time. In practice, that means you must take special care to have all your tensors in your inputs of the same
-shape (so no dynamic padding for instance if you are in an NLP problem) and should not use layer with for loops that
+shape (so no dynamic padding for instance if you are in an NLP problem) and should not use layers with for loops that
 have different lengths depending on the inputs (such as an LSTM) or the training will be excruciatingly slow.
 
 To introduce special behavior in your script for TPUs you can check the `distributed_type` of your
@@ -257,10 +257,10 @@ else:
     # go crazy and be dynamic
 ```
 
-The [NLP example](https://github.com/huggingface/accelerate/blob/main/examples/nlp_example.py) shows an example in
+The [NLP example](https://github.com/huggingface/accelerate/blob/main/examples/nlp_example.py) shows an example in a 
 situation with dynamic padding.
 
-One last thing to pay close attnetion to: if your model has tied weights (such as language models which tie the weights
+One last thing to pay close attention to: if your model has tied weights (such as language models which tie the weights
 of the embedding matrix with the weights of the decoder), moving this model to the TPU (either yourself or after you
 passed your model to [`~Accelerator.prepare`]) will break the tying. You will need to retie the weights
 after. You can find an example of this in the [run_clm_no_trainer](https://github.com/huggingface/transformers/blob/master/examples/pytorch/language-modeling/run_clm.py) script in
@@ -317,8 +317,8 @@ following line in your code:
 accelerator.wait_for_everyone()
 ```
 
-This instruction will block all the processes that arrive them first until all the other processes have reached that
-point (if you run your script on just one GPU or CPU, this wont' do anything).
+This instruction will block all the processes that arrive first until all the other processes have reached that
+point (if you run your script on just one GPU or CPU, this won't do anything).
 
 
 ### Saving/loading a model
@@ -338,7 +338,7 @@ unwrapped_model = accelerator.unwrap_model(model)
 accelerator.save(unwrapped_model.state_dict(), filename)
 ```
 
-If your script contains a logic to load checkpoint, we also recommend you load your weights in the unwrapped model
+If your script contains logic to load a checkpoint, we also recommend you load your weights in the unwrapped model
 (this is only useful if you use the load function after making your model go through
 [`~Accelerator.prepare`]). Here is an example:
 
@@ -368,7 +368,7 @@ and `accelerator.clip_grad_value_` respectively.
 
 ### Mixed Precision training
 
-If you are running your training in Mixed Precision with Accelerate, you will get the best result with your loss being
+If you are running your training in Mixed Precision with 🤗 Accelerate, you will get the best result with your loss being
 computed inside your model (like in Transformer models for instance). Every computation outside of the model will be
 executed in full precision (which is generally what you want for loss computation, expecially if it involves a
 softmax). However you might want to put your loss computation inside the *accelerator.autocast* context manager:
@@ -438,14 +438,14 @@ The random number generator synchronization will by default synchronize:
 - the main random number generator in PyTorch <=1.5.1
 
 You can choose which random number generator(s) to synchronize with the `rng_types` argument of the main
-[`Accelerator`]. In PyTorch >= 1.6, it is recommended to rely on local `generator` to avoid
+[`Accelerator`]. In PyTorch >= 1.6, it is recommended to rely on a local `generator` to avoid
 setting the same seed in the main random number generator in all processes.
 
 <Tip warning={true}>
 
-Synchronization the main torch (or CUDA or XLA) random number generator will affect any other potential random
-artifacts you could have in your dataset (like random data augmentation) in the sense all processes will get the
-same random numbers from the torch random modules (so will apply the same random data augmentation if it's
+Synchronization of the main torch (or CUDA or XLA) random number generator will affect any other potential random
+artifacts you could have in your dataset (like random data augmentation) in the sense that all processes will get
+the same random numbers from the torch random modules (so will apply the same random data augmentation if it's
 controlled by torch).
 
 </Tip>
@@ -457,4 +457,4 @@ The randomization part of your custom sampler, batch sampler or iterable dataset
 
 </Tip>
 
-See more details about the internal in the [Internals page](internal).
+For more details about the internals, see the [Internals page](internal).
diff --git a/docs/source/sagemaker.mdx b/docs/source/sagemaker.mdx
index d55ccd28b..563f0c742 100644
--- a/docs/source/sagemaker.mdx
+++ b/docs/source/sagemaker.mdx
@@ -23,7 +23,7 @@ make it easier than ever to train Hugging Face Transformer models in [Amazon Sag
 Before you can run your 🤗 Accelerate scripts on Amazon SageMaker you need to sign up for an AWS account. If you do not
 have an AWS account yet learn more [here](https://docs.aws.amazon.com/sagemaker/latest/dg/gs-set-up.html).
 
-After you have your AWS Account you need to install the `sagemaker` sdk for 🤗 Accelerate with.
+After you have your AWS Account you need to install the `sagemaker` sdk for 🤗 Accelerate with:
 
 ```bash
 pip install "accelerate[sagemaker]" --upgrade
@@ -31,7 +31,7 @@ pip install "accelerate[sagemaker]" --upgrade
 
 🤗 Accelerate currently uses the 🤗 DLCs, with `transformers`, `datasets` and `tokenizers` pre-installed. 🤗
 Accelerate is not in the DLC yet (will soon be added!) so to use it within Amazon SageMaker you need to create a
-`requirements.txt` in the same directory where your training script is located and add it as dependency.
+`requirements.txt` in the same directory where your training script is located and add it as dependency:
 
 ```
 accelerate
@@ -43,7 +43,7 @@ You should also add any other dependencies you have to this `requirements.txt`.
 ### Configure 🤗 Accelerate
 
 You can configure the launch configuration for Amazon SageMaker the same as you do for non SageMaker training jobs with
-the 🤗 Accelerate CLI.
+the 🤗 Accelerate CLI:
 
 ```bash
 accelerate config
@@ -62,7 +62,7 @@ accelerate config
 
 The training script is very similar to a training script you might run outside of SageMaker, but to save your model
 after training you need to specify either `/opt/ml/model` or use `os.environ["SM_MODEL_DIR"]` as your save
-directory. After training, artifacts in this directory are uploaded to S3.
+directory. After training, artifacts in this directory are uploaded to S3:
 
 
 ```diff
@@ -79,7 +79,7 @@ specify type as bool in your script and provide an explicit True or False value
 
 ### Launch Training
 
-You can launch your training with 🤗 Accelerate CLI with
+You can launch your training with 🤗 Accelerate CLI with:
 
 ```
 accelerate launch path_to_script.py --args_to_the_script
diff --git a/docs/source/tracking.mdx b/docs/source/tracking.mdx
index 6d4ca437c..25e44b3c9 100644
--- a/docs/source/tracking.mdx
+++ b/docs/source/tracking.mdx
@@ -13,7 +13,7 @@ specific language governing permissions and limitations under the License.
 # Tracking
 
 There are a large number of experiment tracking API's available, however getting them all to work with in a multi-processing environment can oftentimes be complex.
-Accelerate provides a general tracking API that can be used to log useful items during your script through [`~Accelerator.log`]
+🤗 Accelerate provides a general tracking API that can be used to log useful items during your script through [`~Accelerator.log`]
 
 ## Integrated Trackers
 
diff --git a/docs/source/utilities.mdx b/docs/source/utilities.mdx
index f5756a262..bfcaeb60e 100644
--- a/docs/source/utilities.mdx
+++ b/docs/source/utilities.mdx
@@ -12,11 +12,11 @@ specific language governing permissions and limitations under the License.
 
 # Helpful Utilities
 
-Below are a variety of utility functions that Accelerate provides, broken down by use-case. 
+Below are a variety of utility functions that 🤗 Accelerate provides, broken down by use-case. 
 
 ## Data Classes
 
-These are basic dataclasses used throughout Accelerate and can be passed in as parameters.
+These are basic dataclasses used throughout 🤗 Accelerate and they can be passed in as parameters.
 
 [[autodoc]] utils.DistributedType
 
@@ -56,7 +56,7 @@ These functionalities check the state of the current working environment includi
 
 [[autodoc]] utils.write_basic_config
 
-When setting up Accelerate for the first time, rather than running `accelerate config` [~utils.write_basic_config] can be used as alternative for quick configuration.
+When setting up 🤗 Accelerate for the first time, rather than running `accelerate config` [~utils.write_basic_config] can be used as an alternative for quick configuration.
 
 ## Modeling
 
@@ -88,4 +88,4 @@ These utilities relate to setting and synchronizing of all the random states.
 
 [[autodoc]] utils.synchronize_rng_state
 
-[[autodoc]] utils.synchronize_rng_states
\ No newline at end of file
+[[autodoc]] utils.synchronize_rng_states

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/472
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 2d2a0bd56..1b3fa27e9 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -68,7 +68,7 @@
         DummyScheduler,
     )
 
-if is_tpu_available():
+if is_tpu_available(check_device=False):
     import torch_xla.distributed.xla_multiprocessing as xmp
 
 logger = get_logger(__name__)
diff --git a/src/accelerate/checkpointing.py b/src/accelerate/checkpointing.py
index 8b41c48f8..d5e816aab 100644
--- a/src/accelerate/checkpointing.py
+++ b/src/accelerate/checkpointing.py
@@ -33,7 +33,7 @@
 )
 
 
-if is_tpu_available():
+if is_tpu_available(check_device=False):
     import torch_xla.core.xla_model as xm
 
 from .logging import get_logger
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index eb26c3a60..806fb7ddc 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -34,7 +34,7 @@
 )
 
 
-if is_tpu_available():
+if is_tpu_available(check_device=False):
     import torch_xla.distributed.parallel_loader as xpl
 
 
diff --git a/src/accelerate/optimizer.py b/src/accelerate/optimizer.py
index 3d5977968..046d5ce5b 100644
--- a/src/accelerate/optimizer.py
+++ b/src/accelerate/optimizer.py
@@ -21,7 +21,7 @@
 from .utils import DistributedType, honor_type, is_torch_version, is_tpu_available
 
 
-if is_tpu_available():
+if is_tpu_available(check_device=False):
     import torch_xla.core.xla_model as xm
 
 
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index 36b11be56..63527124c 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -20,7 +20,7 @@
 from .utils import DistributedType, is_ccl_available, is_deepspeed_available, is_tpu_available
 
 
-if is_tpu_available():
+if is_tpu_available(check_device=False):
     import torch_xla.core.xla_model as xm
 
 
@@ -88,7 +88,7 @@ def __init__(
                     "Please make sure to properly initialize your accelerator via `accelerator = Accelerator()` "
                     "before using any functionality from the `accelerate` library."
                 )
-            elif is_tpu_available() and not cpu:
+            if is_tpu_available() and not cpu:
                 self.distributed_type = DistributedType.TPU
                 self.num_processes = xm.xrt_world_size()
                 self.process_index = xm.get_ordinal()
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index 594f7aa1c..a660db858 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -38,12 +38,7 @@
 try:
     import torch_xla.core.xla_model as xm  # noqa: F401
 
-    try:
-        # Will raise a RuntimeError if no XLA configuration is found
-        _ = xm.xla_device()
-        _tpu_available = True
-    except RuntimeError:
-        _tpu_available = False
+    _tpu_available = True
 except ImportError:
     _tpu_available = False
 
@@ -56,8 +51,15 @@ def is_apex_available():
     return importlib.util.find_spec("apex") is not None
 
 
-def is_tpu_available():
-    "Checks if `torch_xla` is installed and if a TPU is in the environment"
+def is_tpu_available(check_device=True):
+    "Checks if `torch_xla` is installed and potentially if a TPU is in the environment"
+    if _tpu_available and check_device:
+        try:
+            # Will raise a RuntimeError if no XLA configuration is found
+            _ = xm.xla_device()
+            return True
+        except RuntimeError:
+            return False
     return _tpu_available
 
 
diff --git a/src/accelerate/utils/operations.py b/src/accelerate/utils/operations.py
index 954bb5dfd..46753513e 100644
--- a/src/accelerate/utils/operations.py
+++ b/src/accelerate/utils/operations.py
@@ -29,7 +29,7 @@
 from .versions import is_torch_version
 
 
-if is_tpu_available():
+if is_tpu_available(check_device=False):
     import torch_xla.core.xla_model as xm
 
 
diff --git a/src/accelerate/utils/other.py b/src/accelerate/utils/other.py
index 0d4f5b7e2..206c70589 100644
--- a/src/accelerate/utils/other.py
+++ b/src/accelerate/utils/other.py
@@ -28,7 +28,7 @@
 if is_deepspeed_available():
     from deepspeed import DeepSpeedEngine
 
-if is_tpu_available():
+if is_tpu_available(check_device=False):
     import torch_xla.core.xla_model as xm
 
 
diff --git a/src/accelerate/utils/random.py b/src/accelerate/utils/random.py
index 3958022a4..51b7c3719 100644
--- a/src/accelerate/utils/random.py
+++ b/src/accelerate/utils/random.py
@@ -23,7 +23,7 @@
 from .imports import is_tpu_available
 
 
-if is_tpu_available():
+if is_tpu_available(check_device=False):
     import torch_xla.core.xla_model as xm
 
 

Mmm, removing this will break the other places we use `is_tpu_available`
@@ -38,12 +38,7 @@
 try:
     import torch_xla.core.xla_model as xm  # noqa: F401
 
-    try:
-        # Will raise a RuntimeError if no XLA configuration is found
-        _ = xm.xla_device()
-        _tpu_available = True
-    except RuntimeError:
-        _tpu_available = False
The added reason was specifically inside `AcceleratorState`, at the proposed location. But having it as an argument instead works as well. Will refactor
@@ -38,12 +38,7 @@
 try:
     import torch_xla.core.xla_model as xm  # noqa: F401
 
-    try:
-        # Will raise a RuntimeError if no XLA configuration is found
-        _ = xm.xla_device()
-        _tpu_available = True
-    except RuntimeError:
-        _tpu_available = False
This needs to be `True` by default, and `False` when we don't want to check for the device (before launching multiprocessing for instance).
@@ -56,8 +51,15 @@ def is_apex_available():
     return importlib.util.find_spec("apex") is not None
 
 
-def is_tpu_available():
-    "Checks if `torch_xla` is installed and if a TPU is in the environment"
+def is_tpu_available(check_device=False):
It needs to be `False` by default, because otherwise it also does this on import checks that are scattered around the library.
@@ -56,8 +51,15 @@ def is_apex_available():
     return importlib.util.find_spec("apex") is not None
 
 
-def is_tpu_available():
-    "Checks if `torch_xla` is installed and if a TPU is in the environment"
+def is_tpu_available(check_device=False):
Or made into a separate function if we don't want `False` behavior. (I know we're not fans of that, but this is one case where it should be `False`)
@@ -56,8 +51,15 @@ def is_apex_available():
     return importlib.util.find_spec("apex") is not None
 
 
-def is_tpu_available():
-    "Checks if `torch_xla` is installed and if a TPU is in the environment"
+def is_tpu_available(check_device=False):
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/469
diff --git a/setup.py b/setup.py
index 94516436b..e548a3940 100644
--- a/setup.py
+++ b/setup.py
@@ -16,7 +16,7 @@
 from setuptools import find_packages
 
 extras = {}
-extras["quality"] = ["black ~= 22.0", "isort >= 5.5.4", "flake8 >= 3.8.3"]
+extras["quality"] = ["black ~= 22.0", "isort >= 5.5.4", "flake8 >= 3.8.3", "hf-doc-builder >= 0.3.0"]
 extras["docs"] = []
 extras["test"] = [
     "pytest",

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/468
diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml
index 4b2dfac18..04571e50e 100644
--- a/docs/source/_toctree.yml
+++ b/docs/source/_toctree.yml
@@ -19,10 +19,10 @@
     title: Notebook Launcher
   - local: kwargs
     title: Kwargs Handlers
-  - local: internal
-    title: Internals
   - local: checkpoint
     title: Checkpointing
+  - local: internal
+    title: Internals
   - local: tracking
     title: Experiment Tracking
   - local: fsdp
@@ -31,4 +31,6 @@
     title: Memory Utilities
   - local: deepspeed
     title: DeepSpeed
+  - local: utilities
+    title: General Utilities
   title: API Reference
diff --git a/docs/source/internal.mdx b/docs/source/internal.mdx
index 4f7f59500..e1eec36be 100644
--- a/docs/source/internal.mdx
+++ b/docs/source/internal.mdx
@@ -44,34 +44,6 @@ The main work on your PyTorch `DataLoader` is done by the following function:
 
 [[autodoc]] state.AcceleratorState
 
-### DistributedType
-
-[[autodoc]] state.DistributedType
-
 ## Tracking
 
 [[autodoc]] tracking.GeneralTracker
-
-## Utilities
-
-[[autodoc]] utils.extract_model_from_parallel
-
-[[autodoc]] utils.is_bf16_available
-
-[[autodoc]] utils.is_torch_version
-
-[[autodoc]] utils.is_tpu_available
-
-[[autodoc]] utils.gather
-
-[[autodoc]] utils.send_to_device
-
-[[autodoc]] utils.set_seed
-
-[[autodoc]] utils.synchronize_rng_state
-
-[[autodoc]] utils.synchronize_rng_states
-
-[[autodoc]] utils.wait_for_everyone
-
-[[autodoc]] utils.write_basic_config
diff --git a/docs/source/utilities.mdx b/docs/source/utilities.mdx
new file mode 100644
index 000000000..f5756a262
--- /dev/null
+++ b/docs/source/utilities.mdx
@@ -0,0 +1,91 @@
+<!--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.
+-->
+
+# Helpful Utilities
+
+Below are a variety of utility functions that Accelerate provides, broken down by use-case. 
+
+## Data Classes
+
+These are basic dataclasses used throughout Accelerate and can be passed in as parameters.
+
+[[autodoc]] utils.DistributedType
+
+[[autodoc]] utils.LoggerType
+
+[[autodoc]] utils.PrecisionType
+
+## Data Manipulation and Operations
+
+These include data operations that mimic the same `torch` ops but can be used on distributed processes.
+
+[[autodoc]] utils.broadcast
+
+[[autodoc]] utils.concatenate
+
+[[autodoc]] utils.gather
+
+[[autodoc]] utils.pad_across_processes
+
+[[autodoc]] utils.reduce
+
+[[autodoc]] utils.send_to_device
+
+## Environment Checks
+
+These functionalities check the state of the current working environment including information about the operating system itself, what it can support, and if particular dependencies are installed. 
+
+[[autodoc]] utils.get_max_memory
+
+[[autodoc]] utils.is_bf16_available
+
+[[autodoc]] utils.is_torch_version
+
+[[autodoc]] utils.is_tpu_available
+
+## Environment Configuration
+
+[[autodoc]] utils.write_basic_config
+
+When setting up Accelerate for the first time, rather than running `accelerate config` [~utils.write_basic_config] can be used as alternative for quick configuration.
+
+## Modeling
+
+These utilities relate to interacting with PyTorch models
+
+[[autodoc]] utils.extract_model_from_parallel
+
+[[autodoc]] utils.get_max_layer_size
+
+[[autodoc]] utils.offload_state_dict
+
+
+## Parallel
+
+These include general utilities that should be used when working in parallel.
+
+[[autodoc]] utils.extract_model_from_parallel
+
+[[autodoc]] utils.save
+
+[[autodoc]] utils.wait_for_everyone
+
+
+## Random
+
+These utilities relate to setting and synchronizing of all the random states.
+
+[[autodoc]] utils.set_seed
+
+[[autodoc]] utils.synchronize_rng_state
+
+[[autodoc]] utils.synchronize_rng_states
\ No newline at end of file
diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py
index 8b8e30304..3ccf2012b 100644
--- a/src/accelerate/utils/dataclasses.py
+++ b/src/accelerate/utils/dataclasses.py
@@ -179,6 +179,16 @@ def list(cls):
 
 
 class LoggerType(BaseEnum):
+    """Represents a type of supported experiment tracker
+
+    Values:
+
+        - **ALL** -- all available trackers in the environment that are supported
+        - **TENSORBOARD** -- TensorBoard as an experiment tracker
+        - **WANDB** -- wandb as an experiment tracker
+        - **COMETML** -- comet_ml as an experiment tracker
+    """
+
     ALL = "all"
     TENSORBOARD = "tensorboard"
     WANDB = "wandb"
@@ -186,6 +196,15 @@ class LoggerType(BaseEnum):
 
 
 class PrecisionType(BaseEnum):
+    """Represents a type of precision used on floating point values
+
+    Values:
+
+        - **NO** -- using full precision (FP32)
+        - **FP16** -- using half precision
+        - **BF16** -- using brain floating point precision
+    """
+
     NO = "no"
     FP16 = "fp16"
     BF16 = "bf16"

Would be nice to add more content to this page introducing each section!
@@ -0,0 +1,77 @@
+<!--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.
+-->
+
+# Helpful Utilities
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/467
diff --git a/docs/source/internal.mdx b/docs/source/internal.mdx
index 9b496b30e..4f7f59500 100644
--- a/docs/source/internal.mdx
+++ b/docs/source/internal.mdx
@@ -56,6 +56,12 @@ The main work on your PyTorch `DataLoader` is done by the following function:
 
 [[autodoc]] utils.extract_model_from_parallel
 
+[[autodoc]] utils.is_bf16_available
+
+[[autodoc]] utils.is_torch_version
+
+[[autodoc]] utils.is_tpu_available
+
 [[autodoc]] utils.gather
 
 [[autodoc]] utils.send_to_device
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index d1d634355..594f7aa1c 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -57,6 +57,7 @@ def is_apex_available():
 
 
 def is_tpu_available():
+    "Checks if `torch_xla` is installed and if a TPU is in the environment"
     return _tpu_available
 
 

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/466
diff --git a/examples/by_feature/checkpointing.py b/examples/by_feature/checkpointing.py
index c06d06ee9..3b8714911 100644
--- a/examples/by_feature/checkpointing.py
+++ b/examples/by_feature/checkpointing.py
@@ -16,18 +16,13 @@
 import os
 
 import torch
+from torch.optim import AdamW
 from torch.utils.data import DataLoader
 
 import evaluate
 from accelerate import Accelerator, DistributedType
 from datasets import load_dataset
-from transformers import (
-    AdamW,
-    AutoModelForSequenceClassification,
-    AutoTokenizer,
-    get_linear_schedule_with_warmup,
-    set_seed,
-)
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
 
 
 ########################################################################
@@ -117,7 +112,6 @@ def training_function(config, args):
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
     lr = config["lr"]
     num_epochs = int(config["num_epochs"])
-    correct_bias = config["correct_bias"]
     seed = int(config["seed"])
     batch_size = int(config["batch_size"])
 
@@ -155,7 +149,7 @@ def training_function(config, args):
     model = model.to(accelerator.device)
 
     # Instantiate optimizer
-    optimizer = AdamW(params=model.parameters(), lr=lr, correct_bias=correct_bias)
+    optimizer = AdamW(params=model.parameters(), lr=lr)
 
     # Instantiate scheduler
     lr_scheduler = get_linear_schedule_with_warmup(
@@ -297,7 +291,7 @@ def main():
         help="If the training should continue from a checkpoint folder.",
     )
     args = parser.parse_args()
-    config = {"lr": 2e-5, "num_epochs": 3, "correct_bias": True, "seed": 42, "batch_size": 16}
+    config = {"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
     training_function(config, args)
 
 
diff --git a/examples/by_feature/cross_validation.py b/examples/by_feature/cross_validation.py
index 53d582c67..f1f54d78d 100644
--- a/examples/by_feature/cross_validation.py
+++ b/examples/by_feature/cross_validation.py
@@ -17,6 +17,7 @@
 
 import numpy as np
 import torch
+from torch.optim import AdamW
 from torch.utils.data import DataLoader
 
 import evaluate
@@ -26,13 +27,7 @@
 # New Code #
 # We'll be using StratifiedKFold for this example
 from sklearn.model_selection import StratifiedKFold
-from transformers import (
-    AdamW,
-    AutoModelForSequenceClassification,
-    AutoTokenizer,
-    get_linear_schedule_with_warmup,
-    set_seed,
-)
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
 
 
 ########################################################################
@@ -141,7 +136,6 @@ def training_function(config, args):
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
     lr = config["lr"]
     num_epochs = int(config["num_epochs"])
-    correct_bias = config["correct_bias"]
     seed = int(config["seed"])
     batch_size = int(config["batch_size"])
 
@@ -178,7 +172,7 @@ def training_function(config, args):
         model = model.to(accelerator.device)
 
         # Instantiate optimizer
-        optimizer = AdamW(params=model.parameters(), lr=lr, correct_bias=correct_bias)
+        optimizer = AdamW(params=model.parameters(), lr=lr)
 
         # Instantiate scheduler
         lr_scheduler = get_linear_schedule_with_warmup(
@@ -268,7 +262,7 @@ def main():
     # New Code #
     parser.add_argument("--num_folds", type=int, default=3, help="The number of splits to perform across the dataset")
     args = parser.parse_args()
-    config = {"lr": 2e-5, "num_epochs": 3, "correct_bias": True, "seed": 42, "batch_size": 16}
+    config = {"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
     training_function(config, args)
 
 
diff --git a/examples/by_feature/gradient_accumulation.py b/examples/by_feature/gradient_accumulation.py
index 38da07aa1..d0a044bcd 100644
--- a/examples/by_feature/gradient_accumulation.py
+++ b/examples/by_feature/gradient_accumulation.py
@@ -16,18 +16,13 @@
 import os
 
 import torch
+from torch.optim import AdamW
 from torch.utils.data import DataLoader
 
 import evaluate
 from accelerate import Accelerator, DistributedType
 from datasets import load_dataset
-from transformers import (
-    AdamW,
-    AutoModelForSequenceClassification,
-    AutoTokenizer,
-    get_linear_schedule_with_warmup,
-    set_seed,
-)
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
 
 
 ########################################################################
@@ -112,7 +107,6 @@ def training_function(config, args):
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
     lr = config["lr"]
     num_epochs = int(config["num_epochs"])
-    correct_bias = config["correct_bias"]
     seed = int(config["seed"])
     batch_size = int(config["batch_size"])
     # New Code #
@@ -131,7 +125,7 @@ def training_function(config, args):
     model = model.to(accelerator.device)
 
     # Instantiate optimizer
-    optimizer = AdamW(params=model.parameters(), lr=lr, correct_bias=correct_bias)
+    optimizer = AdamW(params=model.parameters(), lr=lr)
 
     # Instantiate scheduler
     lr_scheduler = get_linear_schedule_with_warmup(
@@ -210,7 +204,7 @@ def main():
     )
     parser.add_argument("--cpu", action="store_true", help="If passed, will train on the CPU.")
     args = parser.parse_args()
-    config = {"lr": 2e-5, "num_epochs": 3, "correct_bias": True, "seed": 42, "batch_size": 16}
+    config = {"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
     training_function(config, args)
 
 
diff --git a/examples/by_feature/memory.py b/examples/by_feature/memory.py
index ac4b45866..7aef8e1ff 100644
--- a/examples/by_feature/memory.py
+++ b/examples/by_feature/memory.py
@@ -15,6 +15,7 @@
 import os
 
 import torch
+from torch.optim import AdamW
 from torch.utils.data import DataLoader
 
 # New Code #
@@ -22,13 +23,7 @@
 from accelerate import Accelerator, DistributedType
 from accelerate.utils import find_executable_batch_size
 from datasets import load_dataset
-from transformers import (
-    AdamW,
-    AutoModelForSequenceClassification,
-    AutoTokenizer,
-    get_linear_schedule_with_warmup,
-    set_seed,
-)
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
 
 
 ########################################################################
@@ -117,7 +112,6 @@ def training_function(config, args):
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
     lr = config["lr"]
     num_epochs = int(config["num_epochs"])
-    correct_bias = config["correct_bias"]
     seed = int(config["seed"])
     batch_size = int(config["batch_size"])
 
@@ -139,7 +133,7 @@ def training_function(config, args):
     model = model.to(accelerator.device)
 
     # Instantiate optimizer
-    optimizer = AdamW(params=model.parameters(), lr=lr, correct_bias=correct_bias)
+    optimizer = AdamW(params=model.parameters(), lr=lr)
 
     # New Code #
     # We now can define an inner training loop function. It should take a batch size as the only parameter,
@@ -218,7 +212,7 @@ def main():
     )
     parser.add_argument("--cpu", action="store_true", help="If passed, will train on the CPU.")
     args = parser.parse_args()
-    config = {"lr": 2e-5, "num_epochs": 3, "correct_bias": True, "seed": 42, "batch_size": 16}
+    config = {"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
     training_function(config, args)
 
 
diff --git a/examples/by_feature/multi_process_metrics.py b/examples/by_feature/multi_process_metrics.py
index 0307d251f..788c54f50 100644
--- a/examples/by_feature/multi_process_metrics.py
+++ b/examples/by_feature/multi_process_metrics.py
@@ -16,18 +16,13 @@
 import os
 
 import torch
+from torch.optim import AdamW
 from torch.utils.data import DataLoader
 
 import evaluate
 from accelerate import Accelerator, DistributedType
 from datasets import load_dataset
-from transformers import (
-    AdamW,
-    AutoModelForSequenceClassification,
-    AutoTokenizer,
-    get_linear_schedule_with_warmup,
-    set_seed,
-)
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
 
 
 ########################################################################
@@ -119,7 +114,6 @@ def training_function(config, args):
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
     lr = config["lr"]
     num_epochs = int(config["num_epochs"])
-    correct_bias = config["correct_bias"]
     seed = int(config["seed"])
     batch_size = int(config["batch_size"])
 
@@ -142,7 +136,7 @@ def training_function(config, args):
     model = model.to(accelerator.device)
 
     # Instantiate optimizer
-    optimizer = AdamW(params=model.parameters(), lr=lr, correct_bias=correct_bias)
+    optimizer = AdamW(params=model.parameters(), lr=lr)
 
     # Instantiate scheduler
     lr_scheduler = get_linear_schedule_with_warmup(
@@ -216,7 +210,7 @@ def main():
     )
     parser.add_argument("--cpu", action="store_true", help="If passed, will train on the CPU.")
     args = parser.parse_args()
-    config = {"lr": 2e-5, "num_epochs": 3, "correct_bias": True, "seed": 42, "batch_size": 16}
+    config = {"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
     training_function(config, args)
 
 
diff --git a/examples/by_feature/tracking.py b/examples/by_feature/tracking.py
index 4020b3382..3b7b4c8d4 100644
--- a/examples/by_feature/tracking.py
+++ b/examples/by_feature/tracking.py
@@ -16,18 +16,13 @@
 import os
 
 import torch
+from torch.optim import AdamW
 from torch.utils.data import DataLoader
 
 import evaluate
 from accelerate import Accelerator, DistributedType
 from datasets import load_dataset
-from transformers import (
-    AdamW,
-    AutoModelForSequenceClassification,
-    AutoTokenizer,
-    get_linear_schedule_with_warmup,
-    set_seed,
-)
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
 
 
 ########################################################################
@@ -127,7 +122,6 @@ def training_function(config, args):
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
     lr = config["lr"]
     num_epochs = int(config["num_epochs"])
-    correct_bias = config["correct_bias"]
     seed = int(config["seed"])
     batch_size = int(config["batch_size"])
     set_seed(seed)
@@ -150,7 +144,7 @@ def training_function(config, args):
     model = model.to(accelerator.device)
 
     # Instantiate optimizer
-    optimizer = AdamW(params=model.parameters(), lr=lr, correct_bias=correct_bias)
+    optimizer = AdamW(params=model.parameters(), lr=lr)
 
     # Instantiate scheduler
     lr_scheduler = get_linear_schedule_with_warmup(
@@ -260,7 +254,7 @@ def main():
         help="Location on where to store experiment tracking logs`",
     )
     args = parser.parse_args()
-    config = {"lr": 2e-5, "num_epochs": 3, "correct_bias": True, "seed": 42, "batch_size": 16}
+    config = {"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
     training_function(config, args)
 
 
diff --git a/examples/complete_nlp_example.py b/examples/complete_nlp_example.py
index 8544d1279..f13fdfd27 100644
--- a/examples/complete_nlp_example.py
+++ b/examples/complete_nlp_example.py
@@ -16,18 +16,13 @@
 import os
 
 import torch
+from torch.optim import AdamW
 from torch.utils.data import DataLoader
 
 import evaluate
 from accelerate import Accelerator, DistributedType
 from datasets import load_dataset
-from transformers import (
-    AdamW,
-    AutoModelForSequenceClassification,
-    AutoTokenizer,
-    get_linear_schedule_with_warmup,
-    set_seed,
-)
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
 
 
 ########################################################################
@@ -76,7 +71,6 @@ def training_function(config, args):
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
     lr = config["lr"]
     num_epochs = int(config["num_epochs"])
-    correct_bias = config["correct_bias"]
     seed = int(config["seed"])
     batch_size = int(config["batch_size"])
 
@@ -139,7 +133,7 @@ def collate_fn(examples):
     model = model.to(accelerator.device)
 
     # Instantiate optimizer
-    optimizer = AdamW(params=model.parameters(), lr=lr, correct_bias=correct_bias)
+    optimizer = AdamW(params=model.parameters(), lr=lr)
 
     # Instantiate scheduler
     lr_scheduler = get_linear_schedule_with_warmup(
@@ -305,7 +299,7 @@ def main():
         help="Location on where to store experiment tracking logs`",
     )
     args = parser.parse_args()
-    config = {"lr": 2e-5, "num_epochs": 3, "correct_bias": True, "seed": 42, "batch_size": 16}
+    config = {"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
     training_function(config, args)
 
 
diff --git a/examples/nlp_example.py b/examples/nlp_example.py
index 84c24fd11..00713e886 100644
--- a/examples/nlp_example.py
+++ b/examples/nlp_example.py
@@ -15,18 +15,13 @@
 import argparse
 
 import torch
+from torch.optim import AdamW
 from torch.utils.data import DataLoader
 
 import evaluate
 from accelerate import Accelerator, DistributedType
 from datasets import load_dataset
-from transformers import (
-    AdamW,
-    AutoModelForSequenceClassification,
-    AutoTokenizer,
-    get_linear_schedule_with_warmup,
-    set_seed,
-)
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
 
 
 ########################################################################
@@ -103,7 +98,6 @@ def training_function(config, args):
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
     lr = config["lr"]
     num_epochs = int(config["num_epochs"])
-    correct_bias = config["correct_bias"]
     seed = int(config["seed"])
     batch_size = int(config["batch_size"])
 
@@ -126,7 +120,7 @@ def training_function(config, args):
     model = model.to(accelerator.device)
 
     # Instantiate optimizer
-    optimizer = AdamW(params=model.parameters(), lr=lr, correct_bias=correct_bias)
+    optimizer = AdamW(params=model.parameters(), lr=lr)
 
     # Instantiate scheduler
     lr_scheduler = get_linear_schedule_with_warmup(
@@ -188,7 +182,7 @@ def main():
     )
     parser.add_argument("--cpu", action="store_true", help="If passed, will train on the CPU.")
     args = parser.parse_args()
-    config = {"lr": 2e-5, "num_epochs": 3, "correct_bias": True, "seed": 42, "batch_size": 16}
+    config = {"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
     training_function(config, args)
 
 
diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
index 0329f8786..62c4f9761 100644
--- a/src/accelerate/utils/modeling.py
+++ b/src/accelerate/utils/modeling.py
@@ -78,7 +78,7 @@ def dtype_byte_size(dtype: torch.dtype):
     """
     if dtype == torch.bool:
         return 1 / 8
-    bit_search = re.search("[^\d](\d+)$", str(dtype))
+    bit_search = re.search(r"[^\d](\d+)$", str(dtype))
     if bit_search is None:
         raise ValueError(f"`dtype` is not a valid dtype: {dtype}.")
     bit_size = int(bit_search.groups()[0])
diff --git a/tests/test_utils.py b/tests/test_utils.py
index f2414dfe5..1e9d18c17 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -23,7 +23,7 @@
 from accelerate.utils import convert_outputs_to_fp32, find_device, patch_environment, send_to_device
 
 
-TestNamedTuple = namedtuple("TestNamedTuple", "a b c")
+ExampleNamedTuple = namedtuple("ExampleNamedTuple", "a b c")
 
 
 class UtilsTester(unittest.TestCase):
@@ -50,8 +50,8 @@ def test_send_to_device(self):
         self.assertTrue(torch.equal(result2["b"][1].cpu(), tensor))
         self.assertEqual(result2["c"], 1)
 
-        result3 = send_to_device(TestNamedTuple(a=tensor, b=[tensor, tensor], c=1), device)
-        self.assertIsInstance(result3, TestNamedTuple)
+        result3 = send_to_device(ExampleNamedTuple(a=tensor, b=[tensor, tensor], c=1), device)
+        self.assertIsInstance(result3, ExampleNamedTuple)
         self.assertTrue(torch.equal(result3.a.cpu(), tensor))
         self.assertIsInstance(result3.b, list)
         self.assertTrue(torch.equal(result3.b[0].cpu(), tensor))

Switch to the torch `AdamW`? The HF ones as small bugs that can hurt larger trainings.
@@ -155,7 +155,7 @@ def training_function(config, args):
     model = model.to(accelerator.device)
 
     # Instantiate optimizer
-    optimizer = AdamW(params=model.parameters(), lr=lr, correct_bias=correct_bias)
+    optimizer = AdamW(params=model.parameters(), lr=lr, correct_bias=correct_bias, no_deprecation_warning=True)
We don't set `correct_bias=False`, so makes sense. Will update
@@ -155,7 +155,7 @@ def training_function(config, args):
     model = model.to(accelerator.device)
 
     # Instantiate optimizer
-    optimizer = AdamW(params=model.parameters(), lr=lr, correct_bias=correct_bias)
+    optimizer = AdamW(params=model.parameters(), lr=lr, correct_bias=correct_bias, no_deprecation_warning=True)
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/465
diff --git a/.github/workflows/check_dependencies.yml b/.github/workflows/check_dependencies.yml
index c4818ba85..1e8a4e1bb 100644
--- a/.github/workflows/check_dependencies.yml
+++ b/.github/workflows/check_dependencies.yml
@@ -1,4 +1,4 @@
-name: Check for dependency modification
+name: Trigger docker images and run slow tests
 
 on:
   push:
@@ -37,4 +37,9 @@ jobs:
     needs: check-for-setup
     if: (github.event_name == 'push') && (needs.check-for-setup.outputs.changed == '1')
     uses: ./.github/workflows/build-docker-images.yml
-    secrets: inherit
\ No newline at end of file
+    secrets: inherit
+
+  run-tests:
+    needs: build-docker-containers
+    if: always()
+    uses: ./.github/workflows/on-merge.yml
\ No newline at end of file
diff --git a/.github/workflows/on-merge.yml b/.github/workflows/on-merge.yml
index f8b9ccddb..b95474758 100644
--- a/.github/workflows/on-merge.yml
+++ b/.github/workflows/on-merge.yml
@@ -1,10 +1,7 @@
-name: Self-hosted runner (push to "main")
+name: Self-hosted runner tests (push to "main")
 
 on:
-  workflow_run:
-    workflows: ["Check for dependency modification"]
-    branches: ["main"]
-    types: [completed]
+  workflow_call:
   workflow_dispatch:
 
 env:

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/464
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index c4487f6b6..2d2a0bd56 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -254,7 +254,7 @@ def __init__(
             self.scaler = torch.cuda.amp.GradScaler(**kwargs)
         elif self.state.mixed_precision == "bf16":
             self.native_amp = is_bf16_available(True)
-            if mixed_precision == "bf16" and not self.native_amp:
+            if mixed_precision == "bf16" and not self.native_amp and not is_tpu_available():
                 raise ValueError(err.format(mode="bf16", requirement="PyTorch >= 1.10 and a supported device."))
 
             # Only on the GPU do we care about scaling the gradients
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index 95b9eb30a..36b11be56 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -80,7 +80,9 @@ def __init__(
         if not getattr(self, "initialized", False):
             self.backend = None
             self.deepspeed_plugin = None
-            mixed_precision = mixed_precision.lower() if mixed_precision else None
+            mixed_precision = (
+                parse_choice_from_env("MIXED_PRECISION", "no") if mixed_precision is None else mixed_precision.lower()
+            )
             if not _from_accelerator:
                 raise ValueError(
                     "Please make sure to properly initialize your accelerator via `accelerator = Accelerator()` "

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/463
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 7b6c54608..c4487f6b6 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -46,6 +46,7 @@
     extract_model_from_parallel,
     gather,
     get_pretty_name,
+    is_bf16_available,
     is_deepspeed_available,
     is_torch_version,
     is_tpu_available,
@@ -242,20 +243,24 @@ def __init__(
         # Mixed precision attributes
         self.scaler = None
         self.native_amp = False
+        err = "{mode} mixed precision requires {requirement}"
         if self.state.mixed_precision == "fp16":
             self.native_amp = is_torch_version(">=", "1.6")
             if not self.native_amp:
-                raise ValueError("fp16 mixed precision requires PyTorch >= 1.6")
-
+                raise ValueError(err.format(mode="fp16", requirement="PyTorch >= 1.6"))
+            if not torch.cuda.is_available():
+                raise ValueError(err.format(mode="fp16", requirement="a GPU"))
             kwargs = self.scaler_handler.to_kwargs() if self.scaler_handler is not None else {}
             self.scaler = torch.cuda.amp.GradScaler(**kwargs)
         elif self.state.mixed_precision == "bf16":
-            self.native_amp = is_torch_version(">=", "1.10")
+            self.native_amp = is_bf16_available(True)
             if mixed_precision == "bf16" and not self.native_amp:
-                raise ValueError("bf16 mixed precision requires PyTorch >= 1.10")
+                raise ValueError(err.format(mode="bf16", requirement="PyTorch >= 1.10 and a supported device."))
 
-            kwargs = self.scaler_handler.to_kwargs() if self.scaler_handler is not None else {}
-            self.scaler = torch.cuda.amp.GradScaler(**kwargs)
+            # Only on the GPU do we care about scaling the gradients
+            if torch.cuda.is_available():
+                kwargs = self.scaler_handler.to_kwargs() if self.scaler_handler is not None else {}
+                self.scaler = torch.cuda.amp.GradScaler(**kwargs)
 
         # Internal references to the training objects
         self._optimizers = []
@@ -528,8 +533,9 @@ def prepare_model(self, model):
         if self.native_amp:
             if self.mixed_precision == "fp16" and is_torch_version(">=", "1.10"):
                 model.forward = torch.cuda.amp.autocast(dtype=torch.float16)(model.forward)
-            elif self.mixed_precision == "bf16":
-                model.forward = torch.cuda.amp.autocast(dtype=torch.bfloat16)(model.forward)
+            elif self.mixed_precision == "bf16" and self.distributed_type != DistributedType.TPU:
+                device_type = "cuda" if torch.cuda.is_available() else "cpu"
+                model.forward = torch.autocast(device_type=device_type, dtype=torch.bfloat16)(model.forward)
             else:
                 model.forward = torch.cuda.amp.autocast()(model.forward)
             model.forward = convert_outputs_to_fp32(model.forward)
@@ -1054,8 +1060,10 @@ def autocast(self):
         if self.native_amp:
             if self.mixed_precision == "fp16" and is_torch_version(">=", "1.10"):
                 autocast_context = torch.cuda.amp.autocast(dtype=torch.float16)
-            elif self.mixed_precision == "bf16":
-                autocast_context = torch.cuda.amp.autocast(dtype=torch.bfloat16)
+            elif self.mixed_precision == "bf16" and is_bf16_available():
+                if self.distributed_type in [DistributedType.NO, DistributedType.MULTI_CPU, DistributedType.MULTI_GPU]:
+                    device_type = "cpu" if not torch.cuda.is_available() else "cuda"
+                    autocast_context = torch.autocast(dtype=torch.bfloat16, device_type=device_type)
             else:
                 autocast_context = torch.cuda.amp.autocast()
 
diff --git a/src/accelerate/launchers.py b/src/accelerate/launchers.py
index 298db439b..8ddc9af34 100644
--- a/src/accelerate/launchers.py
+++ b/src/accelerate/launchers.py
@@ -50,6 +50,13 @@ def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, mix
     else:
         in_colab_or_kaggle = False
 
+    try:
+        mixed_precision = PrecisionType(mixed_precision.lower())
+    except ValueError:
+        raise ValueError(
+            f"Unknown mixed_precision mode: {args.mixed_precision.lower()}. Choose between {PrecisionType.list()}."
+        )
+
     if in_colab_or_kaggle:
         if os.environ.get("TPU_NAME", None) is not None:
             # TPU launch
@@ -72,7 +79,7 @@ def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, mix
             if torch.cuda.is_available():
                 print("Launching training on one GPU.")
             else:
-                print("Launching training on CPU.")
+                print("Launching training on one CPU.")
             function(*args)
 
     else:
@@ -105,13 +112,6 @@ def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, mix
                     "function."
                 )
 
-            try:
-                mixed_precision = PrecisionType(mixed_precision.lower())
-            except ValueError:
-                raise ValueError(
-                    f"Unknown mixed_precision mode: {args.mixed_precision.lower()}. Choose between {PrecisionType.list()}."
-                )
-
             if use_fp16:
                 warnings.warn('use_fp16=True is deprecated. Use mixed_precision="fp16" instead.', DeprecationWarning)
                 mixed_precision = "fp16"
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index 06e9cb8fa..95b9eb30a 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -92,9 +92,9 @@ def __init__(
                 self.process_index = xm.get_ordinal()
                 self.local_process_index = xm.get_local_ordinal()
                 self.device = xm.xla_device()
-                self.mixed_precision = (
-                    parse_choice_from_env("MIXED_PRECISION", "no") if mixed_precision is None else mixed_precision
-                )
+                if mixed_precision == "bf16":
+                    os.environ["XLA_USE_BF16"] = str(1)
+                self.mixed_precision = mixed_precision
             elif os.environ.get("USE_DEEPSPEED", "false") == "true" and not cpu:
                 assert (
                     is_deepspeed_available()
@@ -120,9 +120,7 @@ def __init__(
                 self.local_process_index = int(os.environ.get("LOCAL_RANK", -1))
                 self.device = torch.device("cuda", self.local_process_index)
                 torch.cuda.set_device(self.device)
-                self.mixed_precision = (
-                    parse_choice_from_env("MIXED_PRECISION", "no") if mixed_precision is None else mixed_precision
-                )
+                self.mixed_precision = mixed_precision
                 if os.environ.get("USE_FSDP", "false") == "true":
                     self.distributed_type = DistributedType.FSDP
                     if self.mixed_precision != "no":
@@ -166,15 +164,13 @@ def __init__(
                 self.process_index = torch.distributed.get_rank()
                 self.local_process_index = local_rank
                 self.device = torch.device("cpu")
-                self.mixed_precision = "no"
+                self.mixed_precision = mixed_precision
             else:
                 self.distributed_type = DistributedType.NO
                 self.num_processes = 1
                 self.process_index = self.local_process_index = 0
                 self.device = torch.device("cuda" if torch.cuda.is_available() and not cpu else "cpu")
-                self.mixed_precision = (
-                    parse_choice_from_env("MIXED_PRECISION", "no") if mixed_precision is None else mixed_precision
-                )
+                self.mixed_precision = mixed_precision
             self.initialized = True
 
     def __repr__(self):
diff --git a/src/accelerate/test_utils/scripts/test_script.py b/src/accelerate/test_utils/scripts/test_script.py
index 018564fa0..07d794b42 100644
--- a/src/accelerate/test_utils/scripts/test_script.py
+++ b/src/accelerate/test_utils/scripts/test_script.py
@@ -21,7 +21,14 @@
 from accelerate.data_loader import prepare_data_loader
 from accelerate.state import AcceleratorState
 from accelerate.test_utils import RegressionDataset, RegressionModel, are_the_same_tensors
-from accelerate.utils import DistributedType, gather, is_torch_version, set_seed, synchronize_rng_states
+from accelerate.utils import (
+    DistributedType,
+    gather,
+    is_bf16_available,
+    is_torch_version,
+    set_seed,
+    synchronize_rng_states,
+)
 
 
 def init_state_check():
@@ -245,74 +252,77 @@ def training_check():
 
     accelerator.print("Training yielded the same results on one CPU or distributes setup with batch split.")
 
-    # Mostly a test that FP16 doesn't crash as the operation inside the model is not converted to FP16
-    print("FP16 training check.")
-    AcceleratorState._reset_state()
-    accelerator = Accelerator(mixed_precision="fp16")
-    train_dl = DataLoader(train_set, batch_size=batch_size, shuffle=True, generator=generator)
-    model = RegressionModel()
-    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
-
-    train_dl, model, optimizer = accelerator.prepare(train_dl, model, optimizer)
-    set_seed(42)
-    generator.manual_seed(42)
-    for _ in range(3):
-        for batch in train_dl:
-            model.zero_grad()
-            output = model(batch["x"])
-            loss = torch.nn.functional.mse_loss(output, batch["y"])
-            accelerator.backward(loss)
-            optimizer.step()
-
-    model = accelerator.unwrap_model(model).cpu()
-    assert torch.allclose(old_model.a, model.a), "Did not obtain the same model on CPU or distributed training."
-    assert torch.allclose(old_model.b, model.b), "Did not obtain the same model on CPU or distributed training."
-
-    # TEST that previous fp16 flag still works
-    print("Legacy FP16 training check.")
-    AcceleratorState._reset_state()
-    accelerator = Accelerator(fp16=True)
-    train_dl = DataLoader(train_set, batch_size=batch_size, shuffle=True, generator=generator)
-    model = RegressionModel()
-    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
-
-    train_dl, model, optimizer = accelerator.prepare(train_dl, model, optimizer)
-    set_seed(42)
-    generator.manual_seed(42)
-    for _ in range(3):
-        for batch in train_dl:
-            model.zero_grad()
-            output = model(batch["x"])
-            loss = torch.nn.functional.mse_loss(output, batch["y"])
-            accelerator.backward(loss)
-            optimizer.step()
-
-    model = accelerator.unwrap_model(model).cpu()
-    assert torch.allclose(old_model.a, model.a), "Did not obtain the same model on CPU or distributed training."
-    assert torch.allclose(old_model.b, model.b), "Did not obtain the same model on CPU or distributed training."
-
-    # Mostly a test that BF16 doesn't crash as the operation inside the model is not converted to BF16
-    print("BF16 training check.")
-    AcceleratorState._reset_state()
-    accelerator = Accelerator(mixed_precision="bf16")
-    train_dl = DataLoader(train_set, batch_size=batch_size, shuffle=True, generator=generator)
-    model = RegressionModel()
-    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
-
-    train_dl, model, optimizer = accelerator.prepare(train_dl, model, optimizer)
-    set_seed(42)
-    generator.manual_seed(42)
-    for _ in range(3):
-        for batch in train_dl:
-            model.zero_grad()
-            output = model(batch["x"])
-            loss = torch.nn.functional.mse_loss(output, batch["y"])
-            accelerator.backward(loss)
-            optimizer.step()
-
-    model = accelerator.unwrap_model(model).cpu()
-    assert torch.allclose(old_model.a, model.a), "Did not obtain the same model on CPU or distributed training."
-    assert torch.allclose(old_model.b, model.b), "Did not obtain the same model on CPU or distributed training."
+    if torch.cuda.is_available():
+        # Mostly a test that FP16 doesn't crash as the operation inside the model is not converted to FP16
+        print("FP16 training check.")
+        AcceleratorState._reset_state()
+        accelerator = Accelerator(mixed_precision="fp16")
+        train_dl = DataLoader(train_set, batch_size=batch_size, shuffle=True, generator=generator)
+        model = RegressionModel()
+        optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
+
+        train_dl, model, optimizer = accelerator.prepare(train_dl, model, optimizer)
+        set_seed(42)
+        generator.manual_seed(42)
+        for _ in range(3):
+            for batch in train_dl:
+                model.zero_grad()
+                output = model(batch["x"])
+                loss = torch.nn.functional.mse_loss(output, batch["y"])
+                accelerator.backward(loss)
+                optimizer.step()
+
+        model = accelerator.unwrap_model(model).cpu()
+        assert torch.allclose(old_model.a, model.a), "Did not obtain the same model on CPU or distributed training."
+        assert torch.allclose(old_model.b, model.b), "Did not obtain the same model on CPU or distributed training."
+
+        # TEST that previous fp16 flag still works
+        print("Legacy FP16 training check.")
+        AcceleratorState._reset_state()
+        accelerator = Accelerator(fp16=True)
+        train_dl = DataLoader(train_set, batch_size=batch_size, shuffle=True, generator=generator)
+        model = RegressionModel()
+        optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
+
+        train_dl, model, optimizer = accelerator.prepare(train_dl, model, optimizer)
+        set_seed(42)
+        generator.manual_seed(42)
+        for _ in range(3):
+            for batch in train_dl:
+                model.zero_grad()
+                output = model(batch["x"])
+                loss = torch.nn.functional.mse_loss(output, batch["y"])
+                accelerator.backward(loss)
+                optimizer.step()
+
+        model = accelerator.unwrap_model(model).cpu()
+        assert torch.allclose(old_model.a, model.a), "Did not obtain the same model on CPU or distributed training."
+        assert torch.allclose(old_model.b, model.b), "Did not obtain the same model on CPU or distributed training."
+
+    # BF16 support is only for CPU + TPU, and some GPU
+    if is_bf16_available():
+        # Mostly a test that BF16 doesn't crash as the operation inside the model is not converted to BF16
+        print("BF16 training check.")
+        AcceleratorState._reset_state()
+        accelerator = Accelerator(mixed_precision="bf16")
+        train_dl = DataLoader(train_set, batch_size=batch_size, shuffle=True, generator=generator)
+        model = RegressionModel()
+        optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
+
+        train_dl, model, optimizer = accelerator.prepare(train_dl, model, optimizer)
+        set_seed(42)
+        generator.manual_seed(42)
+        for _ in range(3):
+            for batch in train_dl:
+                model.zero_grad()
+                output = model(batch["x"])
+                loss = torch.nn.functional.mse_loss(output, batch["y"])
+                accelerator.backward(loss)
+                optimizer.step()
+
+        model = accelerator.unwrap_model(model).cpu()
+        assert torch.allclose(old_model.a, model.a), "Did not obtain the same model on CPU or distributed training."
+        assert torch.allclose(old_model.b, model.b), "Did not obtain the same model on CPU or distributed training."
 
 
 def main():
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
index 309b49023..350efc899 100644
--- a/src/accelerate/utils/__init__.py
+++ b/src/accelerate/utils/__init__.py
@@ -20,6 +20,7 @@
 )
 from .imports import (
     is_apex_available,
+    is_bf16_available,
     is_boto3_available,
     is_ccl_available,
     is_comet_ml_available,
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index 0a44a093f..d1d634355 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -15,6 +15,10 @@
 import importlib
 import sys
 
+import torch
+
+from .versions import is_torch_version
+
 
 # The package importlib_metadata is in a different place, depending on the Python version.
 if sys.version_info < (3, 8):
@@ -68,6 +72,17 @@ def is_deepspeed_available():
             return False
 
 
+def is_bf16_available(ignore_tpu=False):
+    "Checks if bf16 is supported, optionally ignoring the TPU"
+    if is_tpu_available():
+        return not ignore_tpu
+    if is_torch_version(">=", "1.10"):
+        if torch.cuda.is_available():
+            return torch.cuda.is_bf16_supported()
+        return True
+    return False
+
+
 def is_transformers_available():
     return importlib.util.find_spec("transformers") is not None
 
diff --git a/tests/test_hooks.py b/tests/test_hooks.py
index 51c434ab1..2e0519668 100644
--- a/tests/test_hooks.py
+++ b/tests/test_hooks.py
@@ -77,20 +77,20 @@ def test_pre_forward_hook_is_executed(self):
         test_hook = PreForwardHook()
         add_hook_to_module(test_model, test_hook)
         output1 = test_model(x)
-        self.assertTrue(torch.allclose(output1, expected))
+        self.assertTrue(torch.allclose(output1, expected, atol=1e-5))
 
         # Attaching a hook to a model when it already has one replaces, does not chain
         test_hook = PreForwardHook()
         add_hook_to_module(test_model, test_hook)
         output1 = test_model(x)
-        self.assertTrue(torch.allclose(output1, expected))
+        self.assertTrue(torch.allclose(output1, expected, atol=1e-5))
 
         # You need to use the sequential hook to chain two or more hooks
         test_hook = SequentialHook(PreForwardHook(), PreForwardHook())
         add_hook_to_module(test_model, test_hook)
 
         output2 = test_model(x)
-        assert torch.allclose(output2, expected2)
+        assert torch.allclose(output2, expected2, atol=1e-5)
 
     def test_post_forward_hook_is_executed(self):
         test_model = ModelForTest()
@@ -100,20 +100,20 @@ def test_post_forward_hook_is_executed(self):
         test_hook = PostForwardHook()
         add_hook_to_module(test_model, test_hook)
         output1 = test_model(x)
-        self.assertTrue(torch.allclose(output1, output + 1))
+        self.assertTrue(torch.allclose(output1, output + 1, atol=1e-5))
 
         # Attaching a hook to a model when it already has one replaces, does not chain
         test_hook = PostForwardHook()
         add_hook_to_module(test_model, test_hook)
         output1 = test_model(x)
-        self.assertTrue(torch.allclose(output1, output + 1))
+        self.assertTrue(torch.allclose(output1, output + 1, atol=1e-5))
 
         # You need to use the sequential hook to chain two or more hooks
         test_hook = SequentialHook(PostForwardHook(), PostForwardHook())
         add_hook_to_module(test_model, test_hook)
 
         output2 = test_model(x)
-        assert torch.allclose(output2, output + 2)
+        assert torch.allclose(output2, output + 2, atol=1e-5)
 
     def test_no_grad_in_hook(self):
         test_model = ModelForTest()

Need to be extra sure that this always exists for PyTorch version for which `is_bf16_available()`
@@ -1054,8 +1060,10 @@ def autocast(self):
         if self.native_amp:
             if self.mixed_precision == "fp16" and is_torch_version(">=", "1.10"):
                 autocast_context = torch.cuda.amp.autocast(dtype=torch.float16)
-            elif self.mixed_precision == "bf16":
-                autocast_context = torch.cuda.amp.autocast(dtype=torch.bfloat16)
+            elif self.mixed_precision == "bf16" and is_bf16_available():
+                if self.distributed_type in [DistributedType.NO, DistributedType.MULTI_CPU, DistributedType.MULTI_GPU]:
+                    device_type = "cpu" if not torch.cuda.is_available() else "cuda"
+                    autocast_context = torch.autocast(dtype=torch.bfloat16, device_type=device_type)
Not sure why the period is removed here, but I like them at the end of sentences :-p
@@ -65,14 +72,14 @@ def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, mix
                 num_processes = 8
 
             launcher = PrepareForLaunch(function, distributed_type="TPU")
-            print(f"Launching a training on {num_processes} TPU cores.")
+            print(f"Launching a training on {num_processes} TPU cores")
Do we want to set anything if `mixed_precision != "bf16"`? I'd remove the else statement. But if you really want it, then use one line for such a small test.
@@ -92,9 +92,11 @@ def __init__(
                 self.process_index = xm.get_ordinal()
                 self.local_process_index = xm.get_local_ordinal()
                 self.device = xm.xla_device()
-                self.mixed_precision = (
-                    parse_choice_from_env("MIXED_PRECISION", "no") if mixed_precision is None else mixed_precision
-                )
+                if mixed_precision == "bf16":
+                    os.environ["XLA_USE_BF16"] = 1
+                else:
+                    os.environ["XLA_USE_BF16"] = 0
Not sure we need to keep this. We have done a couple of releases since we deprecated it, so it's okay if we stop testing it IMO.
@@ -245,74 +252,77 @@ def training_check():
 
     accelerator.print("Training yielded the same results on one CPU or distributes setup with batch split.")
 
-    # Mostly a test that FP16 doesn't crash as the operation inside the model is not converted to FP16
-    print("FP16 training check.")
-    AcceleratorState._reset_state()
-    accelerator = Accelerator(mixed_precision="fp16")
-    train_dl = DataLoader(train_set, batch_size=batch_size, shuffle=True, generator=generator)
-    model = RegressionModel()
-    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
-
-    train_dl, model, optimizer = accelerator.prepare(train_dl, model, optimizer)
-    set_seed(42)
-    generator.manual_seed(42)
-    for _ in range(3):
-        for batch in train_dl:
-            model.zero_grad()
-            output = model(batch["x"])
-            loss = torch.nn.functional.mse_loss(output, batch["y"])
-            accelerator.backward(loss)
-            optimizer.step()
-
-    model = accelerator.unwrap_model(model).cpu()
-    assert torch.allclose(old_model.a, model.a), "Did not obtain the same model on CPU or distributed training."
-    assert torch.allclose(old_model.b, model.b), "Did not obtain the same model on CPU or distributed training."
-
-    # TEST that previous fp16 flag still works
-    print("Legacy FP16 training check.")
-    AcceleratorState._reset_state()
-    accelerator = Accelerator(fp16=True)
-    train_dl = DataLoader(train_set, batch_size=batch_size, shuffle=True, generator=generator)
-    model = RegressionModel()
-    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
-
-    train_dl, model, optimizer = accelerator.prepare(train_dl, model, optimizer)
-    set_seed(42)
-    generator.manual_seed(42)
-    for _ in range(3):
-        for batch in train_dl:
-            model.zero_grad()
-            output = model(batch["x"])
-            loss = torch.nn.functional.mse_loss(output, batch["y"])
-            accelerator.backward(loss)
-            optimizer.step()
-
-    model = accelerator.unwrap_model(model).cpu()
-    assert torch.allclose(old_model.a, model.a), "Did not obtain the same model on CPU or distributed training."
-    assert torch.allclose(old_model.b, model.b), "Did not obtain the same model on CPU or distributed training."
-
-    # Mostly a test that BF16 doesn't crash as the operation inside the model is not converted to BF16
-    print("BF16 training check.")
-    AcceleratorState._reset_state()
-    accelerator = Accelerator(mixed_precision="bf16")
-    train_dl = DataLoader(train_set, batch_size=batch_size, shuffle=True, generator=generator)
-    model = RegressionModel()
-    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
-
-    train_dl, model, optimizer = accelerator.prepare(train_dl, model, optimizer)
-    set_seed(42)
-    generator.manual_seed(42)
-    for _ in range(3):
-        for batch in train_dl:
-            model.zero_grad()
-            output = model(batch["x"])
-            loss = torch.nn.functional.mse_loss(output, batch["y"])
-            accelerator.backward(loss)
-            optimizer.step()
-
-    model = accelerator.unwrap_model(model).cpu()
-    assert torch.allclose(old_model.a, model.a), "Did not obtain the same model on CPU or distributed training."
-    assert torch.allclose(old_model.b, model.b), "Did not obtain the same model on CPU or distributed training."
+    if torch.cuda.is_available():
+        # Mostly a test that FP16 doesn't crash as the operation inside the model is not converted to FP16
+        print("FP16 training check.")
+        AcceleratorState._reset_state()
+        accelerator = Accelerator(mixed_precision="fp16")
+        train_dl = DataLoader(train_set, batch_size=batch_size, shuffle=True, generator=generator)
+        model = RegressionModel()
+        optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
+
+        train_dl, model, optimizer = accelerator.prepare(train_dl, model, optimizer)
+        set_seed(42)
+        generator.manual_seed(42)
+        for _ in range(3):
+            for batch in train_dl:
+                model.zero_grad()
+                output = model(batch["x"])
+                loss = torch.nn.functional.mse_loss(output, batch["y"])
+                accelerator.backward(loss)
+                optimizer.step()
+
+        model = accelerator.unwrap_model(model).cpu()
+        assert torch.allclose(old_model.a, model.a), "Did not obtain the same model on CPU or distributed training."
+        assert torch.allclose(old_model.b, model.b), "Did not obtain the same model on CPU or distributed training."
+
+        # TEST that previous fp16 flag still works
+        print("Legacy FP16 training check.")
That's from me messing with some things, put it back!
@@ -65,14 +72,14 @@ def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, mix
                 num_processes = 8
 
             launcher = PrepareForLaunch(function, distributed_type="TPU")
-            print(f"Launching a training on {num_processes} TPU cores.")
+            print(f"Launching a training on {num_processes} TPU cores")
Thought about this for a while, and I'm just going to keep it as assigning if `==`, so users can still set the env var themselves if they're more comfortable/familiar with it. 
@@ -92,9 +92,11 @@ def __init__(
                 self.process_index = xm.get_ordinal()
                 self.local_process_index = xm.get_local_ordinal()
                 self.device = xm.xla_device()
-                self.mixed_precision = (
-                    parse_choice_from_env("MIXED_PRECISION", "no") if mixed_precision is None else mixed_precision
-                )
+                if mixed_precision == "bf16":
+                    os.environ["XLA_USE_BF16"] = 1
+                else:
+                    os.environ["XLA_USE_BF16"] = 0
Fixed by adding a torch check for >= 1.10
@@ -1054,8 +1060,10 @@ def autocast(self):
         if self.native_amp:
             if self.mixed_precision == "fp16" and is_torch_version(">=", "1.10"):
                 autocast_context = torch.cuda.amp.autocast(dtype=torch.float16)
-            elif self.mixed_precision == "bf16":
-                autocast_context = torch.cuda.amp.autocast(dtype=torch.bfloat16)
+            elif self.mixed_precision == "bf16" and is_bf16_available():
+                if self.distributed_type in [DistributedType.NO, DistributedType.MULTI_CPU, DistributedType.MULTI_GPU]:
+                    device_type = "cpu" if not torch.cuda.is_available() else "cuda"
+                    autocast_context = torch.autocast(dtype=torch.bfloat16, device_type=device_type)
I'd feel more comfortable dropping the test once we've removed entirely the legacy param (whenever that may be)
@@ -245,74 +252,77 @@ def training_check():
 
     accelerator.print("Training yielded the same results on one CPU or distributes setup with batch split.")
 
-    # Mostly a test that FP16 doesn't crash as the operation inside the model is not converted to FP16
-    print("FP16 training check.")
-    AcceleratorState._reset_state()
-    accelerator = Accelerator(mixed_precision="fp16")
-    train_dl = DataLoader(train_set, batch_size=batch_size, shuffle=True, generator=generator)
-    model = RegressionModel()
-    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
-
-    train_dl, model, optimizer = accelerator.prepare(train_dl, model, optimizer)
-    set_seed(42)
-    generator.manual_seed(42)
-    for _ in range(3):
-        for batch in train_dl:
-            model.zero_grad()
-            output = model(batch["x"])
-            loss = torch.nn.functional.mse_loss(output, batch["y"])
-            accelerator.backward(loss)
-            optimizer.step()
-
-    model = accelerator.unwrap_model(model).cpu()
-    assert torch.allclose(old_model.a, model.a), "Did not obtain the same model on CPU or distributed training."
-    assert torch.allclose(old_model.b, model.b), "Did not obtain the same model on CPU or distributed training."
-
-    # TEST that previous fp16 flag still works
-    print("Legacy FP16 training check.")
-    AcceleratorState._reset_state()
-    accelerator = Accelerator(fp16=True)
-    train_dl = DataLoader(train_set, batch_size=batch_size, shuffle=True, generator=generator)
-    model = RegressionModel()
-    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
-
-    train_dl, model, optimizer = accelerator.prepare(train_dl, model, optimizer)
-    set_seed(42)
-    generator.manual_seed(42)
-    for _ in range(3):
-        for batch in train_dl:
-            model.zero_grad()
-            output = model(batch["x"])
-            loss = torch.nn.functional.mse_loss(output, batch["y"])
-            accelerator.backward(loss)
-            optimizer.step()
-
-    model = accelerator.unwrap_model(model).cpu()
-    assert torch.allclose(old_model.a, model.a), "Did not obtain the same model on CPU or distributed training."
-    assert torch.allclose(old_model.b, model.b), "Did not obtain the same model on CPU or distributed training."
-
-    # Mostly a test that BF16 doesn't crash as the operation inside the model is not converted to BF16
-    print("BF16 training check.")
-    AcceleratorState._reset_state()
-    accelerator = Accelerator(mixed_precision="bf16")
-    train_dl = DataLoader(train_set, batch_size=batch_size, shuffle=True, generator=generator)
-    model = RegressionModel()
-    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
-
-    train_dl, model, optimizer = accelerator.prepare(train_dl, model, optimizer)
-    set_seed(42)
-    generator.manual_seed(42)
-    for _ in range(3):
-        for batch in train_dl:
-            model.zero_grad()
-            output = model(batch["x"])
-            loss = torch.nn.functional.mse_loss(output, batch["y"])
-            accelerator.backward(loss)
-            optimizer.step()
-
-    model = accelerator.unwrap_model(model).cpu()
-    assert torch.allclose(old_model.a, model.a), "Did not obtain the same model on CPU or distributed training."
-    assert torch.allclose(old_model.b, model.b), "Did not obtain the same model on CPU or distributed training."
+    if torch.cuda.is_available():
+        # Mostly a test that FP16 doesn't crash as the operation inside the model is not converted to FP16
+        print("FP16 training check.")
+        AcceleratorState._reset_state()
+        accelerator = Accelerator(mixed_precision="fp16")
+        train_dl = DataLoader(train_set, batch_size=batch_size, shuffle=True, generator=generator)
+        model = RegressionModel()
+        optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
+
+        train_dl, model, optimizer = accelerator.prepare(train_dl, model, optimizer)
+        set_seed(42)
+        generator.manual_seed(42)
+        for _ in range(3):
+            for batch in train_dl:
+                model.zero_grad()
+                output = model(batch["x"])
+                loss = torch.nn.functional.mse_loss(output, batch["y"])
+                accelerator.backward(loss)
+                optimizer.step()
+
+        model = accelerator.unwrap_model(model).cpu()
+        assert torch.allclose(old_model.a, model.a), "Did not obtain the same model on CPU or distributed training."
+        assert torch.allclose(old_model.b, model.b), "Did not obtain the same model on CPU or distributed training."
+
+        # TEST that previous fp16 flag still works
+        print("Legacy FP16 training check.")
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/462
diff --git a/src/accelerate/utils/offload.py b/src/accelerate/utils/offload.py
index 6e8b34bbc..750ff9d57 100644
--- a/src/accelerate/utils/offload.py
+++ b/src/accelerate/utils/offload.py
@@ -25,8 +25,8 @@ def offload_weight(weight, weight_name, offload_folder, index=None):
     dtype = None
     # Check the string instead of the dtype to be compatible with versions of PyTorch that don't have bfloat16.
     if str(weight.dtype) == "torch.bfloat16":
-        # Need to convert to FP32 since NumPy does not handle bfloat16s.
-        weight = weight.float()
+        # Need to reinterpret the underlined data as int16 since NumPy does not handle bfloat16s.
+        weight = weight.view(torch.int16)
         dtype = "bfloat16"
     array = weight.numpy()
     tensor_file = os.path.join(offload_folder, f"{weight_name}.dat")
@@ -50,8 +50,8 @@ def load_offloaded_weight(weight_file, weight_info):
 
     dtype = weight_info["dtype"]
     if dtype == "bfloat16":
-        # NumPy does not support bfloat16 so this was saved as a float32
-        dtype = "float32"
+        # NumPy does not support bfloat16 so this was saved as a int16
+        dtype = "int16"
 
     weight = np.memmap(weight_file, dtype=dtype, shape=shape, mode="r")
 
@@ -59,7 +59,7 @@ def load_offloaded_weight(weight_file, weight_info):
         weight = weight[0]
     weight = torch.tensor(weight)
     if weight_info["dtype"] == "bfloat16":
-        weight = weight.to(torch.bfloat16)
+        weight = weight.view(torch.bfloat16)
 
     return weight
 

This PR handles a bug reported in https://github.com/huggingface/accelerate/issues/454 concerning bfloat16 and disk offload.
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/461
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
index 60f39c90a..309b49023 100644
--- a/src/accelerate/utils/__init__.py
+++ b/src/accelerate/utils/__init__.py
@@ -48,6 +48,7 @@
     OffloadedWeightsLoader,
     PrefixedDataset,
     extract_submodules_state_dict,
+    load_offloaded_weight,
     offload_state_dict,
     offload_weight,
     save_offload_index,
diff --git a/src/accelerate/utils/offload.py b/src/accelerate/utils/offload.py
index a562e9283..6e8b34bbc 100644
--- a/src/accelerate/utils/offload.py
+++ b/src/accelerate/utils/offload.py
@@ -22,10 +22,18 @@
 
 
 def offload_weight(weight, weight_name, offload_folder, index=None):
+    dtype = None
+    # Check the string instead of the dtype to be compatible with versions of PyTorch that don't have bfloat16.
+    if str(weight.dtype) == "torch.bfloat16":
+        # Need to convert to FP32 since NumPy does not handle bfloat16s.
+        weight = weight.float()
+        dtype = "bfloat16"
     array = weight.numpy()
     tensor_file = os.path.join(offload_folder, f"{weight_name}.dat")
     if index is not None:
-        index[weight_name] = {"dtype": str(array.dtype), "shape": list(array.shape)}
+        if dtype is None:
+            dtype = str(array.dtype)
+        index[weight_name] = {"dtype": dtype, "shape": list(array.shape)}
     if array.ndim == 0:
         array = array[None]
     file_array = np.memmap(tensor_file, dtype=array.dtype, mode="w+", shape=array.shape)
@@ -34,6 +42,28 @@ def offload_weight(weight, weight_name, offload_folder, index=None):
     return index
 
 
+def load_offloaded_weight(weight_file, weight_info):
+    shape = tuple(weight_info["shape"])
+    if shape == ():
+        # NumPy memory-mapped arrays can't have 0 dims so it was saved as 1d tensor
+        shape = (1,)
+
+    dtype = weight_info["dtype"]
+    if dtype == "bfloat16":
+        # NumPy does not support bfloat16 so this was saved as a float32
+        dtype = "float32"
+
+    weight = np.memmap(weight_file, dtype=dtype, shape=shape, mode="r")
+
+    if len(weight_info["shape"]) == 0:
+        weight = weight[0]
+    weight = torch.tensor(weight)
+    if weight_info["dtype"] == "bfloat16":
+        weight = weight.to(torch.bfloat16)
+
+    return weight
+
+
 def save_offload_index(index, offload_folder):
     if index is None or len(index) == 0:
         # Nothing to save
@@ -129,12 +159,7 @@ def __getitem__(self, key: str):
             return self.state_dict[key]
         weight_info = self.index[key]
         weight_file = os.path.join(self.save_folder, f"{key}.dat")
-        shape = tuple(weight_info["shape"])
-        if shape == ():
-            weight = np.memmap(weight_file, dtype=weight_info["dtype"], shape=(1,), mode="r")[0]
-        else:
-            weight = np.memmap(weight_file, dtype=weight_info["dtype"], shape=shape, mode="r")
-        return torch.tensor(weight)
+        return load_offloaded_weight(weight_file, weight_info)
 
     def __iter__(self):
         return iter(self.all_keys)
diff --git a/tests/test_offload.py b/tests/test_offload.py
index 515d5094d..765a9685d 100644
--- a/tests/test_offload.py
+++ b/tests/test_offload.py
@@ -19,7 +19,13 @@
 import torch
 import torch.nn as nn
 
-from accelerate.utils import OffloadedWeightsLoader, offload_state_dict
+from accelerate.utils import (
+    OffloadedWeightsLoader,
+    is_torch_version,
+    load_offloaded_weight,
+    offload_state_dict,
+    offload_weight,
+)
 
 
 class ModelForTest(nn.Module):
@@ -35,8 +41,6 @@ def forward(self, x):
 
 class OffloadTester(unittest.TestCase):
     def test_offload_state_dict(self):
-        from tempfile import TemporaryDirectory
-
         model = ModelForTest()
         with TemporaryDirectory() as tmp_dir:
             offload_state_dict(tmp_dir, model.state_dict())
@@ -49,6 +53,22 @@ def test_offload_state_dict(self):
                 self.assertTrue(os.path.isfile(weight_file))
                 # TODO: add tests on the fact weights are properly loaded
 
+    def test_offload_weight(self):
+        dtypes = [torch.float16, torch.float32]
+        if is_torch_version(">=", "1.10"):
+            dtypes.append(torch.bfloat16)
+
+        for dtype in dtypes:
+            weight = torch.randn(2, 3, dtype=dtype)
+            with TemporaryDirectory() as tmp_dir:
+                index = offload_weight(weight, "weight", tmp_dir, {})
+                weight_file = os.path.join(tmp_dir, "weight.dat")
+                self.assertTrue(os.path.isfile(weight_file))
+                self.assertDictEqual(index, {"weight": {"shape": [2, 3], "dtype": str(dtype).split(".")[1]}})
+
+                new_weight = load_offloaded_weight(weight_file, index["weight"])
+                self.assertTrue(torch.equal(weight, new_weight))
+
     def test_offload_weights_loader(self):
         model = ModelForTest()
         state_dict = model.state_dict()

Since we care about the prevision here, would it be a good idea to test these with floats that marginally get close to the float16 reach to ensure they're lossless? 

E.g.:
```python
weight = torch.rand((2,3), dtype=dtype)
```


@@ -49,6 +53,22 @@ def test_offload_state_dict(self):
                 self.assertTrue(os.path.isfile(weight_file))
                 # TODO: add tests on the fact weights are properly loaded
 
+    def test_offload_weight(self):
+        dtypes = [torch.float16, torch.float32]
+        if is_torch_version(">=", "1.10"):
+            dtypes.append(torch.bfloat16)
+
+        for dtype in dtypes:
+            weight = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=dtype)
Probably yes 😅 
@@ -49,6 +53,22 @@ def test_offload_state_dict(self):
                 self.assertTrue(os.path.isfile(weight_file))
                 # TODO: add tests on the fact weights are properly loaded
 
+    def test_offload_weight(self):
+        dtypes = [torch.float16, torch.float32]
+        if is_torch_version(">=", "1.10"):
+            dtypes.append(torch.bfloat16)
+
+        for dtype in dtypes:
+            weight = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=dtype)
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/460
diff --git a/examples/by_feature/gradient_accumulation.py b/examples/by_feature/gradient_accumulation.py
index bb8aba326..38da07aa1 100644
--- a/examples/by_feature/gradient_accumulation.py
+++ b/examples/by_feature/gradient_accumulation.py
@@ -157,7 +157,7 @@ def training_function(config, args):
             # We use the new `no_sync` context manager to prevent gradient averaging
             # until we want to at the proper step if we happen to be in a distributed setup
             # otherwise it does nothing
-            if step % gradient_accumulation_steps == 0:
+            if step % gradient_accumulation_steps != 0:
                 # Accumulate gradients locally
                 with accelerator.no_sync(model):
                     output = model(**batch)

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/459
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index e341ad93f..06e9cb8fa 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -75,6 +75,7 @@ def __init__(
         self.__dict__ = self._shared_state
         if parse_flag_from_env("USE_CPU"):
             cpu = True
+        self._check_initialized(mixed_precision, cpu)
         self.fork_launched = parse_flag_from_env("FORK_LAUNCHED", 0)
         if not getattr(self, "initialized", False):
             self.backend = None
@@ -196,3 +197,17 @@ def __repr__(self):
     @property
     def use_fp16(self):
         return self.mixed_precision != "no"
+
+    @staticmethod
+    def _reset_state():
+        "Resets `_shared_state`, is used internally and should not be called"
+        AcceleratorState._shared_state = {}
+
+    def _check_initialized(self, mixed_precision=None, cpu=None):
+        "Checks if a modification is trying to be made and the `AcceleratorState` has already been initialized"
+        if getattr(self, "initialized", False):
+            err = "AcceleratorState has already been initialized and cannot be changed, restart your runtime completely and pass `{flag}` to `Accelerate()`."
+            if cpu and self.device.type != "cpu":
+                raise ValueError(err.format(flag="cpu=True"))
+            if mixed_precision is not None and mixed_precision != self.mixed_precision:
+                raise ValueError(err.format(flag=f"mixed_precision='{mixed_precision}'"))
diff --git a/src/accelerate/test_utils/scripts/test_script.py b/src/accelerate/test_utils/scripts/test_script.py
index 8ed869edc..018564fa0 100644
--- a/src/accelerate/test_utils/scripts/test_script.py
+++ b/src/accelerate/test_utils/scripts/test_script.py
@@ -247,6 +247,7 @@ def training_check():
 
     # Mostly a test that FP16 doesn't crash as the operation inside the model is not converted to FP16
     print("FP16 training check.")
+    AcceleratorState._reset_state()
     accelerator = Accelerator(mixed_precision="fp16")
     train_dl = DataLoader(train_set, batch_size=batch_size, shuffle=True, generator=generator)
     model = RegressionModel()
@@ -269,6 +270,7 @@ def training_check():
 
     # TEST that previous fp16 flag still works
     print("Legacy FP16 training check.")
+    AcceleratorState._reset_state()
     accelerator = Accelerator(fp16=True)
     train_dl = DataLoader(train_set, batch_size=batch_size, shuffle=True, generator=generator)
     model = RegressionModel()
@@ -291,6 +293,7 @@ def training_check():
 
     # Mostly a test that BF16 doesn't crash as the operation inside the model is not converted to BF16
     print("BF16 training check.")
+    AcceleratorState._reset_state()
     accelerator = Accelerator(mixed_precision="bf16")
     train_dl = DataLoader(train_set, batch_size=batch_size, shuffle=True, generator=generator)
     model = RegressionModel()

I don't understand why we are using this class, when you can just do a regular string with flag between curly brackets and then use the format method.
@@ -196,3 +198,14 @@ def __repr__(self):
     @property
     def use_fp16(self):
         return self.mixed_precision != "no"
+
+    def _check_initialized(self, mixed_precision=None, cpu=None):
+        "Checks if a modification is trying to be made and the `AcceleratorState` has already been initialized"
+        if getattr(self, "initialized", False):
+            template = Template(
+                "AcceleratorState has already been initialized and cannot be changed, restart your runtime completely and pass `$flag` to `Accelerate()`."
+            )
Definitely not because I forgot strings could do this... 👀 
@@ -196,3 +198,14 @@ def __repr__(self):
     @property
     def use_fp16(self):
         return self.mixed_precision != "no"
+
+    def _check_initialized(self, mixed_precision=None, cpu=None):
+        "Checks if a modification is trying to be made and the `AcceleratorState` has already been initialized"
+        if getattr(self, "initialized", False):
+            template = Template(
+                "AcceleratorState has already been initialized and cannot be changed, restart your runtime completely and pass `$flag` to `Accelerate()`."
+            )
Should be a `staticmethod` no?
@@ -196,3 +197,17 @@ def __repr__(self):
     @property
     def use_fp16(self):
         return self.mixed_precision != "no"
+
+    @classmethod
And here you should change the value of `AcceleratorState._shared_state`
@@ -196,3 +197,17 @@ def __repr__(self):
     @property
     def use_fp16(self):
         return self.mixed_precision != "no"
+
+    @classmethod
+    def _reset_state(cls):
+        "Resets `_shared_state`, is used internally and should not be called"
+        cls._shared_state = {}
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/458
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index c8cce9587..0a44a093f 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -34,7 +34,12 @@
 try:
     import torch_xla.core.xla_model as xm  # noqa: F401
 
-    _tpu_available = True
+    try:
+        # Will raise a RuntimeError if no XLA configuration is found
+        _ = xm.xla_device()
+        _tpu_available = True
+    except RuntimeError:
+        _tpu_available = False
 except ImportError:
     _tpu_available = False
 

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/456
diff --git a/src/accelerate/checkpointing.py b/src/accelerate/checkpointing.py
index b043d5d20..8b41c48f8 100644
--- a/src/accelerate/checkpointing.py
+++ b/src/accelerate/checkpointing.py
@@ -116,7 +116,7 @@ def load_accelerator_state(input_dir, models, optimizers, schedulers, process_in
     Args:
         input_dir (`str` or `os.PathLike`):
             The name of the folder to load all relevant weights and states.
-        model_stmodelsates (`List[torch.nn.Module]`):
+        models (`List[torch.nn.Module]`):
             A list of model instances
         optimizers (`List[torch.optim.Optimizer]`):
             A list of optimizer instances

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/447
diff --git a/setup.py b/setup.py
index 45c4ae912..93916bd9f 100644
--- a/setup.py
+++ b/setup.py
@@ -19,7 +19,6 @@
 extras["quality"] = ["black ~= 22.0", "isort >= 5.5.4", "flake8 >= 3.8.3"]
 extras["docs"] = []
 extras["test"] = [
-    "psutil",
     "pytest",
     "pytest-xdist",
     "pytest-subtests",
@@ -60,7 +59,7 @@
         ]
     },
     python_requires=">=3.7.0",
-    install_requires=["numpy>=1.17", "packaging>=20.0", "pyyaml", "torch>=1.4.0"],
+    install_requires=["numpy>=1.17", "packaging>=20.0", "psutil", "pyyaml", "torch>=1.4.0"],
     extras_require=extras,
     classifiers=[
         "Development Status :: 5 - Production/Stable",

Thanks! LG2M!
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/445
diff --git a/src/accelerate/big_modeling.py b/src/accelerate/big_modeling.py
index 9db015427..4014e3972 100644
--- a/src/accelerate/big_modeling.py
+++ b/src/accelerate/big_modeling.py
@@ -88,6 +88,7 @@ def cpu_offload(
     execution_device: Optional[torch.device] = None,
     offload_buffers: bool = False,
     state_dict: Optional[Dict[str, torch.Tensor]] = None,
+    preload_module_classes: Optional[List[str]] = None,
 ):
     """
     Activates full CPU offload for a model. As a result, all parameters of the model will be offloaded and only one
@@ -104,13 +105,23 @@ def cpu_offload(
             Whether or not to offload the buffers with the model parameters.
         state_dict (`Dict[str, torch.Tensor]`, *optional*):
             The state dict of the model that will be kept on CPU.
+        preload_module_classes (`List[str]`, *optional*):
+            A list of classes whose instances should load all their weights (even in the submodules) at the beginning
+            of the forward. This should only be used for classes that have submodules which are registered but not
+            called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
+            `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
     """
     if execution_device is None:
         execution_device = next(iter(model.parameters())).device
     if state_dict is None:
         state_dict = {n: p.to("cpu") for n, p in model.state_dict().items()}
     attach_align_device_hook(
-        model, execution_device=execution_device, offload=True, offload_buffers=offload_buffers, weights_map=state_dict
+        model,
+        execution_device=execution_device,
+        offload=True,
+        offload_buffers=offload_buffers,
+        weights_map=state_dict,
+        preload_module_classes=preload_module_classes,
     )
     add_hook_to_module(model, AlignDevicesHook(io_same_device=True))
     return model
@@ -121,6 +132,7 @@ def disk_offload(
     offload_dir: Union[str, os.PathLike],
     execution_device: Optional[torch.device] = None,
     offload_buffers: bool = False,
+    preload_module_classes: Optional[List[str]] = None,
 ):
     """
     Activates full disk offload for a model. As a result, all parameters of the model will be offloaded as
@@ -136,6 +148,11 @@ def disk_offload(
             model's first parameter device.
         offload_buffers (`bool`, *optional*, defaults to `False`):
             Whether or not to offload the buffers with the model parameters.
+        preload_module_classes (`List[str]`, *optional*):
+            A list of classes whose instances should load all their weights (even in the submodules) at the beginning
+            of the forward. This should only be used for classes that have submodules which are registered but not
+            called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
+            `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
     """
     if not os.path.isdir(offload_dir) or not os.path.isfile(os.path.join(offload_dir, "index.json")):
         offload_state_dict(offload_dir, model.state_dict())
@@ -148,6 +165,7 @@ def disk_offload(
         offload=True,
         offload_buffers=offload_buffers,
         weights_map=weights_map,
+        preload_module_classes=preload_module_classes,
     )
     add_hook_to_module(model, AlignDevicesHook(io_same_device=True))
     return model
@@ -160,6 +178,7 @@ def dispatch_model(
     state_dict: Optional[Dict[str, torch.Tensor]] = None,
     offload_dir: Union[str, os.PathLike] = None,
     offload_buffers: bool = False,
+    preload_module_classes: Optional[List[str]] = None,
 ):
     """
     Dispatches a model according to a given device map. Layers of the model might be spread across GPUs, offloaded on
@@ -180,6 +199,11 @@ def dispatch_model(
             The folder in which to offload the model weights (or where the model weights are already offloaded).
         offload_buffers (`bool`, *optional*, defaults to `False`):
             Whether or not to offload the buffers with the model parameters.
+        preload_module_classes (`List[str]`, *optional*):
+            A list of classes whose instances should load all their weights (even in the submodules) at the beginning
+            of the forward. This should only be used for classes that have submodules which are registered but not
+            called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
+            `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
     """
     # Error early if the device map is incomplete.
     check_device_map(model, device_map)
@@ -219,6 +243,7 @@ def dispatch_model(
         offload=offload,
         offload_buffers=offload_buffers,
         weights_map=weights_map,
+        preload_module_classes=preload_module_classes,
     )
     model.hf_device_map = device_map
     return model
@@ -234,6 +259,7 @@ def load_checkpoint_and_dispatch(
     offload_buffers: bool = False,
     dtype: Optional[Union[str, torch.dtype]] = None,
     offload_state_dict: bool = False,
+    preload_module_classes: Optional[List[str]] = None,
 ):
     """
     Loads a (potentially sharded) checkpoint inside a model, potentially sending weights to a given device as they are
@@ -267,6 +293,11 @@ def load_checkpoint_and_dispatch(
         offload_state_dict (`bool`, *optional*, defaults to `False`):
             If `True`, will temporarily offload the CPU state dict on the hard drive to avoig getting out of CPU RAM if
             the weight of the CPU state dict + the biggest shard does not fit.
+        preload_module_classes (`List[str]`, *optional*):
+            A list of classes whose instances should load all their weights (even in the submodules) at the beginning
+            of the forward. This should only be used for classes that have submodules which are registered but not
+            called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
+            `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
     """
     if device_map == "auto":
         device_map = infer_auto_device_map(
@@ -282,4 +313,10 @@ def load_checkpoint_and_dispatch(
     )
     if device_map is None:
         return model
-    return dispatch_model(model, device_map=device_map, offload_dir=offload_folder, offload_buffers=offload_buffers)
+    return dispatch_model(
+        model,
+        device_map=device_map,
+        offload_dir=offload_folder,
+        offload_buffers=offload_buffers,
+        preload_module_classes=preload_module_classes,
+    )
diff --git a/src/accelerate/hooks.py b/src/accelerate/hooks.py
index dea10ba50..5c04b8621 100644
--- a/src/accelerate/hooks.py
+++ b/src/accelerate/hooks.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import functools
-from typing import Dict, Mapping, Optional, Union
+from typing import Dict, List, Mapping, Optional, Union
 
 import torch
 import torch.nn as nn
@@ -220,17 +220,23 @@ def init_hook(self, module):
             for name, _ in named_module_tensors(module, recurse=self.place_submodules):
                 set_module_tensor_to_device(module, name, self.execution_device)
         elif self.offload:
-            self.original_devices = {name: param.device for name, param in named_module_tensors(module)}
+            self.original_devices = {
+                name: param.device for name, param in named_module_tensors(module, recurse=self.place_submodules)
+            }
             if self.weights_map is None:
                 self.weights_map = {
                     name: param.to("cpu")
-                    for name, param in named_module_tensors(module, include_buffers=self.offload_buffers)
+                    for name, param in named_module_tensors(
+                        module, include_buffers=self.offload_buffers, recurse=self.place_submodules
+                    )
                 }
 
-            for name, _ in named_module_tensors(module, include_buffers=self.offload_buffers):
+            for name, _ in named_module_tensors(
+                module, include_buffers=self.offload_buffers, recurse=self.place_submodules
+            ):
                 set_module_tensor_to_device(module, name, "meta")
             if not self.offload_buffers and self.execution_device is not None:
-                for name, _ in module.named_buffers(recurse=False):
+                for name, _ in module.named_buffers(recurse=self.place_submodules):
                     set_module_tensor_to_device(module, name, self.execution_device)
         return module
 
@@ -238,14 +244,18 @@ def pre_forward(self, module, *args, **kwargs):
         if self.io_same_device:
             self.input_device = find_device([args, kwargs])
         if self.offload:
-            for name, _ in named_module_tensors(module, include_buffers=self.offload_buffers):
+            for name, _ in named_module_tensors(
+                module, include_buffers=self.offload_buffers, recurse=self.place_submodules
+            ):
                 set_module_tensor_to_device(module, name, self.execution_device, value=self.weights_map[name])
 
         return send_to_device(args, self.execution_device), send_to_device(kwargs, self.execution_device)
 
     def post_forward(self, module, output):
         if self.offload:
-            for name, _ in named_module_tensors(module, include_buffers=self.offload_buffers):
+            for name, _ in named_module_tensors(
+                module, include_buffers=self.offload_buffers, recurse=self.place_submodules
+            ):
                 set_module_tensor_to_device(module, name, "meta")
 
         if self.io_same_device and self.input_device is not None:
@@ -260,7 +270,11 @@ def detach_hook(self, module):
                     set_module_tensor_to_device(module, name, device, value=self.weights_map.get(name, None))
 
 
-def attach_execution_device_hook(module: torch.nn.Module, execution_device: Union[int, str, torch.device]):
+def attach_execution_device_hook(
+    module: torch.nn.Module,
+    execution_device: Union[int, str, torch.device],
+    preload_module_classes: Optional[List[str]] = None,
+):
     """
     Recursively attaches `AlignDevicesHook` to all submodules of a given model to make sure they have the right
     execution device
@@ -270,10 +284,19 @@ def attach_execution_device_hook(module: torch.nn.Module, execution_device: Unio
             The module where we want to attach the hooks.
         execution_device (`int`, `str` or `torch.device`):
             The device on which inputs and model weights should be placed before the forward pass.
+        preload_module_classes (`List[str]`, *optional*):
+            A list of classes whose instances should load all their weights (even in the submodules) at the beginning
+            of the forward. This should only be used for classes that have submodules which are registered but not
+            called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
+            `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
     """
     if not hasattr(module, "_hf_hook") and len(module.state_dict()) > 0:
         add_hook_to_module(module, AlignDevicesHook(execution_device))
 
+    # Break the recursion if we get to a preload module.
+    if preload_module_classes is not None and module.__class__.__name__ in preload_module_classes:
+        return
+
     for child in module.children():
         attach_execution_device_hook(child, execution_device)
 
@@ -285,6 +308,7 @@ def attach_align_device_hook(
     weights_map: Optional[Mapping] = None,
     offload_buffers: bool = False,
     module_name: str = "",
+    preload_module_classes: Optional[List[str]] = None,
 ):
     """
     Recursively attaches `AlignDevicesHook` to all submodules of a given model that have direct parameters and/or
@@ -303,10 +327,19 @@ def attach_align_device_hook(
             Whether or not to include the associated module's buffers when offloading.
         module_name (`str`, *optional*, defaults to `""`):
             The name of the module.
+        preload_module_classes (`List[str]`, *optional*):
+            A list of classes whose instances should load all their weights (even in the submodules) at the beginning
+            of the forward. This should only be used for classes that have submodules which are registered but not
+            called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
+            `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
     """
     # Attach the hook on this module if it has any direct tensor.
     directs = named_module_tensors(module)
-    if len(list(directs)) > 0:
+    full_offload = (
+        offload and preload_module_classes is not None and module.__class__.__name__ in preload_module_classes
+    )
+
+    if len(list(directs)) > 0 or full_offload:
         if weights_map is not None:
             prefix = f"{module_name}." if len(module_name) > 0 else ""
             prefixed_weights_map = PrefixedDataset(weights_map, prefix)
@@ -317,9 +350,14 @@ def attach_align_device_hook(
             offload=offload,
             weights_map=prefixed_weights_map,
             offload_buffers=offload_buffers,
+            place_submodules=full_offload,
         )
         add_hook_to_module(module, hook)
 
+    # We stop the recursion in case we hit the full offload.
+    if full_offload:
+        return
+
     # Recurse on all children of the module.
     for child_name, child in module.named_children():
         child_name = f"{module_name}.{child_name}" if len(module_name) > 0 else child_name
@@ -330,6 +368,7 @@ def attach_align_device_hook(
             weights_map=weights_map,
             offload_buffers=offload_buffers,
             module_name=child_name,
+            preload_module_classes=preload_module_classes,
         )
 
 
@@ -352,6 +391,7 @@ def attach_align_device_hook_on_blocks(
     weights_map: Mapping = None,
     offload_buffers: bool = False,
     module_name: str = "",
+    preload_module_classes: Optional[List[str]] = None,
 ):
     """
     Attaches `AlignDevicesHook` to all blocks of a given model as needed.
@@ -371,6 +411,11 @@ def attach_align_device_hook_on_blocks(
             Whether or not to include the associated module's buffers when offloading.
         module_name (`str`, *optional*, defaults to `""`):
             The name of the module.
+        preload_module_classes (`List[str]`, *optional*):
+            A list of classes whose instances should load all their weights (even in the submodules) at the beginning
+            of the forward. This should only be used for classes that have submodules which are registered but not
+            called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
+            `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
     """
     # If one device and one offload, we've got one hook.
     if not isinstance(execution_device, Mapping) and not isinstance(offload, dict):
@@ -389,7 +434,7 @@ def attach_align_device_hook_on_blocks(
         return
 
     if not isinstance(execution_device, Mapping):
-        execution_device = {key: offload for key in offload.keys()}
+        execution_device = {key: execution_device for key in offload.keys()}
     if not isinstance(offload, Mapping):
         offload = {key: offload for key in execution_device.keys()}
 
@@ -410,11 +455,14 @@ def attach_align_device_hook_on_blocks(
             weights_map=weights_map,
             offload_buffers=offload_buffers,
             module_name=module_name,
+            preload_module_classes=preload_module_classes,
         )
         if not hasattr(module, "_hf_hook"):
             hook = AlignDevicesHook(execution_device=execution_device[module_name], io_same_device=(module_name == ""))
             add_hook_to_module(module, hook)
-        attach_execution_device_hook(module, execution_device[module_name])
+        attach_execution_device_hook(
+            module, execution_device[module_name], preload_module_classes=preload_module_classes
+        )
     elif module_name == "":
         hook = AlignDevicesHook(io_same_device=True)
         add_hook_to_module(module, hook)
@@ -428,4 +476,5 @@ def attach_align_device_hook_on_blocks(
             weights_map=weights_map,
             offload_buffers=offload_buffers,
             module_name=child_name,
+            preload_module_classes=preload_module_classes,
         )
diff --git a/tests/test_big_modeling.py b/tests/test_big_modeling.py
index 852b10c3b..4f3054b73 100644
--- a/tests/test_big_modeling.py
+++ b/tests/test_big_modeling.py
@@ -56,6 +56,29 @@ def forward(self, x):
         return self.linear4(self.linear3(self.batchnorm(self.linear2(self.linear1(x)))))
 
 
+# To test preload_module_classes
+class ModuleWithUnusedSubModules(nn.Module):
+    def __init__(self, input_dim, output_dim):
+        super().__init__()
+        self.linear = nn.Linear(input_dim, output_dim)
+
+    def forward(self, x):
+        return x @ self.linear.weight.t() + self.linear.bias
+
+
+class ModelWithUnusedSubModulesForTest(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear1 = ModuleWithUnusedSubModules(3, 4)
+        self.linear2 = ModuleWithUnusedSubModules(4, 5)
+        self.batchnorm = nn.BatchNorm1d(5)
+        self.linear3 = ModuleWithUnusedSubModules(5, 6)
+        self.linear4 = ModuleWithUnusedSubModules(6, 5)
+
+    def forward(self, x):
+        return self.linear4(self.linear3(self.batchnorm(self.linear2(self.linear1(x)))))
+
+
 class BigModelingTester(unittest.TestCase):
     def test_init_empty_weights(self):
         # base use
@@ -107,6 +130,33 @@ def test_cpu_offload(self):
             torch.allclose(expected, output.cpu(), 1e-4, 1e-5), msg=f"Expected: {expected}\nActual: {output.cpu()}"
         )
 
+    def test_cpu_offload_with_unused_submodules(self):
+        model = ModelWithUnusedSubModulesForTest()
+        x = torch.randn(2, 3)
+        expected = model(x)
+
+        device = torch.device(0 if torch.cuda.is_available() else "cpu")
+
+        cpu_offload(model, execution_device=device, preload_module_classes=["ModuleWithUnusedSubModules"])
+        output = model(x)
+        self.assertTrue(
+            torch.allclose(expected, output.cpu(), 1e-4, 1e-5), msg=f"Expected: {expected}\nActual: {output.cpu()}"
+        )
+
+        # Clean up for next test.
+        remove_hook_from_submodules(model)
+
+        cpu_offload(
+            model,
+            execution_device=device,
+            offload_buffers=True,
+            preload_module_classes=["ModuleWithUnusedSubModules"],
+        )
+        output = model(x)
+        self.assertTrue(
+            torch.allclose(expected, output.cpu(), 1e-4, 1e-5), msg=f"Expected: {expected}\nActual: {output.cpu()}"
+        )
+
     @slow
     @require_cuda
     def test_cpu_offload_gpt2(self):
@@ -145,6 +195,38 @@ def test_disk_offload(self):
                 torch.allclose(expected, output.cpu(), 1e-4, 1e-5), msg=f"Expected: {expected}\nActual: {output.cpu()}"
             )
 
+    def test_disk_offload_with_unused_submodules(self):
+        model = ModelWithUnusedSubModulesForTest()
+        x = torch.randn(2, 3)
+        expected = model(x)
+
+        device = torch.device(0 if torch.cuda.is_available() else "cpu")
+
+        with TemporaryDirectory() as tmp_dir:
+            disk_offload(
+                model, tmp_dir, execution_device=device, preload_module_classes=["ModuleWithUnusedSubModules"]
+            )
+            output = model(x)
+            self.assertTrue(
+                torch.allclose(expected, output.cpu(), 1e-4, 1e-5), msg=f"Expected: {expected}\nActual: {output.cpu()}"
+            )
+
+            # Clean up for next test.
+            remove_hook_from_submodules(model)
+
+        with TemporaryDirectory() as tmp_dir:
+            disk_offload(
+                model,
+                tmp_dir,
+                execution_device=device,
+                offload_buffers=True,
+                preload_module_classes=["ModuleWithUnusedSubModules"],
+            )
+            output = model(x)
+            self.assertTrue(
+                torch.allclose(expected, output.cpu(), 1e-4, 1e-5), msg=f"Expected: {expected}\nActual: {output.cpu()}"
+            )
+
     @slow
     @require_cuda
     def test_disk_offload_gpt2(self):
@@ -237,6 +319,36 @@ def test_dispatch_model_gpt2_on_two_gpus(self):
                 "Hello world! My name is Kiyoshi, and I'm a student at the University of Tokyo",
             )
 
+    @require_cuda
+    def test_dispatch_model_with_unused_submodules(self):
+        model = ModelWithUnusedSubModulesForTest()
+        device_map = {"linear1": "cpu", "linear2": "disk", "batchnorm": "cpu", "linear3": 0, "linear4": 0}
+
+        x = torch.randn(2, 3)
+        expected = model(x)
+
+        with TemporaryDirectory() as tmp_dir:
+            dispatch_model(
+                model, device_map, offload_dir=tmp_dir, preload_module_classes=["ModuleWithUnusedSubModules"]
+            )
+            output = model(x)
+            self.assertTrue(torch.allclose(expected, output.cpu(), atol=1e-5))
+
+    @require_multi_gpu
+    def test_dispatch_model_with_unused_submodules_multi_gpu(self):
+        model = ModelWithUnusedSubModulesForTest()
+        device_map = {"linear1": "cpu", "linear2": "disk", "batchnorm": "cpu", "linear3": 0, "linear4": 1}
+
+        x = torch.randn(2, 3)
+        expected = model(x)
+
+        with TemporaryDirectory() as tmp_dir:
+            dispatch_model(
+                model, device_map, offload_dir=tmp_dir, preload_module_classes=["ModuleWithUnusedSubModules"]
+            )
+            output = model(x)
+            self.assertTrue(torch.allclose(expected, output.cpu(), atol=1e-5))
+
     @require_cuda
     def test_load_checkpoint_and_dispatch(self):
         model = ModelForTest()
@@ -282,3 +394,55 @@ def test_load_checkpoint_and_dispatch_multi_gpu(self):
 
         output = new_model(x)
         self.assertTrue(torch.allclose(expected, output.cpu(), atol=1e-5))
+
+    @require_cuda
+    def test_load_checkpoint_and_dispatch_with_unused_submodules(self):
+        model = ModelWithUnusedSubModulesForTest()
+        device_map = {"linear1": "cpu", "linear2": "cpu", "batchnorm": 0, "linear3": 0, "linear4": 0}
+
+        x = torch.randn(2, 3)
+        expected = model(x)
+
+        with TemporaryDirectory() as tmp_dir:
+            checkpoint = os.path.join(tmp_dir, "pt_model.bin")
+            torch.save(model.state_dict(), checkpoint)
+
+            new_model = ModelWithUnusedSubModulesForTest()
+            new_model = load_checkpoint_and_dispatch(
+                new_model, checkpoint, device_map=device_map, preload_module_classes=["ModuleWithUnusedSubModules"]
+            )
+
+        # CPU-offloaded weights are on the meta device while waiting for the forward pass.
+        self.assertEqual(new_model.linear1.linear.weight.device, torch.device("meta"))
+        self.assertEqual(new_model.linear2.linear.weight.device, torch.device("meta"))
+        self.assertEqual(new_model.linear3.linear.weight.device, torch.device(0))
+        self.assertEqual(new_model.linear4.linear.weight.device, torch.device(0))
+
+        output = new_model(x)
+        self.assertTrue(torch.allclose(expected, output.cpu(), atol=1e-5))
+
+    @require_multi_gpu
+    def test_load_checkpoint_and_dispatch_multi_gpu_with_unused_submodules(self):
+        model = ModelWithUnusedSubModulesForTest()
+        device_map = {"linear1": "cpu", "linear2": "cpu", "batchnorm": 0, "linear3": 0, "linear4": 1}
+
+        x = torch.randn(2, 3)
+        expected = model(x)
+
+        with TemporaryDirectory() as tmp_dir:
+            checkpoint = os.path.join(tmp_dir, "pt_model.bin")
+            torch.save(model.state_dict(), checkpoint)
+
+            new_model = ModelWithUnusedSubModulesForTest()
+            new_model = load_checkpoint_and_dispatch(
+                new_model, checkpoint, device_map=device_map, preload_module_classes=["ModuleWithUnusedSubModules"]
+            )
+
+        # CPU-offloaded weights are on the meta device while waiting for the forward pass.
+        self.assertEqual(new_model.linear1.linear.weight.device, torch.device("meta"))
+        self.assertEqual(new_model.linear2.linear.weight.device, torch.device("meta"))
+        self.assertEqual(new_model.linear3.linear.weight.device, torch.device(0))
+        self.assertEqual(new_model.linear4.linear.weight.device, torch.device(1))
+
+        output = new_model(x)
+        self.assertTrue(torch.allclose(expected, output.cpu(), atol=1e-5))

```suggestion
        preloaded_modules (`List[str]`, *optional*):
```

What about something such as `preloaded_modules`?
@@ -371,6 +411,11 @@ def attach_align_device_hook_on_blocks(
             Whether or not to include the associated module's buffers when offloading.
         module_name (`str`, *optional*, defaults to `""`):
             The name of the module.
+        load_all_weights_classes (`List[str]`, *optional*):
Maybe `preloaded_module_classes` to be similar to `no_split_module_classes`? I like `preloaded` a lot, thanks for the suggestion!
@@ -371,6 +411,11 @@ def attach_align_device_hook_on_blocks(
             Whether or not to include the associated module's buffers when offloading.
         module_name (`str`, *optional*, defaults to `""`):
             The name of the module.
+        load_all_weights_classes (`List[str]`, *optional*):
Works with me!
@@ -371,6 +411,11 @@ def attach_align_device_hook_on_blocks(
             Whether or not to include the associated module's buffers when offloading.
         module_name (`str`, *optional*, defaults to `""`):
             The name of the module.
+        load_all_weights_classes (`List[str]`, *optional*):
```suggestion
    # Break the recursion if we get to a preload module.
```
@@ -270,10 +284,19 @@ def attach_execution_device_hook(module: torch.nn.Module, execution_device: Unio
             The module where we want to attach the hooks.
         execution_device (`int`, `str` or `torch.device`):
             The device on which inputs and model weights should be placed before the forward pass.
+        preload_module_classes (`List[str]`, *optional*):
+            A list of classes whose instances should load all their weights (even in the submodules) at the beginning
+            of the forward. This should only be used for classes that have submodules which are registered but not
+            called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
+            `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
     """
     if not hasattr(module, "_hf_hook") and len(module.state_dict()) > 0:
         add_hook_to_module(module, AlignDevicesHook(execution_device))
 
+    # Break the recursion if we get to a load all weights module.
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/442
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 407b90184..0ead08c81 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -49,6 +49,7 @@
     get_pretty_name,
     is_deepspeed_available,
     is_torch_version,
+    is_tpu_available,
     is_transformers_available,
     pad_across_processes,
     reduce,
@@ -68,6 +69,9 @@
         DummyScheduler,
     )
 
+if is_tpu_available():
+    import torch_xla.distributed.xla_multiprocessing as xmp
+
 logger = get_logger(__name__)
 
 
@@ -546,6 +550,8 @@ def prepare_model(self, model):
             else:
                 model.forward = torch.cuda.amp.autocast()(model.forward)
             model.forward = convert_outputs_to_fp32(model.forward)
+        if self.distributed_type == DistributedType.TPU and self.state.fork_launched:
+            model = xmp.MpModelWrapper(model).to(self.device)
         return model
 
     def _prepare_deepspeed(self, *args):
@@ -712,7 +718,7 @@ def prepare_data_loader(self, data_loader):
             num_processes=self.num_processes,
             process_index=self.process_index,
             split_batches=self.split_batches,
-            put_on_device=self.device_placement,
+            put_on_device=self.device_placement if self.distributed_type != DistributedType.TPU else False,
             rng_types=self.rng_types.copy(),
             dispatch_batches=self.dispatch_batches,
         )
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index 79ba6b3f3..eb26c3a60 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -35,7 +35,7 @@
 
 
 if is_tpu_available():
-    import torch_xla.core.xla_model as xm
+    import torch_xla.distributed.parallel_loader as xpl
 
 
 # kwargs of the DataLoader in min version 1.4.0.
@@ -298,10 +298,7 @@ def __init__(self, dataset, device=None, rng_types=None, generator=None, **kwarg
     def __iter__(self):
         if self.rng_types is not None:
             synchronize_rng_states(self.rng_types, self.generator)
-        state = AcceleratorState()
         for batch in super().__iter__():
-            if state.distributed_type == DistributedType.TPU:
-                xm.mark_step()
             yield batch if self.device is None else send_to_device(batch, self.device)
 
 
@@ -408,9 +405,6 @@ def __iter__(self):
                 batch_size += 1
 
             data_slice = slice(state.process_index * batch_size, (state.process_index + 1) * batch_size)
-
-            if state.distributed_type == DistributedType.TPU:
-                xm.mark_step()
             yield slice_tensors(batch, data_slice)
 
     def __len__(self):
@@ -565,15 +559,19 @@ def prepare_data_loader(
         kwargs["batch_size"] = dataloader.batch_size // num_processes if split_batches else dataloader.batch_size
 
     if dispatch_batches:
-        return DataLoaderDispatcher(
+        dataloader = DataLoaderDispatcher(
             new_dataset, split_batches=split_batches, batch_sampler=new_batch_sampler, **kwargs
         )
+    else:
+        dataloader = DataLoaderShard(
+            new_dataset,
+            device=device if put_on_device and state.distributed_type != DistributedType.TPU else None,
+            batch_sampler=new_batch_sampler,
+            rng_types=rng_types,
+            generator=generator,
+            **kwargs,
+        )
 
-    return DataLoaderShard(
-        new_dataset,
-        device=device if put_on_device else None,
-        batch_sampler=new_batch_sampler,
-        rng_types=rng_types,
-        generator=generator,
-        **kwargs,
-    )
+    if state.distributed_type == DistributedType.TPU:
+        return xpl.MpDeviceLoader(dataloader, device)
+    return dataloader
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index 6e09fdae4..891b0762b 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -75,6 +75,7 @@ def __init__(
         self.__dict__ = self._shared_state
         if parse_flag_from_env("USE_CPU"):
             cpu = True
+        self.fork_launched = parse_flag_from_env("FORK_LAUNCHED", 0)
         if not getattr(self, "initialized", False):
             self.backend = None
             self.deepspeed_plugin = None
@@ -90,7 +91,9 @@ def __init__(
                 self.process_index = xm.get_ordinal()
                 self.local_process_index = xm.get_local_ordinal()
                 self.device = xm.xla_device()
-                self.mixed_precision = "no"
+                self.mixed_precision = (
+                    parse_choice_from_env("MIXED_PRECISION", "no") if mixed_precision is None else mixed_precision
+                )
             elif os.environ.get("USE_DEEPSPEED", "false") == "true" and not cpu:
                 assert (
                     is_deepspeed_available()
diff --git a/src/accelerate/test_utils/scripts/test_script.py b/src/accelerate/test_utils/scripts/test_script.py
index a23179bd7..8ed869edc 100644
--- a/src/accelerate/test_utils/scripts/test_script.py
+++ b/src/accelerate/test_utils/scripts/test_script.py
@@ -326,7 +326,8 @@ def main():
     if state.local_process_index == 0:
         print("\n**DataLoader integration test**")
     dl_preparation_check()
-    central_dl_preparation_check()
+    if state.distributed_type != DistributedType.TPU:
+        central_dl_preparation_check()
 
     # Trainings are not exactly the same in DeepSpeed and CPU mode
     if state.distributed_type == DistributedType.DEEPSPEED:
@@ -337,5 +338,10 @@ def main():
     training_check()
 
 
+def _mp_fn(index):
+    # For xla_spawn (TPUs)
+    main()
+
+
 if __name__ == "__main__":
     main()
diff --git a/src/accelerate/test_utils/scripts/test_sync.py b/src/accelerate/test_utils/scripts/test_sync.py
index 2f5cc3722..b58a8c092 100644
--- a/src/accelerate/test_utils/scripts/test_sync.py
+++ b/src/accelerate/test_utils/scripts/test_sync.py
@@ -126,5 +126,10 @@ def main():
         test_distributed_sync(accelerator)
 
 
+def _mp_fn(index):
+    # For xla_spawn (TPUs)
+    main()
+
+
 if __name__ == "__main__":
     main()
diff --git a/src/accelerate/utils/launch.py b/src/accelerate/utils/launch.py
index d4f0042fb..c74185c5d 100644
--- a/src/accelerate/utils/launch.py
+++ b/src/accelerate/utils/launch.py
@@ -68,4 +68,5 @@ def __call__(self, index, *args):
             os.environ["LOCAL_RANK"] = str(index)
             os.environ["RANK"] = str(index)
 
+        os.environ["FORK_LAUNCHED"] = str(1)
         self.launcher(*args)
diff --git a/tests/test_tpu.py b/tests/test_tpu.py
index d0b2f0b72..bffa8b8b4 100644
--- a/tests/test_tpu.py
+++ b/tests/test_tpu.py
@@ -24,7 +24,7 @@
 class MultiTPUTester(unittest.TestCase):
     def setUp(self):
         mod_file = inspect.getfile(accelerate.test_utils)
-        self.test_file_path = os.path.sep.join(mod_file.split(os.path.sep)[:-1] + ["test_script.py"])
+        self.test_file_path = os.path.sep.join(mod_file.split(os.path.sep)[:-1] + ["scripts", "test_script.py"])
         self.test_dir = os.path.sep.join(inspect.getfile(self.__class__).split(os.path.sep)[:-1])
 
     @require_tpu
diff --git a/tests/xla_spawn.py b/tests/xla_spawn.py
index c65878056..1a07af29c 100644
--- a/tests/xla_spawn.py
+++ b/tests/xla_spawn.py
@@ -78,7 +78,6 @@ def main():
 
     # Patch sys.argv
     sys.argv = [args.training_script] + args.training_script_args + ["--tpu_num_cores", str(args.num_cores)]
-
     xmp.spawn(mod._mp_fn, args=(), nprocs=args.num_cores)
 
 

No, this shouldn't always be done IMO. This is an optimization for Colab/notebooks so we should only do it when we arrive at this code after a `notebook_launcher`.

Or be ready to benchmark the scripts on TPU to show that this does not affect performance in any way, but usually optimizations for memory come at a cost for speed.
@@ -546,6 +553,8 @@ def prepare_model(self, model):
             else:
                 model.forward = torch.cuda.amp.autocast()(model.forward)
             model.forward = convert_outputs_to_fp32(model.forward)
+        if self.distributed_type == DistributedType.TPU:
+            model = xmp.MpModelWrapper(model).to(self.device)
Same, would like to see some benchmarks on launching a script on a TPU VM to make sure this does not affect speed.
@@ -565,15 +563,19 @@ def prepare_data_loader(
         kwargs["batch_size"] = dataloader.batch_size // num_processes if split_batches else dataloader.batch_size
 
     if dispatch_batches:
-        return DataLoaderDispatcher(
+        dataloader = DataLoaderDispatcher(
             new_dataset, split_batches=split_batches, batch_sampler=new_batch_sampler, **kwargs
         )
+    else:
+        dataloader = DataLoaderShard(
+            new_dataset,
+            device=device if put_on_device else None,
+            batch_sampler=new_batch_sampler,
+            rng_types=rng_types,
+            generator=generator,
+            **kwargs,
+        )
 
-    return DataLoaderShard(
-        new_dataset,
-        device=device if put_on_device else None,
-        batch_sampler=new_batch_sampler,
-        rng_types=rng_types,
-        generator=generator,
-        **kwargs,
-    )
+    if state.distributed_type == DistributedType.TPU:
+        return xpl.MpDeviceLoader(dataloader, device)
Here's the docs on [MpModelWrapper](https://pytorch.org/xla/release/1.11/index.html#torch_xla.distributed.xla_multiprocessing.MpModelWrapper), but here's their comment on it:

> This class should be used together with the spawn(…, start_method=’fork’) API to minimize the use of host memory. Instead of creating models on each multiprocessing process, hence replicating the model’s initial host memory, the model is created once at global scope, and then moved into each device inside the spawn() target function. This method has two advantages. First it uses only one copy of the memory pages to host the original model weights, and second it serializes the move of the wrapped model into each device, by lowering the load onto the system memory during the process.

Will perform benchmarks today and tomorrow
@@ -546,6 +553,8 @@ def prepare_model(self, model):
             else:
                 model.forward = torch.cuda.amp.autocast()(model.forward)
             model.forward = convert_outputs_to_fp32(model.forward)
+        if self.distributed_type == DistributedType.TPU:
+            model = xmp.MpModelWrapper(model).to(self.device)
Yes, and we don't use that spawn method when launching a script on TPU, hence why I don't think we should use it if we don't come from a colab/notebook here :-)
@@ -546,6 +553,8 @@ def prepare_model(self, model):
             else:
                 model.forward = torch.cuda.amp.autocast()(model.forward)
             model.forward = convert_outputs_to_fp32(model.forward)
+        if self.distributed_type == DistributedType.TPU:
+            model = xmp.MpModelWrapper(model).to(self.device)
Makes sense! Will test and report back if there's a loss at all :) 
@@ -546,6 +553,8 @@ def prepare_model(self, model):
             else:
                 model.forward = torch.cuda.amp.autocast()(model.forward)
             model.forward = convert_outputs_to_fp32(model.forward)
+        if self.distributed_type == DistributedType.TPU:
+            model = xmp.MpModelWrapper(model).to(self.device)
```suggestion
        self.fork_launched = parse_flag_from_env("FORK_LAUNCHED", 0)
```
Might be clearer this way.
@@ -75,6 +75,7 @@ def __init__(
         self.__dict__ = self._shared_state
         if parse_flag_from_env("USE_CPU"):
             cpu = True
+        self.use_fork = parse_flag_from_env("USE_LAUNCHER", 0)
Makes sense!
@@ -75,6 +75,7 @@ def __init__(
         self.__dict__ = self._shared_state
         if parse_flag_from_env("USE_CPU"):
             cpu = True
+        self.use_fork = parse_flag_from_env("USE_LAUNCHER", 0)
_The documentation is not available anymore as the PR was closed or merged._
Thanks for fixing up the TPU support with best practices! Looks great to me! 🔥 
@sgugger this PR also modifies the TPU selection in `AcceleratorState` to actually use mixed precision types.

- [FP16 Support](https://github.com/pytorch/xla/issues/1936)
- [BF16 Support](https://pytorch.org/xla/release/1.11/index.html#xla-tensors-and-bfloat16)
We're passing 😄 
![image](https://user-images.githubusercontent.com/7831895/173685310-82547d2a-2219-4236-b899-dd5b7599a9ef.png)

https://api.github.com/repos/huggingface/accelerate/pulls/441
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 8801011ec..407b90184 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -199,6 +199,10 @@ def __init__(
                     ds_config.update({"bf16": {"enabled": True}})
                 self.dschf = HfDeepSpeedConfig(ds_config)  # keep this object alive # noqa
 
+        if os.environ.get("USE_FSDP", "false") == "true" or isinstance(fsdp_plugin, FullyShardedDataParallelPlugin):
+            if is_torch_version("<", "1.12.0.dev20220418+cu113"):
+                raise ValueError("FSDP requires PyTorch >= 1.12.0.dev20220418+cu113")
+
         if fsdp_plugin is None:  # init from env variables
             fsdp_plugin = FullyShardedDataParallelPlugin() if os.environ.get("USE_FSDP", "false") == "true" else None
         else:
@@ -206,10 +210,6 @@ def __init__(
                 raise TypeError("`fsdp_plugin` must be a FullyShardedDataParallelPlugin object.")
             os.environ["USE_FSDP"] = "true"  # use FSDP if plugin is provided
 
-        if os.environ.get("USE_FSDP", "false") == "true":
-            if is_torch_version("<", "1.12.0.dev20220418+cu113"):
-                raise ValueError("FSDP requires PyTorch >= 1.12.0.dev20220418+cu113")
-
         # Kwargs handlers
         self.ddp_handler = None
         self.scaler_handler = None

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/437
diff --git a/README.md b/README.md
index 44b25778e..c67d5dd50 100644
--- a/README.md
+++ b/README.md
@@ -241,4 +241,5 @@ pip install accelerate
 - multi-GPU on several nodes (machines)
 - TPU
 - FP16 with native AMP (apex on the roadmap)
-- DeepSpeed support (experimental)
+- DeepSpeed support (Experimental)
+- PyTorch Fully Sharded Data Parallel (FSDP) support (Experimental)
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 0ead08c81..7b6c54608 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -19,7 +19,6 @@
 import sys
 import warnings
 from contextlib import contextmanager
-from copy import deepcopy
 from typing import List, Optional, Union
 
 import torch
@@ -50,7 +49,6 @@
     is_deepspeed_available,
     is_torch_version,
     is_tpu_available,
-    is_transformers_available,
     pad_across_processes,
     reduce,
     save,
@@ -183,25 +181,10 @@ def __init__(
                 raise ImportError("DeepSpeed is not installed => run `pip install deepspeed` or build it from source.")
             if compare_versions("deepspeed", "<", "0.6.5"):
                 raise ImportError("DeepSpeed version must be >= 0.6.5. Please update DeepSpeed.")
-            if os.environ.get("DEEPSPEED_ZERO3_INIT", "false") == "true" or deepspeed_plugin.zero3_init_flag:
-                if not is_transformers_available():
-                    raise Exception(
-                        "When `zero3_init_flag` is set, it requires Transformers to be installed. "
-                        "Please run `pip install transformers`."
-                    )
-                from transformers.deepspeed import HfDeepSpeedConfig
 
-                ds_config = deepcopy(deepspeed_plugin.deepspeed_config)
-                del ds_config["train_batch_size"]
-                ds_config.update({"train_micro_batch_size_per_gpu": 1, "gradient_accumulation_steps": 1})
-                mixed_precision = (
-                    os.environ.get("MIXED_PRECISION", "no") if mixed_precision is None else mixed_precision
-                )
-                if mixed_precision == "fp16":
-                    ds_config.update({"fp16": {"enabled": True}})
-                elif mixed_precision == "bf16":
-                    ds_config.update({"bf16": {"enabled": True}})
-                self.dschf = HfDeepSpeedConfig(ds_config)  # keep this object alive # noqa
+            mixed_precision = os.environ.get("MIXED_PRECISION", "no") if mixed_precision is None else mixed_precision
+            deepspeed_plugin.set_mixed_precision(mixed_precision)
+            deepspeed_plugin.set_deepspeed_weakref()
 
         if os.environ.get("USE_FSDP", "false") == "true" or isinstance(fsdp_plugin, FullyShardedDataParallelPlugin):
             if is_torch_version("<", "1.12.0.dev20220418+cu113"):
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 2b59d39f6..2bc6fbe4d 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -353,7 +353,7 @@ def deepspeed_launcher(args):
     current_env["USE_DEEPSPEED"] = "true"
     current_env["DEEPSPEED_ZERO_STAGE"] = str(args.zero_stage)
     current_env["GRADIENT_ACCUMULATION_STEPS"] = str(args.gradient_accumulation_steps)
-    current_env["GRADIENT_CLIPPING"] = str(args.gradient_clipping)
+    current_env["GRADIENT_CLIPPING"] = str(args.gradient_clipping).lower()
     current_env["DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE"] = str(args.offload_optimizer_device).lower()
     current_env["DEEPSPEED_OFFLOAD_PARAM_DEVICE"] = str(args.offload_param_device).lower()
     current_env["DEEPSPEED_ZERO3_INIT"] = str(args.zero3_init_flag).lower()
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index 891b0762b..e341ad93f 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -108,21 +108,6 @@ def __init__(
                 self.device = torch.device("cuda", self.local_process_index)
                 torch.cuda.set_device(self.device)
                 self.mixed_precision = "no"  # deepspeed handles mixed_precision using deepspeed_config
-                mixed_precision = (
-                    parse_choice_from_env("MIXED_PRECISION", "no") if mixed_precision is None else mixed_precision
-                )
-                if (
-                    mixed_precision == "fp16"
-                    and "fp16" not in deepspeed_plugin.deepspeed_config
-                    and "bf16" not in deepspeed_plugin.deepspeed_config
-                ):
-                    deepspeed_plugin.deepspeed_config.update({"fp16": {"enabled": True}})
-                elif (
-                    mixed_precision == "bf16"
-                    and "fp16" not in deepspeed_plugin.deepspeed_config
-                    and "bf16" not in deepspeed_plugin.deepspeed_config
-                ):
-                    deepspeed_plugin.deepspeed_config.update({"bf16": {"enabled": True}})
                 self.deepspeed_plugin = deepspeed_plugin
             elif int(os.environ.get("LOCAL_RANK", -1)) != -1 and not cpu:
                 self.distributed_type = DistributedType.MULTI_GPU
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
index a3e193915..60f39c90a 100644
--- a/src/accelerate/utils/__init__.py
+++ b/src/accelerate/utils/__init__.py
@@ -83,6 +83,7 @@
         DeepSpeedSchedulerWrapper,
         DummyOptim,
         DummyScheduler,
+        HfDeepSpeedConfig,
     )
 
 from .launch import PrepareForLaunch, get_launch_prefix
diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py
index b731ca12a..8b8e30304 100644
--- a/src/accelerate/utils/dataclasses.py
+++ b/src/accelerate/utils/dataclasses.py
@@ -19,14 +19,12 @@
 import copy
 import enum
 import functools
-import io
-import json
 import os
 import typing
 import warnings
 from dataclasses import dataclass, field
 from datetime import timedelta
-from typing import Callable, Iterable, Optional
+from typing import Any, Callable, Iterable, Optional
 
 import torch
 
@@ -215,7 +213,12 @@ class DeepSpeedPlugin:
     This plugin is used to integrate DeepSpeed.
     """
 
-    config_file: str = field(default=None, metadata={"help": "Path to the DeepSpeed config file."})
+    hf_ds_config: Any = field(
+        default=None,
+        metadata={
+            "help": "path to DeepSpeed config file or dict or an object of class `accelerate.utils.deepspeed.HfDeepSpeedConfig`."
+        },
+    )
     gradient_accumulation_steps: int = field(
         default=None, metadata={"help": "Number of steps to accumulate gradients before updating optimizer states"}
     )
@@ -249,17 +252,23 @@ class DeepSpeedPlugin:
     )
 
     def __post_init__(self):
-        if self.config_file is None:
-            self.config_file = os.environ.get("DEEPSPEED_CONFIG_FILE", "none")
-        if self.config_file != "none":
-            with io.open(self.config_file, "r", encoding="utf-8") as f:
-                self.deepspeed_config = json.load(f)
-            if "gradient_accumulation_steps" not in self.deepspeed_config:
-                self.deepspeed_config["gradient_accumulation_steps"] = 1
-            elif self.deepspeed_config["gradient_accumulation_steps"] == "auto":
-                raise ValueError("gradient_accumulation_steps cannot be set to 'auto' in the DeepSpeed config file.")
-            if "zero_optimization" not in self.deepspeed_config:
-                raise ValueError("Please specify the ZeRO optimization config in the DeepSpeed config file.")
+        from .deepspeed import HfDeepSpeedConfig
+
+        if self.hf_ds_config is None:
+            self.hf_ds_config = os.environ.get("DEEPSPEED_CONFIG_FILE", "none")
+        if (
+            isinstance(self.hf_ds_config, dict)
+            or (isinstance(self.hf_ds_config, str) and self.hf_ds_config != "none")
+            or isinstance(self.hf_ds_config, HfDeepSpeedConfig)
+        ):
+            if not isinstance(self.hf_ds_config, HfDeepSpeedConfig):
+                self.hf_ds_config = HfDeepSpeedConfig(self.hf_ds_config)
+            if "gradient_accumulation_steps" not in self.hf_ds_config.config:
+                self.hf_ds_config.config["gradient_accumulation_steps"] = 1
+            elif self.hf_ds_config.config["gradient_accumulation_steps"] == "auto":
+                raise ValueError("gradient_accumulation_steps cannot be set to 'auto' in the DeepSpeed config.")
+            if "zero_optimization" not in self.hf_ds_config.config:
+                raise ValueError("Please specify the ZeRO optimization config in the DeepSpeed config.")
         else:
             if self.gradient_accumulation_steps is None:
                 self.gradient_accumulation_steps = int(os.environ.get("GRADIENT_ACCUMULATION_STEPS", 1))
@@ -281,8 +290,9 @@ def __post_init__(self):
             if self.zero3_save_16bit_model is None:
                 self.zero3_save_16bit_model = os.environ.get("DEEPSPEED_ZERO3_SAVE_16BIT_MODEL", "false") == "true"
 
-            self.deepspeed_config = {
+            config = {
                 "train_batch_size": "auto",
+                "train_micro_batch_size_per_gpu": "auto",
                 "gradient_accumulation_steps": self.gradient_accumulation_steps,
                 "zero_optimization": {
                     "stage": self.zero_stage,
@@ -296,29 +306,18 @@ def __post_init__(self):
                 },
             }
             if self.gradient_clipping:
-                self.deepspeed_config["gradient_clipping"] = self.gradient_clipping
+                config["gradient_clipping"] = self.gradient_clipping
+            self.hf_ds_config = HfDeepSpeedConfig(config)
+        self.deepspeed_config = self.hf_ds_config.config
         self.deepspeed_config["steps_per_print"] = float("inf")  # this will stop deepspeed from logging @ stdout
         if self.zero3_init_flag is None:
             self.zero3_init_flag = os.environ.get("DEEPSPEED_ZERO3_INIT", "false") == "true"
-        if self.zero3_init_flag and self.deepspeed_config["zero_optimization"]["stage"] != 3:
+        if self.zero3_init_flag and not self.hf_ds_config.is_zero3():
             warnings.warn("DeepSpeed Zero3 Init flag is only applicable for ZeRO Stage 3. Setting it to False.")
             self.zero3_init_flag = False
 
-    def find_config_node(self, ds_key_long):
-        config = self.deepspeed_config
-
-        # find the config node of interest if it exists
-        nodes = ds_key_long.split(".")
-        ds_key = nodes.pop()
-        for node in nodes:
-            config = config.get(node)
-            if config is None:
-                return None, ds_key
-
-        return config, ds_key
-
     def fill_match(self, ds_key_long, mismatches, must_match=True, **kwargs):
-        config, ds_key = self.find_config_node(ds_key_long)
+        config, ds_key = self.hf_ds_config.find_config_node(ds_key_long)
         if config is None:
             return
 
@@ -360,6 +359,37 @@ def deepspeed_config_process(self, prefix="", mismatches=None, config=None, must
                 f" values:\n{mismatches_msg}\nThe easiest method is to set these DeepSpeed config values to 'auto'."
             )
 
+    def set_mixed_precision(self, mixed_precision):
+        ds_config = self.deepspeed_config
+        if mixed_precision == "fp16" and "fp16" not in ds_config and "bf16" not in ds_config:
+            ds_config.update({"fp16": {"enabled": True}})
+        elif mixed_precision == "bf16" and "fp16" not in ds_config and "bf16" not in ds_config:
+            ds_config.update({"bf16": {"enabled": True}})
+
+    def set_deepspeed_weakref(self):
+        from .imports import is_transformers_available
+
+        if self.zero3_init_flag:
+            if not is_transformers_available():
+                raise Exception(
+                    "When `zero3_init_flag` is set, it requires Transformers to be installed. "
+                    "Please run `pip install transformers`."
+                )
+            ds_config = copy.deepcopy(self.deepspeed_config)
+            if "gradient_accumulation_steps" not in ds_config or ds_config["gradient_accumulation_steps"] == "auto":
+                ds_config["gradient_accumulation_steps"] = 1
+            if (
+                "train_micro_batch_size_per_gpu" not in ds_config
+                or ds_config["train_micro_batch_size_per_gpu"] == "auto"
+            ):
+                ds_config["train_micro_batch_size_per_gpu"] = 1
+            if ds_config["train_batch_size"] == "auto":
+                del ds_config["train_batch_size"]
+
+            from transformers.deepspeed import HfDeepSpeedConfig
+
+            self.dschf = HfDeepSpeedConfig(ds_config)  # keep this object alive # noqa
+
 
 @dataclass
 class FullyShardedDataParallelPlugin:
diff --git a/src/accelerate/utils/deepspeed.py b/src/accelerate/utils/deepspeed.py
index dde1b9c63..02d1ab8bc 100644
--- a/src/accelerate/utils/deepspeed.py
+++ b/src/accelerate/utils/deepspeed.py
@@ -12,9 +12,132 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from accelerate.scheduler import AcceleratedScheduler
+import io
+import json
+from copy import deepcopy
 
 from ..optimizer import AcceleratedOptimizer
+from ..scheduler import AcceleratedScheduler
+
+
+class HfDeepSpeedConfig:
+    """
+    This object contains a DeepSpeed configuration dictionary and can be quickly queried for things like zero stage.
+
+    A `weakref` of this object is stored in the module's globals to be able to access the config from areas where
+    things like the Trainer object is not available (e.g. `from_pretrained` and `_get_resized_embeddings`). Therefore
+    it's important that this object remains alive while the program is still running.
+
+    [`Trainer`] uses the `HfTrainerDeepSpeedConfig` subclass instead. That subclass has logic to sync the configuration
+    with values of [`TrainingArguments`] by replacing special placeholder values: `"auto"`. Without this special logic
+    the DeepSpeed configuration is not modified in any way.
+
+    Args:
+        config_file_or_dict (`Union[str, Dict]`): path to DeepSpeed config file or dict.
+
+    """
+
+    def __init__(self, config_file_or_dict):
+
+        if isinstance(config_file_or_dict, dict):
+            # Don't modify user's data should they want to reuse it (e.g. in tests), because once we
+            # modified it, it will not be accepted here again, since `auto` values would have been overridden
+            config = deepcopy(config_file_or_dict)
+        elif isinstance(config_file_or_dict, str):
+            with io.open(config_file_or_dict, "r", encoding="utf-8") as f:
+                config = json.load(f)
+        else:
+            raise ValueError("expecting either a path to a DeepSpeed config file or a pre-populated dict")
+        self.config = config
+
+        # zero stage - this is done as early as possible, before model is created, to allow
+        # ``is_deepspeed_zero3_enabled`` query and getting to the early deepspeed config object
+        # during ``zero.Init()`` which needs to know the dtype, and some other hparams.
+        self._stage = self.get_value("zero_optimization.stage", -1)
+
+        # offload
+        self._offload = False
+        if self.is_zero2() or self.is_zero3():
+            offload_devices_valid = set(["cpu", "nvme"])
+            offload_devices = set(
+                [
+                    self.get_value("zero_optimization.offload_optimizer.device"),
+                    self.get_value("zero_optimization.offload_param.device"),
+                ]
+            )
+            if len(offload_devices & offload_devices_valid) > 0:
+                self._offload = True
+
+    def find_config_node(self, ds_key_long):
+        config = self.config
+
+        # find the config node of interest if it exists
+        nodes = ds_key_long.split(".")
+        ds_key = nodes.pop()
+        for node in nodes:
+            config = config.get(node)
+            if config is None:
+                return None, ds_key
+
+        return config, ds_key
+
+    def get_value(self, ds_key_long, default=None):
+        """
+        Returns the set value or `default` if no value is set
+        """
+        config, ds_key = self.find_config_node(ds_key_long)
+        if config is None:
+            return default
+        return config.get(ds_key, default)
+
+    def del_config_sub_tree(self, ds_key_long, must_exist=False):
+        """
+        Deletes a sub-section of the config file if it's found.
+
+        Unless `must_exist` is `True` the section doesn't have to exist.
+        """
+        config = self.config
+
+        # find the config node of interest if it exists
+        nodes = ds_key_long.split(".")
+        for node in nodes:
+            parent_config = config
+            config = config.get(node)
+            if config is None:
+                if must_exist:
+                    raise ValueError(f"Can't find {ds_key_long} entry in the config: {self.config}")
+                else:
+                    return
+
+        # if found remove it
+        if parent_config is not None:
+            parent_config.pop(node)
+
+    def is_true(self, ds_key_long):
+        """
+        Returns `True`/``False` only if the value is set, always `False` otherwise. So use this method to ask the very
+        specific question of whether the value is set to `True` (and it's not set to `False`` or isn't set).
+
+        """
+        value = self.get_value(ds_key_long)
+        return False if value is None else bool(value)
+
+    def is_false(self, ds_key_long):
+        """
+        Returns `True`/``False` only if the value is set, always `False` otherwise. So use this method to ask the very
+        specific question of whether the value is set to `False` (and it's not set to `True`` or isn't set).
+        """
+        value = self.get_value(ds_key_long)
+        return False if value is None else not bool(value)
+
+    def is_zero2(self):
+        return self._stage == 2
+
+    def is_zero3(self):
+        return self._stage == 3
+
+    def is_offload(self):
+        return self._offload
 
 
 class DeepSpeedEngineWrapper:
diff --git a/tests/deepspeed/test_deepspeed.py b/tests/deepspeed/test_deepspeed.py
index 9ece7b0f7..12cca415c 100644
--- a/tests/deepspeed/test_deepspeed.py
+++ b/tests/deepspeed/test_deepspeed.py
@@ -40,7 +40,6 @@
 )
 from parameterized import parameterized
 from transformers import AutoModel, AutoModelForCausalLM, get_scheduler
-from transformers.deepspeed import HfDeepSpeedConfig
 from transformers.testing_utils import mockenv_context
 from transformers.trainer_utils import set_seed
 from transformers.utils import is_torch_bf16_available
@@ -153,7 +152,7 @@ def test_deepspeed_plugin(self, stage):
         deepspeed_plugin.deepspeed_config = None
 
         # Test config files are loaded correctly
-        deepspeed_plugin = DeepSpeedPlugin(config_file=self.ds_config_file[stage], zero3_init_flag=True)
+        deepspeed_plugin = DeepSpeedPlugin(hf_ds_config=self.ds_config_file[stage], zero3_init_flag=True)
         if stage == ZERO2:
             self.assertFalse(deepspeed_plugin.zero3_init_flag)
         elif stage == ZERO3:
@@ -165,7 +164,7 @@ def test_deepspeed_plugin(self, stage):
             del ds_config["gradient_accumulation_steps"]
             with open(os.path.join(dirpath, "ds_config.json"), "w") as out_file:
                 json.dump(ds_config, out_file)
-            deepspeed_plugin = DeepSpeedPlugin(config_file=os.path.join(dirpath, "ds_config.json"))
+            deepspeed_plugin = DeepSpeedPlugin(hf_ds_config=os.path.join(dirpath, "ds_config.json"))
             self.assertEqual(deepspeed_plugin.deepspeed_config["gradient_accumulation_steps"], 1)
             deepspeed_plugin.deepspeed_config = None
 
@@ -176,14 +175,14 @@ def test_deepspeed_plugin(self, stage):
             with open(os.path.join(dirpath, "ds_config.json"), "w") as out_file:
                 json.dump(ds_config, out_file)
             with self.assertRaises(ValueError) as cm:
-                deepspeed_plugin = DeepSpeedPlugin(config_file=os.path.join(dirpath, "ds_config.json"))
+                deepspeed_plugin = DeepSpeedPlugin(hf_ds_config=os.path.join(dirpath, "ds_config.json"))
             self.assertTrue(
-                "Please specify the ZeRO optimization config in the DeepSpeed config file." in str(cm.exception)
+                "Please specify the ZeRO optimization config in the DeepSpeed config." in str(cm.exception)
             )
             deepspeed_plugin.deepspeed_config = None
 
         # Test `deepspeed_config_process`
-        deepspeed_plugin = DeepSpeedPlugin(config_file=self.ds_config_file[stage])
+        deepspeed_plugin = DeepSpeedPlugin(hf_ds_config=self.ds_config_file[stage])
         kwargs = {
             "fp16.enabled": True,
             "bf16.enabled": False,
@@ -202,7 +201,7 @@ def test_deepspeed_plugin(self, stage):
         }
         deepspeed_plugin.deepspeed_config_process(**kwargs)
         for ds_key_long, value in kwargs.items():
-            config, ds_key = deepspeed_plugin.find_config_node(ds_key_long)
+            config, ds_key = deepspeed_plugin.hf_ds_config.find_config_node(ds_key_long)
             if config.get(ds_key) is not None:
                 self.assertEqual(config.get(ds_key), value)
 
@@ -245,7 +244,7 @@ def test_accelerate_state_deepspeed(self, dtype):
             zero3_init_flag=True,
         )
         with mockenv_context(**self.dist_env):
-            state = AcceleratorState(mixed_precision=dtype, deepspeed_plugin=deepspeed_plugin, _from_accelerator=True)
+            state = Accelerator(mixed_precision=dtype, deepspeed_plugin=deepspeed_plugin).state
             self.assertTrue(state.deepspeed_plugin.deepspeed_config[dtype]["enabled"])
             state.initialized = False
 
@@ -262,8 +261,10 @@ def test_init_zero3(self):
 
         with mockenv_context(**self.dist_env):
             accelerator = Accelerator(deepspeed_plugin=deepspeed_plugin)
-            self.assertTrue("dschf" in accelerator.__dict__)
-            self.assertTrue(type(accelerator.dschf) == HfDeepSpeedConfig)
+            from transformers.deepspeed import is_deepspeed_zero3_enabled
+
+            self.assertTrue(is_deepspeed_zero3_enabled())
+            accelerator.state.initialized = False
 
     @parameterized.expand(optim_scheduler_params, name_func=parameterized_custom_name_func)
     def test_prepare_deepspeed(self, optim_type, scheduler_type):
@@ -301,7 +302,6 @@ def test_prepare_deepspeed(self, optim_type, scheduler_type):
             )
             with mockenv_context(**self.dist_env):
                 accelerator = Accelerator(mixed_precision="fp16", deepspeed_plugin=deepspeed_plugin)
-                self.assertEqual(accelerator.state.deepspeed_plugin.config_file, "none")
 
                 train_set = RegressionDataset(length=80)
                 eval_set = RegressionDataset(length=20)
@@ -354,7 +354,7 @@ def test_prepare_deepspeed(self, optim_type, scheduler_type):
 
         elif optim_type == DS_OPTIMIZER and scheduler_type == DS_SCHEDULER:
             # Test DeepSpeed optimizer + DeepSpeed scheduler
-            deepspeed_plugin = DeepSpeedPlugin(config_file=self.ds_config_file[ZERO2])
+            deepspeed_plugin = DeepSpeedPlugin(hf_ds_config=self.ds_config_file[ZERO2])
             with mockenv_context(**self.dist_env):
                 accelerator = Accelerator(deepspeed_plugin=deepspeed_plugin)
                 train_set = RegressionDataset(length=80)
@@ -414,7 +414,7 @@ def test_prepare_deepspeed(self, optim_type, scheduler_type):
 
         elif optim_type == CUSTOM_OPTIMIZER and scheduler_type == DS_SCHEDULER:
             # Test custom optimizer + DeepSpeed scheduler
-            deepspeed_plugin = DeepSpeedPlugin(config_file=self.ds_config_file[ZERO2])
+            deepspeed_plugin = DeepSpeedPlugin(hf_ds_config=self.ds_config_file[ZERO2])
             with mockenv_context(**self.dist_env):
                 accelerator = Accelerator(deepspeed_plugin=deepspeed_plugin)
                 train_set = RegressionDataset(length=80)
@@ -447,7 +447,7 @@ def test_prepare_deepspeed(self, optim_type, scheduler_type):
                 self.assertTrue(type(accelerator.deepspeed_engine_wrapped) == DeepSpeedEngineWrapper)
         elif optim_type == DS_OPTIMIZER and scheduler_type == CUSTOM_SCHEDULER:
             # Test deepspeed optimizer + custom scheduler
-            deepspeed_plugin = DeepSpeedPlugin(config_file=self.ds_config_file[ZERO2])
+            deepspeed_plugin = DeepSpeedPlugin(hf_ds_config=self.ds_config_file[ZERO2])
             with mockenv_context(**self.dist_env):
                 accelerator = Accelerator(deepspeed_plugin=deepspeed_plugin)
                 train_set = RegressionDataset(length=80)
@@ -483,7 +483,7 @@ def test_prepare_deepspeed(self, optim_type, scheduler_type):
 
     def test_save_checkpoints(self):
         deepspeed_plugin = DeepSpeedPlugin(
-            config_file=self.ds_config_file[ZERO3],
+            hf_ds_config=self.ds_config_file[ZERO3],
             zero3_init_flag=True,
         )
         del deepspeed_plugin.deepspeed_config["bf16"]
@@ -537,7 +537,7 @@ def test_save_checkpoints(self):
 
     def test_autofill_dsconfig(self):
         deepspeed_plugin = DeepSpeedPlugin(
-            config_file=self.ds_config_file[ZERO3],
+            hf_ds_config=self.ds_config_file[ZERO3],
             zero3_init_flag=True,
         )
         del deepspeed_plugin.deepspeed_config["bf16"]

This is breaking. While the API was marked as experimental and we can do such a change, I'd like to understand why we can't keep the old argument name here.
@@ -215,7 +213,12 @@ class DeepSpeedPlugin:
     This plugin is used to integrate DeepSpeed.
     """
 
-    config_file: str = field(default=None, metadata={"help": "Path to the DeepSpeed config file."})
+    hf_ds_config: Any = field(
Hello, this isn't breaking change as this feature was added in PR from couple days ago #405 . It isn't part of the previous release (0.9.0) and as such won't be a breaking change. The renaming was done to clearly explain that it isn't just limited to file and can accept either file, dict or `HFDeepSpeedConfig` Object.
@@ -215,7 +213,12 @@ class DeepSpeedPlugin:
     This plugin is used to integrate DeepSpeed.
     """
 
-    config_file: str = field(default=None, metadata={"help": "Path to the DeepSpeed config file."})
+    hf_ds_config: Any = field(
Oh okay then, I thought it had been there for longer.
@@ -215,7 +213,12 @@ class DeepSpeedPlugin:
     This plugin is used to integrate DeepSpeed.
     """
 
-    config_file: str = field(default=None, metadata={"help": "Path to the DeepSpeed config file."})
+    hf_ds_config: Any = field(
Let's check if the key exists before setting it a None?
@@ -193,3 +194,47 @@ def __repr__(self):
     @property
     def use_fp16(self):
         return self.mixed_precision != "no"
+
+    @staticmethod
+    def is_deepspeed_zero3_enabled():
+        if (
+            "deepspeed_plugin" in AcceleratorState._shared_state
+            and AcceleratorState._shared_state["deepspeed_plugin"] is not None
+        ):
+            deepspeed_plugin = AcceleratorState._shared_state["deepspeed_plugin"]
+            return deepspeed_plugin.hf_ds_config.is_zero3() and deepspeed_plugin.zero3_init_flag
+        return False
+
+    @staticmethod
+    def get_deepspeed_config():
+        ds_config = None
+        if (
+            "deepspeed_plugin" in AcceleratorState._shared_state
+            and AcceleratorState._shared_state["deepspeed_plugin"] is not None
+        ):
+            ds_config = deepcopy(AcceleratorState._shared_state["deepspeed_plugin"].hf_ds_config.config)
+            if ds_config["gradient_accumulation_steps"] == "auto":
+                ds_config["gradient_accumulation_steps"] = 1
+            if (
+                "train_micro_batch_size_per_gpu" not in ds_config
+                or ds_config["train_micro_batch_size_per_gpu"] == "auto"
+            ):
+                ds_config["train_micro_batch_size_per_gpu"] = 1
+            if ds_config["train_batch_size"] == "auto":
+                del ds_config["train_batch_size"]
+        return ds_config
+
+    @staticmethod
+    def set_deepspeed_config(hf_ds_config):
+        from .utils import DeepSpeedPlugin, HfDeepSpeedConfig
+
+        if isinstance(hf_ds_config, HfDeepSpeedConfig):
+            AcceleratorState._shared_state["deepspeed_plugin"] = DeepSpeedPlugin(
+                hf_ds_config=hf_ds_config, zero3_init_flag=True
+            )
+        else:
+            raise ValueError("`hf_ds_config` must be an instance of `accelerate.utils.HfDeepSpeedConfig`")
+
+    @staticmethod
+    def unset_deepspeed_config():
+        AcceleratorState._shared_state["deepspeed_plugin"] = None
Let's error first if DeepSpeed is not available.
@@ -193,3 +194,47 @@ def __repr__(self):
     @property
     def use_fp16(self):
         return self.mixed_precision != "no"
+
+    @staticmethod
+    def is_deepspeed_zero3_enabled():
+        if (
+            "deepspeed_plugin" in AcceleratorState._shared_state
+            and AcceleratorState._shared_state["deepspeed_plugin"] is not None
+        ):
+            deepspeed_plugin = AcceleratorState._shared_state["deepspeed_plugin"]
+            return deepspeed_plugin.hf_ds_config.is_zero3() and deepspeed_plugin.zero3_init_flag
+        return False
+
+    @staticmethod
+    def get_deepspeed_config():
+        ds_config = None
+        if (
+            "deepspeed_plugin" in AcceleratorState._shared_state
+            and AcceleratorState._shared_state["deepspeed_plugin"] is not None
+        ):
+            ds_config = deepcopy(AcceleratorState._shared_state["deepspeed_plugin"].hf_ds_config.config)
+            if ds_config["gradient_accumulation_steps"] == "auto":
+                ds_config["gradient_accumulation_steps"] = 1
+            if (
+                "train_micro_batch_size_per_gpu" not in ds_config
+                or ds_config["train_micro_batch_size_per_gpu"] == "auto"
+            ):
+                ds_config["train_micro_batch_size_per_gpu"] = 1
+            if ds_config["train_batch_size"] == "auto":
+                del ds_config["train_batch_size"]
+        return ds_config
+
+    @staticmethod
+    def set_deepspeed_config(hf_ds_config):
+        from .utils import DeepSpeedPlugin, HfDeepSpeedConfig
setting deepspeed_plugin doesn't use any deepspeed functionality; Accelerate and Trainer check DeepSpeed availability at different places. Should we still check deepspeed availability here?
@@ -193,3 +194,47 @@ def __repr__(self):
     @property
     def use_fp16(self):
         return self.mixed_precision != "no"
+
+    @staticmethod
+    def is_deepspeed_zero3_enabled():
+        if (
+            "deepspeed_plugin" in AcceleratorState._shared_state
+            and AcceleratorState._shared_state["deepspeed_plugin"] is not None
+        ):
+            deepspeed_plugin = AcceleratorState._shared_state["deepspeed_plugin"]
+            return deepspeed_plugin.hf_ds_config.is_zero3() and deepspeed_plugin.zero3_init_flag
+        return False
+
+    @staticmethod
+    def get_deepspeed_config():
+        ds_config = None
+        if (
+            "deepspeed_plugin" in AcceleratorState._shared_state
+            and AcceleratorState._shared_state["deepspeed_plugin"] is not None
+        ):
+            ds_config = deepcopy(AcceleratorState._shared_state["deepspeed_plugin"].hf_ds_config.config)
+            if ds_config["gradient_accumulation_steps"] == "auto":
+                ds_config["gradient_accumulation_steps"] = 1
+            if (
+                "train_micro_batch_size_per_gpu" not in ds_config
+                or ds_config["train_micro_batch_size_per_gpu"] == "auto"
+            ):
+                ds_config["train_micro_batch_size_per_gpu"] = 1
+            if ds_config["train_batch_size"] == "auto":
+                del ds_config["train_batch_size"]
+        return ds_config
+
+    @staticmethod
+    def set_deepspeed_config(hf_ds_config):
+        from .utils import DeepSpeedPlugin, HfDeepSpeedConfig
As we are directly setting the key in the dictionary, if it isn't present, it will be set to None which should not make a difference. I don't think it is required. 
@@ -193,3 +194,47 @@ def __repr__(self):
     @property
     def use_fp16(self):
         return self.mixed_precision != "no"
+
+    @staticmethod
+    def is_deepspeed_zero3_enabled():
+        if (
+            "deepspeed_plugin" in AcceleratorState._shared_state
+            and AcceleratorState._shared_state["deepspeed_plugin"] is not None
+        ):
+            deepspeed_plugin = AcceleratorState._shared_state["deepspeed_plugin"]
+            return deepspeed_plugin.hf_ds_config.is_zero3() and deepspeed_plugin.zero3_init_flag
+        return False
+
+    @staticmethod
+    def get_deepspeed_config():
+        ds_config = None
+        if (
+            "deepspeed_plugin" in AcceleratorState._shared_state
+            and AcceleratorState._shared_state["deepspeed_plugin"] is not None
+        ):
+            ds_config = deepcopy(AcceleratorState._shared_state["deepspeed_plugin"].hf_ds_config.config)
+            if ds_config["gradient_accumulation_steps"] == "auto":
+                ds_config["gradient_accumulation_steps"] = 1
+            if (
+                "train_micro_batch_size_per_gpu" not in ds_config
+                or ds_config["train_micro_batch_size_per_gpu"] == "auto"
+            ):
+                ds_config["train_micro_batch_size_per_gpu"] = 1
+            if ds_config["train_batch_size"] == "auto":
+                del ds_config["train_batch_size"]
+        return ds_config
+
+    @staticmethod
+    def set_deepspeed_config(hf_ds_config):
+        from .utils import DeepSpeedPlugin, HfDeepSpeedConfig
+
+        if isinstance(hf_ds_config, HfDeepSpeedConfig):
+            AcceleratorState._shared_state["deepspeed_plugin"] = DeepSpeedPlugin(
+                hf_ds_config=hf_ds_config, zero3_init_flag=True
+            )
+        else:
+            raise ValueError("`hf_ds_config` must be an instance of `accelerate.utils.HfDeepSpeedConfig`")
+
+    @staticmethod
+    def unset_deepspeed_config():
+        AcceleratorState._shared_state["deepspeed_plugin"] = None
Oh I thought there was some DeepSpeed import in that file. Okay to leave as is.
@@ -193,3 +194,47 @@ def __repr__(self):
     @property
     def use_fp16(self):
         return self.mixed_precision != "no"
+
+    @staticmethod
+    def is_deepspeed_zero3_enabled():
+        if (
+            "deepspeed_plugin" in AcceleratorState._shared_state
+            and AcceleratorState._shared_state["deepspeed_plugin"] is not None
+        ):
+            deepspeed_plugin = AcceleratorState._shared_state["deepspeed_plugin"]
+            return deepspeed_plugin.hf_ds_config.is_zero3() and deepspeed_plugin.zero3_init_flag
+        return False
+
+    @staticmethod
+    def get_deepspeed_config():
+        ds_config = None
+        if (
+            "deepspeed_plugin" in AcceleratorState._shared_state
+            and AcceleratorState._shared_state["deepspeed_plugin"] is not None
+        ):
+            ds_config = deepcopy(AcceleratorState._shared_state["deepspeed_plugin"].hf_ds_config.config)
+            if ds_config["gradient_accumulation_steps"] == "auto":
+                ds_config["gradient_accumulation_steps"] = 1
+            if (
+                "train_micro_batch_size_per_gpu" not in ds_config
+                or ds_config["train_micro_batch_size_per_gpu"] == "auto"
+            ):
+                ds_config["train_micro_batch_size_per_gpu"] = 1
+            if ds_config["train_batch_size"] == "auto":
+                del ds_config["train_batch_size"]
+        return ds_config
+
+    @staticmethod
+    def set_deepspeed_config(hf_ds_config):
+        from .utils import DeepSpeedPlugin, HfDeepSpeedConfig
As long as you're sure it won't make any deepspeed related stuff break, fine by me.
@@ -193,3 +194,47 @@ def __repr__(self):
     @property
     def use_fp16(self):
         return self.mixed_precision != "no"
+
+    @staticmethod
+    def is_deepspeed_zero3_enabled():
+        if (
+            "deepspeed_plugin" in AcceleratorState._shared_state
+            and AcceleratorState._shared_state["deepspeed_plugin"] is not None
+        ):
+            deepspeed_plugin = AcceleratorState._shared_state["deepspeed_plugin"]
+            return deepspeed_plugin.hf_ds_config.is_zero3() and deepspeed_plugin.zero3_init_flag
+        return False
+
+    @staticmethod
+    def get_deepspeed_config():
+        ds_config = None
+        if (
+            "deepspeed_plugin" in AcceleratorState._shared_state
+            and AcceleratorState._shared_state["deepspeed_plugin"] is not None
+        ):
+            ds_config = deepcopy(AcceleratorState._shared_state["deepspeed_plugin"].hf_ds_config.config)
+            if ds_config["gradient_accumulation_steps"] == "auto":
+                ds_config["gradient_accumulation_steps"] = 1
+            if (
+                "train_micro_batch_size_per_gpu" not in ds_config
+                or ds_config["train_micro_batch_size_per_gpu"] == "auto"
+            ):
+                ds_config["train_micro_batch_size_per_gpu"] = 1
+            if ds_config["train_batch_size"] == "auto":
+                del ds_config["train_batch_size"]
+        return ds_config
+
+    @staticmethod
+    def set_deepspeed_config(hf_ds_config):
+        from .utils import DeepSpeedPlugin, HfDeepSpeedConfig
+
+        if isinstance(hf_ds_config, HfDeepSpeedConfig):
+            AcceleratorState._shared_state["deepspeed_plugin"] = DeepSpeedPlugin(
+                hf_ds_config=hf_ds_config, zero3_init_flag=True
+            )
+        else:
+            raise ValueError("`hf_ds_config` must be an instance of `accelerate.utils.HfDeepSpeedConfig`")
+
+    @staticmethod
+    def unset_deepspeed_config():
+        AcceleratorState._shared_state["deepspeed_plugin"] = None
Reverted this change as per latest discussions on transformers PR
@@ -193,3 +194,47 @@ def __repr__(self):
     @property
     def use_fp16(self):
         return self.mixed_precision != "no"
+
+    @staticmethod
+    def is_deepspeed_zero3_enabled():
+        if (
+            "deepspeed_plugin" in AcceleratorState._shared_state
+            and AcceleratorState._shared_state["deepspeed_plugin"] is not None
+        ):
+            deepspeed_plugin = AcceleratorState._shared_state["deepspeed_plugin"]
+            return deepspeed_plugin.hf_ds_config.is_zero3() and deepspeed_plugin.zero3_init_flag
+        return False
+
+    @staticmethod
+    def get_deepspeed_config():
+        ds_config = None
+        if (
+            "deepspeed_plugin" in AcceleratorState._shared_state
+            and AcceleratorState._shared_state["deepspeed_plugin"] is not None
+        ):
+            ds_config = deepcopy(AcceleratorState._shared_state["deepspeed_plugin"].hf_ds_config.config)
+            if ds_config["gradient_accumulation_steps"] == "auto":
+                ds_config["gradient_accumulation_steps"] = 1
+            if (
+                "train_micro_batch_size_per_gpu" not in ds_config
+                or ds_config["train_micro_batch_size_per_gpu"] == "auto"
+            ):
+                ds_config["train_micro_batch_size_per_gpu"] = 1
+            if ds_config["train_batch_size"] == "auto":
+                del ds_config["train_batch_size"]
+        return ds_config
+
+    @staticmethod
+    def set_deepspeed_config(hf_ds_config):
+        from .utils import DeepSpeedPlugin, HfDeepSpeedConfig
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/432
diff --git a/examples/by_feature/README.md b/examples/by_feature/README.md
index f0db09937..689127a35 100644
--- a/examples/by_feature/README.md
+++ b/examples/by_feature/README.md
@@ -42,6 +42,18 @@ These arguments should be added at the end of any method for starting the python
 accelerate launch ./checkpointing.py --checkpointing_steps epoch output_dir "checkpointing_tutorial" --resume_from_checkpoint "checkpointing_tutorial/epoch_0"
 ```
 
+### Cross Validation (`cross_validation.py`)
+
+- Shows how to use `Accelerator.free_memory` and run cross validation efficiently with `datasets`.
+- Arguments available:
+  - `num_folds`, the number of folds the training dataset should be split into.
+
+These arguments should be added at the end of any method for starting the python script (such as `python`, `accelerate launch`, `python -m torch.distributed.launch`), such as:
+
+```bash
+accelerate launch ./cross_validation.py --num_folds 2
+```
+
 ### Experiment Tracking (`tracking.py`)
 
 - Shows how to use `Accelerate.init_trackers` and `Accelerator.log`
@@ -55,14 +67,14 @@ These arguments should be added at the end of any method for starting the python
 accelerate launch ./tracking.py --with_tracking
 ```
 
-### Cross Validation (`cross_validation.py`)
+### Gradient Accumulation (`gradient_accumulation.py`)
 
-- Shows how to use `Accelerator.free_memory` and run cross validation efficiently with `datasets`.
+- Shows how to use `Accelerator.no_sync` to prevent gradient averaging in a distributed setup.
 - Arguments available:
-  - `num_folds`, the number of folds the training dataset should be split into.
+  - `gradient_accumulation_steps`, the number of steps to perform before the gradients are accumulated and the optimizer and scheduler are stepped + zero_grad
 
 These arguments should be added at the end of any method for starting the python script (such as `python`, `accelerate launch`, `python -m torch.distributed.launch`), such as:
 
 ```bash
-accelerate launch ./cross_validation.py --num_folds 2
-```
+accelerate launch ./gradient_accumulation.py --gradient_accumulation_steps 5
+```
\ No newline at end of file
diff --git a/examples/by_feature/gradient_accumulation.py b/examples/by_feature/gradient_accumulation.py
new file mode 100644
index 000000000..bb8aba326
--- /dev/null
+++ b/examples/by_feature/gradient_accumulation.py
@@ -0,0 +1,218 @@
+# coding=utf-8
+# Copyright 2021 The HuggingFace Inc. 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 torch
+from torch.utils.data import DataLoader
+
+import evaluate
+from accelerate import Accelerator, DistributedType
+from datasets import load_dataset
+from transformers import (
+    AdamW,
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    get_linear_schedule_with_warmup,
+    set_seed,
+)
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate
+# and perform gradient accumulation
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+def get_dataloaders(accelerator: Accelerator, batch_size: int = 16):
+    """
+    Creates a set of `DataLoader`s for the `glue` dataset,
+    using "bert-base-cased" as the tokenizer.
+
+    Args:
+        accelerator (`Accelerator`):
+            An `Accelerator` object
+        batch_size (`int`, *optional*):
+            The batch size for the train and validation DataLoaders.
+    """
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    datasets = load_dataset("glue", "mrpc")
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["idx", "sentence1", "sentence2"],
+    )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    return train_dataloader, eval_dataloader
+
+
+# For testing only
+if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+    from accelerate.test_utils.training import mocked_dataloaders
+
+    get_dataloaders = mocked_dataloaders  # noqa: F811
+
+
+def training_function(config, args):
+    # Initialize accelerator
+    accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    correct_bias = config["correct_bias"]
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+    # New Code #
+    gradient_accumulation_steps = int(args.gradient_accumulation_steps)
+
+    metric = evaluate.load("glue", "mrpc")
+
+    set_seed(seed)
+    train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
+
+    # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
+    # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
+    # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
+    model = model.to(accelerator.device)
+
+    # Instantiate optimizer
+    optimizer = AdamW(params=model.parameters(), lr=lr, correct_bias=correct_bias)
+
+    # Instantiate scheduler
+    lr_scheduler = get_linear_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=100,
+        num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,
+    )
+
+    # Prepare everything
+    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
+    # prepare method.
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    # Now we train the model
+    for epoch in range(num_epochs):
+        model.train()
+        for step, batch in enumerate(train_dataloader):
+            # We could avoid this line since we set the accelerator with `device_placement=True`.
+            batch.to(accelerator.device)
+            # New code #
+            # We use the new `no_sync` context manager to prevent gradient averaging
+            # until we want to at the proper step if we happen to be in a distributed setup
+            # otherwise it does nothing
+            if step % gradient_accumulation_steps == 0:
+                # Accumulate gradients locally
+                with accelerator.no_sync(model):
+                    output = model(**batch)
+                    loss = output.loss / gradient_accumulation_steps
+                    accelerator.backward(loss)
+            else:
+                # Sync gradients and step
+                output = model(**batch)
+                loss = output.loss / gradient_accumulation_steps
+                accelerator.backward(loss)
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+
+        model.eval()
+        for step, batch in enumerate(eval_dataloader):
+            # We could avoid this line since we set the accelerator with `device_placement=True`.
+            batch.to(accelerator.device)
+            with torch.no_grad():
+                outputs = model(**batch)
+            predictions = outputs.logits.argmax(dim=-1)
+            predictions, references = accelerator.gather((predictions, batch["labels"]))
+            metric.add_batch(
+                predictions=predictions,
+                references=references,
+            )
+
+        eval_metric = metric.compute()
+        # Use accelerator.print to print only on the main process.
+        accelerator.print(f"epoch {epoch}:", eval_metric)
+
+
+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.",
+    )
+    # New Code #
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help="The number of minibatches to be ran before gradients are accumulated.",
+    )
+    parser.add_argument("--cpu", action="store_true", help="If passed, will train on the CPU.")
+    args = parser.parse_args()
+    config = {"lr": 2e-5, "num_epochs": 3, "correct_bias": True, "seed": 42, "batch_size": 16}
+    training_function(config, args)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/accelerate/test_utils/scripts/__init__.py b/src/accelerate/test_utils/scripts/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/test_examples.py b/tests/test_examples.py
index 6ad4a3ca8..4b298d253 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -34,6 +34,7 @@
 
 EXCLUDE_EXAMPLES = [
     "cross_validation.py",
+    "gradient_accumulation.py",
     "multi_process_metrics.py",
     "memory.py",
     "fsdp_with_peak_mem_tracking.py",
@@ -215,3 +216,7 @@ def test_tracking(self):
             """.split()
             _ = subprocess.run(self._launch_args + testargs, stdout=subprocess.PIPE)
             self.assertTrue(os.path.exists(os.path.join(tmpdir, "tracking")))
+
+    def test_gradient_accumulation(self):
+        testargs = ["examples/by_feature/gradient_accumulation.py"]
+        _ = subprocess.run(self._launch_args + testargs, stdout=subprocess.PIPE)

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/431
diff --git a/src/accelerate/test_utils/scripts/__init__.py b/src/accelerate/test_utils/scripts/__init__.py
new file mode 100644
index 000000000..e69de29bb

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/430
diff --git a/src/accelerate/test_utils/scripts/__init__.py b/src/accelerate/test_utils/scripts/__init__.py
new file mode 100644
index 000000000..e69de29bb

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/429
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 9205f0854..8801011ec 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import contextlib
 import gc
 import math
 import os
@@ -355,6 +356,25 @@ def _goes_first(self, is_main):
         if is_main:
             self.wait_for_everyone()
 
+    @contextmanager
+    def no_sync(self, model):
+        """
+        A context manager to disable gradient synchronizations across DDP processes by calling
+        `torch.nn.parallel.DistributedDataParallel.no_sync`.
+
+        If `model` is not in DDP, this context manager does nothing
+
+        Args:
+            model (`torch.nn.Module`):
+                PyTorch Module that was prepared with `Accelerator.prepare`
+        """
+        context = contextlib.nullcontext
+        if self.num_processes > 1:
+            context = getattr(model, "no_sync", context)
+
+        with context():
+            yield
+
     def print(self, *args, **kwargs):
         """
         Use in replacement of `print()` to only print once per server.
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 93b04d6ab..2b59d39f6 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -31,10 +31,10 @@
     DistributedType,
     PrecisionType,
     PrepareForLaunch,
+    get_launch_prefix,
     is_deepspeed_available,
     is_sagemaker_available,
 )
-from accelerate.utils.versions import is_torch_version
 
 
 def launch_command_parser(subparsers=None):
@@ -251,12 +251,7 @@ def simple_launcher(args):
 
 
 def multi_gpu_launcher(args):
-    if is_torch_version(">=", "1.10.0"):
-        cmd = ["torchrun"]
-    elif is_torch_version(">=", "1.9.0"):
-        cmd = [sys.executable, "-m", "torch.distributed.run"]
-    else:
-        cmd = [sys.executable, "-m", "torch.distributed.launch", "--use_env"]
+    cmd = get_launch_prefix()
     if args.num_machines > 1:
         cmd.extend(
             [
diff --git a/src/accelerate/test_utils/__init__.py b/src/accelerate/test_utils/__init__.py
index e8dc99c8c..2d8527be2 100644
--- a/src/accelerate/test_utils/__init__.py
+++ b/src/accelerate/test_utils/__init__.py
@@ -8,7 +8,11 @@
     require_cpu,
     require_cuda,
     require_multi_gpu,
+    require_single_gpu,
     require_tpu,
     slow,
 )
 from .training import RegressionDataset, RegressionModel
+
+
+from .scripts import test_script, test_sync  # isort:skip
diff --git a/src/accelerate/test_utils/test_script.py b/src/accelerate/test_utils/scripts/test_script.py
similarity index 100%
rename from src/accelerate/test_utils/test_script.py
rename to src/accelerate/test_utils/scripts/test_script.py
diff --git a/src/accelerate/test_utils/scripts/test_sync.py b/src/accelerate/test_utils/scripts/test_sync.py
new file mode 100644
index 000000000..2f5cc3722
--- /dev/null
+++ b/src/accelerate/test_utils/scripts/test_sync.py
@@ -0,0 +1,130 @@
+# 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.
+
+from copy import deepcopy
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+
+from accelerate import Accelerator
+from accelerate.test_utils import RegressionDataset, RegressionModel
+from accelerate.utils import DistributedType, set_seed
+
+
+def step_model(model, input, target, accelerator):
+    model.train()
+    output = model(input)
+    loss = F.mse_loss(output, target.to(output.device))
+    accelerator.backward(loss)
+
+
+def get_training_setup(accelerator):
+    "Returns everything needed to perform basic training"
+    set_seed(42)
+    model = RegressionModel()
+    model.to(accelerator.device)
+    dset = RegressionDataset()
+    dataloader = DataLoader(dset, batch_size=16)
+    # Make a copy of `model`
+    ddp_model, dataloader = accelerator.prepare(deepcopy(model), dataloader)
+    # Use a single batch for all of the tests
+    ddp_input, ddp_target = next(iter(dataloader)).values()
+    return model, ddp_model, ddp_input, ddp_target
+
+
+def test_noop_sync(accelerator):
+    # Test when on a single CPU or GPU that the context manager does nothing
+    model, ddp_model, ddp_input, ddp_target = get_training_setup(accelerator)
+    for iteration in range(3):
+        # Gather the distributed inputs and targs for the base model
+        input, target = accelerator.gather((ddp_input, ddp_target))
+        input, target = input.to(accelerator.device), target.to(accelerator.device)
+        # Perform our initial ground truth step in non "DDP"
+        step_model(model, input, target, accelerator)
+        # Do "gradient accumulation" (noop)
+        if iteration % 2 == 0:
+            # Accumulate grads locally
+            with accelerator.no_sync(ddp_model):
+                step_model(ddp_model, ddp_input, ddp_target, accelerator)
+        else:
+            # Sync grads
+            step_model(ddp_model, ddp_input, ddp_target, accelerator)
+
+        # Since `no_sync` is a noop, `ddp_model` and `model` grads should always be in sync
+        for param, ddp_param in zip(model.parameters(), ddp_model.parameters()):
+            if not param.requires_grad:
+                continue
+            assert torch.allclose(
+                param.grad, ddp_param.grad
+            ), f"Gradients not in sync when they should be:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})"
+
+        # Shuffle ddp_input on each iteration
+        torch.manual_seed(1337 + iteration)
+        ddp_input = ddp_input[torch.randperm(16)]
+
+
+def test_distributed_sync(accelerator):
+    # Test on distributed setup that context manager behaves properly
+    model, ddp_model, ddp_input, ddp_target = get_training_setup(accelerator)
+    for iteration in range(3):
+        # Gather the distributed inputs and targs for the base model
+        input, target = accelerator.gather((ddp_input, ddp_target))
+        input, target = input.to(accelerator.device), target.to(accelerator.device)
+        # Perform our initial ground truth step in non "DDP"
+        step_model(model, input, target, accelerator)
+        # Do "gradient accumulation" (noop)
+        if iteration % 2 == 0:
+            # Accumulate grads locally
+            with accelerator.no_sync(ddp_model):
+                step_model(ddp_model, ddp_input, ddp_target, accelerator)
+        else:
+            # Sync grads
+            step_model(ddp_model, ddp_input, ddp_target, accelerator)
+
+        # DDP model and model should only be in sync when not (iteration % 2 == 0)
+        for param, ddp_param in zip(model.parameters(), ddp_model.parameters()):
+            if not param.requires_grad:
+                continue
+            if iteration % 2 == 0:
+                # Grads should not be in sync
+                assert (
+                    torch.allclose(param.grad, ddp_param.grad) is False
+                ), f"Gradients in sync when they should not be:\nModel grad ({param.grad}) == DDP grad ({ddp_param.grad})"
+            else:
+                # Grads should be in sync
+                assert (
+                    torch.allclose(param.grad, ddp_param.grad) is True
+                ), f"Gradients not in sync when they should be:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})"
+
+        # Shuffle ddp_input on each iteration
+        torch.manual_seed(1337 + iteration)
+        ddp_input = ddp_input[torch.randperm(16)]
+
+
+def main():
+    accelerator = Accelerator()
+    state = accelerator.state
+    if state.distributed_type == DistributedType.NO:
+        if state.local_process_index == 0:
+            print("**NOOP `no_sync` gradient accumulation**")
+        test_noop_sync(accelerator)
+    if state.distributed_type in (DistributedType.MULTI_GPU, DistributedType.MULTI_CPU):
+        if state.local_process_index == 0:
+            print("**Distributed `no_sync` gradient accumulation**")
+        test_distributed_sync(accelerator)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/accelerate/test_utils/testing.py b/src/accelerate/test_utils/testing.py
index 89d104a49..017f66d06 100644
--- a/src/accelerate/test_utils/testing.py
+++ b/src/accelerate/test_utils/testing.py
@@ -84,6 +84,14 @@ def require_tpu(test_case):
     return unittest.skipUnless(is_tpu_available(), "test requires TPU")(test_case)
 
 
+def require_single_gpu(test_case):
+    """
+    Decorator marking a test that requires CUDA on a single GPU. These tests are skipped when there are no GPU
+    available or number of GPUs is more than one.
+    """
+    return unittest.skipUnless(torch.cuda.device_count() == 1, "test requires a GPU")(test_case)
+
+
 def require_multi_gpu(test_case):
     """
     Decorator marking a test that requires a multi-GPU setup. These tests are skipped on a machine without multiple
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
index b66b6fd1d..a3e193915 100644
--- a/src/accelerate/utils/__init__.py
+++ b/src/accelerate/utils/__init__.py
@@ -85,7 +85,7 @@
         DummyScheduler,
     )
 
-from .launch import PrepareForLaunch
+from .launch import PrepareForLaunch, get_launch_prefix
 from .memory import find_executable_batch_size
 from .other import (
     extract_model_from_parallel,
diff --git a/src/accelerate/utils/launch.py b/src/accelerate/utils/launch.py
index 8ff624250..d4f0042fb 100644
--- a/src/accelerate/utils/launch.py
+++ b/src/accelerate/utils/launch.py
@@ -13,12 +13,28 @@
 # limitations under the License.
 
 import os
+import sys
 
 import torch
 
+from ..utils import is_torch_version
 from .dataclasses import DistributedType
 
 
+def get_launch_prefix():
+    """
+    Grabs the correct launcher for starting a distributed command, such as either `torchrun`, `python -m
+    torch.distributed.run`, etc
+    """
+    if is_torch_version(">=", "1.10.0"):
+        cmd = ["torchrun"]
+    elif is_torch_version(">=", "1.9.0"):
+        cmd = [sys.executable, "-m", "torch.distributed.run"]
+    else:
+        cmd = [sys.executable, "-m", "torch.distributed.launch", "--use_env"]
+    return cmd
+
+
 class PrepareForLaunch:
     """
     Prepare a function that will launched in a distributed setup.
diff --git a/tests/test_grad_sync.py b/tests/test_grad_sync.py
new file mode 100644
index 000000000..c066bd41a
--- /dev/null
+++ b/tests/test_grad_sync.py
@@ -0,0 +1,56 @@
+# 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 inspect
+import os
+import unittest
+
+import torch
+
+import accelerate
+from accelerate import debug_launcher
+from accelerate.test_utils import (
+    execute_subprocess_async,
+    require_cpu,
+    require_multi_gpu,
+    require_single_gpu,
+    test_sync,
+)
+from accelerate.utils import get_launch_prefix, patch_environment
+
+
+class SyncScheduler(unittest.TestCase):
+    def setUp(self):
+        mod_file = inspect.getfile(accelerate.test_utils)
+        self.test_file_path = os.path.sep.join(mod_file.split(os.path.sep)[:-1] + ["scripts", "test_sync.py"])
+
+    @require_cpu
+    def test_gradient_sync_single_cpu_noop(self):
+        debug_launcher(test_sync.main)
+        debug_launcher(test_sync.main, num_processes=1)
+
+    @require_cpu
+    def test_gradient_sync_multi_cpu(self):
+        debug_launcher(test_sync.main)
+
+    @require_single_gpu
+    def test_gradient_sync_single_gpu(self):
+        debug_launcher(test_sync.main, num_processes=1)
+
+    @require_multi_gpu
+    def test_gradient_sync_multi_gpu(self):
+        print(f"Found {torch.cuda.device_count()} devices.")
+        cmd = get_launch_prefix() + [f"--nproc_per_node={torch.cuda.device_count()}", self.test_file_path]
+        with patch_environment(omp_num_threads=1):
+            execute_subprocess_async(cmd, env=os.environ.copy())
diff --git a/tests/test_multigpu.py b/tests/test_multigpu.py
index 7ffbd94f8..2c0403ed5 100644
--- a/tests/test_multigpu.py
+++ b/tests/test_multigpu.py
@@ -14,7 +14,6 @@
 
 import inspect
 import os
-import sys
 import unittest
 
 import torch
@@ -22,35 +21,26 @@
 import accelerate
 from accelerate import Accelerator
 from accelerate.test_utils import execute_subprocess_async, require_multi_gpu
+from accelerate.utils import get_launch_prefix, patch_environment
 
 
 class MultiGPUTester(unittest.TestCase):
     def setUp(self):
         mod_file = inspect.getfile(accelerate.test_utils)
-        self.test_file_path = os.path.sep.join(mod_file.split(os.path.sep)[:-1] + ["test_script.py"])
+        self.test_file_path = os.path.sep.join(mod_file.split(os.path.sep)[:-1] + ["scripts", "test_script.py"])
 
     @require_multi_gpu
     def test_multi_gpu(self):
         print(f"Found {torch.cuda.device_count()} devices.")
-        distributed_args = f"""
-            -m torch.distributed.launch
-            --nproc_per_node={torch.cuda.device_count()}
-            --use_env
-            {self.test_file_path}
-        """.split()
-        cmd = [sys.executable] + distributed_args
-        execute_subprocess_async(cmd, env=os.environ.copy())
+        cmd = get_launch_prefix() + [self.test_file_path]
+        with patch_environment(omp_num_threads=1):
+            execute_subprocess_async(cmd, env=os.environ.copy())
 
     @require_multi_gpu
     def test_pad_across_processes(self):
-        distributed_args = f"""
-            -m torch.distributed.launch
-            --nproc_per_node={torch.cuda.device_count()}
-            --use_env
-            {inspect.getfile(self.__class__)}
-        """.split()
-        cmd = [sys.executable] + distributed_args
-        execute_subprocess_async(cmd, env=os.environ.copy())
+        cmd = get_launch_prefix() + [inspect.getfile(self.__class__)]
+        with patch_environment(omp_num_threads=1):
+            execute_subprocess_async(cmd, env=os.environ.copy())
 
 
 if __name__ == "__main__":

I would double check we are in a distributed state on top of this.
@@ -355,6 +356,22 @@ def _goes_first(self, is_main):
         if is_main:
             self.wait_for_everyone()
 
+    @contextmanager
+    def no_sync(self, model):
+        """
+        A context manager to disable gradient synchronizations across DDP processes by calling
+        `torch.nn.parallel.DistributedDataParallel.no_sync`.
+
+        If `model` is not in DDP, this context manager does nothing
+
+        Args:
+            model (`torch.nn.Module`):
+                PyTorch Module that was prepared with `Accelerator.prepare`
+        """
+        context = getattr(model, "no_sync", contextlib.nullcontext)
This is not a core feature, I don't see why it should go in `accelerate test`, which is just a quick test that everything is going well. It should be tested on its own in the tests that's all.
@@ -312,6 +315,55 @@ def training_check():
     assert torch.allclose(old_model.b, model.b), "Did not obtain the same model on CPU or distributed training."
 
 
+def sync_test():
Better variable names would help understand what this is testing. Also, those should be equal at the proper `iteration`.
@@ -0,0 +1,63 @@
+# 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 unittest
+from copy import deepcopy
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+
+from accelerate import Accelerator
+from accelerate.test_utils import RegressionDataset, RegressionModel, require_cpu
+from accelerate.utils import set_seed
+
+
+def step_model(model, input, target, accelerator):
+    model.train()
+    output = model(input)
+    loss = F.mse_loss(output, target.to(output.device))
+    accelerator.backward(loss)
+
+
+@require_cpu
+class SyncTest(unittest.TestCase):
+    def test_noop_wrapper(self):
+        accelerator = Accelerator()
+        device = accelerator.device
+        set_seed(42)
+        model = RegressionModel()
+        dset = RegressionDataset()
+        dl = DataLoader(dset, batch_size=16)
+        ddp_model, dl = accelerator.prepare(deepcopy(model), dl)
+        model.to(device)
+        ddp_input, ddp_target = next(iter(dl)).values()
+
+        for iteration in range(2):
+            input, target = accelerator.gather((ddp_input, ddp_target))
+            input = input.to(accelerator.device)
+            target = target.to(accelerator.device)
+            step_model(model, input, target, accelerator)
+            if iteration % 2 == 0:
+                # Accumulate grads locally
+                with accelerator.no_sync(ddp_model):
+                    step_model(ddp_model, ddp_input, ddp_target, accelerator)
+            else:
+                # Sync grads
+                step_model(ddp_model, ddp_input, ddp_target, accelerator)
+
+            for i, j in zip(model.parameters(), ddp_model.parameters()):
+                if not i.requires_grad:
+                    continue
+                assert torch.allclose(i.grad, j.grad), f"{i.grad} != {j.grad}"
_The documentation is not available anymore as the PR was closed or merged._
@sgugger good for a rereview, slow test logs can be seen here passing: https://github.com/huggingface/accelerate/runs/6794400970?check_suite_focus=true
https://api.github.com/repos/huggingface/accelerate/pulls/428
diff --git a/tests/deepspeed/test_deepspeed.py b/tests/deepspeed/test_deepspeed.py
index dbe64b5c7..9ece7b0f7 100644
--- a/tests/deepspeed/test_deepspeed.py
+++ b/tests/deepspeed/test_deepspeed.py
@@ -158,7 +158,6 @@ def test_deepspeed_plugin(self, stage):
             self.assertFalse(deepspeed_plugin.zero3_init_flag)
         elif stage == ZERO3:
             self.assertTrue(deepspeed_plugin.zero3_init_flag)
-            deepspeed_plugin.deepspeed_config = None
 
         # Test `gradient_accumulation_steps` is set to 1 if unavailable in config file
         with tempfile.TemporaryDirectory() as dirpath:
@@ -232,6 +231,10 @@ def test_deepspeed_plugin(self, stage):
 
     @parameterized.expand([FP16, BF16], name_func=parameterized_custom_name_func)
     def test_accelerate_state_deepspeed(self, dtype):
+        state = AcceleratorState(_from_accelerator=True)
+        if state.initialized:
+            state.initialized = False
+
         deepspeed_plugin = DeepSpeedPlugin(
             gradient_accumulation_steps=1,
             gradient_clipping=1.0,

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/427
diff --git a/.github/workflows/check_dependencies.yml b/.github/workflows/check_dependencies.yml
index 65afdd6c9..c4818ba85 100644
--- a/.github/workflows/check_dependencies.yml
+++ b/.github/workflows/check_dependencies.yml
@@ -36,4 +36,5 @@ jobs:
   build-docker-containers:
     needs: check-for-setup
     if: (github.event_name == 'push') && (needs.check-for-setup.outputs.changed == '1')
-    uses: ./.github/workflows/build-docker-images.yml
\ No newline at end of file
+    uses: ./.github/workflows/build-docker-images.yml
+    secrets: inherit
\ No newline at end of file

_The documentation is not available anymore as the PR was closed or merged._
@muellerzr Thank you pining for me. I read (part of) that page previously, but not in detail.

So it looks like to me that the called (reusable) workflow can't access `secrets`, and need to get it from the caller.
(Only the workflows triggered by events like `push/pull_request/schedule`, etc.  can access `secrets`)
Yup exactly. And we can do so by using `secrets`, by specifying inherit it can have all of the secrets the caller has. Keeping it to be what's needed only isn't a bad idea either, rather than giving them everything. But for this one I decided all wouldn't be a terrible idea. 
https://api.github.com/repos/huggingface/accelerate/pulls/426
diff --git a/setup.py b/setup.py
index 45c4ae912..eb0f50ab8 100644
--- a/setup.py
+++ b/setup.py
@@ -14,7 +14,7 @@
 
 from setuptools import setup
 from setuptools import find_packages
-
+# A comment
 extras = {}
 extras["quality"] = ["black ~= 22.0", "isort >= 5.5.4", "flake8 >= 3.8.3"]
 extras["docs"] = []

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/425
diff --git a/.github/workflows/build-docker-images.yml b/.github/workflows/build-docker-images.yml
index de8542359..c207f5515 100644
--- a/.github/workflows/build-docker-images.yml
+++ b/.github/workflows/build-docker-images.yml
@@ -2,6 +2,7 @@ name: Build Docker images (scheduled)
 
 on:
   workflow_dispatch:
+  workflow_call:
   schedule:
     - cron: "0 1 * * *"
 
diff --git a/.github/workflows/check_dependencies.yml b/.github/workflows/check_dependencies.yml
new file mode 100644
index 000000000..65afdd6c9
--- /dev/null
+++ b/.github/workflows/check_dependencies.yml
@@ -0,0 +1,39 @@
+name: Check for dependency modification
+
+on:
+  push:
+    branches:
+      - main
+  workflow_dispatch:
+
+env:
+  GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+
+jobs:
+  check-for-setup:
+    runs-on: ubuntu-latest
+    name: Check if setup was changed
+    outputs:
+      changed: ${{ steps.was_changed.outputs.changed }}
+    steps:
+      - uses: actions/checkout@v3
+        with: 
+          fetch-depth: "2"
+      
+      - name: Get changed files
+        id: changed-files
+        uses: tj-actions/changed-files@v22.2
+      
+      - name: Was setup changed 
+        id: was_changed
+        run: |
+          for file in ${{ steps.changed-files.outputs.all_changed_files }}; do
+            if [ `basename "${file}"` = "setup.py" ]; then
+              echo ::set-output name=changed::"1"
+            fi
+          done
+          
+  build-docker-containers:
+    needs: check-for-setup
+    if: (github.event_name == 'push') && (needs.check-for-setup.outputs.changed == '1')
+    uses: ./.github/workflows/build-docker-images.yml
\ No newline at end of file
diff --git a/.github/workflows/on-merge.yml b/.github/workflows/on-merge.yml
index 4d603bf8e..f8b9ccddb 100644
--- a/.github/workflows/on-merge.yml
+++ b/.github/workflows/on-merge.yml
@@ -1,10 +1,11 @@
 name: Self-hosted runner (push to "main")
 
 on:
+  workflow_run:
+    workflows: ["Check for dependency modification"]
+    branches: ["main"]
+    types: [completed]
   workflow_dispatch:
-  push:
-    branches:
-      - "main"
 
 env:
   TESTING_MOCKED_DATALOADERS: "1"

This is not set anywhere, so I don't think this will work (also pashing the github token on a PR is not a trivial thing as you don't get secrets on pull requests. I'd just remove this whole writing comments thingy.
@@ -0,0 +1,97 @@
+# Copyright 2022 The HuggingFace Team, the AllenNLP library authors. 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.
+"""
+Script to check if the setup.py was modified during a Pull Request,
+and if so will write a comment stating that Docker Images will be 
+rebuilt upon a merge. 
+
+Environment variables available:
+- PR_NUMBER: The number of the currently open pull request
+- REPO: The full name of the repository (such as 'huggingface/accelerate')
+"""
+
+import os
+import re
+from github import Github, PullRequest
+
+def get_setup_diff(pull_request:PullRequest):
+    """Checks whether `setup.py` was changed during this pull request, and will 
+    return the diff if so
+
+    Args:
+        pull_request (`PullRequest`):
+            A pull request returned from `Github.repo.get_pull_request`
+    """
+    files = pull_request.get_files()
+    for f in files:
+        if f.filename == "setup.py":
+            return f'''```diff\n{f.patch}\n```'''
+    return None
+
+def does_comment_exist(pull_request:PullRequest):
+    """Checks whether the bot has already commented on this pull request
+
+    Args:
+        pull_request (`PullRequest`):
+            A pull request returned from `Github.repo.get_pull_request`
+    """
+    comments = pull_request.get_issue_comments()
+    for c in comments:
+        if c.user.login == "github-actions[bot]" and 'This PR modifies `setup.py`.' in c.body:
+            return True
+    return False
+
+def write_comment(pull_request:PullRequest, diff:str):
+    """Writes a comment stating that the pr modified setup.py, and that new Docker images will be built
+
+    Args:
+        pull_request (`PullRequest`):
+            A pull request returned from `Github.repo.get_pull_request`
+        diff (`str`):
+            The diff of the modified setup.py
+    """
+    s = f'This PR modifies `setup.py`. New latest Docker images will be built and deployed before running tests once this has been merged:\n\n{diff}'
+    pull_request.create_issue_comment(s)
+
+def update_diff(pull_request:PullRequest, diff:str):
+    """Updates the diff of the setup.py in the existing comment
+
+    Args:
+        pull_request (`PullRequest`):
+            A pull request returned from `Github.repo.get_pull_request`
+        diff (`str`):
+            The diff of the modified setup.py
+    """
+    comments = pull_request.get_issue_comments()
+    for c in comments:
+        if c.user.login == "github-actions[bot]" and 'This PR modifies `setup.py`.' in c.body:
+            comment = c
+            break
+    original_diff = re.search(r'```([^`]*)```', comment.body).group(0)
+    new_body = comment.body.replace(original_diff, diff)
+    comment.edit(new_body)
+
+def main():
+    g = Github(os.environ["GITHUB_TOKEN"])
I'd remove this job personally.
@@ -0,0 +1,68 @@
+name: Check for dependency modification
+
+on:
+  pull_request:
+    types: [opened, synchronize, reopened]
+  push:
+    branches:
+      - main
+  workflow_dispatch:
+
+env:
+  GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+
+jobs:
+  check-for-setup:
+    runs-on: ubuntu-latest
+    name: Check if setup was changed
+    outputs:
+      changed: ${{ steps.was_changed.outputs.changed }}
+    steps:
+      - uses: actions/checkout@v3
+        with: 
+          fetch-depth: "2"
+      
+      - name: Get changed files
+        id: changed-files
+        uses: tj-actions/changed-files@v22.2
+      
+      - name: Was setup changed 
+        id: was_changed
+        run: |
+          for file in ${{ steps.changed-files.outputs.all_changed_files }}; do
+            if [ `basename "${file}"` = "setup.py" ]; then
+              echo ::set-output name=changed::"1"
+            fi
+          done
+
+  write-message:
Is there any reason that we should run this workflow `on pull_request`?
@@ -0,0 +1,41 @@
+name: Check for dependency modification
+
+on:
+  pull_request:
+    types: [opened, synchronize, reopened]
Very nice - we should build the image only for the push event (to `main`).
@@ -0,0 +1,41 @@
+name: Check for dependency modification
+
+on:
+  pull_request:
+    types: [opened, synchronize, reopened]
+  push:
+    branches:
+      - main
+  workflow_dispatch:
+
+env:
+  GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+
+jobs:
+  check-for-setup:
+    runs-on: ubuntu-latest
+    name: Check if setup was changed
+    outputs:
+      changed: ${{ steps.was_changed.outputs.changed }}
+    steps:
+      - uses: actions/checkout@v3
+        with: 
+          fetch-depth: "2"
+      
+      - name: Get changed files
+        id: changed-files
+        uses: tj-actions/changed-files@v22.2
+      
+      - name: Was setup changed 
+        id: was_changed
+        run: |
+          for file in ${{ steps.changed-files.outputs.all_changed_files }}; do
+            if [ `basename "${file}"` = "setup.py" ]; then
+              echo ::set-output name=changed::"1"
+            fi
+          done
+          
+  build-docker-containers:
+    needs: check-for-setup
+    if: (github.event_name == 'push') && (needs.check-for-setup.outputs.changed == '1')
I was thinking that `completed` means the CI will be run only if the workflow `Check for dependency modification` is successful.

After reading this [doc](https://docs.github.com/en/actions/using-workflows/events-that-trigger-workflows#running-a-workflow-based-on-the-conclusion-of-another-workflow), it is not the case. And this is good, as we want to run the push CI anyway (even if something wrong in the setup check).


@@ -1,10 +1,11 @@
 name: Self-hosted runner (push to "main")
 
 on:
+  workflow_run:
+    workflows: ["Check for dependency modification"]
+    branches: ["main"]
+    types: [completed]
Yup exactly
@@ -1,10 +1,11 @@
 name: Self-hosted runner (push to "main")
 
 on:
+  workflow_run:
+    workflows: ["Check for dependency modification"]
+    branches: ["main"]
+    types: [completed]
This was due to the overengineering part. We definitely could just trigger it on a push to main. 
@@ -0,0 +1,41 @@
+name: Check for dependency modification
+
+on:
+  pull_request:
+    types: [opened, synchronize, reopened]
Yeah, I meanly mean `push` in this remark, and you are doing it as what I would expect. (`to main` is to make it more clear only 😄 )
@@ -0,0 +1,41 @@
+name: Check for dependency modification
+
+on:
+  pull_request:
+    types: [opened, synchronize, reopened]
+  push:
+    branches:
+      - main
+  workflow_dispatch:
+
+env:
+  GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+
+jobs:
+  check-for-setup:
+    runs-on: ubuntu-latest
+    name: Check if setup was changed
+    outputs:
+      changed: ${{ steps.was_changed.outputs.changed }}
+    steps:
+      - uses: actions/checkout@v3
+        with: 
+          fetch-depth: "2"
+      
+      - name: Get changed files
+        id: changed-files
+        uses: tj-actions/changed-files@v22.2
+      
+      - name: Was setup changed 
+        id: was_changed
+        run: |
+          for file in ${{ steps.changed-files.outputs.all_changed_files }}; do
+            if [ `basename "${file}"` = "setup.py" ]; then
+              echo ::set-output name=changed::"1"
+            fi
+          done
+          
+  build-docker-containers:
+    needs: check-for-setup
+    if: (github.event_name == 'push') && (needs.check-for-setup.outputs.changed == '1')
_The documentation is not available anymore as the PR was closed or merged._
@ydshieh you can see my almost 250 debug logs here: https://github.com/muellerzr/accelerate/actions

**Very** *very* comfortable with this working
https://api.github.com/repos/huggingface/accelerate/pulls/424
diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml
index 1764e2381..d52cf5e1e 100644
--- a/.github/workflows/nightly.yml
+++ b/.github/workflows/nightly.yml
@@ -5,12 +5,14 @@ on:
   schedule:
     - cron: "0 2 * * *"
 
+env:
+  RUN_SLOW: "yes"
+
 jobs:
   run_all_tests_single_gpu:
     runs-on: [self-hosted, docker-gpu, multi-gpu]
     env:
       CUDA_VISIBLE_DEVICES: "0"
-      ALL_SLOW: "yes"
     container:
       image: huggingface/accelerate-gpu:latest
       options: --gpus all --shm-size "16gb"
@@ -38,7 +40,6 @@ jobs:
   run_all_tests_multi_gpu:
     runs-on: [self-hosted, docker-gpu, multi-gpu]
     env:
-      ALL_SLOW: "yes"
       CUDA_VISIBLE_DEVICES: "0,1"
     container:
       image: huggingface/accelerate-gpu:latest

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/421
diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml
index d50aab329..1764e2381 100644
--- a/.github/workflows/nightly.yml
+++ b/.github/workflows/nightly.yml
@@ -23,12 +23,9 @@ jobs:
         run: |
           source activate accelerate
           git config --global --add safe.directory '*'
-          git fetch && git checkout ${{ github.sha }}
-          pip install tensorflow -U 
+          git fetch && git checkout ${{ github.sha }} 
           pip install -e . --no-deps
 
-# Note: tensorflow upgrade is needed until dropped py 3.6 support
-
       - name: Run test on GPUs
         run: |
           source activate accelerate
@@ -56,11 +53,8 @@ jobs:
           source activate accelerate
           git config --global --add safe.directory '*'
           git fetch && git checkout ${{ github.sha }}
-          pip install tensorflow -U
           pip install -e . --no-deps
 
-        # Note: tensorflow upgrade is needed until dropped py 3.6 support
-
       - name: Run test on GPUs
         run: |
           source activate accelerate
diff --git a/.github/workflows/on-merge.yml b/.github/workflows/on-merge.yml
index 8df455563..4d603bf8e 100644
--- a/.github/workflows/on-merge.yml
+++ b/.github/workflows/on-merge.yml
@@ -6,6 +6,9 @@ on:
     branches:
       - "main"
 
+env:
+  TESTING_MOCKED_DATALOADERS: "1"
+
 jobs:
   run_all_tests_single_gpu:
     runs-on: [self-hosted, docker-gpu, multi-gpu]
@@ -24,10 +27,7 @@ jobs:
           source activate accelerate
           git config --global --add safe.directory '*'
           git fetch && git checkout ${{ github.sha }}
-          pip install tensorflow -U 
-          pip install -e . -U
-
-# Note: tensorflow upgrade is needed until dropped py 3.6 support
+          pip install -e .[test,test_trackers]
 
       - name: Run test on GPUs
         run: |
@@ -53,10 +53,7 @@ jobs:
           source activate accelerate
           git config --global --add safe.directory '*'
           git fetch && git checkout ${{ github.sha }}
-          pip install tensorflow -U
-          pip install -e . -U
-
-        # Note: tensorflow upgrade is needed until dropped py 3.6 support
+          pip install -e .[test,test_trackers]
 
       - name: Run test on GPUs
         run: |
diff --git a/.github/workflows/quality.yml b/.github/workflows/quality.yml
index 9b54a269f..5d4707b4e 100644
--- a/.github/workflows/quality.yml
+++ b/.github/workflows/quality.yml
@@ -7,10 +7,10 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v2
-    - name: Set up Python 3.6
-      uses: actions/setup-python@v2
+    - name: Set up Python 3.7
+      uses: actions/setup-python@v3
       with:
-        python-version: 3.6
+        python-version: 3.7
     - name: Install Python dependencies
       run: pip install -e .[quality]
     - name: Run Quality check
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 2bd407d85..d2de25a79 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -4,6 +4,7 @@ on: [pull_request]
 
 env:
   HF_HOME: ~/hf_cache
+  TESTING_MOCKED_DATALOADERS: "1"
 
 jobs:
   run-tests:
@@ -35,7 +36,7 @@ jobs:
     - name: Install the library
       run: |
         pip install --upgrade pip
-        pip install -e .[test,test_trackers,tensorboard]
+        pip install -e .[test,test_trackers]
     
     - name: Run Tests
       run: |
diff --git a/docker/accelerate-cpu/Dockerfile b/docker/accelerate-cpu/Dockerfile
index 1d99c4035..5d1286772 100644
--- a/docker/accelerate-cpu/Dockerfile
+++ b/docker/accelerate-cpu/Dockerfile
@@ -1,7 +1,7 @@
 # Builds CPU-only Docker image of PyTorch
 # Uses multi-staged approach to reduce size
 # Stage 1
-FROM python:3.6-slim as compile-image
+FROM python:3.7-slim as compile-image
 
 ARG DEBIAN_FRONTEND=noninteractive
 
@@ -25,7 +25,7 @@ RUN python3 -m pip install --no-cache-dir \
     --extra-index-url https://download.pytorch.org/whl/cpu
     
 # Stage 2
-FROM python:3.6-slim AS build-image
+FROM python:3.7-slim AS build-image
 COPY --from=compile-image /opt/venv /opt/venv
 RUN useradd -ms /bin/bash user
 USER user
diff --git a/docker/accelerate-gpu/Dockerfile b/docker/accelerate-gpu/Dockerfile
index f495d9063..746e65350 100644
--- a/docker/accelerate-gpu/Dockerfile
+++ b/docker/accelerate-gpu/Dockerfile
@@ -4,7 +4,7 @@
 # Use base conda image to reduce time
 FROM continuumio/miniconda3:latest AS compile-image
 # Specify py version
-ENV PYTHON_VERSION=3.6
+ENV PYTHON_VERSION=3.7.3
 # Install apt libs
 RUN apt-get update && \
     apt-get install -y curl git wget && \
diff --git a/setup.py b/setup.py
index 51a6c1a01..a28b9254c 100644
--- a/setup.py
+++ b/setup.py
@@ -29,8 +29,8 @@
     "scipy",
     "sklearn"
 ]
-extras["tensorboard"] = ["protobuf<=3.20.1", "tensorflow>=2.6.2", "tensorboard"]
-extras["test_trackers"] = extras["tensorboard"] + ["wandb", "comet-ml"]
+
+extras["test_trackers"] = ["wandb", "comet-ml", "tensorboard"]
 extras["dev"] = extras["quality"] + extras["test"]
 
 extras["sagemaker"] = [
@@ -57,7 +57,7 @@
             "accelerate-launch=accelerate.commands.launch:main",
         ]
     },
-    python_requires=">=3.6.0",
+    python_requires=">=3.7.0",
     install_requires=["numpy>=1.17", "packaging>=20.0", "pyyaml", "torch>=1.4.0"],
     extras_require=extras,
     classifiers=[
@@ -68,7 +68,6 @@
         "License :: OSI Approved :: Apache Software License",
         "Operating System :: OS Independent",
         "Programming Language :: Python :: 3",
-        "Programming Language :: Python :: 3.6",
         "Programming Language :: Python :: 3.7",
         "Topic :: Scientific/Engineering :: Artificial Intelligence",
     ],
diff --git a/src/accelerate/test_utils/testing.py b/src/accelerate/test_utils/testing.py
index 013dfad44..91dc0f7ad 100644
--- a/src/accelerate/test_utils/testing.py
+++ b/src/accelerate/test_utils/testing.py
@@ -26,7 +26,7 @@
 import torch
 
 from ..state import AcceleratorState
-from ..utils import gather, is_comet_ml_available, is_tensorflow_available, is_tpu_available, is_wandb_available
+from ..utils import gather, is_comet_ml_available, is_tensorboard_available, is_tpu_available, is_wandb_available
 
 
 def parse_flag_from_env(key, default=False):
@@ -85,12 +85,12 @@ def require_multi_gpu(test_case):
     return unittest.skipUnless(torch.cuda.device_count() > 1, "test requires multiple GPUs")(test_case)
 
 
-def require_tensorflow(test_case):
+def require_tensorboard(test_case):
     """
-    Decorator marking a test that requires TensorFlow installed. These tests are skipped when TensorFlow isn't
+    Decorator marking a test that requires tensorboard installed. These tests are skipped when tensorboard isn't
     installed
     """
-    return unittest.skipUnless(is_tensorflow_available(), "test requires TensorFlow")(test_case)
+    return unittest.skipUnless(is_tensorboard_available(), "test requires Tensorboard")(test_case)
 
 
 def require_wandb(test_case):
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
index 186f27bb7..a39368dd3 100644
--- a/src/accelerate/utils/__init__.py
+++ b/src/accelerate/utils/__init__.py
@@ -26,7 +26,6 @@
     is_deepspeed_available,
     is_sagemaker_available,
     is_tensorboard_available,
-    is_tensorflow_available,
     is_tpu_available,
     is_wandb_available,
 )
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index b492f2f9c..947508726 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -63,10 +63,6 @@ def is_deepspeed_available():
             return False
 
 
-def is_tensorflow_available():
-    return importlib.util.find_spec("tensorflow") is not None
-
-
 def is_tensorboard_available():
     return importlib.util.find_spec("tensorboard") is not None or importlib.util.find_spec("tensorboardX") is not None
 
diff --git a/tests/test_big_modeling.py b/tests/test_big_modeling.py
index be2238794..852b10c3b 100644
--- a/tests/test_big_modeling.py
+++ b/tests/test_big_modeling.py
@@ -94,14 +94,18 @@ def test_cpu_offload(self):
 
         cpu_offload(model, execution_device=device)
         output = model(x)
-        self.assertTrue(torch.allclose(expected, output.cpu()))
+        self.assertTrue(
+            torch.allclose(expected, output.cpu(), 1e-4, 1e-5), msg=f"Expected: {expected}\nActual: {output.cpu()}"
+        )
 
         # Clean up for next test.
         remove_hook_from_submodules(model)
 
         cpu_offload(model, execution_device=device, offload_buffers=True)
         output = model(x)
-        self.assertTrue(torch.allclose(expected, output.cpu()))
+        self.assertTrue(
+            torch.allclose(expected, output.cpu(), 1e-4, 1e-5), msg=f"Expected: {expected}\nActual: {output.cpu()}"
+        )
 
     @slow
     @require_cuda
@@ -127,7 +131,9 @@ def test_disk_offload(self):
         with TemporaryDirectory() as tmp_dir:
             disk_offload(model, tmp_dir, execution_device=device)
             output = model(x)
-            self.assertTrue(torch.allclose(expected, output.cpu()))
+            self.assertTrue(
+                torch.allclose(expected, output.cpu(), 1e-4, 1e-5), msg=f"Expected: {expected}\nActual: {output.cpu()}"
+            )
 
             # Clean up for next test.
             remove_hook_from_submodules(model)
@@ -135,7 +141,9 @@ def test_disk_offload(self):
         with TemporaryDirectory() as tmp_dir:
             disk_offload(model, tmp_dir, execution_device=device, offload_buffers=True)
             output = model(x)
-            self.assertTrue(torch.allclose(expected, output.cpu()))
+            self.assertTrue(
+                torch.allclose(expected, output.cpu(), 1e-4, 1e-5), msg=f"Expected: {expected}\nActual: {output.cpu()}"
+            )
 
     @slow
     @require_cuda
diff --git a/tests/test_tracking.py b/tests/test_tracking.py
index f02db8d21..a37ae7128 100644
--- a/tests/test_tracking.py
+++ b/tests/test_tracking.py
@@ -30,31 +30,22 @@
     MockingTestCase,
     TempDirTestCase,
     require_comet_ml,
-    require_tensorflow,
+    require_tensorboard,
     require_wandb,
 )
 from accelerate.tracking import CometMLTracker, GeneralTracker
-from accelerate.utils import is_comet_ml_available, is_tensorflow_available
+from accelerate.utils import is_comet_ml_available
 
 
 if is_comet_ml_available():
     from comet_ml import OfflineExperiment
 
-
-if is_tensorflow_available():
-    import tensorflow as tf
-    from tensorboard.plugins.hparams import plugin_data_pb2
-    from tensorflow.core.util import event_pb2
-    from tensorflow.python.summary.summary_iterator import summary_iterator
-
-
 logger = logging.getLogger(__name__)
 
 
+@require_tensorboard
 class TensorBoardTrackingTest(unittest.TestCase):
-    @require_tensorflow
     def test_init_trackers(self):
-        hps = None
         project_name = "test_project_with_config"
         with tempfile.TemporaryDirectory() as dirpath:
             accelerator = Accelerator(log_with="tensorboard", logging_dir=dirpath)
@@ -63,29 +54,9 @@ def test_init_trackers(self):
             accelerator.end_training()
             for child in Path(f"{dirpath}/{project_name}").glob("*/**"):
                 log = list(filter(lambda x: x.is_file(), child.iterdir()))[0]
-                # The config log is stored one layer deeper in the logged directory
-                # And names are randomly generated each time
-            si = summary_iterator(str(log))
-            # Pull HPS through careful parsing
-            for event in si:
-                for value in event.summary.value:
-                    proto_bytes = value.metadata.plugin_data.content
-                    plugin_data = plugin_data_pb2.HParamsPluginData.FromString(proto_bytes)
-                    if plugin_data.HasField("session_start_info"):
-                        hps = dict(plugin_data.session_start_info.hparams)
-
-        self.assertTrue(isinstance(hps, dict))
-        keys = list(hps.keys())
-        keys.sort()
-        self.assertEqual(keys, ["learning_rate", "num_iterations", "some_boolean", "some_string"])
-        self.assertEqual(hps["num_iterations"].number_value, 12)
-        self.assertEqual(hps["learning_rate"].number_value, 0.01)
-        self.assertEqual(hps["some_boolean"].bool_value, False)
-        self.assertEqual(hps["some_string"].string_value, "some_value")
+            self.assertNotEqual(str(log), "")
 
-    @require_tensorflow
     def test_log(self):
-        step = None
         project_name = "test_project_with_log"
         with tempfile.TemporaryDirectory() as dirpath:
             accelerator = Accelerator(log_with="tensorboard", logging_dir=dirpath)
@@ -96,21 +67,7 @@ def test_log(self):
             # Logged values are stored in the outermost-tfevents file and can be read in as a TFRecord
             # Names are randomly generated each time
             log = list(filter(lambda x: x.is_file(), Path(f"{dirpath}/{project_name}").iterdir()))[0]
-            serialized_examples = tf.data.TFRecordDataset(log)
-            for e in serialized_examples:
-                event = event_pb2.Event.FromString(e.numpy())
-                if step is None:
-                    step = event.step
-                for value in event.summary.value:
-                    if value.tag == "total_loss":
-                        total_loss = value.simple_value
-                    elif value.tag == "iteration":
-                        iteration = value.simple_value
-                    elif value.tag == "my_text/text_summary":  # Append /text_summary to the key
-                        my_text = value.tensor.string_val[0].decode()
-        self.assertAlmostEqual(total_loss, values["total_loss"])
-        self.assertEqual(iteration, values["iteration"])
-        self.assertEqual(my_text, values["my_text"])
+            self.assertNotEqual(str(log), "")
 
     def test_logging_dir(self):
         with self.assertRaisesRegex(ValueError, "Logging with `tensorboard` requires a `logging_dir`"):

```suggestion
    return unittest.skipUnless(is_tensorboard_available(), "test requires Tensorboard")(test_case)
```
@@ -85,12 +85,11 @@ def require_multi_gpu(test_case):
     return unittest.skipUnless(torch.cuda.device_count() > 1, "test requires multiple GPUs")(test_case)
 
 
-def require_tensorflow(test_case):
+def require_tensorboard(test_case):
     """
-    Decorator marking a test that requires TensorFlow installed. These tests are skipped when TensorFlow isn't
-    installed
+    Decorator marking a test that requires tensorboard installed. These tests are skipped when tensorboard isn't installed
     """
-    return unittest.skipUnless(is_tensorflow_available(), "test requires TensorFlow")(test_case)
+    return unittest.skipUnless(is_tensorboard_available(), "test requires wandb")(test_case)
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/420
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index d764d3927..2bd407d85 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -2,33 +2,41 @@ name: Run Tests
 
 on: [pull_request]
 
+env:
+  HF_HOME: ~/hf_cache
+
 jobs:
-  test:
+  run-tests:
     runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        test-kind: [
+          test,
+          test_example_differences,
+          test_checkpoint_step,
+          test_checkpoint_epoch,
+          test_rest
+        ]
     steps:
-    - uses: actions/checkout@v2
-    - name: Set up Python 3.7
-      uses: actions/setup-python@v2
+    - uses: actions/checkout@v3
+    - name: Set up python 3.7
+      uses: actions/setup-python@v3
       with:
         python-version: 3.7
-    - name: Install Python dependencies
-      run: |
-        pip install --upgrade pip
-        pip install -e .[test,test_trackers]
-    - name: Run Tests
-      run: make test
-      
-  test_examples:
-    runs-on: ubuntu-latest
-    steps:
-    - uses: actions/checkout@v2
-    - name: Set up Python 3.7
-      uses: actions/setup-python@v2
+    
+    - name: Activate python cache
+      uses: actions/cache@v3
       with:
-        python-version: 3.7
-    - name: Install Python dependencies
+        path: |
+          ${{ env.pythonLocation }}
+          ${{ env.HF_HOME }}
+        key: ${{ env.pythonLocation }}-${{ matrix.test-kind }}-${{ hashFiles('setup.py') }}
+    
+    - name: Install the library
       run: |
         pip install --upgrade pip
-        pip install -e .[test,tensorboard]
+        pip install -e .[test,test_trackers,tensorboard]
+    
     - name: Run Tests
-      run: make test_examples
+      run: |
+        make ${{ matrix.test-kind }}
\ No newline at end of file
diff --git a/Makefile b/Makefile
index f4534e12c..38cb16cfc 100644
--- a/Makefile
+++ b/Makefile
@@ -29,3 +29,16 @@ test:
 
 test_examples:
 	python -m pytest -s -v ./tests/test_examples.py
+
+# Broken down example tests for the CI runners
+test_example_differences:
+	python -m pytest -s -v ./tests/test_examples.py::ExampleDifferenceTests
+
+test_checkpoint_epoch:
+	python -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k "by_epoch"
+
+test_checkpoint_step:
+	python -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k "by_step"
+
+test_rest:
+	python -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k "not by_step and not by_epoch"
\ No newline at end of file

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/419
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index e9d797083..d764d3927 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -7,12 +7,14 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v2
-    - name: Set up Python 3.6
+    - name: Set up Python 3.7
       uses: actions/setup-python@v2
       with:
-        python-version: 3.6
+        python-version: 3.7
     - name: Install Python dependencies
-      run: pip install setuptools==59.5.0; pip install -e .[test,test_trackers]
+      run: |
+        pip install --upgrade pip
+        pip install -e .[test,test_trackers]
     - name: Run Tests
       run: make test
       
@@ -20,11 +22,13 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v2
-    - name: Set up Python 3.6
+    - name: Set up Python 3.7
       uses: actions/setup-python@v2
       with:
-        python-version: 3.6
+        python-version: 3.7
     - name: Install Python dependencies
-      run: pip install setuptools==59.5.0; pip install -e .[test] tensorboard
+      run: |
+        pip install --upgrade pip
+        pip install -e .[test,tensorboard]
     - name: Run Tests
       run: make test_examples
diff --git a/examples/README.md b/examples/README.md
index bb248e9d1..271ffd3c5 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -23,7 +23,7 @@ The [nlp_example.py](./nlp_example.py) script is a simple example to train a Ber
 Prior to running it you should install 🤗 Dataset and 🤗 Transformers:
 
 ```bash
-pip install datasets transformers
+pip install datasets evaluate transformers
 ```
 
 The same script can be run in any of the following configurations:
diff --git a/examples/by_feature/checkpointing.py b/examples/by_feature/checkpointing.py
index 06edd8540..c06d06ee9 100644
--- a/examples/by_feature/checkpointing.py
+++ b/examples/by_feature/checkpointing.py
@@ -18,8 +18,9 @@
 import torch
 from torch.utils.data import DataLoader
 
+import evaluate
 from accelerate import Accelerator, DistributedType
-from datasets import load_dataset, load_metric
+from datasets import load_dataset
 from transformers import (
     AdamW,
     AutoModelForSequenceClassification,
@@ -137,7 +138,7 @@ def training_function(config, args):
     set_seed(seed)
 
     train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)
-    metric = load_metric("glue", "mrpc")
+    metric = evaluate.load("glue", "mrpc")
 
     # If the batch size is too big we use gradient accumulation
     gradient_accumulation_steps = 1
diff --git a/examples/by_feature/cross_validation.py b/examples/by_feature/cross_validation.py
index 9fdc3fe72..53d582c67 100644
--- a/examples/by_feature/cross_validation.py
+++ b/examples/by_feature/cross_validation.py
@@ -19,8 +19,9 @@
 import torch
 from torch.utils.data import DataLoader
 
+import evaluate
 from accelerate import Accelerator, DistributedType
-from datasets import DatasetDict, load_dataset, load_metric
+from datasets import DatasetDict, load_dataset
 
 # New Code #
 # We'll be using StratifiedKFold for this example
@@ -144,7 +145,7 @@ def training_function(config, args):
     seed = int(config["seed"])
     batch_size = int(config["batch_size"])
 
-    metric = load_metric("glue", "mrpc")
+    metric = evaluate.load("glue", "mrpc")
 
     # If the batch size is too big we use gradient accumulation
     gradient_accumulation_steps = 1
diff --git a/examples/by_feature/fsdp_with_peak_mem_tracking.py b/examples/by_feature/fsdp_with_peak_mem_tracking.py
index 8304bcefe..4f0311785 100644
--- a/examples/by_feature/fsdp_with_peak_mem_tracking.py
+++ b/examples/by_feature/fsdp_with_peak_mem_tracking.py
@@ -19,8 +19,9 @@
 import torch
 from torch.utils.data import DataLoader
 
+import evaluate
 from accelerate import Accelerator, DistributedType
-from datasets import load_dataset, load_metric
+from datasets import load_dataset
 from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
 
 
@@ -119,7 +120,7 @@ def training_function(config, args):
 
     tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
     datasets = load_dataset("glue", "mrpc")
-    metric = load_metric("glue", "mrpc")
+    metric = evaluate.load("glue", "mrpc")
 
     def tokenize_function(examples):
         # max_length=None => use the model max length (it's actually the default)
diff --git a/examples/by_feature/memory.py b/examples/by_feature/memory.py
index 140f14d75..ac4b45866 100644
--- a/examples/by_feature/memory.py
+++ b/examples/by_feature/memory.py
@@ -17,11 +17,11 @@
 import torch
 from torch.utils.data import DataLoader
 
-from accelerate import Accelerator, DistributedType
-
 # New Code #
+import evaluate
+from accelerate import Accelerator, DistributedType
 from accelerate.utils import find_executable_batch_size
-from datasets import load_dataset, load_metric
+from datasets import load_dataset
 from transformers import (
     AdamW,
     AutoModelForSequenceClassification,
@@ -121,7 +121,7 @@ def training_function(config, args):
     seed = int(config["seed"])
     batch_size = int(config["batch_size"])
 
-    metric = load_metric("glue", "mrpc")
+    metric = evaluate.load("glue", "mrpc")
 
     # If the batch size is too big we use gradient accumulation
     gradient_accumulation_steps = 1
diff --git a/examples/by_feature/multi_process_metrics.py b/examples/by_feature/multi_process_metrics.py
index 6babf879a..0307d251f 100644
--- a/examples/by_feature/multi_process_metrics.py
+++ b/examples/by_feature/multi_process_metrics.py
@@ -18,8 +18,9 @@
 import torch
 from torch.utils.data import DataLoader
 
+import evaluate
 from accelerate import Accelerator, DistributedType
-from datasets import load_dataset, load_metric
+from datasets import load_dataset
 from transformers import (
     AdamW,
     AutoModelForSequenceClassification,
@@ -122,7 +123,7 @@ def training_function(config, args):
     seed = int(config["seed"])
     batch_size = int(config["batch_size"])
 
-    metric = load_metric("glue", "mrpc")
+    metric = evaluate.load("glue", "mrpc")
 
     # If the batch size is too big we use gradient accumulation
     gradient_accumulation_steps = 1
diff --git a/examples/by_feature/tracking.py b/examples/by_feature/tracking.py
index 9538b21b1..4020b3382 100644
--- a/examples/by_feature/tracking.py
+++ b/examples/by_feature/tracking.py
@@ -18,8 +18,9 @@
 import torch
 from torch.utils.data import DataLoader
 
+import evaluate
 from accelerate import Accelerator, DistributedType
-from datasets import load_dataset, load_metric
+from datasets import load_dataset
 from transformers import (
     AdamW,
     AutoModelForSequenceClassification,
@@ -132,7 +133,7 @@ def training_function(config, args):
     set_seed(seed)
 
     train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)
-    metric = load_metric("glue", "mrpc")
+    metric = evaluate.load("glue", "mrpc")
 
     # If the batch size is too big we use gradient accumulation
     gradient_accumulation_steps = 1
diff --git a/examples/complete_nlp_example.py b/examples/complete_nlp_example.py
index 9c8b89706..8544d1279 100644
--- a/examples/complete_nlp_example.py
+++ b/examples/complete_nlp_example.py
@@ -18,8 +18,9 @@
 import torch
 from torch.utils.data import DataLoader
 
+import evaluate
 from accelerate import Accelerator, DistributedType
-from datasets import load_dataset, load_metric
+from datasets import load_dataset
 from transformers import (
     AdamW,
     AutoModelForSequenceClassification,
@@ -89,7 +90,7 @@ def training_function(config, args):
 
     tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
     datasets = load_dataset("glue", "mrpc")
-    metric = load_metric("glue", "mrpc")
+    metric = evaluate.load("glue", "mrpc")
 
     def tokenize_function(examples):
         # max_length=None => use the model max length (it's actually the default)
diff --git a/examples/nlp_example.py b/examples/nlp_example.py
index 87cfd7698..84c24fd11 100644
--- a/examples/nlp_example.py
+++ b/examples/nlp_example.py
@@ -17,8 +17,9 @@
 import torch
 from torch.utils.data import DataLoader
 
+import evaluate
 from accelerate import Accelerator, DistributedType
-from datasets import load_dataset, load_metric
+from datasets import load_dataset
 from transformers import (
     AdamW,
     AutoModelForSequenceClassification,
@@ -106,7 +107,7 @@ def training_function(config, args):
     seed = int(config["seed"])
     batch_size = int(config["batch_size"])
 
-    metric = load_metric("glue", "mrpc")
+    metric = evaluate.load("glue", "mrpc")
 
     # If the batch size is too big we use gradient accumulation
     gradient_accumulation_steps = 1
diff --git a/setup.py b/setup.py
index 5f6e128c1..51a6c1a01 100644
--- a/setup.py
+++ b/setup.py
@@ -24,11 +24,13 @@
     "pytest-xdist",
     "pytest-subtests",
     "datasets",
+    "evaluate",
     "transformers",
     "scipy",
     "sklearn"
 ]
-extras["test_trackers"] = ["wandb", "comet-ml", "tensorflow>=2.6.2", "tensorboard"]
+extras["tensorboard"] = ["protobuf<=3.20.1", "tensorflow>=2.6.2", "tensorboard"]
+extras["test_trackers"] = extras["tensorboard"] + ["wandb", "comet-ml"]
 extras["dev"] = extras["quality"] + extras["test"]
 
 extras["sagemaker"] = [

There's some micro optimization that can be done I found that can save us a number of minutes on each CI, and will actually let this no longer need to be a thing. I'll make a PR after this with that, it'll simplify some things quite a bit :) 
@@ -7,24 +7,28 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v2
-    - name: Set up Python 3.6
+    - name: Set up Python 3.7
       uses: actions/setup-python@v2
       with:
-        python-version: 3.6
+        python-version: 3.7
     - name: Install Python dependencies
-      run: pip install setuptools==59.5.0; pip install -e .[test,test_trackers]
+      run: |
+        pip install --upgrade pip
+        pip install -e .[test,test_trackers]
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/417
diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml
index 11946b544..d50aab329 100644
--- a/.github/workflows/nightly.yml
+++ b/.github/workflows/nightly.yml
@@ -32,7 +32,7 @@ jobs:
       - name: Run test on GPUs
         run: |
           source activate accelerate
-          make test_cuda
+          make test
       - name: Run examples on GPUs
         run: |
           source activate accelerate
@@ -64,7 +64,7 @@ jobs:
       - name: Run test on GPUs
         run: |
           source activate accelerate
-          make test_cuda
+          make test
 
       - name: Run examples on GPUs
         run: |
diff --git a/.github/workflows/on-merge.yml b/.github/workflows/on-merge.yml
new file mode 100644
index 000000000..8df455563
--- /dev/null
+++ b/.github/workflows/on-merge.yml
@@ -0,0 +1,69 @@
+name: Self-hosted runner (push to "main")
+
+on:
+  workflow_dispatch:
+  push:
+    branches:
+      - "main"
+
+jobs:
+  run_all_tests_single_gpu:
+    runs-on: [self-hosted, docker-gpu, multi-gpu]
+    env:
+      CUDA_VISIBLE_DEVICES: "0"
+    container:
+      image: huggingface/accelerate-gpu:latest
+      options: --gpus all --shm-size "16gb"
+    defaults:
+      run:
+        working-directory: accelerate/
+        shell: bash
+    steps:
+      - name: Update clone & pip install
+        run: |
+          source activate accelerate
+          git config --global --add safe.directory '*'
+          git fetch && git checkout ${{ github.sha }}
+          pip install tensorflow -U 
+          pip install -e . -U
+
+# Note: tensorflow upgrade is needed until dropped py 3.6 support
+
+      - name: Run test on GPUs
+        run: |
+          source activate accelerate
+          make test
+      - name: Run examples on GPUs
+        run: |
+          source activate accelerate
+          make test_examples
+
+  run_all_tests_multi_gpu:
+    runs-on: [self-hosted, docker-gpu, multi-gpu]
+    container:
+      image: huggingface/accelerate-gpu:latest
+      options: --gpus all --shm-size "16gb"
+    defaults:
+      run:
+        working-directory: accelerate/
+        shell: bash
+    steps:
+      - name: Update clone
+        run: |
+          source activate accelerate
+          git config --global --add safe.directory '*'
+          git fetch && git checkout ${{ github.sha }}
+          pip install tensorflow -U
+          pip install -e . -U
+
+        # Note: tensorflow upgrade is needed until dropped py 3.6 support
+
+      - name: Run test on GPUs
+        run: |
+          source activate accelerate
+          make test
+
+      - name: Run examples on GPUs
+        run: |
+          source activate accelerate
+          make test_examples
\ No newline at end of file

Thinking on one of your frustrations @sgugger, I decided to add a `pip install -U` instead of `--no-deps` so that if a dep was added on a push, we can test it (and not instead have all those failing tests on a push to main and freak out 😉 )
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/416
diff --git a/.github/ISSUE_TEMPLATE/bug-report.yml b/.github/ISSUE_TEMPLATE/bug-report.yml
new file mode 100644
index 000000000..69a101307
--- /dev/null
+++ b/.github/ISSUE_TEMPLATE/bug-report.yml
@@ -0,0 +1,59 @@
+name: "\U0001F41B Bug Report"
+description: Submit a bug report to help us improve Accelerate
+labels: [ "bug" ]
+body:
+  - type: textarea
+    id: system-info
+    attributes:
+      label: System Info
+      description: Please share your accelerate configuration with us. You can run the command `accelerate env` and copy-paste its outputs below
+      render: Shell
+      placeholder: accelerate version, OS, python version, numpy version, torch version, and accelerate's configuration
+    validations:
+      required: true
+  
+  - type: checkboxes
+    id: information-scripts-examples
+    attributes:
+      label: Information
+      description: 'The problem arises when using:'
+      options:
+        - label: "The official example scripts"
+        - label: "My own modified scripts"
+  
+  - type: checkboxes
+    id: information-tasks
+    attributes:
+      label: Tasks
+      description: "The tasks I am working on are:"
+      options:
+        - label: "One of the scripts in the examples/ folder of Accelerate or an officially supported `no_trainer` script in the `examples` folder of the `transformers` repo (such as `run_no_trainer_glue.py`)"
+        - label: "My own task or dataset (give details below)"
+  
+  - type: textarea
+    id: reproduction
+    validations:
+      required: true
+    attributes:
+      label: Reproduction
+      description: |
+        Please provide a code sample that reproduces the problem you ran into. It can be a Colab link or just a code snippet.
+        If you have code snippets, error messages, stack traces please provide them here as well.
+        Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
+        Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.
+
+      placeholder: |
+        Steps to reproduce the behavior:
+          
+          1.
+          2.
+          3.
+
+  - type: textarea
+    id: expected-behavior
+    validations:
+      required: true
+    attributes:
+      label: Expected behavior
+      description: "A clear and concise description of what you would expect to happen."
+      render: Shell

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/415
diff --git a/docker/accelerate-gpu/Dockerfile b/docker/accelerate-gpu/Dockerfile
index 6195ae2ae..f495d9063 100644
--- a/docker/accelerate-gpu/Dockerfile
+++ b/docker/accelerate-gpu/Dockerfile
@@ -22,7 +22,7 @@ SHELL ["/bin/bash", "-c"]
 # Activate the conda env and install torch + accelerate
 RUN source activate accelerate && \
     python3 -m pip install --no-cache-dir \
-    git+https://github.com/huggingface/accelerate@nightly-runner#egg=accelerate[test,test_trackers] \
+    git+https://github.com/huggingface/accelerate#egg=accelerate[test,test_trackers] \
     --extra-index-url https://download.pytorch.org/whl/cu113
 
 # Stage 2

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/414
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 4613e1e8d..e9d797083 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -14,7 +14,7 @@ jobs:
     - name: Install Python dependencies
       run: pip install setuptools==59.5.0; pip install -e .[test,test_trackers]
     - name: Run Tests
-      run: make test_cpu
+      run: make test
       
   test_examples:
     runs-on: ubuntu-latest
diff --git a/Makefile b/Makefile
index d2822153b..f4534e12c 100644
--- a/Makefile
+++ b/Makefile
@@ -24,12 +24,8 @@ style:
 	python utils/style_doc.py src/accelerate docs/source --max_len 119
 	
 # Run tests for the library
-test_cpu:
+test:
 	python -m pytest -s -v ./tests/ --ignore=./tests/test_examples.py
 
-test_cuda:
-	python -m pytest -s -v ./tests/ --ignore=./tests/test_examples.py --ignore=./tests/test_scheduler.py --ignore=./tests/test_cpu.py
-	python -m pytest -s -v ./tests/test_cpu.py ./tests/test_scheduler.py
-
 test_examples:
 	python -m pytest -s -v ./tests/test_examples.py
diff --git a/src/accelerate/test_utils/__init__.py b/src/accelerate/test_utils/__init__.py
index 46723ebfa..e8dc99c8c 100644
--- a/src/accelerate/test_utils/__init__.py
+++ b/src/accelerate/test_utils/__init__.py
@@ -2,5 +2,13 @@
 # There's no way to ignore "F401 '...' imported but unused" warnings in this
 # module, but to preserve other warnings. So, don't check this module at all.
 
-from .testing import are_the_same_tensors, execute_subprocess_async, require_cuda, require_multi_gpu, require_tpu, slow
+from .testing import (
+    are_the_same_tensors,
+    execute_subprocess_async,
+    require_cpu,
+    require_cuda,
+    require_multi_gpu,
+    require_tpu,
+    slow,
+)
 from .training import RegressionDataset, RegressionModel
diff --git a/src/accelerate/test_utils/testing.py b/src/accelerate/test_utils/testing.py
index 9379eb0b9..013dfad44 100644
--- a/src/accelerate/test_utils/testing.py
+++ b/src/accelerate/test_utils/testing.py
@@ -56,6 +56,13 @@ def slow(test_case):
     return unittest.skipUnless(_run_slow_tests, "test is slow")(test_case)
 
 
+def require_cpu(test_case):
+    """
+    Decorator marking a test that must be only ran on the CPU. These tests are skipped when a GPU is available.
+    """
+    return unittest.skipUnless(not torch.cuda.is_available(), "test requires only a CPU")(test_case)
+
+
 def require_cuda(test_case):
     """
     Decorator marking a test that requires CUDA. These tests are skipped when there are no GPU available.
diff --git a/tests/test_cpu.py b/tests/test_cpu.py
index 63ee69a97..ab73058c1 100644
--- a/tests/test_cpu.py
+++ b/tests/test_cpu.py
@@ -15,9 +15,10 @@
 import unittest
 
 from accelerate import debug_launcher
-from accelerate.test_utils import test_script
+from accelerate.test_utils import require_cpu, test_script
 
 
+@require_cpu
 class MultiCPUTester(unittest.TestCase):
     def test_cpu(self):
         debug_launcher(test_script.main)
diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py
index 8ae9e56ae..be4f975fb 100644
--- a/tests/test_scheduler.py
+++ b/tests/test_scheduler.py
@@ -18,6 +18,7 @@
 import torch
 
 from accelerate import Accelerator, debug_launcher
+from accelerate.test_utils import require_cpu
 
 
 def scheduler_test(num_processes=2, step_scheduler_with_optimizer=True, split_batches=False):
@@ -46,6 +47,7 @@ def scheduler_test(num_processes=2, step_scheduler_with_optimizer=True, split_ba
     ), f"Wrong lr found at second step, expected {expected_lr}, got {scheduler.get_last_lr()[0]}"
 
 
+@require_cpu
 class SchedulerTester(unittest.TestCase):
     def test_scheduler_steps_with_optimizer_single_process(self):
         debug_launcher(partial(scheduler_test, num_processes=1), num_processes=1)

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/413
diff --git a/src/accelerate/test_utils/training.py b/src/accelerate/test_utils/training.py
index c26587abe..7345b93cb 100644
--- a/src/accelerate/test_utils/training.py
+++ b/src/accelerate/test_utils/training.py
@@ -17,8 +17,6 @@
 from torch.utils.data import DataLoader
 
 from accelerate.utils.dataclasses import DistributedType
-from datasets import load_dataset
-from transformers import AutoTokenizer
 
 
 class RegressionDataset:
@@ -51,6 +49,9 @@ def forward(self, x=None):
 
 
 def mocked_dataloaders(accelerator, batch_size: int = 16):
+    from datasets import load_dataset
+    from transformers import AutoTokenizer
+
     tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
     data_files = {"train": "tests/test_samples/MRPC/train.csv", "validation": "tests/test_samples/MRPC/dev.csv"}
     datasets = load_dataset("csv", data_files=data_files)

_The documentation is not available anymore as the PR was closed or merged._
Thanks folks, for the library and the fix!
https://api.github.com/repos/huggingface/accelerate/pulls/411
diff --git a/.github/workflows/build-docker-images.yml b/.github/workflows/build-docker-images.yml
index 6b7987e51..de8542359 100644
--- a/.github/workflows/build-docker-images.yml
+++ b/.github/workflows/build-docker-images.yml
@@ -1,7 +1,7 @@
 name: Build Docker images (scheduled)
 
 on:
-  repository_dispatch:
+  workflow_dispatch:
   schedule:
     - cron: "0 1 * * *"
 
diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml
new file mode 100644
index 000000000..11946b544
--- /dev/null
+++ b/.github/workflows/nightly.yml
@@ -0,0 +1,72 @@
+name: Self-hosted runner (scheduled)
+
+on:
+  workflow_dispatch:
+  schedule:
+    - cron: "0 2 * * *"
+
+jobs:
+  run_all_tests_single_gpu:
+    runs-on: [self-hosted, docker-gpu, multi-gpu]
+    env:
+      CUDA_VISIBLE_DEVICES: "0"
+      ALL_SLOW: "yes"
+    container:
+      image: huggingface/accelerate-gpu:latest
+      options: --gpus all --shm-size "16gb"
+    defaults:
+      run:
+        working-directory: accelerate/
+        shell: bash
+    steps:
+      - name: Update clone & pip install
+        run: |
+          source activate accelerate
+          git config --global --add safe.directory '*'
+          git fetch && git checkout ${{ github.sha }}
+          pip install tensorflow -U 
+          pip install -e . --no-deps
+
+# Note: tensorflow upgrade is needed until dropped py 3.6 support
+
+      - name: Run test on GPUs
+        run: |
+          source activate accelerate
+          make test_cuda
+      - name: Run examples on GPUs
+        run: |
+          source activate accelerate
+          make test_examples
+
+  run_all_tests_multi_gpu:
+    runs-on: [self-hosted, docker-gpu, multi-gpu]
+    env:
+      ALL_SLOW: "yes"
+      CUDA_VISIBLE_DEVICES: "0,1"
+    container:
+      image: huggingface/accelerate-gpu:latest
+      options: --gpus all --shm-size "16gb"
+    defaults:
+      run:
+        working-directory: accelerate/
+        shell: bash
+    steps:
+      - name: Update clone
+        run: |
+          source activate accelerate
+          git config --global --add safe.directory '*'
+          git fetch && git checkout ${{ github.sha }}
+          pip install tensorflow -U
+          pip install -e . --no-deps
+
+        # Note: tensorflow upgrade is needed until dropped py 3.6 support
+
+      - name: Run test on GPUs
+        run: |
+          source activate accelerate
+          make test_cuda
+
+      - name: Run examples on GPUs
+        run: |
+          source activate accelerate
+          make test_examples
\ No newline at end of file
diff --git a/docker/accelerate-cpu/Dockerfile b/docker/accelerate-cpu/Dockerfile
index c536fb70c..1d99c4035 100644
--- a/docker/accelerate-cpu/Dockerfile
+++ b/docker/accelerate-cpu/Dockerfile
@@ -21,8 +21,8 @@ WORKDIR /workspace
 RUN python3 -m pip install --upgrade --no-cache-dir pip
 RUN python3 -m pip install --no-cache-dir \
     jupyter \
-    torch --extra-index-url https://download.pytorch.org/whl/cpu \
-    git+https://github.com/huggingface/accelerate#egg=accelerate[dev]
+    git+https://github.com/huggingface/accelerate#egg=accelerate[test,test_trackers] \
+    --extra-index-url https://download.pytorch.org/whl/cpu
     
 # Stage 2
 FROM python:3.6-slim AS build-image
diff --git a/docker/accelerate-gpu/Dockerfile b/docker/accelerate-gpu/Dockerfile
index 299f573d8..6195ae2ae 100644
--- a/docker/accelerate-gpu/Dockerfile
+++ b/docker/accelerate-gpu/Dockerfile
@@ -4,7 +4,7 @@
 # Use base conda image to reduce time
 FROM continuumio/miniconda3:latest AS compile-image
 # Specify py version
-ENV PYTHON_VERSION=3.6 
+ENV PYTHON_VERSION=3.6
 # Install apt libs
 RUN apt-get update && \
     apt-get install -y curl git wget && \
@@ -22,17 +22,21 @@ SHELL ["/bin/bash", "-c"]
 # Activate the conda env and install torch + accelerate
 RUN source activate accelerate && \
     python3 -m pip install --no-cache-dir \
-    torch --extra-index-url https://download.pytorch.org/whl/cu113 \
-    git+https://github.com/huggingface/accelerate#egg=accelerate[dev]
+    git+https://github.com/huggingface/accelerate@nightly-runner#egg=accelerate[test,test_trackers] \
+    --extra-index-url https://download.pytorch.org/whl/cu113
 
 # Stage 2
 FROM nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04 AS build-image
 COPY --from=compile-image /opt/conda /opt/conda
 ENV PATH /opt/conda/bin:$PATH
 
-RUN echo "source activate accelerate" >> /.bashrc
+# Install apt libs
+RUN apt-get update && \
+    apt-get install -y curl git wget && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists*
+
+RUN echo "source activate accelerate" >> ~/.profile
 
-RUN useradd -ms /bin/bash user
-USER user
 # Activate the virtualenv
 CMD ["/bin/bash"]
\ No newline at end of file
diff --git a/setup.py b/setup.py
index dcfe3fd8e..5f6e128c1 100644
--- a/setup.py
+++ b/setup.py
@@ -19,6 +19,7 @@
 extras["quality"] = ["black ~= 22.0", "isort >= 5.5.4", "flake8 >= 3.8.3"]
 extras["docs"] = []
 extras["test"] = [
+    "psutil",
     "pytest",
     "pytest-xdist",
     "pytest-subtests",
@@ -27,7 +28,7 @@
     "scipy",
     "sklearn"
 ]
-extras["test_trackers"] = ["wandb", "comet-ml", "tensorflow"]
+extras["test_trackers"] = ["wandb", "comet-ml", "tensorflow>=2.6.2", "tensorboard"]
 extras["dev"] = extras["quality"] + extras["test"]
 
 extras["sagemaker"] = [
diff --git a/src/accelerate/test_utils/training.py b/src/accelerate/test_utils/training.py
index c26587abe..7345b93cb 100644
--- a/src/accelerate/test_utils/training.py
+++ b/src/accelerate/test_utils/training.py
@@ -17,8 +17,6 @@
 from torch.utils.data import DataLoader
 
 from accelerate.utils.dataclasses import DistributedType
-from datasets import load_dataset
-from transformers import AutoTokenizer
 
 
 class RegressionDataset:
@@ -51,6 +49,9 @@ def forward(self, x=None):
 
 
 def mocked_dataloaders(accelerator, batch_size: int = 16):
+    from datasets import load_dataset
+    from transformers import AutoTokenizer
+
     tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
     data_files = {"train": "tests/test_samples/MRPC/train.csv", "validation": "tests/test_samples/MRPC/dev.csv"}
     datasets = load_dataset("csv", data_files=data_files)
diff --git a/tests/test_examples.py b/tests/test_examples.py
index 6ec2a143a..8f277b932 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -21,6 +21,8 @@
 import unittest
 from unittest import mock
 
+import torch
+
 from accelerate.test_utils.examples import compare_against_test
 from accelerate.test_utils.testing import TempDirTestCase, slow
 from accelerate.utils import write_basic_config
@@ -142,11 +144,11 @@ def test_checkpointing_by_epoch(self):
     def test_checkpointing_by_steps(self):
         testargs = f"""
         examples/by_feature/checkpointing.py
-        --checkpointing_steps 2
+        --checkpointing_steps 1
         --output_dir {self.tmpdir}
         """.split()
         _ = subprocess.run(self._launch_args + testargs, stdout=subprocess.PIPE, env=os.environ)
-        self.assertTrue(os.path.exists(os.path.join(self.tmpdir, "step_4")))
+        self.assertTrue(os.path.exists(os.path.join(self.tmpdir, "step_5")))
 
     def test_load_states_by_epoch(self):
         testargs = f"""
@@ -163,14 +165,23 @@ def test_load_states_by_epoch(self):
     def test_load_states_by_steps(self):
         testargs = f"""
         examples/by_feature/checkpointing.py
-        --resume_from_checkpoint {os.path.join(self.tmpdir, "step_4")}
+        --resume_from_checkpoint {os.path.join(self.tmpdir, "step_5")}
         """.split()
         output = subprocess.run(
             self._launch_args + testargs, universal_newlines=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE
         ).stdout
-        self.assertNotIn("epoch 0:", output)
-        self.assertIn("epoch 1:", output)
-        self.assertIn("epoch 2:", output)
+        if torch.cuda.is_available():
+            num_processes = torch.cuda.device_count()
+        else:
+            num_processes = 1
+        if num_processes > 1:
+            self.assertNotIn("epoch 0:", output)
+            self.assertNotIn("epoch 1:", output)
+            self.assertIn("epoch 2:", output)
+        else:
+            self.assertNotIn("epoch 0:", output)
+            self.assertIn("epoch 1:", output)
+            self.assertIn("epoch 2:", output)
 
     @slow
     def test_cross_validation(self):

Why was this necessary?
@@ -0,0 +1,65 @@
+name: Self-hosted runner (scheduled)
+
+on:
+  workflow_dispatch:
+  schedule:
+    - cron: "0 2 * * *"
+
+jobs:
+  run_all_tests_single_gpu:
+    runs-on: [self-hosted, docker-gpu, multi-gpu]
+    env:
+      CUDA_VISIBLE_DEVICES: "0"
+      ALL_SLOW: "yes"
+    container:
+      image: huggingface/accelerate-gpu:latest
+      options: --gpus all --shm-size "16gb"
+    defaults:
+      run:
+        working-directory: accelerate/
+        shell: bash
+    steps:
+      - name: Update clone & pip install
+        run: |
+          source activate accelerate
+          git config --global --add safe.directory '*'
+          git fetch && git checkout ${{ github.sha }}
+          pip install -e . --no-deps
+          pip install tensorflow -U
+      - name: Run test on GPUs
+        run: |
+          source activate accelerate
+          make test_cuda
+      - name: Run examples on GPUs
+        run: |
+          source activate accelerate
+          make test_examples
+
+  run_all_tests_multi_gpu:
+    runs-on: [self-hosted, docker-gpu, multi-gpu]
+    env:
+      ALL_SLOW: "yes"
+    container:
+      image: huggingface/accelerate-gpu:latest
+      options: --gpus all --shm-size "16gb"
+    defaults:
+      run:
+        working-directory: accelerate/
+        shell: bash
+    steps:
+      - name: Update clone
+        run: |
+          source activate accelerate
+          git config --global --add safe.directory '*'
Is `tensorflow` required?
@@ -0,0 +1,65 @@
+name: Self-hosted runner (scheduled)
+
+on:
+  workflow_dispatch:
+  schedule:
+    - cron: "0 2 * * *"
+
+jobs:
+  run_all_tests_single_gpu:
+    runs-on: [self-hosted, docker-gpu, multi-gpu]
+    env:
+      CUDA_VISIBLE_DEVICES: "0"
+      ALL_SLOW: "yes"
+    container:
+      image: huggingface/accelerate-gpu:latest
+      options: --gpus all --shm-size "16gb"
+    defaults:
+      run:
+        working-directory: accelerate/
+        shell: bash
+    steps:
+      - name: Update clone & pip install
+        run: |
+          source activate accelerate
+          git config --global --add safe.directory '*'
+          git fetch && git checkout ${{ github.sha }}
+          pip install -e . --no-deps
+          pip install tensorflow -U
I was getting odd permission errors, and this was the documented solution I found. (Will post the link in the AM)
@@ -0,0 +1,65 @@
+name: Self-hosted runner (scheduled)
+
+on:
+  workflow_dispatch:
+  schedule:
+    - cron: "0 2 * * *"
+
+jobs:
+  run_all_tests_single_gpu:
+    runs-on: [self-hosted, docker-gpu, multi-gpu]
+    env:
+      CUDA_VISIBLE_DEVICES: "0"
+      ALL_SLOW: "yes"
+    container:
+      image: huggingface/accelerate-gpu:latest
+      options: --gpus all --shm-size "16gb"
+    defaults:
+      run:
+        working-directory: accelerate/
+        shell: bash
+    steps:
+      - name: Update clone & pip install
+        run: |
+          source activate accelerate
+          git config --global --add safe.directory '*'
+          git fetch && git checkout ${{ github.sha }}
+          pip install -e . --no-deps
+          pip install tensorflow -U
+      - name: Run test on GPUs
+        run: |
+          source activate accelerate
+          make test_cuda
+      - name: Run examples on GPUs
+        run: |
+          source activate accelerate
+          make test_examples
+
+  run_all_tests_multi_gpu:
+    runs-on: [self-hosted, docker-gpu, multi-gpu]
+    env:
+      ALL_SLOW: "yes"
+    container:
+      image: huggingface/accelerate-gpu:latest
+      options: --gpus all --shm-size "16gb"
+    defaults:
+      run:
+        working-directory: accelerate/
+        shell: bash
+    steps:
+      - name: Update clone
+        run: |
+          source activate accelerate
+          git config --global --add safe.directory '*'
It won't be after the prior prs are merged, nice catch!
@@ -0,0 +1,65 @@
+name: Self-hosted runner (scheduled)
+
+on:
+  workflow_dispatch:
+  schedule:
+    - cron: "0 2 * * *"
+
+jobs:
+  run_all_tests_single_gpu:
+    runs-on: [self-hosted, docker-gpu, multi-gpu]
+    env:
+      CUDA_VISIBLE_DEVICES: "0"
+      ALL_SLOW: "yes"
+    container:
+      image: huggingface/accelerate-gpu:latest
+      options: --gpus all --shm-size "16gb"
+    defaults:
+      run:
+        working-directory: accelerate/
+        shell: bash
+    steps:
+      - name: Update clone & pip install
+        run: |
+          source activate accelerate
+          git config --global --add safe.directory '*'
+          git fetch && git checkout ${{ github.sha }}
+          pip install -e . --no-deps
+          pip install tensorflow -U
@LysandreJik here's the specific thread: https://github.com/actions/checkout/issues/760#issuecomment-1097519290

How we got here is the runner has a clone of accelerate inside the working directory for the action, and that clone is then updated, checked out, etc. But we run into this permission denied unsafe repository from github, so github needs to be told we can safely adjust the repository. 
@@ -0,0 +1,65 @@
+name: Self-hosted runner (scheduled)
+
+on:
+  workflow_dispatch:
+  schedule:
+    - cron: "0 2 * * *"
+
+jobs:
+  run_all_tests_single_gpu:
+    runs-on: [self-hosted, docker-gpu, multi-gpu]
+    env:
+      CUDA_VISIBLE_DEVICES: "0"
+      ALL_SLOW: "yes"
+    container:
+      image: huggingface/accelerate-gpu:latest
+      options: --gpus all --shm-size "16gb"
+    defaults:
+      run:
+        working-directory: accelerate/
+        shell: bash
+    steps:
+      - name: Update clone & pip install
+        run: |
+          source activate accelerate
+          git config --global --add safe.directory '*'
+          git fetch && git checkout ${{ github.sha }}
+          pip install -e . --no-deps
+          pip install tensorflow -U
+      - name: Run test on GPUs
+        run: |
+          source activate accelerate
+          make test_cuda
+      - name: Run examples on GPUs
+        run: |
+          source activate accelerate
+          make test_examples
+
+  run_all_tests_multi_gpu:
+    runs-on: [self-hosted, docker-gpu, multi-gpu]
+    env:
+      ALL_SLOW: "yes"
+    container:
+      image: huggingface/accelerate-gpu:latest
+      options: --gpus all --shm-size "16gb"
+    defaults:
+      run:
+        working-directory: accelerate/
+        shell: bash
+    steps:
+      - name: Update clone
+        run: |
+          source activate accelerate
+          git config --global --add safe.directory '*'
Okay, interesting. Thanks!
@@ -0,0 +1,65 @@
+name: Self-hosted runner (scheduled)
+
+on:
+  workflow_dispatch:
+  schedule:
+    - cron: "0 2 * * *"
+
+jobs:
+  run_all_tests_single_gpu:
+    runs-on: [self-hosted, docker-gpu, multi-gpu]
+    env:
+      CUDA_VISIBLE_DEVICES: "0"
+      ALL_SLOW: "yes"
+    container:
+      image: huggingface/accelerate-gpu:latest
+      options: --gpus all --shm-size "16gb"
+    defaults:
+      run:
+        working-directory: accelerate/
+        shell: bash
+    steps:
+      - name: Update clone & pip install
+        run: |
+          source activate accelerate
+          git config --global --add safe.directory '*'
+          git fetch && git checkout ${{ github.sha }}
+          pip install -e . --no-deps
+          pip install tensorflow -U
+      - name: Run test on GPUs
+        run: |
+          source activate accelerate
+          make test_cuda
+      - name: Run examples on GPUs
+        run: |
+          source activate accelerate
+          make test_examples
+
+  run_all_tests_multi_gpu:
+    runs-on: [self-hosted, docker-gpu, multi-gpu]
+    env:
+      ALL_SLOW: "yes"
+    container:
+      image: huggingface/accelerate-gpu:latest
+      options: --gpus all --shm-size "16gb"
+    defaults:
+      run:
+        working-directory: accelerate/
+        shell: bash
+    steps:
+      - name: Update clone
+        run: |
+          source activate accelerate
+          git config --global --add safe.directory '*'
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/410
diff --git a/docker/accelerate-cpu/Dockerfile b/docker/accelerate-cpu/Dockerfile
index c536fb70c..0c7073f7c 100644
--- a/docker/accelerate-cpu/Dockerfile
+++ b/docker/accelerate-cpu/Dockerfile
@@ -27,8 +27,11 @@ RUN python3 -m pip install --no-cache-dir \
 # Stage 2
 FROM python:3.6-slim AS build-image
 COPY --from=compile-image /opt/venv /opt/venv
-RUN useradd -ms /bin/bash user
-USER user
+# Install apt libs
+RUN apt-get update && \
+    apt-get install -y curl git wget && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists*
 
 # Make sure we use the virtualenv
 ENV PATH="/opt/venv/bin:$PATH"
diff --git a/docker/accelerate-gpu/Dockerfile b/docker/accelerate-gpu/Dockerfile
index 299f573d8..5eabb00f7 100644
--- a/docker/accelerate-gpu/Dockerfile
+++ b/docker/accelerate-gpu/Dockerfile
@@ -30,9 +30,13 @@ FROM nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04 AS build-image
 COPY --from=compile-image /opt/conda /opt/conda
 ENV PATH /opt/conda/bin:$PATH
 
-RUN echo "source activate accelerate" >> /.bashrc
+# Install apt libs
+RUN apt-get update && \
+    apt-get install -y curl git wget && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists*
+
+RUN echo "source activate accelerate" >> ~/.profile
 
-RUN useradd -ms /bin/bash user
-USER user
 # Activate the virtualenv
 CMD ["/bin/bash"]
\ No newline at end of file

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/409
diff --git a/setup.py b/setup.py
index dcfe3fd8e..b53519525 100644
--- a/setup.py
+++ b/setup.py
@@ -19,6 +19,7 @@
 extras["quality"] = ["black ~= 22.0", "isort >= 5.5.4", "flake8 >= 3.8.3"]
 extras["docs"] = []
 extras["test"] = [
+    "psutil", 
     "pytest",
     "pytest-xdist",
     "pytest-subtests",
@@ -27,7 +28,7 @@
     "scipy",
     "sklearn"
 ]
-extras["test_trackers"] = ["wandb", "comet-ml", "tensorflow"]
+extras["test_trackers"] = ["wandb", "comet-ml", "tensorflow>=2.6.2", "tensorboard"]
 extras["dev"] = extras["quality"] + extras["test"]
 
 extras["sagemaker"] = [

No, `psutil` should be an optional dependency as it's only used for big model inference when `device_map="auto"`.
@@ -55,7 +55,7 @@
         ]
     },
     python_requires=">=3.6.0",
-    install_requires=["numpy>=1.17", "packaging>=20.0", "pyyaml", "torch>=1.4.0"],
+    install_requires=["numpy>=1.17", "packaging>=20.0", "psutil", "pyyaml", "torch>=1.4.0"],
Makes sense, will move it over there. Thanks!
@@ -55,7 +55,7 @@
         ]
     },
     python_requires=">=3.6.0",
-    install_requires=["numpy>=1.17", "packaging>=20.0", "pyyaml", "torch>=1.4.0"],
+    install_requires=["numpy>=1.17", "packaging>=20.0", "psutil", "pyyaml", "torch>=1.4.0"],
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/408
diff --git a/tests/test_examples.py b/tests/test_examples.py
index 6ec2a143a..3c875e02d 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -21,6 +21,7 @@
 import unittest
 from unittest import mock
 
+from accelerate import Accelerator
 from accelerate.test_utils.examples import compare_against_test
 from accelerate.test_utils.testing import TempDirTestCase, slow
 from accelerate.utils import write_basic_config
@@ -142,11 +143,11 @@ def test_checkpointing_by_epoch(self):
     def test_checkpointing_by_steps(self):
         testargs = f"""
         examples/by_feature/checkpointing.py
-        --checkpointing_steps 2
+        --checkpointing_steps 1
         --output_dir {self.tmpdir}
         """.split()
         _ = subprocess.run(self._launch_args + testargs, stdout=subprocess.PIPE, env=os.environ)
-        self.assertTrue(os.path.exists(os.path.join(self.tmpdir, "step_4")))
+        self.assertTrue(os.path.exists(os.path.join(self.tmpdir, "step_5")))
 
     def test_load_states_by_epoch(self):
         testargs = f"""
@@ -163,14 +164,20 @@ def test_load_states_by_epoch(self):
     def test_load_states_by_steps(self):
         testargs = f"""
         examples/by_feature/checkpointing.py
-        --resume_from_checkpoint {os.path.join(self.tmpdir, "step_4")}
+        --resume_from_checkpoint {os.path.join(self.tmpdir, "step_5")}
         """.split()
         output = subprocess.run(
             self._launch_args + testargs, universal_newlines=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE
         ).stdout
-        self.assertNotIn("epoch 0:", output)
-        self.assertIn("epoch 1:", output)
-        self.assertIn("epoch 2:", output)
+        num_processes = Accelerator().num_processes
+        if num_processes > 1:
+            self.assertNotIn("epoch 0:", output)
+            self.assertNotIn("epoch 1:", output)
+            self.assertIn("epoch 2:", output)
+        else:
+            self.assertNotIn("epoch 0:", output)
+            self.assertIn("epoch 1:", output)
+            self.assertIn("epoch 2:", output)
 
     @slow
     def test_cross_validation(self):

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/407
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index d2de25a79..c585e742a 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -13,6 +13,7 @@ jobs:
       matrix:
         test-kind: [
           test,
+          test_deepspeed,
           test_example_differences,
           test_checkpoint_step,
           test_checkpoint_epoch,
diff --git a/Makefile b/Makefile
index 38cb16cfc..58cb89ab9 100644
--- a/Makefile
+++ b/Makefile
@@ -27,6 +27,9 @@ style:
 test:
 	python -m pytest -s -v ./tests/ --ignore=./tests/test_examples.py
 
+test_deepspeed:
+	python -m pytest -s -v ./tests/deepspeed
+
 test_examples:
 	python -m pytest -s -v ./tests/test_examples.py
 
diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml
index 130659196..4b2dfac18 100644
--- a/docs/source/_toctree.yml
+++ b/docs/source/_toctree.yml
@@ -29,4 +29,6 @@
     title: Fully Sharded Data Parallel
   - local: memory
     title: Memory Utilities
+  - local: deepspeed
+    title: DeepSpeed
   title: API Reference
diff --git a/docs/source/deepspeed.mdx b/docs/source/deepspeed.mdx
new file mode 100644
index 000000000..c0385c352
--- /dev/null
+++ b/docs/source/deepspeed.mdx
@@ -0,0 +1,508 @@
+<!--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.
+-->
+
+# DeepSpeed 
+
+[DeepSpeed](https://github.com/microsoft/DeepSpeed) implements everything described in the [ZeRO paper](https://arxiv.org/abs/1910.02054). Currently it provides full support for:
+
+1. Optimizer state partitioning (ZeRO stage 1)
+2. Gradient partitioning (ZeRO stage 2)
+3. Parameter partitioning (ZeRO stage 3)
+4. Custom mixed precision training handling
+5. A range of fast CUDA-extension-based optimizers
+6. ZeRO-Offload to CPU and Disk/NVMe
+
+ZeRO-Offload has its own dedicated paper: [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://arxiv.org/abs/2101.06840). And NVMe-support is described in the paper [ZeRO-Infinity: Breaking the GPU
+Memory Wall for Extreme Scale Deep Learning](https://arxiv.org/abs/2104.07857).
+
+DeepSpeed ZeRO-2 is primarily used only for training, as its features are of no use to inference.
+
+DeepSpeed ZeRO-3 can be used for inference as well, since it allows huge models to be loaded on multiple GPUs, which
+won't be possible on a single GPU.
+
+🤗 Accelerate integrates [DeepSpeed](https://github.com/microsoft/DeepSpeed) via 2 options:
+
+1. Integration of the DeepSpeed features via `deepspeed config file` specification in `accelerate config` . You just supply your custom config file or use our template. Most of
+   this document is focused on this feature. This supports all the core features of DeepSpeed and gives user a lot of flexibility. 
+   User may have to change few lines of code depending on the config.
+2. Integration via `deepspeed_plugin`.This supports subset of the DeepSpeed features and uses default options for the rest of the configurations. 
+   User need not change any code and is good for those who are fine with most of the default settings of DeepSpeed.
+
+## What is integrated?
+
+Training:
+
+1. DeepSpeed ZeRO training supports the full ZeRO stages 1, 2 and 3 as well as CPU/Disk offload of optimizer states, gradients and parameters. 
+Below is a short description of Data Parallelism using ZeRO - Zero Redundancy Optimizer along with diagram from this [blog post](https://www.microsoft.com/en-us/research/blog/zero-deepspeed-new-system-optimizations-enable-training-models-with-over-100-billion-parameters/)
+![ZeRO Data Parallelism](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/parallelism-zero.png)
+
+(Source: [link](https://www.microsoft.com/en-us/research/blog/zero-deepspeed-new-system-optimizations-enable-training-models-with-over-100-billion-parameters/))
+
+ a. **Stage 1** : Shards optimizer states across data parallel workers/GPUs
+
+ b. **Stage 2** : Shards optimizer states + gradients across data parallel workers/GPUs
+
+ c. **Stage 3**: Shards optimizer states + gradients + model parameters across data parallel workers/GPUs
+
+ d. **Optimizer Offload**: Offloads the gradients + optimizer states to CPU/Disk building on top of ZERO Stage 2
+
+ e. **Param Offload**: Offloads the model parameters to CPU/Disk building on top of ZERO Stage 3
+
+<u>Note</u>: With respect to Disk Offload, the disk should be an NVME for decent speed but it technically work on any Disk
+
+Inference:
+
+1. DeepSpeed ZeRO Inference supports ZeRO stage 3 with ZeRO-Infinity. It uses the same ZeRO protocol as training, but
+   it doesn't use an optimizer and a lr scheduler and only stage 3 is relevant. For more details see:
+   [deepspeed-zero-inference](#deepspeed-zero-inference).
+
+
+## How it works?
+
+**Pre-Requisites**: Install DeepSpeed version >=0.6.5. Please refer to the [DeepSpeed Insallation details](https://github.com/microsoft/DeepSpeed#installation)
+for more information.
+
+We will first look at easy to use integration via `accelerate config`. 
+Followed by more flexible and feature rich `deepspeed config file` integration. 
+
+### Accelerate DeepSpeed Plugin
+On your machine(s) just run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked. It will ask whether you want to use a config file for DeepSpeed to which you should answer no. Then answer the following questions to generate a basic DeepSpeed config.
+This will generate a config file that will be used automatically to properly set the
+default options when doing
+
+```bash
+accelerate launch my_script.py --args_to_my_script
+```
+
+For instance, here is how you would run the NLP example `examples/nlp_example.py` (from the root of the repo) with DeepSpeed Plugin:
+
+**ZeRO Stage-2 DeepSpeed Plugin Example**
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+ gradient_accumulation_steps: 1
+ gradient_clipping: 1.0
+ offload_optimizer_device: none
+ offload_param_device: none
+ zero3_init_flag: true
+ zero_stage: 2
+distributed_type: DEEPSPEED
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 2
+use_cpu: false
+```
+
+```bash
+accelerate launch examples/nlp_example.py --mixed_precision fp16
+```
+
+**ZeRO Stage-3 with CPU Offload DeepSpeed Plugin Example**
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+  gradient_accumulation_steps: 1
+  gradient_clipping: 1.0
+  offload_optimizer_device: cpu
+  offload_param_device: cpu
+  zero3_init_flag: true
+  zero3_save_16bit_model: true
+  zero_stage: 3
+distributed_type: DEEPSPEED
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 2
+use_cpu: false
+```
+
+```bash
+accelerate launch examples/nlp_example.py --mixed_precision fp16
+```
+
+Currently, `Accelerate` supports following config through the CLI:
+
+```bash
+`zero_stage`: [0] Disabled, [1] optimizer state partitioning, [2] optimizer+gradient state partitioning and [3] optimizer+gradient+parameter partitioning
+`gradient_accumulation_steps`: Number of training steps to accumulate gradients before averaging and applying them.
+`gradient_clipping`: Enable gradient clipping with value.
+`offload_optimizer_device`: [none] Disable optimizer offloading, [cpu] offload optimizer to CPU, [nvme] offload optimizer to NVMe SSD. Only applicable with ZeRO >= Stage-2.
+`offload_param_device`: [none] Disable parameter offloading, [cpu] offload parameters to CPU, [nvme] offload parameters to NVMe SSD. Only applicable with ZeRO Stage-3.
+`zero3_init_flag`: Decides whether to enable `deepspeed.zero.Init` for constructing massive models. Only applicable with ZeRO Stage-3.
+`zero3_save_16bit_model`: Decides whether to save 16-bit model weights when using ZeRO Stage-3.
+`mixed_precision`: `no` for FP32 training, `fp16` for FP16 mixed-precision training and `bf16` for BF16 mixed-precision training. 
+```
+To be able to tweak more options, you will need to use a DeepSpeed config file.
+
+### DeepSpeed Config File
+On your machine(s) just run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked. It will ask whether you want to use a config file for deepspeed to which you answer yes 
+and provide the path to the deepspeed config file. 
+This will generate a config file that will be used automatically to properly set the
+default options when doing
+
+```bash
+accelerate launch my_script.py --args_to_my_script
+```
+
+For instance, here is how you would run the NLP example `examples/by_feature/deepspeed_with_config_support.py` (from the root of the repo) with DeepSpeed Config File:
+
+**ZeRO Stage-2 DeepSpeed Config File Example**
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+ deepspeed_config_file: /home/ubuntu/accelerate/examples/configs/deepspeed_config_templates/zero_stage2_config.json
+ zero3_init_flag: true
+distributed_type: DEEPSPEED
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 2
+use_cpu: false
+```
+
+with the contents of `zero_stage2_config.json` being:
+```json
+{
+    "fp16": {
+        "enabled": true,
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "optimizer": {
+        "type": "AdamW",
+        "params": {
+            "lr": "auto",
+            "weight_decay": "auto",
+            "torch_adam": true,
+            "adam_w_mode": true
+        }
+    },
+    "scheduler": {
+        "type": "WarmupDecayLR",
+        "params": {
+            "warmup_min_lr": "auto",
+            "warmup_max_lr": "auto",
+            "warmup_num_steps": "auto",
+            "total_num_steps": "auto"
+        }
+    },
+    "zero_optimization": {
+        "stage": 2,
+        "allgather_partitions": true,
+        "allgather_bucket_size": 2e8,
+        "overlap_comm": true,
+        "reduce_scatter": true,
+        "reduce_bucket_size": "auto",
+        "contiguous_gradients": true
+    },
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}
+```
+
+```bash
+accelerate launch examples/by_feature/deepspeed_with_config_support.py \
+--config_name "gpt2-large" \
+--tokenizer_name "gpt2-large" \
+--dataset_name "wikitext" \
+--dataset_config_name "wikitext-2-raw-v1" \
+--block_size 128 \
+--output_dir "./clm/clm_deepspeed_stage2_accelerate" \
+--learning_rate 5e-4 \
+--per_device_train_batch_size 24 \
+--per_device_eval_batch_size 24 \
+--num_train_epochs 3 \
+--with_tracking \
+--report_to "wandb"\
+```
+
+**ZeRO Stage-3 with CPU offload DeepSpeed Config File Example**
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+ deepspeed_config_file: /home/ubuntu/accelerate/examples/configs/deepspeed_config_templates/zero_stage3_offload_config.json
+ zero3_init_flag: true
+distributed_type: DEEPSPEED
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 2
+use_cpu: false
+```
+with the contents of `zero_stage3_offload_config.json` being:
+```json
+{
+    "fp16": {
+        "enabled": true,
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "optimizer": {
+        "type": "AdamW",
+        "params": {
+            "lr": "auto",
+            "weight_decay": "auto"
+        }
+    },
+    "scheduler": {
+        "type": "WarmupDecayLR",
+        "params": {
+            "warmup_min_lr": "auto",
+            "warmup_max_lr": "auto",
+            "warmup_num_steps": "auto",
+            "total_num_steps": "auto"
+        }
+    },
+    "zero_optimization": {
+        "stage": 3,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "offload_param": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "reduce_bucket_size": "auto",
+        "stage3_prefetch_bucket_size": "auto",
+        "stage3_param_persistence_threshold": "auto",
+        "sub_group_size": 1e9,
+        "stage3_max_live_parameters": 1e9,
+        "stage3_max_reuse_distance": 1e9,
+        "stage3_gather_16bit_weights_on_model_save": "auto"
+    },
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}
+```
+
+```bash
+accelerate launch examples/by_feature/deepspeed_with_config_support.py \
+--config_name "gpt2-large" \
+--tokenizer_name "gpt2-large" \
+--dataset_name "wikitext" \
+--dataset_config_name "wikitext-2-raw-v1" \
+--block_size 128 \
+--output_dir "./clm/clm_deepspeed_stage3_offload_accelerate" \
+--learning_rate 5e-4 \
+--per_device_train_batch_size 32 \
+--per_device_eval_batch_size 32 \
+--num_train_epochs 3 \
+--with_tracking \
+--report_to "wandb"\
+```
+
+**Important code changes when using DeepSpeed Config File**
+
+1. DeepSpeed Optimizers and Schedulers. For more information on these, 
+see the [DeepSpeed Optimizers](https://deepspeed.readthedocs.io/en/latest/optimizers.html) and [DeepSpeed Schedulers](https://deepspeed.readthedocs.io/en/latest/schedulers.html) documentation.
+We will look at the changes needed in the code when using these.
+   
+   a. DS Optim + DS Scheduler: The case when both `optimizer` and `scheduler` keys present in the DeepSpeed config file.
+   In this situation, those will be used and user has to use `accelerate.utils.DummyOptim` and `accelerate.utils.DummyScheduler` to replace the PyTorch/Custom optimizers and schedulers in their code.
+   Below is the snippet from `examples/by_feature/deepspeed_with_config_support.py` showing this:
+   ```python
+    # Creates Dummy Optimizer if `optimizer` was spcified in the config file else creates Adam Optimizer
+    optimizer_cls = (
+        torch.optim.AdamW
+        if accelerator.state.deepspeed_plugin is None
+        or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config
+        else DummyOptim
+    )
+    optimizer = optimizer_cls(optimizer_grouped_parameters, lr=args.learning_rate)
+
+    # Creates Dummy Scheduler if `scheduler` was spcified in the config file else creates `args.lr_scheduler_type` Scheduler
+    if (
+        accelerator.state.deepspeed_plugin is None
+        or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
+    ):
+        lr_scheduler = get_scheduler(
+            name=args.lr_scheduler_type,
+            optimizer=optimizer,
+            num_warmup_steps=args.num_warmup_steps,
+            num_training_steps=args.max_train_steps,
+        )
+    else:
+        lr_scheduler = DummyScheduler(
+            optimizer, total_num_steps=args.max_train_steps, warmup_num_steps=args.num_warmup_steps
+        )
+   ```
+   b. Custom Optim + Custom Scheduler: The case when both `optimizer` and `scheduler` keys are absent in the DeepSpeed config file.
+   In this situation, no code changes are needed from the user and this is the case when using integration via DeepSpeed Plugin.
+   In the above example we can see that the code reamins unchanged if the `optimizer` and `scheduler` keys are absent in the DeepSpeed config file.
+
+   c. Custom Optim + DS Scheduler: The case when only `scheduler` key is present in the DeepSpeed config file. 
+   In this situation, user has to use `accelerate.utils.DummyScheduler` to replace the PyTorch/Custom scheduler in their code. 
+
+   d. DS Optim + Custom Scheduler: The case when only `optimizer` key is present in the DeepSpeed config file. 
+   This will result in an error because one can only use DS Scheduler when using DS Optim.
+
+2. Notice the `auto` values in the above example DeepSpeed config files. These are automatically handled by `prepare` method 
+based on model, dataloaders, dummy optimizer and dummy schedulers provided to `prepare` method. 
+Only the `auto` fields specified in above examples are handled by `prepare` method and the rest have to be explicitly specified by the user.
+
+## Saving and loading
+
+1. Saving and loading of models is unchanged for ZeRO Stage-1 and Stage-2.
+
+2. under ZeRO Stage-3, `state_dict` contains just the placeholders since the model weights are partitioned across multiple GPUs.
+ZeRO Stage-3 has 2 options:
+
+   a. Saving the entire 16bit model weights to directly load later on using `model.load_state_dict(torch.load(pytorch_model.bin))`.
+   For this, either set `zero_optimization.stage3_gather_16bit_weights_on_model_save` to True in DeepSpeed Config file or set
+   `zero3_save_16bit_model` to True in DeepSpeed Plugin. 
+   **Note that this option requires consolidation of the weights on one GPU it can be slow and memory demanding, so only use this feature when needed.**
+   Below is the snippet from `examples/by_feature/deepspeed_with_config_support.py` showing this:
+   ```python
+   unwrapped_model = accelerator.unwrap_model(model)
+
+   # New Code #
+   # Saves the whole/unpartitioned fp16 model when in ZeRO Stage-3 to the output directory if
+   # `stage3_gather_16bit_weights_on_model_save` is True in DeepSpeed Config file or
+   # `zero3_save_16bit_model` is True in DeepSpeed Plugin.
+   # For Zero Stages 1 and 2, models are saved as usual in the output directory.
+   # The model name saved is `pytorch_model.bin`
+   unwrapped_model.save_pretrained(
+       args.output_dir,
+       is_main_process=accelerator.is_main_process,
+       save_function=accelerator.save,
+       state_dict=accelerator.get_state_dict(model),
+   )
+   ```
+
+   b. To get 32bit weights, first save the model using `model.save_checkpoint()`.
+   Below is the snippet from `examples/by_feature/deepspeed_with_config_support.py` showing this:
+   ```python
+   success = model.save_checkpoint(PATH, ckpt_id, checkpoint_state_dict)
+   status_msg = "checkpointing: PATH={}, ckpt_id={}".format(PATH, ckpt_id)
+   if success:
+       logging.info(f"Success {status_msg}")
+   else:
+       logging.warning(f"Failure {status_msg}")
+   ``` 
+   This will create ZeRO model and optimizer partitions along with `zero_to_fp32.py` script in checkpoint directory.
+   One can use this script to do offline consolidation.  
+   It requires no configuration files or GPUs. Here is an example of its usage:  
+   ```bash
+   $ cd /path/to/checkpoint_dir
+   $ ./zero_to_fp32.py . pytorch_model.bin
+   Processing zero checkpoint at global_step1
+   Detected checkpoint of type zero stage 3, world_size: 2
+   Saving fp32 state dict to pytorch_model.bin (total_numel=60506624)
+   ```
+   To get 32bit model for saving/inference, one can do the following:
+   ```python
+   from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
+
+   unwrapped_model = accelerator.unwrap_model(model)
+   fp32_model = load_state_dict_from_zero_checkpoint(unwrapped_model, checkpoint_dir)
+   ```
+   If only interested in state_dict, one can do the following:
+   ```python
+   from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+
+   state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir)
+   ```
+   Note that all these functions require ~2x memory (general RAM) of the size of the final checkpoint.
+
+## ZeRO Inference
+DeepSpeed ZeRO Inference supports ZeRO stage 3 with ZeRO-Infinity. 
+It uses the same ZeRO protocol as training, but it doesn't use an optimizer and a lr scheduler and only stage 3 is relevant.
+With accelerate integration, one has to just prepare model and dataloader as shown below:
+
+```python
+model, eval_dataloader = accelerator.prepare(model, eval_dataloader)
+```
+
+## Few caveats to be aware of 
+
+1. Current integration doesn’t support Pipeline Parallelism of DeepSpeed.
+2. Current integration doesn’t support `mpu`, limiting the tensor parallelism which is supported in Megatron-LM. 
+3. Current integration doesn’t support multiple models for a given `accelerator` object. 
+
+
+## Internals
+
+[[autodoc]] utils.DeepSpeedPlugin
+
+[[autodoc]] utils.DummyOptim
+
+[[autodoc]] utils.DummyScheduler
+
+[[autodoc]] utils.DeepSpeedEngineWrapper
+
+[[autodoc]] utils.DeepSpeedOptimizerWrapper
+
+[[autodoc]] utils.DeepSpeedSchedulerWrapper
+
+
+## Main DeepSpeed Resources
+
+- [Project's github](https://github.com/microsoft/deepspeed)
+- [Usage docs](https://www.deepspeed.ai/getting-started/)
+- [API docs](https://deepspeed.readthedocs.io/en/latest/index.html)
+- [Blog posts](https://www.microsoft.com/en-us/research/search/?q=deepspeed)
+
+Papers:
+
+- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://arxiv.org/abs/1910.02054)
+- [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://arxiv.org/abs/2101.06840)
+- [ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning](https://arxiv.org/abs/2104.07857)
+
+Finally, please, remember that, 🤗 `Accelerate` only integrates DeepSpeed, therefore if you
+have any problems or questions with regards to DeepSpeed usage, please, file an issue with [DeepSpeed GitHub](https://github.com/microsoft/DeepSpeed/issues).
+
diff --git a/examples/by_feature/deepspeed_with_config_support.py b/examples/by_feature/deepspeed_with_config_support.py
new file mode 100755
index 000000000..fd8a8fa82
--- /dev/null
+++ b/examples/by_feature/deepspeed_with_config_support.py
@@ -0,0 +1,736 @@
+#!/usr/bin/env python
+# coding=utf-8
+# Copyright 2022 The HuggingFace Inc. 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.
+"""
+Fine-tuning the library models for causal language modeling (GPT, GPT-2, CTRL, ...)
+on a text file or a dataset without using HuggingFace Trainer.
+
+Here is the full list of checkpoints on the hub that can be fine-tuned by this script:
+https://huggingface.co/models?filter=text-generation
+"""
+# You can also adapt this script on your own causal language modeling task. Pointers for this are left as comments.
+
+import argparse
+import json
+import logging
+import math
+import os
+import random
+from itertools import chain
+from pathlib import Path
+
+import torch
+from torch.utils.data import DataLoader
+
+import datasets
+import transformers
+from accelerate import Accelerator, DistributedType
+from accelerate.logging import get_logger
+from accelerate.utils import DummyOptim, DummyScheduler, set_seed
+from datasets import load_dataset
+from huggingface_hub import Repository
+from tqdm.auto import tqdm
+from transformers import (
+    CONFIG_MAPPING,
+    MODEL_MAPPING,
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    SchedulerType,
+    default_data_collator,
+    get_scheduler,
+)
+from transformers.utils import get_full_repo_name
+from transformers.utils.versions import require_version
+
+
+logger = get_logger(__name__)
+
+require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt")
+
+MODEL_CONFIG_CLASSES = list(MODEL_MAPPING.keys())
+MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Finetune a transformers model on a causal language modeling task")
+    parser.add_argument(
+        "--dataset_name",
+        type=str,
+        default=None,
+        help="The name of the dataset to use (via the datasets library).",
+    )
+    parser.add_argument(
+        "--dataset_config_name",
+        type=str,
+        default=None,
+        help="The configuration name of the dataset to use (via the datasets library).",
+    )
+    parser.add_argument(
+        "--train_file", type=str, default=None, help="A csv or a json file containing the training data."
+    )
+    parser.add_argument(
+        "--validation_file", type=str, default=None, help="A csv or a json file containing the validation data."
+    )
+    parser.add_argument(
+        "--validation_split_percentage",
+        default=5,
+        help="The percentage of the train set used as validation set in case there's no validation split",
+    )
+    parser.add_argument(
+        "--model_name_or_path",
+        type=str,
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+        required=False,
+    )
+    parser.add_argument(
+        "--config_name",
+        type=str,
+        default=None,
+        help="Pretrained config name or path if not the same as model_name",
+    )
+    parser.add_argument(
+        "--tokenizer_name",
+        type=str,
+        default=None,
+        help="Pretrained tokenizer name or path if not the same as model_name",
+    )
+    parser.add_argument(
+        "--use_slow_tokenizer",
+        action="store_true",
+        help="If passed, will use a slow tokenizer (not backed by the 🤗 Tokenizers library).",
+    )
+    parser.add_argument(
+        "--per_device_train_batch_size",
+        type=int,
+        default=8,
+        help="Batch size (per device) for the training dataloader.",
+    )
+    parser.add_argument(
+        "--per_device_eval_batch_size",
+        type=int,
+        default=8,
+        help="Batch size (per device) for the evaluation dataloader.",
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=5e-5,
+        help="Initial learning rate (after the potential warmup period) to use.",
+    )
+    parser.add_argument("--weight_decay", type=float, default=0.0, help="Weight decay to use.")
+    parser.add_argument("--num_train_epochs", type=int, default=3, help="Total number of training epochs to perform.")
+    parser.add_argument(
+        "--max_train_steps",
+        type=int,
+        default=None,
+        help="Total number of training steps to perform. If provided, overrides num_train_epochs.",
+    )
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help="Number of updates steps to accumulate before performing a backward/update pass.",
+    )
+    parser.add_argument(
+        "--lr_scheduler_type",
+        type=SchedulerType,
+        default="linear",
+        help="The scheduler type to use.",
+        choices=["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"],
+    )
+    parser.add_argument(
+        "--num_warmup_steps", type=int, default=0, help="Number of steps for the warmup in the lr scheduler."
+    )
+    parser.add_argument("--output_dir", type=str, default=None, help="Where to store the final model.")
+    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
+    parser.add_argument(
+        "--model_type",
+        type=str,
+        default=None,
+        help="Model type to use if training from scratch.",
+        choices=MODEL_TYPES,
+    )
+    parser.add_argument(
+        "--block_size",
+        type=int,
+        default=None,
+        help=(
+            "Optional input sequence length after tokenization. The training dataset will be truncated in block of"
+            " this size for training. Default to the model max input length for single sentence inputs (take into"
+            " account special tokens)."
+        ),
+    )
+    parser.add_argument(
+        "--preprocessing_num_workers",
+        type=int,
+        default=None,
+        help="The number of processes to use for the preprocessing.",
+    )
+    parser.add_argument(
+        "--overwrite_cache", type=bool, default=False, help="Overwrite the cached training and evaluation sets"
+    )
+    parser.add_argument(
+        "--no_keep_linebreaks", action="store_true", help="Do not keep line breaks when using TXT files."
+    )
+    parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")
+    parser.add_argument(
+        "--hub_model_id", type=str, help="The name of the repository to keep in sync with the local `output_dir`."
+    )
+    parser.add_argument("--hub_token", type=str, help="The token to use to push to the Model Hub.")
+    parser.add_argument(
+        "--checkpointing_steps",
+        type=str,
+        default=None,
+        help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.",
+    )
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        default=None,
+        help="If the training should continue from a checkpoint folder.",
+    )
+    # New Code #
+    # Whether to load the best model at the end of training
+    parser.add_argument(
+        "--load_best_model",
+        action="store_true",
+        help="Whether to load the best model at the end of training",
+    )
+    parser.add_argument(
+        "--with_tracking",
+        action="store_true",
+        help="Whether to enable experiment trackers for logging.",
+    )
+    parser.add_argument(
+        "--report_to",
+        type=str,
+        default="all",
+        help=(
+            'The integration to report the results and logs to. Supported platforms are `"tensorboard"`,'
+            ' `"wandb"` and `"comet_ml"`. Use `"all"` (default) to report to all integrations.'
+            "Only applicable when `--with_tracking` is passed."
+        ),
+    )
+    args = parser.parse_args()
+
+    # Sanity checks
+    if args.dataset_name is None and args.train_file is None and args.validation_file is None:
+        raise ValueError("Need either a dataset name or a training/validation file.")
+    else:
+        if args.train_file is not None:
+            extension = args.train_file.split(".")[-1]
+            assert extension in ["csv", "json", "txt"], "`train_file` should be a csv, json or txt file."
+        if args.validation_file is not None:
+            extension = args.validation_file.split(".")[-1]
+            assert extension in ["csv", "json", "txt"], "`validation_file` should be a csv, json or txt file."
+
+    if args.push_to_hub:
+        assert args.output_dir is not None, "Need an `output_dir` to create a repo when `--push_to_hub` is passed."
+
+    return args
+
+
+# New Code #
+def checkpoint_model(checkpoint_folder, ckpt_id, model, epoch, last_global_step, **kwargs):
+    """Utility function for checkpointing model + optimizer dictionaries
+    The main purpose for this is to be able to resume training from that instant again
+    """
+    checkpoint_state_dict = {
+        "epoch": epoch,
+        "last_global_step": last_global_step,
+    }
+    # Add extra kwargs too
+    checkpoint_state_dict.update(kwargs)
+
+    success = model.save_checkpoint(checkpoint_folder, ckpt_id, checkpoint_state_dict)
+    status_msg = f"checkpointing: checkpoint_folder={checkpoint_folder}, ckpt_id={ckpt_id}"
+    if success:
+        logging.info(f"Success {status_msg}")
+    else:
+        logging.warning(f"Failure {status_msg}")
+    return
+
+
+# New Code #
+def load_training_checkpoint(model, load_dir, tag=None, **kwargs):
+    """Utility function for checkpointing model + optimizer dictionaries
+    The main purpose for this is to be able to resume training from that instant again
+    """
+    _, checkpoint_state_dict = model.load_checkpoint(load_dir, tag=tag, **kwargs)
+    epoch = checkpoint_state_dict["epoch"]
+    last_global_step = checkpoint_state_dict["last_global_step"]
+    del checkpoint_state_dict
+    return (epoch, last_global_step)
+
+
+# New Code #
+def evaluate(args, model, eval_dataloader, accelerator, eval_dataset):
+    model.eval()
+    losses = []
+    for step, batch in enumerate(eval_dataloader):
+        with torch.no_grad():
+            outputs = model(**batch)
+
+        loss = outputs.loss
+        losses.append(accelerator.gather(loss.repeat(args.per_device_eval_batch_size)))
+
+    losses = torch.cat(losses)
+    losses = losses[: len(eval_dataset)]
+    try:
+        eval_loss = torch.mean(losses)
+        perplexity = math.exp(eval_loss)
+    except OverflowError:
+        perplexity = float("inf")
+    return perplexity, eval_loss
+
+
+def main():
+    args = parse_args()
+
+    # Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
+    # If we're using tracking, we also need to initialize it here and it will by default pick up all supported trackers
+    # in the environment
+    accelerator = (
+        Accelerator(log_with=args.report_to, logging_dir=args.output_dir) if args.with_tracking else Accelerator()
+    )
+    # Make one log on every process with the configuration for debugging.
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO,
+    )
+    logger.info(accelerator.state, main_process_only=False)
+    if accelerator.is_local_main_process:
+        datasets.utils.logging.set_verbosity_warning()
+        transformers.utils.logging.set_verbosity_info()
+    else:
+        datasets.utils.logging.set_verbosity_error()
+        transformers.utils.logging.set_verbosity_error()
+
+    # If passed along, set the training seed now.
+    if args.seed is not None:
+        set_seed(args.seed)
+
+    # Handle the repository creation
+    if accelerator.is_main_process:
+        if args.push_to_hub:
+            if args.hub_model_id is None:
+                repo_name = get_full_repo_name(Path(args.output_dir).name, token=args.hub_token)
+            else:
+                repo_name = args.hub_model_id
+            repo = Repository(args.output_dir, clone_from=repo_name)
+
+            with open(os.path.join(args.output_dir, ".gitignore"), "w+") as gitignore:
+                if "step_*" not in gitignore:
+                    gitignore.write("step_*\n")
+                if "epoch_*" not in gitignore:
+                    gitignore.write("epoch_*\n")
+        elif args.output_dir is not None:
+            os.makedirs(args.output_dir, exist_ok=True)
+    accelerator.wait_for_everyone()
+
+    # Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)
+    # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
+    # (the dataset will be downloaded automatically from the datasets Hub).
+    #
+    # For CSV/JSON files, this script will use the column called 'text' or the first column if no column called
+    # 'text' is found. You can easily tweak this behavior (see below).
+    #
+    # In distributed training, the load_dataset function guarantee that only one local process can concurrently
+    # download the dataset.
+    if args.dataset_name is not None:
+        # Downloading and loading a dataset from the hub.
+        raw_datasets = load_dataset(args.dataset_name, args.dataset_config_name)
+        if "validation" not in raw_datasets.keys():
+            raw_datasets["validation"] = load_dataset(
+                args.dataset_name,
+                args.dataset_config_name,
+                split=f"train[:{args.validation_split_percentage}%]",
+            )
+            raw_datasets["train"] = load_dataset(
+                args.dataset_name,
+                args.dataset_config_name,
+                split=f"train[{args.validation_split_percentage}%:]",
+            )
+    else:
+        data_files = {}
+        dataset_args = {}
+        if args.train_file is not None:
+            data_files["train"] = args.train_file
+        if args.validation_file is not None:
+            data_files["validation"] = args.validation_file
+        extension = args.train_file.split(".")[-1]
+        if extension == "txt":
+            extension = "text"
+            dataset_args["keep_linebreaks"] = not args.no_keep_linebreaks
+        raw_datasets = load_dataset(extension, data_files=data_files, **dataset_args)
+        # If no validation data is there, validation_split_percentage will be used to divide the dataset.
+        if "validation" not in raw_datasets.keys():
+            raw_datasets["validation"] = load_dataset(
+                extension,
+                data_files=data_files,
+                split=f"train[:{args.validation_split_percentage}%]",
+                **dataset_args,
+            )
+            raw_datasets["train"] = load_dataset(
+                extension,
+                data_files=data_files,
+                split=f"train[{args.validation_split_percentage}%:]",
+                **dataset_args,
+            )
+
+    # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
+    # https://huggingface.co/docs/datasets/loading_datasets.html.
+
+    # Load pretrained model and tokenizer
+    #
+    # In distributed training, the .from_pretrained methods guarantee that only one local process can concurrently
+    # download model & vocab.
+    if args.config_name:
+        config = AutoConfig.from_pretrained(args.config_name)
+    elif args.model_name_or_path:
+        config = AutoConfig.from_pretrained(args.model_name_or_path)
+    else:
+        config = CONFIG_MAPPING[args.model_type]()
+        logger.warning("You are instantiating a new config instance from scratch.")
+
+    if args.tokenizer_name:
+        tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name, use_fast=not args.use_slow_tokenizer)
+    elif args.model_name_or_path:
+        tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, use_fast=not args.use_slow_tokenizer)
+    else:
+        raise ValueError(
+            "You are instantiating a new tokenizer from scratch. This is not supported by this script."
+            "You can do it from another script, save it, and load it from here, using --tokenizer_name."
+        )
+
+    if args.model_name_or_path:
+        model = AutoModelForCausalLM.from_pretrained(
+            args.model_name_or_path,
+            from_tf=bool(".ckpt" in args.model_name_or_path),
+            config=config,
+        )
+    else:
+        logger.info("Training new model from scratch")
+        model = AutoModelForCausalLM.from_config(config)
+
+    model.resize_token_embeddings(len(tokenizer))
+
+    # Preprocessing the datasets.
+    # First we tokenize all the texts.
+    column_names = raw_datasets["train"].column_names
+    text_column_name = "text" if "text" in column_names else column_names[0]
+
+    def tokenize_function(examples):
+        return tokenizer(examples[text_column_name])
+
+    with accelerator.main_process_first():
+        tokenized_datasets = raw_datasets.map(
+            tokenize_function,
+            batched=True,
+            num_proc=args.preprocessing_num_workers,
+            remove_columns=column_names,
+            load_from_cache_file=not args.overwrite_cache,
+            desc="Running tokenizer on dataset",
+        )
+
+    if args.block_size is None:
+        block_size = tokenizer.model_max_length
+        if block_size > 1024:
+            logger.warning(
+                f"The tokenizer picked seems to have a very large `model_max_length` ({tokenizer.model_max_length}). "
+                "Picking 1024 instead. You can change that default value by passing --block_size xxx."
+            )
+        block_size = 1024
+    else:
+        if args.block_size > tokenizer.model_max_length:
+            logger.warning(
+                f"The block_size passed ({args.block_size}) is larger than the maximum length for the model"
+                f"({tokenizer.model_max_length}). Using block_size={tokenizer.model_max_length}."
+            )
+        block_size = min(args.block_size, tokenizer.model_max_length)
+
+    # Main data processing function that will concatenate all texts from our dataset and generate chunks of block_size.
+    def group_texts(examples):
+        # Concatenate all texts.
+        concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}
+        total_length = len(concatenated_examples[list(examples.keys())[0]])
+        # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
+        # customize this part to your needs.
+        if total_length >= block_size:
+            total_length = (total_length // block_size) * block_size
+        # Split by chunks of max_len.
+        result = {
+            k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
+            for k, t in concatenated_examples.items()
+        }
+        result["labels"] = result["input_ids"].copy()
+        return result
+
+    # Note that with `batched=True`, this map processes 1,000 texts together, so group_texts throws away a remainder
+    # for each of those groups of 1,000 texts. You can adjust that batch_size here but a higher value might be slower
+    # to preprocess.
+    #
+    # To speed up this part, we use multiprocessing. See the documentation of the map method for more information:
+    # https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasets.Dataset.map
+
+    with accelerator.main_process_first():
+        lm_datasets = tokenized_datasets.map(
+            group_texts,
+            batched=True,
+            num_proc=args.preprocessing_num_workers,
+            load_from_cache_file=not args.overwrite_cache,
+            desc=f"Grouping texts in chunks of {block_size}",
+        )
+
+    train_dataset = lm_datasets["train"]
+    eval_dataset = lm_datasets["validation"]
+
+    # Log a few random samples from the training set:
+    for index in random.sample(range(len(train_dataset)), 3):
+        logger.info(f"Sample {index} of the training set: {train_dataset[index]}.")
+
+    # DataLoaders creation:
+    train_dataloader = DataLoader(
+        train_dataset, shuffle=True, collate_fn=default_data_collator, batch_size=args.per_device_train_batch_size
+    )
+    eval_dataloader = DataLoader(
+        eval_dataset, collate_fn=default_data_collator, batch_size=args.per_device_eval_batch_size
+    )
+
+    # Optimizer
+    # Split weights in two groups, one with weight decay and the other not.
+    no_decay = ["bias", "LayerNorm.weight"]
+    optimizer_grouped_parameters = [
+        {
+            "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
+            "weight_decay": args.weight_decay,
+        },
+        {
+            "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
+            "weight_decay": 0.0,
+        },
+    ]
+    # New Code #
+    # Creates Dummy Optimizer if `optimizer` was spcified in the config file else creates Adam Optimizer
+    optimizer_cls = (
+        torch.optim.AdamW
+        if accelerator.state.deepspeed_plugin is None
+        or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config
+        else DummyOptim
+    )
+    optimizer = optimizer_cls(optimizer_grouped_parameters, lr=args.learning_rate)
+
+    # On TPU, the tie weights in our model have been disconnected, so we need to restore the ties.
+    if accelerator.distributed_type == DistributedType.TPU:
+        model.tie_weights()
+
+    # Scheduler and math around the number of training steps.
+
+    # New Code
+    # Get gradient accumulation steps from deepspeed config if available
+    if accelerator.state.deepspeed_plugin is not None:
+        args.gradient_accumulation_steps = accelerator.state.deepspeed_plugin.deepspeed_config[
+            "gradient_accumulation_steps"
+        ]
+
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    if args.max_train_steps is None:
+        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+    else:
+        args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+    # New Code #
+    # Creates Dummy Scheduler if `scheduler` was spcified in the config file else creates `args.lr_scheduler_type` Scheduler
+    if (
+        accelerator.state.deepspeed_plugin is None
+        or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
+    ):
+        lr_scheduler = get_scheduler(
+            name=args.lr_scheduler_type,
+            optimizer=optimizer,
+            num_warmup_steps=args.num_warmup_steps,
+            num_training_steps=args.max_train_steps,
+        )
+    else:
+        lr_scheduler = DummyScheduler(
+            optimizer, total_num_steps=args.max_train_steps, warmup_num_steps=args.num_warmup_steps
+        )
+
+    # Prepare everything with our `accelerator`.
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+
+    # Figure out how many steps we should save the Accelerator states
+    if hasattr(args.checkpointing_steps, "isdigit"):
+        checkpointing_steps = args.checkpointing_steps
+        if args.checkpointing_steps.isdigit():
+            checkpointing_steps = int(args.checkpointing_steps)
+    else:
+        checkpointing_steps = None
+
+    # We need to initialize the trackers we use, and also store our configuration.
+    # We initialize the trackers only on main process because `accelerator.log`
+    # only logs on main process and we don't want empty logs/runs on other processes.
+    if args.with_tracking:
+        if accelerator.is_main_process:
+            experiment_config = vars(args)
+            # TensorBoard cannot log Enums, need the raw value
+            experiment_config["lr_scheduler_type"] = experiment_config["lr_scheduler_type"].value
+            accelerator.init_trackers("clm_no_trainer", experiment_config)
+
+    # Train!
+    total_batch_size = args.per_device_train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+
+    logger.info("***** Running training *****")
+    logger.info(f"  Num examples = {len(train_dataset)}")
+    logger.info(f"  Num Epochs = {args.num_train_epochs}")
+    logger.info(f"  Instantaneous batch size per device = {args.per_device_train_batch_size}")
+    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
+    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
+    logger.info(f"  Total optimization steps = {args.max_train_steps}")
+    # Only show the progress bar once on each machine.
+    progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process)
+    completed_steps = 0
+    starting_epoch = 0
+    best_metric = None
+    best_metric_checkpoint = None
+
+    # Potentially load in the weights and states from a previous save
+    if args.resume_from_checkpoint:
+        # New Code #
+        # Loads the DeepSpeed checkpoint from the specified path
+        _, last_global_step = load_training_checkpoint(
+            model,
+            args.resume_from_checkpoint,
+            **{"load_optimizer_states": True, "load_lr_scheduler_states": True},
+        )
+        accelerator.print(f"Resumed from checkpoint: {args.resume_from_checkpoint}")
+        resume_step = last_global_step
+        starting_epoch = resume_step // len(train_dataloader)
+        resume_step -= starting_epoch * len(train_dataloader)
+
+    for epoch in range(starting_epoch, args.num_train_epochs):
+        model.train()
+        if args.with_tracking:
+            total_loss = 0
+        for step, batch in enumerate(train_dataloader):
+            # We need to skip steps until we reach the resumed step
+            if args.resume_from_checkpoint and epoch == starting_epoch:
+                if resume_step is not None and step < resume_step:
+                    completed_steps += 1
+                    continue
+            outputs = model(**batch)
+            loss = outputs.loss
+            # We keep track of the loss at each epoch
+            if args.with_tracking:
+                total_loss += loss.detach().float()
+            loss = loss / args.gradient_accumulation_steps
+            accelerator.backward(loss)
+            if step % args.gradient_accumulation_steps == 0 or step == len(train_dataloader) - 1:
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+                progress_bar.update(1)
+                completed_steps += 1
+
+            if isinstance(checkpointing_steps, int):
+                if completed_steps % checkpointing_steps == 0:
+                    output_dir = f"step_{completed_steps }"
+                    if args.output_dir is not None:
+                        output_dir = os.path.join(args.output_dir, output_dir)
+                    accelerator.save_state(output_dir)
+            if completed_steps >= args.max_train_steps:
+                break
+
+        perplexity, eval_loss = evaluate(args, model, eval_dataloader, accelerator, eval_dataset)
+        logger.info(f"epoch {epoch}: perplexity: {perplexity} eval_loss: {eval_loss}")
+
+        if args.with_tracking:
+            accelerator.log(
+                {
+                    "perplexity": perplexity,
+                    "eval_loss": eval_loss,
+                    "train_loss": total_loss.item() / len(train_dataloader),
+                    "epoch": epoch,
+                    "step": completed_steps,
+                },
+                step=completed_steps,
+            )
+
+        # New Code #
+        # Save the DeepSpeed checkpoint to the specified path
+        checkpoint_model(args.output_dir, epoch, model, epoch, completed_steps)
+
+        # New Code #
+        # Tracks the best checkpoint and best metric
+        if best_metric is None or best_metric > perplexity:
+            best_metric = perplexity
+            best_metric_checkpoint = os.path.join(args.output_dir, str(epoch))
+            accelerator.print(f"New best metric: {best_metric} at epoch {epoch}")
+            accelerator.print(f"best_metric_checkpoint: {best_metric_checkpoint}")
+
+    # New Code #
+    # Loads the best checkpoint after the training is finished
+    if args.load_best_model:
+        _, last_global_step = load_training_checkpoint(
+            model,
+            "/".join(best_metric_checkpoint.split("/")[:-1]),
+            tag=best_metric_checkpoint.split("/")[-1],
+            **{"load_optimizer_states": True, "load_lr_scheduler_states": True},
+        )
+
+    # New Code #
+    # Evaluates using the best checkpoint
+    perplexity, eval_loss = evaluate(args, model, eval_dataloader, accelerator, eval_dataset)
+    logger.info(f"Best model metrics: perplexity: {perplexity} eval_loss: {eval_loss}")
+    if perplexity != best_metric:
+        raise AssertionError(
+            f"Best metric {best_metric} does not match the metric {perplexity} of the loaded best model."
+        )
+
+    if args.output_dir is not None:
+        accelerator.wait_for_everyone()
+        unwrapped_model = accelerator.unwrap_model(model)
+
+        # New Code #
+        # Saves the whole/unpartitioned fp16 model when in ZeRO Stage-3 to the output directory if
+        # `stage3_gather_16bit_weights_on_model_save` is True in DeepSpeed Config file or
+        # `zero3_save_16bit_model` is True in DeepSpeed Plugin.
+        # For Zero Stages 1 and 2, models are saved as usual in the output directory.
+        # The model name saved is `pytorch_model.bin`
+        unwrapped_model.save_pretrained(
+            args.output_dir,
+            is_main_process=accelerator.is_main_process,
+            save_function=accelerator.save,
+            state_dict=accelerator.get_state_dict(model),
+        )
+        if accelerator.is_main_process:
+            tokenizer.save_pretrained(args.output_dir)
+            if args.push_to_hub:
+                repo.push_to_hub(commit_message="End of training", auto_lfs_prune=True)
+
+        with open(os.path.join(args.output_dir, "all_results.json"), "w") as f:
+            json.dump({"perplexity": perplexity, "eval_loss": eval_loss.item()}, f)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/deepspeed_config_templates/zero_stage1_config.json b/examples/deepspeed_config_templates/zero_stage1_config.json
new file mode 100644
index 000000000..674420e57
--- /dev/null
+++ b/examples/deepspeed_config_templates/zero_stage1_config.json
@@ -0,0 +1,43 @@
+{
+    "fp16": {
+        "enabled": true,
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "optimizer": {
+        "type": "AdamW",
+        "params": {
+            "lr": "auto",
+            "weight_decay": "auto",
+            "torch_adam": true,
+            "adam_w_mode": true
+        }
+    },
+    "scheduler": {
+        "type": "WarmupDecayLR",
+        "params": {
+            "warmup_min_lr": "auto",
+            "warmup_max_lr": "auto",
+            "warmup_num_steps": "auto",
+            "total_num_steps": "auto"
+        }
+    },
+    "zero_optimization": {
+        "stage": 1,
+        "allgather_partitions": true,
+        "allgather_bucket_size": 2e8,
+        "overlap_comm": true,
+        "reduce_scatter": true,
+        "reduce_bucket_size": "auto",
+        "contiguous_gradients": true
+    },
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}
\ No newline at end of file
diff --git a/examples/deepspeed_config_templates/zero_stage2_config.json b/examples/deepspeed_config_templates/zero_stage2_config.json
new file mode 100644
index 000000000..9597f8485
--- /dev/null
+++ b/examples/deepspeed_config_templates/zero_stage2_config.json
@@ -0,0 +1,43 @@
+{
+    "fp16": {
+        "enabled": true,
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "optimizer": {
+        "type": "AdamW",
+        "params": {
+            "lr": "auto",
+            "weight_decay": "auto",
+            "torch_adam": true,
+            "adam_w_mode": true
+        }
+    },
+    "scheduler": {
+        "type": "WarmupDecayLR",
+        "params": {
+            "warmup_min_lr": "auto",
+            "warmup_max_lr": "auto",
+            "warmup_num_steps": "auto",
+            "total_num_steps": "auto"
+        }
+    },
+    "zero_optimization": {
+        "stage": 2,
+        "allgather_partitions": true,
+        "allgather_bucket_size": 2e8,
+        "overlap_comm": true,
+        "reduce_scatter": true,
+        "reduce_bucket_size": "auto",
+        "contiguous_gradients": true
+    },
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}
\ No newline at end of file
diff --git a/examples/deepspeed_config_templates/zero_stage2_offload_config.json b/examples/deepspeed_config_templates/zero_stage2_offload_config.json
new file mode 100644
index 000000000..98baedef3
--- /dev/null
+++ b/examples/deepspeed_config_templates/zero_stage2_offload_config.json
@@ -0,0 +1,47 @@
+{
+    "fp16": {
+        "enabled": true,
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "optimizer": {
+        "type": "AdamW",
+        "params": {
+            "lr": "auto",
+            "weight_decay": "auto",
+            "torch_adam": true,
+            "adam_w_mode": true
+        }
+    },
+    "scheduler": {
+        "type": "WarmupDecayLR",
+        "params": {
+            "warmup_min_lr": "auto",
+            "warmup_max_lr": "auto",
+            "warmup_num_steps": "auto",
+            "total_num_steps": "auto"
+        }
+    },
+    "zero_optimization": {
+        "stage": 2,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "allgather_partitions": true,
+        "allgather_bucket_size": 2e8,
+        "overlap_comm": true,
+        "reduce_scatter": true,
+        "reduce_bucket_size": "auto",
+        "contiguous_gradients": true
+    },
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}
\ No newline at end of file
diff --git a/examples/deepspeed_config_templates/zero_stage3_config.json b/examples/deepspeed_config_templates/zero_stage3_config.json
new file mode 100644
index 000000000..2ec6fff41
--- /dev/null
+++ b/examples/deepspeed_config_templates/zero_stage3_config.json
@@ -0,0 +1,44 @@
+{
+    "fp16": {
+        "enabled": true,
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "optimizer": {
+        "type": "AdamW",
+        "params": {
+            "lr": "auto",
+            "weight_decay": "auto"
+        }
+    },
+    "scheduler": {
+        "type": "WarmupDecayLR",
+        "params": {
+            "warmup_min_lr": "auto",
+            "warmup_max_lr": "auto",
+            "warmup_num_steps": "auto",
+            "total_num_steps": "auto"
+        }
+    },
+    "zero_optimization": {
+        "stage": 3,
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "reduce_bucket_size": "auto",
+        "stage3_prefetch_bucket_size": "auto",
+        "stage3_param_persistence_threshold": "auto",
+        "sub_group_size": 1e9,
+        "stage3_max_live_parameters": 1e9,
+        "stage3_max_reuse_distance": 1e9,
+        "stage3_gather_16bit_weights_on_model_save": "auto"
+    },
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}
\ No newline at end of file
diff --git a/examples/deepspeed_config_templates/zero_stage3_offload_config.json b/examples/deepspeed_config_templates/zero_stage3_offload_config.json
new file mode 100644
index 000000000..edae8e628
--- /dev/null
+++ b/examples/deepspeed_config_templates/zero_stage3_offload_config.json
@@ -0,0 +1,52 @@
+{
+    "fp16": {
+        "enabled": true,
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "optimizer": {
+        "type": "AdamW",
+        "params": {
+            "lr": "auto",
+            "weight_decay": "auto"
+        }
+    },
+    "scheduler": {
+        "type": "WarmupDecayLR",
+        "params": {
+            "warmup_min_lr": "auto",
+            "warmup_max_lr": "auto",
+            "warmup_num_steps": "auto",
+            "total_num_steps": "auto"
+        }
+    },
+    "zero_optimization": {
+        "stage": 3,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "offload_param": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "reduce_bucket_size": "auto",
+        "stage3_prefetch_bucket_size": "auto",
+        "stage3_param_persistence_threshold": "auto",
+        "sub_group_size": 1e9,
+        "stage3_max_live_parameters": 1e9,
+        "stage3_max_reuse_distance": 1e9,
+        "stage3_gather_16bit_weights_on_model_save": "auto"
+    },
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}
\ No newline at end of file
diff --git a/setup.py b/setup.py
index a28b9254c..45c4ae912 100644
--- a/setup.py
+++ b/setup.py
@@ -27,7 +27,9 @@
     "evaluate",
     "transformers",
     "scipy",
-    "sklearn"
+    "sklearn",
+    "parameterized",
+    "deepspeed",
 ]
 
 extras["test_trackers"] = ["wandb", "comet-ml", "tensorboard"]
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 6dc1f847c..9205f0854 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -13,10 +13,12 @@
 # limitations under the License.
 
 import gc
+import math
 import os
 import sys
 import warnings
 from contextlib import contextmanager
+from copy import deepcopy
 from typing import List, Optional, Union
 
 import torch
@@ -39,12 +41,14 @@
     LoggerType,
     PrecisionType,
     RNGType,
+    compare_versions,
     convert_outputs_to_fp32,
     extract_model_from_parallel,
     gather,
     get_pretty_name,
     is_deepspeed_available,
     is_torch_version,
+    is_transformers_available,
     pad_across_processes,
     reduce,
     save,
@@ -55,7 +59,13 @@
 if is_deepspeed_available():
     import deepspeed
 
-    from .utils import DeepSpeedEngineWrapper, DeepSpeedOptimizerWrapper
+    from .utils import (
+        DeepSpeedEngineWrapper,
+        DeepSpeedOptimizerWrapper,
+        DeepSpeedSchedulerWrapper,
+        DummyOptim,
+        DummyScheduler,
+    )
 
 logger = get_logger(__name__)
 
@@ -163,6 +173,30 @@ def __init__(
                 deepspeed_plugin, DeepSpeedPlugin
             ), "`deepspeed_plugin` must be a DeepSpeedPlugin object."
             os.environ["USE_DEEPSPEED"] = "true"  # use DeepSpeed if plugin is provided
+        if deepspeed_plugin:
+            if not is_deepspeed_available():
+                raise ImportError("DeepSpeed is not installed => run `pip install deepspeed` or build it from source.")
+            if compare_versions("deepspeed", "<", "0.6.5"):
+                raise ImportError("DeepSpeed version must be >= 0.6.5. Please update DeepSpeed.")
+            if os.environ.get("DEEPSPEED_ZERO3_INIT", "false") == "true" or deepspeed_plugin.zero3_init_flag:
+                if not is_transformers_available():
+                    raise Exception(
+                        "When `zero3_init_flag` is set, it requires Transformers to be installed. "
+                        "Please run `pip install transformers`."
+                    )
+                from transformers.deepspeed import HfDeepSpeedConfig
+
+                ds_config = deepcopy(deepspeed_plugin.deepspeed_config)
+                del ds_config["train_batch_size"]
+                ds_config.update({"train_micro_batch_size_per_gpu": 1, "gradient_accumulation_steps": 1})
+                mixed_precision = (
+                    os.environ.get("MIXED_PRECISION", "no") if mixed_precision is None else mixed_precision
+                )
+                if mixed_precision == "fp16":
+                    ds_config.update({"fp16": {"enabled": True}})
+                elif mixed_precision == "bf16":
+                    ds_config.update({"bf16": {"enabled": True}})
+                self.dschf = HfDeepSpeedConfig(ds_config)  # keep this object alive # noqa
 
         if fsdp_plugin is None:  # init from env variables
             fsdp_plugin = FullyShardedDataParallelPlugin() if os.environ.get("USE_FSDP", "false") == "true" else None
@@ -497,9 +531,15 @@ def prepare_model(self, model):
     def _prepare_deepspeed(self, *args):
 
         deepspeed_plugin = self.state.deepspeed_plugin
-        self.deepspeed_config = deepspeed_plugin.deepspeed_config
+
+        result = [
+            self._prepare_one(obj, first_pass=True) if isinstance(obj, torch.utils.data.DataLoader) else obj
+            for obj in args
+        ]
 
         batch_sizes = [obj.batch_size for obj in args if hasattr(obj, "batch_size")]
+        if self.split_batches:
+            batch_sizes = [batch_size // self.num_processes for batch_size in batch_sizes]
         if len(batch_sizes) == 0:
             raise ValueError(
                 "You must specify a training or evaluation dataloader in `accelerate.prepare()` when using DeepSpeed."
@@ -508,73 +548,141 @@ def _prepare_deepspeed(self, *args):
         batch_size_per_device = min(batch_sizes) if deepspeed_plugin.is_train_batch_min else max(batch_sizes)
         if len(batch_sizes) > 1:
             logger.info(
-                f"Since you passed both train and evaluation dataloader, `is_train_batch_min` (here \
-                {deepspeed_plugin.is_train_batch_min} will decide the `train_batch_size` ({batch_size_per_device})."
+                "Since you passed both train and evaluation dataloader, `is_train_batch_min` (here "
+                f"{deepspeed_plugin.is_train_batch_min} will decide the `train_batch_size` ({batch_size_per_device})."
             )
 
-        self.deepspeed_config["train_batch_size"] = (
-            batch_size_per_device * deepspeed_plugin.gradient_accumulation_steps * self.num_processes
-        )
-
-        result = [
-            self._prepare_one(obj, first_pass=True) if isinstance(obj, torch.utils.data.DataLoader) else obj
-            for obj in args
-        ]
+        config_kwargs = {
+            "train_micro_batch_size_per_gpu": batch_size_per_device,
+            "train_batch_size": batch_size_per_device
+            * deepspeed_plugin.deepspeed_config["gradient_accumulation_steps"]
+            * self.num_processes,
+            "gradient_clipping": 1.0,
+            "zero_optimization.stage3_gather_16bit_weights_on_model_save": False,
+        }
 
         model = None
         optimizer = None
+        scheduler = None
         for obj in result:
             if isinstance(obj, torch.nn.Module):
                 model = obj
-            elif isinstance(obj, (torch.optim.Optimizer, dict)):
+            elif isinstance(obj, (torch.optim.Optimizer, DummyOptim)):
                 optimizer = obj
+            elif (isinstance(obj, (torch.optim.lr_scheduler._LRScheduler, DummyScheduler))) or (
+                type(obj).__name__ in deepspeed.runtime.lr_schedules.VALID_LR_SCHEDULES
+            ):
+                scheduler = obj
+
+        if optimizer is not None:
+            if "optimizer" in deepspeed_plugin.deepspeed_config and not isinstance(optimizer, (DummyOptim)):
+                raise ValueError(
+                    "You cannot specify an optimizer in the config file and in the code at the same time. "
+                    "Please remove the optimizer from the config file or "
+                    "create `accelerate.utils.DummyOptim` in the code."
+                )
+            elif "optimizer" not in deepspeed_plugin.deepspeed_config and isinstance(optimizer, (DummyOptim)):
+                raise ValueError(
+                    "You cannot create a `DummyOptim` without specifying an optimizer in the config file."
+                )
+
+            if isinstance(optimizer, (torch.optim.Optimizer)):
+                deepspeed_plugin.deepspeed_config["zero_allow_untested_optimizer"] = True
+
+        if scheduler is not None:
+            if "scheduler" in deepspeed_plugin.deepspeed_config and not isinstance(scheduler, (DummyScheduler)):
+                raise ValueError(
+                    "You cannot specify a scheduler in the config file and in the code at the same time. "
+                    "Please remove the scheduler from the config file or "
+                    "create `accelerate.utils.DummyScheduler` in the code."
+                )
+            elif "scheduler" not in deepspeed_plugin.deepspeed_config and isinstance(scheduler, (DummyScheduler)):
+                raise ValueError(
+                    "You cannot create a `DummyScheduler` without specifying a scheduler in the config file."
+                )
 
-        if deepspeed_plugin.auto_opt_mapping:
-            is_adam = isinstance(optimizer, torch.optim.Adam)
-            is_adamw = isinstance(optimizer, torch.optim.AdamW)
-            if (is_adam or is_adamw) and deepspeed_plugin.offload_optimizer_device == "cpu":
-                defaults = optimizer.defaults
-                params = []
-                for group in optimizer.param_groups:
-                    params.extend(group["params"])
-
-                optimizer = deepspeed.ops.adam.DeepSpeedCPUAdam(
-                    params,
-                    lr=defaults["lr"],
-                    bias_correction=True,
-                    betas=defaults["betas"],
-                    eps=defaults["eps"],
-                    weight_decay=defaults["weight_decay"],
-                    amsgrad=defaults["amsgrad"],
-                    adamw_mode=is_adamw,
+        if optimizer is not None and scheduler is not None:
+            if isinstance(optimizer, (DummyOptim)) and not isinstance(scheduler, (DummyScheduler)):
+                raise ValueError(
+                    "You can only specify `accelerate.utils.DummyScheduler` in the code when using "
+                    "`accelerate.utils.DummyOptim`."
                 )
 
-        # useful when only eval_dataloader is given into `accelerator.prepare()`
         if model is not None:
-            engine = DeepSpeedEngineWrapper(
-                args=None,
-                model=model,
-                optimizer=optimizer,
-                config_params=self.deepspeed_config,
-                dist_init_required=False,
-            )
+            if hasattr(model, "config") and hasattr(model.config, "hidden_size"):
+                hidden_size = model.config.hidden_size
+                config_kwargs.update(
+                    {
+                        "zero_optimization.reduce_bucket_size": hidden_size * hidden_size,
+                        "zero_optimization.stage3_prefetch_bucket_size": 0.9 * hidden_size * hidden_size,
+                        "zero_optimization.stage3_param_persistence_threshold": 10 * hidden_size,
+                    }
+                )
+
+            if isinstance(optimizer, (DummyOptim)):
+                config_kwargs.update(
+                    {"optimizer.params.lr": optimizer.lr, "optimizer.params.weight_decay": optimizer.weight_decay}
+                )
+            if isinstance(scheduler, (DummyScheduler)):
+                config_kwargs.update(
+                    {
+                        "scheduler.params.warmup_min_lr": 0,
+                        "scheduler.params.warmup_max_lr": scheduler.optimizer.lr,
+                        "scheduler.params.warmup_num_steps": scheduler.warmup_num_steps,
+                    }
+                )
+                if scheduler.total_num_steps is not None:
+                    config_kwargs["scheduler.params.total_num_steps"] = (
+                        math.ceil(scheduler.total_num_steps / self.num_processes)
+                        if not self.split_batches
+                        else scheduler.total_num_steps
+                    )
+            deepspeed_plugin.deepspeed_config_process(must_match=False, **config_kwargs)
+            self.deepspeed_config = deepspeed_plugin.deepspeed_config
+            kwargs = dict(model=model, config_params=self.deepspeed_config)
+            if optimizer is not None:
+                if isinstance(optimizer, (DummyOptim)):
+                    kwargs["model_parameters"] = optimizer.params
+                else:
+                    kwargs["optimizer"] = optimizer
+                    if scheduler is not None:
+                        if type(scheduler).__name__ in deepspeed.runtime.lr_schedules.VALID_LR_SCHEDULES:
+                            kwargs["lr_scheduler"] = scheduler
+
+            engine, optimizer, _, lr_scheduler = deepspeed.initialize(**kwargs)
+            if optimizer is not None:
+                optimizer = DeepSpeedOptimizerWrapper(optimizer)
+            if scheduler is not None:
+                if lr_scheduler is None:
+                    scheduler = AcceleratedScheduler(
+                        scheduler,
+                        optimizer,
+                        step_with_optimizer=self.step_scheduler_with_optimizer,
+                        split_batches=self.split_batches,
+                    )
+                else:
+                    scheduler = DeepSpeedSchedulerWrapper(lr_scheduler, optimizer)
+
             for i in range(len(result)):
                 if isinstance(result[i], torch.nn.Module):
                     result[i] = engine
-                elif isinstance(result[i], torch.optim.Optimizer):
-                    result[i] = DeepSpeedOptimizerWrapper(engine.optimizer, engine)
-            self.deepspeed_engine = engine  # pointing for deepspeed_engine.backward()
+                elif isinstance(result[i], (torch.optim.Optimizer, DummyOptim)):
+                    result[i] = optimizer
+                elif (isinstance(result[i], (torch.optim.lr_scheduler._LRScheduler, DummyScheduler))) or (
+                    type(result[i]).__name__ in deepspeed.runtime.lr_schedules.VALID_LR_SCHEDULES
+                ):
+                    result[i] = scheduler
+            # pointing for deepspeed_engine_wrapped.backward()
+            self.deepspeed_engine_wrapped = DeepSpeedEngineWrapper(engine)
             self._models.append(engine)
-            self._optimizers.append(engine.optimizer)
-            assert (
-                len(self._models) == 1
-            ), "You can't use same `Accelerator()` instance with 2 models when using DeepSpeed"
-
-        if self.distributed_type == DistributedType.DEEPSPEED:
-            assert hasattr(
-                self, "deepspeed_engine"
-            ), "You need to pass the model along the optimizer when using Deepspeed."
-
+            if optimizer is not None:
+                self._optimizers.append(optimizer)
+            if scheduler is not None:
+                self._schedulers.append(scheduler)
+            if len(self._models) > 1:
+                raise AssertionError(
+                    "You can't use same `Accelerator()` instance with multiple models when using DeepSpeed"
+                )
         return tuple(result)
 
     def prepare_data_loader(self, data_loader):
@@ -612,7 +720,7 @@ def backward(self, loss, **kwargs):
         Use `accelerator.backward(loss)` in lieu of `loss.backward()`.
         """
         if self.distributed_type == DistributedType.DEEPSPEED:
-            self.deepspeed_engine.backward(loss, **kwargs)
+            self.deepspeed_engine_wrapped.backward(loss, **kwargs)
         elif self.scaler is not None:
             self.scaler.scale(loss).backward(**kwargs)
         else:
@@ -648,6 +756,9 @@ def clip_grad_norm_(self, parameters, max_norm, norm_type=2):
                 if parameters == [p for p in model.parameters()]:
                     model.clip_grad_norm_(max_norm, norm_type)
                     return
+        elif self.distributed_type == DistributedType.DEEPSPEED:
+            # `accelerator.backward(loss)` is doing that automatically. Therefore, it's implementation is not needed
+            return
         self.unscale_gradients()
         torch.nn.utils.clip_grad_norm_(parameters, max_norm, norm_type=norm_type)
 
@@ -655,6 +766,8 @@ def clip_grad_value_(self, parameters, clip_value):
         """
         Should be used in place of `torch.nn.utils.clip_grad_value_`.
         """
+        if self.distributed_type in [DistributedType.DEEPSPEED, DistributedType.FSDP]:
+            raise Exception("DeepSpeed and FSDP  do not support `clip_grad_value_`. Use `clip_grad_norm_` instead.")
         self.unscale_gradients()
         torch.nn.utils.clip_grad_value_(parameters, clip_value)
 
@@ -837,7 +950,7 @@ def free_memory(self):
         self._schedulers = []
         self._optimizers = []
         self._models = []
-        self.deepspeed_engine = None
+        self.deepspeed_engine_wrapped = None
         gc.collect()
         torch.cuda.empty_cache()
 
@@ -875,12 +988,19 @@ def _get_devices(self, *args):
 
     def get_state_dict(self, model):
         is_zero_3 = False
-        if is_deepspeed_available():
-            if isinstance(model, DeepSpeedEngineWrapper) and self.distributed_type == DistributedType.DEEPSPEED:
-                is_zero_3 = self.state.deepspeed_plugin.zero_stage == 3
+        if self.distributed_type == DistributedType.DEEPSPEED:
+            is_zero_3 = self.deepspeed_config["zero_optimization"]["stage"] == 3
 
         if is_zero_3:
-            state_dict = model._zero3_consolidated_16bit_state_dict()
+            if model.zero_gather_16bit_weights_on_model_save():
+                state_dict = model._zero3_consolidated_16bit_state_dict()
+            else:
+                raise ValueError(
+                    "Cannot get 16bit model weights because `stage3_gather_16bit_weights_on_model_save` in DeepSpeed config is False. "
+                    "To save the model weights in 16bit, set `stage3_gather_16bit_weights_on_model_save` to True in DeepSpeed config file or "
+                    "set `zero3_save_16bit_model` to True when using `accelerate config`. "
+                    "To save the full checkpoint, run `model.save_checkpoint(save_dir)` and use `zero_to_fp32.py` to recover weights."
+                )
         else:
             model = self.unwrap_model(model)
             state_dict = model.state_dict()
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index 01a55e646..907122036 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -14,7 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ...utils import ComputeEnvironment, DistributedType, is_deepspeed_available
+from ...utils import ComputeEnvironment, DistributedType, is_deepspeed_available, is_transformers_available
 from .config_args import ClusterConfig
 from .config_utils import _ask_field, _convert_distributed_mode, _convert_yes_no_to_bool
 
@@ -77,24 +77,72 @@ def get_cluster_input():
             ), "DeepSpeed is not installed => run `pip3 install deepspeed` or build it from source"
 
         if distributed_type == DistributedType.DEEPSPEED:
-            deepspeed_config["zero_stage"] = _ask_field(
-                "What should be your DeepSpeed's ZeRO optimization stage (0, 1, 2, 3)? [2]: ",
-                lambda x: int(x),
-                default=2,
+            use_deepspeed_config = _ask_field(
+                "Do you want to specify a json file to a DeepSpeed config? [yes/NO]: ",
+                _convert_yes_no_to_bool,
+                default=False,
+                error_message="Please enter yes or no.",
             )
-
-            if deepspeed_config["zero_stage"] >= 2:
-                deepspeed_config["offload_optimizer_device"] = _ask_field(
-                    "Where to offload optimizer states? [NONE/cpu/nvme]: ",
+            if use_deepspeed_config:
+                deepspeed_config["deepspeed_config_file"] = _ask_field(
+                    "Please enter the path to the json DeepSpeed config file: ",
                     lambda x: str(x),
                     default="none",
                 )
+            else:
+                deepspeed_config["zero_stage"] = _ask_field(
+                    "What should be your DeepSpeed's ZeRO optimization stage (0, 1, 2, 3)? [2]: ",
+                    lambda x: int(x),
+                    default=2,
+                )
 
-            deepspeed_config["gradient_accumulation_steps"] = _ask_field(
-                "How many gradient accumulation steps you're passing in your script? [1]: ",
-                lambda x: int(x),
-                default=1,
+                if deepspeed_config["zero_stage"] >= 2:
+                    deepspeed_config["offload_optimizer_device"] = _ask_field(
+                        "Where to offload optimizer states? [none/cpu/nvme]: ",
+                        lambda x: str(x),
+                        default="none",
+                    )
+                    deepspeed_config["offload_param_device"] = _ask_field(
+                        "Where to offload parameters? [none/cpu/nvme]: ",
+                        lambda x: str(x),
+                        default="none",
+                    )
+                deepspeed_config["gradient_accumulation_steps"] = _ask_field(
+                    "How many gradient accumulation steps you're passing in your script? [1]: ",
+                    lambda x: int(x),
+                    default=1,
+                )
+                use_gradient_clipping = _ask_field(
+                    "Do you want to use gradient clipping? [yes/NO]: ",
+                    _convert_yes_no_to_bool,
+                    default=False,
+                    error_message="Please enter yes or no.",
+                )
+                if use_gradient_clipping:
+                    deepspeed_config["gradient_clipping"] = _ask_field(
+                        "What is the gradient clipping value? [1.0]: ",
+                        lambda x: float(x),
+                        default=1.0,
+                    )
+                if deepspeed_config["zero_stage"] == 3:
+                    deepspeed_config["zero3_save_16bit_model"] = _ask_field(
+                        "Do you want to save 16-bit model weights when using ZeRO Stage-3? [yes/NO]: ",
+                        _convert_yes_no_to_bool,
+                        default=False,
+                        error_message="Please enter yes or no.",
+                    )
+            deepspeed_config["zero3_init_flag"] = _ask_field(
+                "Do you want to enable `deepspeed.zero.Init` when using ZeRO Stage-3 for constructing massive models? [yes/NO]: ",
+                _convert_yes_no_to_bool,
+                default=False,
+                error_message="Please enter yes or no.",
             )
+            if deepspeed_config["zero3_init_flag"]:
+                if not is_transformers_available():
+                    raise Exception(
+                        "When `zero3_init_flag` is set, it requires Transformers to be installed. "
+                        "Please run `pip3 install transformers`."
+                    )
 
     fsdp_config = {}
     if distributed_type in [DistributedType.MULTI_GPU]:
@@ -155,11 +203,14 @@ def get_cluster_input():
         num_processes = 1
 
     if distributed_type != DistributedType.TPU:
-        mixed_precision = _ask_field(
-            "Do you wish to use FP16 or BF16 (mixed precision)? [NO/fp16/bf16]: ",
-            lambda x: str(x).lower(),
-            default="no",
-        )
+        if distributed_type == DistributedType.DEEPSPEED and use_deepspeed_config:
+            mixed_precision = "no"
+        else:
+            mixed_precision = _ask_field(
+                "Do you wish to use FP16 or BF16 (mixed precision)? [NO/fp16/bf16]: ",
+                lambda x: str(x).lower(),
+                default="no",
+            )
     else:
         mixed_precision = "no"
 
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 22c8e621c..93b04d6ab 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -31,6 +31,7 @@
     DistributedType,
     PrecisionType,
     PrepareForLaunch,
+    is_deepspeed_available,
     is_sagemaker_available,
 )
 from accelerate.utils.versions import is_torch_version
@@ -57,6 +58,56 @@ def launch_command_parser(subparsers=None):
         action="store_true",
         help="Whether to use deepspeed.",
     )
+    parser.add_argument(
+        "--deepspeed_config_file",
+        default=None,
+        type=str,
+        help="DeepSpeed config file.",
+    )
+    parser.add_argument(
+        "--zero_stage",
+        default=None,
+        type=int,
+        help="DeepSpeed's ZeRO optimization stage (useful only when `use_deepspeed` flag is passed).",
+    )
+    parser.add_argument(
+        "--offload_optimizer_device",
+        default=None,
+        type=str,
+        help="Decides where (none|cpu|nvme) to offload optimizer states (useful only when `use_deepspeed` flag is passed).",
+    )
+    parser.add_argument(
+        "--offload_param_device",
+        default=None,
+        type=str,
+        help="Decides where (none|cpu|nvme) to offload parameters (useful only when `use_deepspeed` flag is passed).",
+    )
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        default=None,
+        type=int,
+        help="No of gradient_accumulation_steps used in your training script (useful only when `use_deepspeed` flag is passed).",
+    )
+    parser.add_argument(
+        "--gradient_clipping",
+        default=None,
+        type=float,
+        help="gradient clipping value used in your training script (useful only when `use_deepspeed` flag is passed).",
+    )
+    parser.add_argument(
+        "--zero3_init_flag",
+        default=None,
+        type=str,
+        help="Decides Whether (true|false) to enable `deepspeed.zero.Init` for constructing massive models. "
+        "Only applicable with DeepSpeed ZeRO Stage-3.",
+    )
+    parser.add_argument(
+        "--zero3_save_16bit_model",
+        default=None,
+        type=str,
+        help="Decides Whether (true|false) to save 16-bit model weights when using ZeRO Stage-3. "
+        "Only applicable with DeepSpeed ZeRO Stage-3.",
+    )
     parser.add_argument(
         "--use_fsdp",
         default=False,
@@ -158,24 +209,6 @@ def launch_command_parser(subparsers=None):
             "script."
         ),
     )
-    parser.add_argument(
-        "--zero_stage",
-        default=None,
-        type=int,
-        help="DeepSpeed's ZeRO optimization stage (useful only when `use_deepspeed` flag is passed).",
-    )
-    parser.add_argument(
-        "--offload_optimizer_device",
-        default=None,
-        type=str,
-        help="Decides where (none|cpu|nvme) to offload optimizer states (useful only when `use_deepspeed` flag is passed).",
-    )
-    parser.add_argument(
-        "--gradient_accumulation_steps",
-        default=None,
-        type=int,
-        help="No of gradient_accumulation_steps used in your training script (useful only when `use_deepspeed` flag is passed).",
-    )
 
     # Other arguments of the training scripts
     parser.add_argument("training_script_args", nargs=argparse.REMAINDER, help="Arguments of the training script.")
@@ -279,6 +312,8 @@ def multi_gpu_launcher(args):
 
 
 def deepspeed_launcher(args):
+    if not is_deepspeed_available():
+        raise ImportError("DeepSpeed is not installed => run `pip3 install deepspeed` or build it from source.")
     cmd = ["deepspeed", "--no_local_rank"]
     if args.num_machines > 1:
         cmd.extend(
@@ -323,7 +358,12 @@ def deepspeed_launcher(args):
     current_env["USE_DEEPSPEED"] = "true"
     current_env["DEEPSPEED_ZERO_STAGE"] = str(args.zero_stage)
     current_env["GRADIENT_ACCUMULATION_STEPS"] = str(args.gradient_accumulation_steps)
+    current_env["GRADIENT_CLIPPING"] = str(args.gradient_clipping)
     current_env["DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE"] = str(args.offload_optimizer_device).lower()
+    current_env["DEEPSPEED_OFFLOAD_PARAM_DEVICE"] = str(args.offload_param_device).lower()
+    current_env["DEEPSPEED_ZERO3_INIT"] = str(args.zero3_init_flag).lower()
+    current_env["DEEPSPEED_ZERO3_SAVE_16BIT_MODEL"] = str(args.zero3_save_16bit_model).lower()
+    current_env["DEEPSPEED_CONFIG_FILE"] = str(args.deepspeed_config_file).lower()
 
     process = subprocess.Popen(cmd, env=current_env)
     process.wait()
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index fbda09754..6e09fdae4 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -108,10 +108,18 @@ def __init__(
                 mixed_precision = (
                     parse_choice_from_env("MIXED_PRECISION", "no") if mixed_precision is None else mixed_precision
                 )
-                if mixed_precision == "fp16":
+                if (
+                    mixed_precision == "fp16"
+                    and "fp16" not in deepspeed_plugin.deepspeed_config
+                    and "bf16" not in deepspeed_plugin.deepspeed_config
+                ):
                     deepspeed_plugin.deepspeed_config.update({"fp16": {"enabled": True}})
-                elif mixed_precision == "bf16":
-                    deepspeed_plugin.deepspeed_config.update({"bfloat16": {"enabled": True}})
+                elif (
+                    mixed_precision == "bf16"
+                    and "fp16" not in deepspeed_plugin.deepspeed_config
+                    and "bf16" not in deepspeed_plugin.deepspeed_config
+                ):
+                    deepspeed_plugin.deepspeed_config.update({"bf16": {"enabled": True}})
                 self.deepspeed_plugin = deepspeed_plugin
             elif int(os.environ.get("LOCAL_RANK", -1)) != -1 and not cpu:
                 self.distributed_type = DistributedType.MULTI_GPU
@@ -189,10 +197,11 @@ def __repr__(self):
             f"Process index: {self.process_index}\n"
             f"Local process index: {self.local_process_index}\n"
             f"Device: {self.device}\n"
-            f"Mixed precision type: {mixed_precision}\n"
         )
         if self.distributed_type == DistributedType.DEEPSPEED:
             repr += f"ds_config: {self.deepspeed_plugin.deepspeed_config}\n"
+        else:
+            f"Mixed precision type: {mixed_precision}\n"
         return repr
 
     # For backward compatibility
diff --git a/src/accelerate/test_utils/testing.py b/src/accelerate/test_utils/testing.py
index 91dc0f7ad..89d104a49 100644
--- a/src/accelerate/test_utils/testing.py
+++ b/src/accelerate/test_utils/testing.py
@@ -26,7 +26,14 @@
 import torch
 
 from ..state import AcceleratorState
-from ..utils import gather, is_comet_ml_available, is_tensorboard_available, is_tpu_available, is_wandb_available
+from ..utils import (
+    gather,
+    is_comet_ml_available,
+    is_deepspeed_available,
+    is_tensorboard_available,
+    is_tpu_available,
+    is_wandb_available,
+)
 
 
 def parse_flag_from_env(key, default=False):
@@ -85,6 +92,13 @@ def require_multi_gpu(test_case):
     return unittest.skipUnless(torch.cuda.device_count() > 1, "test requires multiple GPUs")(test_case)
 
 
+def require_deepspeed(test_case):
+    """
+    Decorator marking a test that requires DeepSpeed installed. These tests are skipped when DeepSpeed isn't installed
+    """
+    return unittest.skipUnless(is_deepspeed_available(), "test requires DeepSpeed")(test_case)
+
+
 def require_tensorboard(test_case):
     """
     Decorator marking a test that requires tensorboard installed. These tests are skipped when tensorboard isn't
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
index a39368dd3..b66b6fd1d 100644
--- a/src/accelerate/utils/__init__.py
+++ b/src/accelerate/utils/__init__.py
@@ -27,6 +27,7 @@
     is_sagemaker_available,
     is_tensorboard_available,
     is_tpu_available,
+    is_transformers_available,
     is_wandb_available,
 )
 from .modeling import (
@@ -76,7 +77,13 @@
 
 
 if is_deepspeed_available():
-    from .deepspeed import DeepSpeedEngineWrapper, DeepSpeedOptimizerWrapper
+    from .deepspeed import (
+        DeepSpeedEngineWrapper,
+        DeepSpeedOptimizerWrapper,
+        DeepSpeedSchedulerWrapper,
+        DummyOptim,
+        DummyScheduler,
+    )
 
 from .launch import PrepareForLaunch
 from .memory import find_executable_batch_size
diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py
index 4a19f8b51..b731ca12a 100644
--- a/src/accelerate/utils/dataclasses.py
+++ b/src/accelerate/utils/dataclasses.py
@@ -19,8 +19,11 @@
 import copy
 import enum
 import functools
+import io
+import json
 import os
 import typing
+import warnings
 from dataclasses import dataclass, field
 from datetime import timedelta
 from typing import Callable, Iterable, Optional
@@ -208,10 +211,15 @@ class TensorInformation:
 
 @dataclass
 class DeepSpeedPlugin:
+    """
+    This plugin is used to integrate DeepSpeed.
+    """
 
+    config_file: str = field(default=None, metadata={"help": "Path to the DeepSpeed config file."})
     gradient_accumulation_steps: int = field(
         default=None, metadata={"help": "Number of steps to accumulate gradients before updating optimizer states"}
     )
+    gradient_clipping: float = field(default=None, metadata={"help": "Enable gradient clipping with value"})
     zero_stage: int = field(
         default=None,
         metadata={"help": "Possible options are 0,1,2,3; Default will be taken from environment variable"},
@@ -220,37 +228,137 @@ class DeepSpeedPlugin:
         default=True,
         metadata={"help": "If both train & eval dataloaders are specified, this will decide the train_batch_size"},
     )
-
-    auto_opt_mapping: bool = field(
-        default=True,
-        metadata={"help": "whether to map torch.adam to deepspeed optimizer version of adam based on config"},
+    offload_optimizer_device: bool = field(
+        default=None,
+        metadata={"help": "Possible options are none|cpu|nvme. Only applicable with ZeRO Stages 2 and 3."},
+    )
+    offload_param_device: bool = field(
+        default=None,
+        metadata={"help": "Possible options are none|cpu|nvme. Only applicable with ZeRO Stage 3."},
+    )
+    zero3_init_flag: bool = field(
+        default=None,
+        metadata={
+            "help": "Flag to indicate whether to enable `deepspeed.zero.Init` for constructing massive models."
+            "Only applicable with ZeRO Stage-3."
+        },
+    )
+    zero3_save_16bit_model: bool = field(
+        default=None,
+        metadata={"help": "Flag to indicate whether to save 16-bit model. Only applicable with ZeRO Stage-3."},
     )
-
-    offload_optimizer_device: bool = field(default=None, metadata={"help": "Possible options are none|cpu|nvme"})
 
     def __post_init__(self):
-
-        if self.gradient_accumulation_steps is None:
-            self.gradient_accumulation_steps = int(os.environ.get("GRADIENT_ACCUMULATION_STEPS", 1))
-
-        if self.zero_stage is None:
-            self.zero_stage = int(os.environ.get("DEEPSPEED_ZERO_STAGE", 2))
-
-        if self.offload_optimizer_device is None:
-            self.offload_optimizer_device = os.environ.get("DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE", "none")
-
-        self.deepspeed_config = {
-            "train_batch_size": None,
-            "gradient_accumulation_steps": self.gradient_accumulation_steps,
-            "zero_optimization": {
-                "stage": self.zero_stage,
-                "offload_optimizer": {
-                    "device": self.offload_optimizer_device,
+        if self.config_file is None:
+            self.config_file = os.environ.get("DEEPSPEED_CONFIG_FILE", "none")
+        if self.config_file != "none":
+            with io.open(self.config_file, "r", encoding="utf-8") as f:
+                self.deepspeed_config = json.load(f)
+            if "gradient_accumulation_steps" not in self.deepspeed_config:
+                self.deepspeed_config["gradient_accumulation_steps"] = 1
+            elif self.deepspeed_config["gradient_accumulation_steps"] == "auto":
+                raise ValueError("gradient_accumulation_steps cannot be set to 'auto' in the DeepSpeed config file.")
+            if "zero_optimization" not in self.deepspeed_config:
+                raise ValueError("Please specify the ZeRO optimization config in the DeepSpeed config file.")
+        else:
+            if self.gradient_accumulation_steps is None:
+                self.gradient_accumulation_steps = int(os.environ.get("GRADIENT_ACCUMULATION_STEPS", 1))
+
+            if self.gradient_clipping is None:
+                gradient_clipping = os.environ.get("GRADIENT_CLIPPING", "none")
+                if gradient_clipping != "none":
+                    self.gradient_clipping = float(gradient_clipping)
+
+            if self.zero_stage is None:
+                self.zero_stage = int(os.environ.get("DEEPSPEED_ZERO_STAGE", 2))
+
+            if self.offload_optimizer_device is None:
+                self.offload_optimizer_device = os.environ.get("DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE", "none")
+
+            if self.offload_param_device is None:
+                self.offload_param_device = os.environ.get("DEEPSPEED_OFFLOAD_PARAM_DEVICE", "none")
+
+            if self.zero3_save_16bit_model is None:
+                self.zero3_save_16bit_model = os.environ.get("DEEPSPEED_ZERO3_SAVE_16BIT_MODEL", "false") == "true"
+
+            self.deepspeed_config = {
+                "train_batch_size": "auto",
+                "gradient_accumulation_steps": self.gradient_accumulation_steps,
+                "zero_optimization": {
+                    "stage": self.zero_stage,
+                    "offload_optimizer": {
+                        "device": self.offload_optimizer_device,
+                    },
+                    "offload_param": {
+                        "device": self.offload_param_device,
+                    },
+                    "stage3_gather_16bit_weights_on_model_save": self.zero3_save_16bit_model,
                 },
-            },
-            "steps_per_print": float("inf"),  # this will stop deepspeed from logging @ stdout
-            "zero_allow_untested_optimizer": True,
-        }
+            }
+            if self.gradient_clipping:
+                self.deepspeed_config["gradient_clipping"] = self.gradient_clipping
+        self.deepspeed_config["steps_per_print"] = float("inf")  # this will stop deepspeed from logging @ stdout
+        if self.zero3_init_flag is None:
+            self.zero3_init_flag = os.environ.get("DEEPSPEED_ZERO3_INIT", "false") == "true"
+        if self.zero3_init_flag and self.deepspeed_config["zero_optimization"]["stage"] != 3:
+            warnings.warn("DeepSpeed Zero3 Init flag is only applicable for ZeRO Stage 3. Setting it to False.")
+            self.zero3_init_flag = False
+
+    def find_config_node(self, ds_key_long):
+        config = self.deepspeed_config
+
+        # find the config node of interest if it exists
+        nodes = ds_key_long.split(".")
+        ds_key = nodes.pop()
+        for node in nodes:
+            config = config.get(node)
+            if config is None:
+                return None, ds_key
+
+        return config, ds_key
+
+    def fill_match(self, ds_key_long, mismatches, must_match=True, **kwargs):
+        config, ds_key = self.find_config_node(ds_key_long)
+        if config is None:
+            return
+
+        if config.get(ds_key) == "auto":
+            if ds_key_long in kwargs:
+                config[ds_key] = kwargs[ds_key_long]
+                return
+            else:
+                raise ValueError(
+                    f"`{ds_key_long}` not found in kwargs. "
+                    f"Please specify `{ds_key_long}` without `auto`(set to correct value) in the DeepSpeed config file or "
+                    "pass it in kwargs."
+                )
+
+        if not must_match:
+            return
+
+        ds_val = config.get(ds_key)
+        if ds_val is not None and ds_key_long in kwargs:
+            if ds_val != kwargs[ds_key_long]:
+                mismatches.append(f"- ds {ds_key_long}={ds_val} vs arg {ds_key_long}={kwargs[ds_key_long]}")
+
+    def deepspeed_config_process(self, prefix="", mismatches=None, config=None, must_match=True, **kwargs):
+        """Process the DeepSpeed config with the values from the kwargs."""
+        mismatches = [] if mismatches is None else mismatches
+        if config is None:
+            config = self.deepspeed_config
+        for key, value in config.items():
+            if isinstance(value, dict):
+                self.deepspeed_config_process(
+                    prefix=prefix + key + ".", mismatches=mismatches, config=value, must_match=must_match, **kwargs
+                )
+            else:
+                self.fill_match(prefix + key, mismatches, must_match=must_match, **kwargs)
+        if len(mismatches) > 0 and prefix == "":
+            mismatches_msg = "\n".join(mismatches)
+            raise ValueError(
+                "Please correct the following DeepSpeed config values that mismatch kwargs "
+                f" values:\n{mismatches_msg}\nThe easiest method is to set these DeepSpeed config values to 'auto'."
+            )
 
 
 @dataclass
diff --git a/src/accelerate/utils/deepspeed.py b/src/accelerate/utils/deepspeed.py
index 891f6fdb8..dde1b9c63 100644
--- a/src/accelerate/utils/deepspeed.py
+++ b/src/accelerate/utils/deepspeed.py
@@ -12,58 +12,34 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ..optimizer import AcceleratedOptimizer
-from .imports import is_apex_available, is_deepspeed_available
-
-
-if is_deepspeed_available():
-    from deepspeed import DeepSpeedEngine
+from accelerate.scheduler import AcceleratedScheduler
 
-if is_apex_available():
-    from apex import amp
+from ..optimizer import AcceleratedOptimizer
 
 
-class DeepSpeedEngineWrapper(DeepSpeedEngine):
-    """
-    Wrapper over deepspeed.DeepSpeedEngine object
+class DeepSpeedEngineWrapper:
     """
+    Internal wrapper for deepspeed.runtime.engine.DeepSpeedEngine. This is used to follow conventional training loop.
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-        # overwriting micro_steps for user's gradient_accumulation
-        self.micro_steps = -1
-
-    def step(self, lr_kwargs=None):
-        """DeepSpeedEngine.step() without `micro_steps` update & no profiling"""
-        if self.is_gradient_accumulation_boundary():  # it shouldn't matter whether we keep this line or not
-            if self.progressive_layer_drop:
-                self.progressive_layer_drop.update_state(self.global_steps)
+    Args:
+        engine (deepspeed.runtime.engine.DeepSpeedEngine): deepspeed engine to wrap
+    """
 
-            self._take_model_step(lr_kwargs)
+    def __init__(self, engine):
+        self.engine = engine
 
     def backward(self, loss):
-        """DeepSpeedEngine.backward() with with no loss scaling; no profiling but with `micro_steps` update"""
-
-        if self.zero_optimization():
-            self.optimizer.is_gradient_accumulation_boundary = self.is_gradient_accumulation_boundary()
-            self.optimizer.backward(loss)
-        elif self.amp_enabled():
-            # AMP requires delaying unscale when inside gradient accumulation boundaries
-            # https://nvidia.github.io/apex/advanced.html#gradient-accumulation-across-iterations
-            delay_unscale = not self.is_gradient_accumulation_boundary()
-            with amp.scale_loss(loss, self.optimizer, delay_unscale=delay_unscale) as scaled_loss:
-                scaled_loss.backward()
-        elif self.fp16_enabled():
-            self.optimizer.backward(loss)
-        else:
-            loss.backward()
+        # runs backpropagation and handles mixed precision
+        self.engine.backward(loss)
 
-        if self.enable_backward_allreduce:
-            self.allreduce_gradients()
-
-        # this will ensure deepspeed gradient_accumulation matches user's accumulation
-        self.micro_steps += 1
+        # deepspeed `engine.step` performs following operations:
+        # gradient accumulation check
+        # gradient clipping
+        # optimizer step
+        # zero grad
+        # checking overflow
+        # lr_scheduler step
+        self.engine.step()
 
 
 class DeepSpeedOptimizerWrapper(AcceleratedOptimizer):
@@ -75,22 +51,79 @@ class DeepSpeedOptimizerWrapper(AcceleratedOptimizer):
             The optimizer to wrap.
     """
 
-    def __init__(self, optimizer, model: DeepSpeedEngineWrapper):
+    def __init__(self, optimizer):
         super().__init__(optimizer, device_placement=False, scaler=None)
 
-        self.model = model
-
     def zero_grad(self, set_to_none=None):
-        pass  # `model.step()` is doing that automatically. Therefore, it's implementation is not needed
+        pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, it's implementation is not needed
 
     def step(self):
-        """This will handle optimizer.step() & optimizer.zero_grad() with gradient_accumulation"""
-        self.model.step()
+        pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, it's implementation is not needed
 
     @property
-    def is_overflow(self):
+    def step_was_skipped(self):
         """Whether or not the optimizer step was done, or skipped because of gradient overflow."""
-        overflow = False
-        if hasattr(self.optimizer, "overflow"):
-            overflow = self.optimizer.overflow
-        return overflow
+        return self.optimizer.overflow
+
+
+class DeepSpeedSchedulerWrapper(AcceleratedScheduler):
+    """
+    Internal wrapper around a deepspeed scheduler.
+
+    Args:
+        scheduler (`torch.optim.lr_scheduler.LambdaLR`):
+            The scheduler to wrap.
+        optimizers (one or a list of `torch.optim.Optimizer`):
+    """
+
+    def __init__(self, scheduler, optimizers):
+        super().__init__(scheduler, optimizers)
+
+    def step(self):
+        pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, it's implementation is not needed
+
+
+class DummyOptim:
+    """
+    Dummy optimizer presents model parameters or param groups, this is primarily used to follow conventional training
+    loop when optimizer config is specified in the deepspeed config file.
+
+    Args:
+        lr (float):
+            Learning rate.
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        weight_decay (float):
+            Weight decay.
+        **kwargs:
+            Other arguments.
+    """
+
+    def __init__(self, params, lr=0.001, weight_decay=0, **kwargs):
+        self.params = params
+        self.lr = lr
+        self.weight_decay = weight_decay
+        self.kwargs = kwargs
+
+
+class DummyScheduler:
+    """
+    Dummy scheduler presents model parameters or param groups, this is primarily used to follow conventional training
+    loop when scheduler config is specified in the deepspeed config file.
+
+    Args:
+        optimizer (`torch.optim.optimizer.Optimizer`):
+            The optimizer to wrap.
+        total_num_steps (int):
+            Total number of steps.
+        warmup_num_steps (int):
+            Number of steps for warmup.
+        **kwargs:
+            Other arguments.
+    """
+
+    def __init__(self, optimizer, total_num_steps=None, warmup_num_steps=0, **kwargs):
+        self.optimizer = optimizer
+        self.total_num_steps = total_num_steps
+        self.warmup_num_steps = warmup_num_steps
+        self.kwargs = kwargs
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index 947508726..c8cce9587 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -63,6 +63,10 @@ def is_deepspeed_available():
             return False
 
 
+def is_transformers_available():
+    return importlib.util.find_spec("transformers") is not None
+
+
 def is_tensorboard_available():
     return importlib.util.find_spec("tensorboard") is not None or importlib.util.find_spec("tensorboardX") is not None
 
diff --git a/tests/deepspeed/ds_config_zero2.json b/tests/deepspeed/ds_config_zero2.json
new file mode 100644
index 000000000..f031969ee
--- /dev/null
+++ b/tests/deepspeed/ds_config_zero2.json
@@ -0,0 +1,49 @@
+{
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "bf16": {
+        "enabled": "auto"
+    },
+    "optimizer": {
+        "type": "AdamW",
+        "params": {
+            "lr": "auto",
+            "weight_decay": "auto",
+            "torch_adam": true,
+            "adam_w_mode": true
+        }
+    },
+    "scheduler": {
+        "type": "WarmupLR",
+        "params": {
+            "warmup_min_lr": "auto",
+            "warmup_max_lr": "auto",
+            "warmup_num_steps": "auto"
+        }
+    },
+    "zero_optimization": {
+        "stage": 2,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "allgather_partitions": true,
+        "allgather_bucket_size": 2e8,
+        "overlap_comm": true,
+        "reduce_scatter": true,
+        "reduce_bucket_size": "auto",
+        "contiguous_gradients": true
+    },
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}
\ No newline at end of file
diff --git a/tests/deepspeed/ds_config_zero3.json b/tests/deepspeed/ds_config_zero3.json
new file mode 100644
index 000000000..846cd7361
--- /dev/null
+++ b/tests/deepspeed/ds_config_zero3.json
@@ -0,0 +1,56 @@
+{
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "bf16": {
+        "enabled": "auto"
+    },
+    "optimizer": {
+        "type": "AdamW",
+        "params": {
+            "lr": "auto",
+            "weight_decay": "auto",
+            "torch_adam": true,
+            "adam_w_mode": true
+        }
+    },
+    "scheduler": {
+        "type": "WarmupLR",
+        "params": {
+            "warmup_min_lr": "auto",
+            "warmup_max_lr": "auto",
+            "warmup_num_steps": "auto"
+        }
+    },
+    "zero_optimization": {
+        "stage": 3,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "offload_param": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "sub_group_size": 1e9,
+        "reduce_bucket_size": "auto",
+        "stage3_prefetch_bucket_size": "auto",
+        "stage3_param_persistence_threshold": "auto",
+        "stage3_max_live_parameters": 1e9,
+        "stage3_max_reuse_distance": 1e9,
+        "stage3_gather_16bit_weights_on_model_save": "auto"
+    },
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}
\ No newline at end of file
diff --git a/tests/deepspeed/test_deepspeed.py b/tests/deepspeed/test_deepspeed.py
new file mode 100644
index 000000000..dbe64b5c7
--- /dev/null
+++ b/tests/deepspeed/test_deepspeed.py
@@ -0,0 +1,581 @@
+# 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 inspect
+import io
+import itertools
+import json
+import os
+import tempfile
+import unittest
+from copy import deepcopy
+from pathlib import Path
+
+import torch
+from torch.utils.data import DataLoader
+
+from accelerate.accelerator import Accelerator
+from accelerate.scheduler import AcceleratedScheduler
+from accelerate.state import AcceleratorState
+from accelerate.test_utils.testing import require_cuda, require_deepspeed
+from accelerate.test_utils.training import RegressionDataset
+from accelerate.utils.dataclasses import DeepSpeedPlugin
+from accelerate.utils.deepspeed import (
+    DeepSpeedEngineWrapper,
+    DeepSpeedOptimizerWrapper,
+    DeepSpeedSchedulerWrapper,
+    DummyOptim,
+    DummyScheduler,
+)
+from parameterized import parameterized
+from transformers import AutoModel, AutoModelForCausalLM, get_scheduler
+from transformers.deepspeed import HfDeepSpeedConfig
+from transformers.testing_utils import mockenv_context
+from transformers.trainer_utils import set_seed
+from transformers.utils import is_torch_bf16_available
+
+
+set_seed(42)
+
+T5_SMALL = "t5-small"
+T5_TINY = "patrickvonplaten/t5-tiny-random"
+GPT2_TINY = "sshleifer/tiny-gpt2"
+
+ZERO2 = "zero2"
+ZERO3 = "zero3"
+
+FP16 = "fp16"
+BF16 = "bf16"
+
+CUSTOM_OPTIMIZER = "custom_optimizer"
+CUSTOM_SCHEDULER = "custom_scheduler"
+DS_OPTIMIZER = "deepspeed_optimizer"
+DS_SCHEDULER = "deepspeed_scheduler"
+
+stages = [ZERO2, ZERO3]
+optims = [CUSTOM_OPTIMIZER, DS_OPTIMIZER]
+schedulers = [CUSTOM_SCHEDULER, DS_SCHEDULER]
+if is_torch_bf16_available():
+    dtypes = [FP16, BF16]
+else:
+    dtypes = [FP16]
+
+
+def parameterized_custom_name_func(func, param_num, param):
+    # customize the test name generator function as we want both params to appear in the sub-test
+    # name, as by default it shows only the first param
+    param_based_name = parameterized.to_safe_name("_".join(str(x) for x in param.args))
+    return f"{func.__name__}_{param_based_name}"
+
+
+# Cartesian-product of zero stages with models to test
+params = list(itertools.product(stages, dtypes))
+optim_scheduler_params = list(itertools.product(optims, schedulers))
+
+
+@require_deepspeed
+@require_cuda
+class DeepSpeedConfigIntegration(unittest.TestCase):
+    def setUp(self):
+        super().setUp()
+
+        self._test_file_path = inspect.getfile(self.__class__)
+        path = Path(self._test_file_path).resolve()
+        self.test_file_dir_str = str(path.parents[0])
+
+        self.ds_config_file = dict(
+            zero2=f"{self.test_file_dir_str}/ds_config_zero2.json",
+            zero3=f"{self.test_file_dir_str}/ds_config_zero3.json",
+        )
+
+        # use self.get_config_dict(stage) to use these to ensure the original is not modified
+        with io.open(self.ds_config_file[ZERO2], "r", encoding="utf-8") as f:
+            config_zero2 = json.load(f)
+        with io.open(self.ds_config_file[ZERO3], "r", encoding="utf-8") as f:
+            config_zero3 = json.load(f)
+            # The following setting slows things down, so don't enable it by default unless needed by a test.
+            # It's in the file as a demo for users since we want everything to work out of the box even if slower.
+            config_zero3["zero_optimization"]["stage3_gather_16bit_weights_on_model_save"] = False
+
+        self.ds_config_dict = dict(zero2=config_zero2, zero3=config_zero3)
+
+        self.dist_env = dict(
+            USE_DEEPSPEED="true",
+            MASTER_ADDR="localhost",
+            MASTER_PORT="10999",
+            RANK="0",
+            LOCAL_RANK="0",
+            WORLD_SIZE="1",
+        )
+
+    def get_config_dict(self, stage):
+        # As some tests modify the dict, always make a copy
+        return deepcopy(self.ds_config_dict[stage])
+
+    @parameterized.expand(stages, name_func=parameterized_custom_name_func)
+    def test_deepspeed_plugin(self, stage):
+
+        # Test zero3_init_flag will be set to False when ZeRO stage != 3
+        deepspeed_plugin = DeepSpeedPlugin(
+            gradient_accumulation_steps=1,
+            gradient_clipping=1.0,
+            zero_stage=2,
+            offload_optimizer_device="cpu",
+            offload_param_device="cpu",
+            zero3_save_16bit_model=True,
+            zero3_init_flag=True,
+        )
+        self.assertFalse(deepspeed_plugin.zero3_init_flag)
+        deepspeed_plugin.deepspeed_config = None
+
+        # Test zero3_init_flag will be set to True only when ZeRO stage == 3
+        deepspeed_plugin = DeepSpeedPlugin(
+            gradient_accumulation_steps=1,
+            gradient_clipping=1.0,
+            zero_stage=3,
+            offload_optimizer_device="cpu",
+            offload_param_device="cpu",
+            zero3_save_16bit_model=True,
+            zero3_init_flag=True,
+        )
+        self.assertTrue(deepspeed_plugin.zero3_init_flag)
+        deepspeed_plugin.deepspeed_config = None
+
+        # Test config files are loaded correctly
+        deepspeed_plugin = DeepSpeedPlugin(config_file=self.ds_config_file[stage], zero3_init_flag=True)
+        if stage == ZERO2:
+            self.assertFalse(deepspeed_plugin.zero3_init_flag)
+        elif stage == ZERO3:
+            self.assertTrue(deepspeed_plugin.zero3_init_flag)
+            deepspeed_plugin.deepspeed_config = None
+
+        # Test `gradient_accumulation_steps` is set to 1 if unavailable in config file
+        with tempfile.TemporaryDirectory() as dirpath:
+            ds_config = self.get_config_dict(stage)
+            del ds_config["gradient_accumulation_steps"]
+            with open(os.path.join(dirpath, "ds_config.json"), "w") as out_file:
+                json.dump(ds_config, out_file)
+            deepspeed_plugin = DeepSpeedPlugin(config_file=os.path.join(dirpath, "ds_config.json"))
+            self.assertEqual(deepspeed_plugin.deepspeed_config["gradient_accumulation_steps"], 1)
+            deepspeed_plugin.deepspeed_config = None
+
+        # Test `ValueError` is raised if `zero_optimization` is unavailable in config file
+        with tempfile.TemporaryDirectory() as dirpath:
+            ds_config = self.get_config_dict(stage)
+            del ds_config["zero_optimization"]
+            with open(os.path.join(dirpath, "ds_config.json"), "w") as out_file:
+                json.dump(ds_config, out_file)
+            with self.assertRaises(ValueError) as cm:
+                deepspeed_plugin = DeepSpeedPlugin(config_file=os.path.join(dirpath, "ds_config.json"))
+            self.assertTrue(
+                "Please specify the ZeRO optimization config in the DeepSpeed config file." in str(cm.exception)
+            )
+            deepspeed_plugin.deepspeed_config = None
+
+        # Test `deepspeed_config_process`
+        deepspeed_plugin = DeepSpeedPlugin(config_file=self.ds_config_file[stage])
+        kwargs = {
+            "fp16.enabled": True,
+            "bf16.enabled": False,
+            "optimizer.params.lr": 5e-5,
+            "optimizer.params.weight_decay": 0.0,
+            "scheduler.params.warmup_min_lr": 0.0,
+            "scheduler.params.warmup_max_lr": 5e-5,
+            "scheduler.params.warmup_num_steps": 0,
+            "train_micro_batch_size_per_gpu": 16,
+            "gradient_clipping": 1.0,
+            "train_batch_size": 16,
+            "zero_optimization.reduce_bucket_size": 5e5,
+            "zero_optimization.stage3_prefetch_bucket_size": 5e5,
+            "zero_optimization.stage3_param_persistence_threshold": 5e5,
+            "zero_optimization.stage3_gather_16bit_weights_on_model_save": False,
+        }
+        deepspeed_plugin.deepspeed_config_process(**kwargs)
+        for ds_key_long, value in kwargs.items():
+            config, ds_key = deepspeed_plugin.find_config_node(ds_key_long)
+            if config.get(ds_key) is not None:
+                self.assertEqual(config.get(ds_key), value)
+
+        # Test mismatches
+        mismatches = {
+            "optimizer.params.lr": 1e-5,
+            "optimizer.params.weight_decay": 1e-5,
+            "gradient_accumulation_steps": 2,
+        }
+        with self.assertRaises(ValueError) as cm:
+            new_kwargs = deepcopy(kwargs)
+            new_kwargs.update(mismatches)
+            deepspeed_plugin.deepspeed_config_process(**new_kwargs)
+        for key in mismatches.keys():
+            self.assertTrue(
+                key in str(cm.exception),
+                f"{key} is not in the exception message:\n{cm.exception}",
+            )
+
+        # Test `ValueError` is raised if some config file fields with `auto` value is missing in `kwargs`
+        deepspeed_plugin.deepspeed_config["optimizer"]["params"]["lr"] = "auto"
+        with self.assertRaises(ValueError) as cm:
+            del kwargs["optimizer.params.lr"]
+            deepspeed_plugin.deepspeed_config_process(**kwargs)
+        self.assertTrue("`optimizer.params.lr` not found in kwargs." in str(cm.exception))
+
+    @parameterized.expand([FP16, BF16], name_func=parameterized_custom_name_func)
+    def test_accelerate_state_deepspeed(self, dtype):
+        deepspeed_plugin = DeepSpeedPlugin(
+            gradient_accumulation_steps=1,
+            gradient_clipping=1.0,
+            zero_stage=ZERO2,
+            offload_optimizer_device="cpu",
+            offload_param_device="cpu",
+            zero3_save_16bit_model=True,
+            zero3_init_flag=True,
+        )
+        with mockenv_context(**self.dist_env):
+            state = AcceleratorState(mixed_precision=dtype, deepspeed_plugin=deepspeed_plugin, _from_accelerator=True)
+            self.assertTrue(state.deepspeed_plugin.deepspeed_config[dtype]["enabled"])
+            state.initialized = False
+
+    def test_init_zero3(self):
+        deepspeed_plugin = DeepSpeedPlugin(
+            gradient_accumulation_steps=1,
+            gradient_clipping=1.0,
+            zero_stage=3,
+            offload_optimizer_device="cpu",
+            offload_param_device="cpu",
+            zero3_save_16bit_model=True,
+            zero3_init_flag=True,
+        )
+
+        with mockenv_context(**self.dist_env):
+            accelerator = Accelerator(deepspeed_plugin=deepspeed_plugin)
+            self.assertTrue("dschf" in accelerator.__dict__)
+            self.assertTrue(type(accelerator.dschf) == HfDeepSpeedConfig)
+
+    @parameterized.expand(optim_scheduler_params, name_func=parameterized_custom_name_func)
+    def test_prepare_deepspeed(self, optim_type, scheduler_type):
+        # 1. Testing with one of the ZeRO Stages is enough to test the `_prepare_deepspeed` function.
+        # Here we test using ZeRO Stage 2 with FP16 enabled.
+        from deepspeed.runtime.engine import DeepSpeedEngine
+
+        kwargs = {
+            "fp16.enabled": True,
+            "bf16.enabled": False,
+            "optimizer.params.lr": 5e-5,
+            "optimizer.params.weight_decay": 0.0,
+            "scheduler.params.warmup_min_lr": 0.0,
+            "scheduler.params.warmup_max_lr": 5e-5,
+            "scheduler.params.warmup_num_steps": 0,
+            "train_micro_batch_size_per_gpu": 16,
+            "gradient_clipping": 1.0,
+            "train_batch_size": 16,
+            "zero_optimization.reduce_bucket_size": 5e5,
+            "zero_optimization.stage3_prefetch_bucket_size": 5e5,
+            "zero_optimization.stage3_param_persistence_threshold": 5e5,
+            "zero_optimization.stage3_gather_16bit_weights_on_model_save": False,
+        }
+
+        if optim_type == CUSTOM_OPTIMIZER and scheduler_type == CUSTOM_SCHEDULER:
+            # Test custom optimizer + custom scheduler
+            deepspeed_plugin = DeepSpeedPlugin(
+                gradient_accumulation_steps=1,
+                gradient_clipping=1.0,
+                zero_stage=2,
+                offload_optimizer_device="cpu",
+                offload_param_device="cpu",
+                zero3_save_16bit_model=False,
+                zero3_init_flag=False,
+            )
+            with mockenv_context(**self.dist_env):
+                accelerator = Accelerator(mixed_precision="fp16", deepspeed_plugin=deepspeed_plugin)
+                self.assertEqual(accelerator.state.deepspeed_plugin.config_file, "none")
+
+                train_set = RegressionDataset(length=80)
+                eval_set = RegressionDataset(length=20)
+                train_dataloader = DataLoader(train_set, batch_size=16, shuffle=True)
+                eval_dataloader = DataLoader(eval_set, batch_size=32, shuffle=False)
+                model = AutoModel.from_pretrained(GPT2_TINY)
+                optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)
+                lr_scheduler = get_scheduler(
+                    name="linear",
+                    optimizer=optimizer,
+                    num_warmup_steps=0,
+                    num_training_steps=1000,
+                )
+                dummy_optimizer = DummyOptim(params=model.parameters())
+                dummy_lr_scheduler = DummyScheduler(dummy_optimizer)
+
+                with self.assertRaises(ValueError) as cm:
+                    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+                        model, dummy_optimizer, train_dataloader, eval_dataloader, lr_scheduler
+                    )
+                self.assertTrue(
+                    "You cannot create a `DummyOptim` without specifying an optimizer in the config file."
+                    in str(cm.exception)
+                )
+                with self.assertRaises(ValueError) as cm:
+                    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+                        model, optimizer, train_dataloader, eval_dataloader, dummy_lr_scheduler
+                    )
+                self.assertTrue(
+                    "You cannot create a `DummyScheduler` without specifying a scheduler in the config file."
+                    in str(cm.exception)
+                )
+
+                with self.assertRaises(ValueError) as cm:
+                    model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler)
+                self.assertTrue(
+                    "You must specify a training or evaluation dataloader in `accelerate.prepare()` when using DeepSpeed."
+                    in str(cm.exception)
+                )
+
+                model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+                    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+                )
+                self.assertTrue(accelerator.deepspeed_config["zero_allow_untested_optimizer"])
+                self.assertTrue(accelerator.deepspeed_config["train_batch_size"], 16)
+                self.assertEqual(type(model), DeepSpeedEngine)
+                self.assertEqual(type(optimizer), DeepSpeedOptimizerWrapper)
+                self.assertEqual(type(lr_scheduler), AcceleratedScheduler)
+                self.assertEqual(type(accelerator.deepspeed_engine_wrapped), DeepSpeedEngineWrapper)
+
+        elif optim_type == DS_OPTIMIZER and scheduler_type == DS_SCHEDULER:
+            # Test DeepSpeed optimizer + DeepSpeed scheduler
+            deepspeed_plugin = DeepSpeedPlugin(config_file=self.ds_config_file[ZERO2])
+            with mockenv_context(**self.dist_env):
+                accelerator = Accelerator(deepspeed_plugin=deepspeed_plugin)
+                train_set = RegressionDataset(length=80)
+                eval_set = RegressionDataset(length=20)
+                train_dataloader = DataLoader(train_set, batch_size=10, shuffle=True)
+                eval_dataloader = DataLoader(eval_set, batch_size=5, shuffle=False)
+                model = AutoModel.from_pretrained(GPT2_TINY)
+                optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)
+                lr_scheduler = get_scheduler(
+                    name="linear",
+                    optimizer=optimizer,
+                    num_warmup_steps=0,
+                    num_training_steps=1000,
+                )
+                dummy_optimizer = DummyOptim(params=model.parameters())
+                dummy_lr_scheduler = DummyScheduler(dummy_optimizer)
+                kwargs["train_batch_size"] = (
+                    kwargs["train_micro_batch_size_per_gpu"]
+                    * deepspeed_plugin.deepspeed_config["gradient_accumulation_steps"]
+                    * accelerator.num_processes
+                )
+                accelerator.state.deepspeed_plugin.deepspeed_config_process(**kwargs)
+                with self.assertRaises(ValueError) as cm:
+                    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+                        model, optimizer, train_dataloader, eval_dataloader, dummy_lr_scheduler
+                    )
+                self.assertTrue(
+                    "You cannot specify an optimizer in the config file and in the code at the same time"
+                    in str(cm.exception)
+                )
+
+                with self.assertRaises(ValueError) as cm:
+                    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+                        model, dummy_optimizer, train_dataloader, eval_dataloader, lr_scheduler
+                    )
+                self.assertTrue(
+                    "You cannot specify a scheduler in the config file and in the code at the same time"
+                    in str(cm.exception)
+                )
+
+                with self.assertRaises(ValueError) as cm:
+                    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+                        model, dummy_optimizer, train_dataloader, eval_dataloader, lr_scheduler
+                    )
+                self.assertTrue(
+                    "You cannot specify a scheduler in the config file and in the code at the same time"
+                    in str(cm.exception)
+                )
+
+                model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+                    model, dummy_optimizer, train_dataloader, eval_dataloader, dummy_lr_scheduler
+                )
+                self.assertTrue(type(model) == DeepSpeedEngine)
+                self.assertTrue(type(optimizer) == DeepSpeedOptimizerWrapper)
+                self.assertTrue(type(lr_scheduler) == DeepSpeedSchedulerWrapper)
+                self.assertTrue(type(accelerator.deepspeed_engine_wrapped) == DeepSpeedEngineWrapper)
+
+        elif optim_type == CUSTOM_OPTIMIZER and scheduler_type == DS_SCHEDULER:
+            # Test custom optimizer + DeepSpeed scheduler
+            deepspeed_plugin = DeepSpeedPlugin(config_file=self.ds_config_file[ZERO2])
+            with mockenv_context(**self.dist_env):
+                accelerator = Accelerator(deepspeed_plugin=deepspeed_plugin)
+                train_set = RegressionDataset(length=80)
+                eval_set = RegressionDataset(length=20)
+                train_dataloader = DataLoader(train_set, batch_size=10, shuffle=True)
+                eval_dataloader = DataLoader(eval_set, batch_size=5, shuffle=False)
+                model = AutoModel.from_pretrained(GPT2_TINY)
+                optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)
+                lr_scheduler = get_scheduler(
+                    name="linear",
+                    optimizer=optimizer,
+                    num_warmup_steps=0,
+                    num_training_steps=1000,
+                )
+                dummy_optimizer = DummyOptim(params=model.parameters())
+                dummy_lr_scheduler = DummyScheduler(dummy_optimizer)
+                kwargs["train_batch_size"] = (
+                    kwargs["train_micro_batch_size_per_gpu"]
+                    * deepspeed_plugin.deepspeed_config["gradient_accumulation_steps"]
+                    * accelerator.num_processes
+                )
+                accelerator.state.deepspeed_plugin.deepspeed_config_process(**kwargs)
+                del accelerator.state.deepspeed_plugin.deepspeed_config["optimizer"]
+                model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+                    model, optimizer, train_dataloader, eval_dataloader, dummy_lr_scheduler
+                )
+                self.assertTrue(type(model) == DeepSpeedEngine)
+                self.assertTrue(type(optimizer) == DeepSpeedOptimizerWrapper)
+                self.assertTrue(type(lr_scheduler) == DeepSpeedSchedulerWrapper)
+                self.assertTrue(type(accelerator.deepspeed_engine_wrapped) == DeepSpeedEngineWrapper)
+        elif optim_type == DS_OPTIMIZER and scheduler_type == CUSTOM_SCHEDULER:
+            # Test deepspeed optimizer + custom scheduler
+            deepspeed_plugin = DeepSpeedPlugin(config_file=self.ds_config_file[ZERO2])
+            with mockenv_context(**self.dist_env):
+                accelerator = Accelerator(deepspeed_plugin=deepspeed_plugin)
+                train_set = RegressionDataset(length=80)
+                eval_set = RegressionDataset(length=20)
+                train_dataloader = DataLoader(train_set, batch_size=10, shuffle=True)
+                eval_dataloader = DataLoader(eval_set, batch_size=5, shuffle=False)
+                model = AutoModel.from_pretrained(GPT2_TINY)
+                optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)
+                lr_scheduler = get_scheduler(
+                    name="linear",
+                    optimizer=optimizer,
+                    num_warmup_steps=0,
+                    num_training_steps=1000,
+                )
+                dummy_optimizer = DummyOptim(params=model.parameters())
+                dummy_lr_scheduler = DummyScheduler(dummy_optimizer)
+                kwargs["train_batch_size"] = (
+                    kwargs["train_micro_batch_size_per_gpu"]
+                    * deepspeed_plugin.deepspeed_config["gradient_accumulation_steps"]
+                    * accelerator.num_processes
+                )
+                accelerator.state.deepspeed_plugin.deepspeed_config_process(**kwargs)
+                del accelerator.state.deepspeed_plugin.deepspeed_config["scheduler"]
+                with self.assertRaises(ValueError) as cm:
+                    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+                        model, dummy_optimizer, train_dataloader, eval_dataloader, lr_scheduler
+                    )
+                self.assertTrue(
+                    "You can only specify `accelerate.utils.DummyScheduler` in the code when using `accelerate.utils.DummyOptim`."
+                    in str(cm.exception)
+                )
+        accelerator.state.initialized = False
+
+    def test_save_checkpoints(self):
+        deepspeed_plugin = DeepSpeedPlugin(
+            config_file=self.ds_config_file[ZERO3],
+            zero3_init_flag=True,
+        )
+        del deepspeed_plugin.deepspeed_config["bf16"]
+        kwargs = {
+            "fp16.enabled": True,
+            "bf16.enabled": False,
+            "optimizer.params.lr": 5e-5,
+            "optimizer.params.weight_decay": 0.0,
+            "scheduler.params.warmup_min_lr": 0.0,
+            "scheduler.params.warmup_max_lr": 5e-5,
+            "scheduler.params.warmup_num_steps": 0,
+            "train_micro_batch_size_per_gpu": 16,
+            "gradient_clipping": 1.0,
+            "train_batch_size": 16,
+            "zero_optimization.reduce_bucket_size": 5e5,
+            "zero_optimization.stage3_prefetch_bucket_size": 5e5,
+            "zero_optimization.stage3_param_persistence_threshold": 5e5,
+            "zero_optimization.stage3_gather_16bit_weights_on_model_save": False,
+        }
+
+        with mockenv_context(**self.dist_env):
+            accelerator = Accelerator(deepspeed_plugin=deepspeed_plugin)
+            kwargs["train_batch_size"] = (
+                kwargs["train_micro_batch_size_per_gpu"]
+                * deepspeed_plugin.deepspeed_config["gradient_accumulation_steps"]
+                * accelerator.num_processes
+            )
+            accelerator.state.deepspeed_plugin.deepspeed_config_process(**kwargs)
+
+            train_set = RegressionDataset(length=80)
+            eval_set = RegressionDataset(length=20)
+            train_dataloader = DataLoader(train_set, batch_size=16, shuffle=True)
+            eval_dataloader = DataLoader(eval_set, batch_size=32, shuffle=False)
+            model = AutoModelForCausalLM.from_pretrained("gpt2")
+            dummy_optimizer = DummyOptim(params=model.parameters())
+            dummy_lr_scheduler = DummyScheduler(dummy_optimizer)
+
+            model, _, train_dataloader, eval_dataloader, _ = accelerator.prepare(
+                model, dummy_optimizer, train_dataloader, eval_dataloader, dummy_lr_scheduler
+            )
+            with self.assertRaises(ValueError) as cm:
+                accelerator.get_state_dict(model)
+            msg = (
+                "Cannot get 16bit model weights because `stage3_gather_16bit_weights_on_model_save` in DeepSpeed config is False. "
+                "To save the model weights in 16bit, set `stage3_gather_16bit_weights_on_model_save` to True in DeepSpeed config file or "
+                "set `zero3_save_16bit_model` to True when using `accelerate config`. "
+                "To save the full checkpoint, run `model.save_checkpoint(save_dir)` and use `zero_to_fp32.py` to recover weights."
+            )
+            self.assertTrue(msg in str(cm.exception))
+        accelerator.state.initialized = False
+
+    def test_autofill_dsconfig(self):
+        deepspeed_plugin = DeepSpeedPlugin(
+            config_file=self.ds_config_file[ZERO3],
+            zero3_init_flag=True,
+        )
+        del deepspeed_plugin.deepspeed_config["bf16"]
+        del deepspeed_plugin.deepspeed_config["fp16"]
+
+        with mockenv_context(**self.dist_env):
+            accelerator = Accelerator(deepspeed_plugin=deepspeed_plugin)
+            train_set = RegressionDataset(length=80)
+            eval_set = RegressionDataset(length=20)
+            train_dataloader = DataLoader(train_set, batch_size=16, shuffle=True)
+            eval_dataloader = DataLoader(eval_set, batch_size=32, shuffle=False)
+            model = AutoModelForCausalLM.from_pretrained("gpt2")
+            dummy_optimizer = DummyOptim(params=model.parameters(), lr=5e-5, weight_decay=1e-4)
+            dummy_lr_scheduler = DummyScheduler(dummy_optimizer, warmup_num_steps=10, total_num_steps=1000)
+            hidden_size = model.config.hidden_size
+            model, _, train_dataloader, eval_dataloader, _ = accelerator.prepare(
+                model, dummy_optimizer, train_dataloader, eval_dataloader, dummy_lr_scheduler
+            )
+            self.assertEqual(accelerator.deepspeed_config["train_micro_batch_size_per_gpu"], 16)
+            self.assertEqual(accelerator.deepspeed_config["train_batch_size"], 16)
+
+            self.assertEqual(accelerator.deepspeed_config["optimizer"]["params"]["lr"], 5e-5)
+            self.assertEqual(accelerator.deepspeed_config["optimizer"]["params"]["weight_decay"], 1e-4)
+
+            self.assertEqual(accelerator.deepspeed_config["scheduler"]["params"]["warmup_min_lr"], 0.0)
+            self.assertEqual(accelerator.deepspeed_config["scheduler"]["params"]["warmup_max_lr"], 5e-5)
+            self.assertEqual(accelerator.deepspeed_config["scheduler"]["params"]["warmup_num_steps"], 10)
+
+            self.assertEqual(accelerator.deepspeed_config["gradient_clipping"], 1.0)
+            self.assertEqual(
+                accelerator.deepspeed_config["zero_optimization"]["reduce_bucket_size"], hidden_size * hidden_size
+            )
+            self.assertEqual(
+                accelerator.deepspeed_config["zero_optimization"]["stage3_prefetch_bucket_size"],
+                0.9 * hidden_size * hidden_size,
+            )
+            self.assertEqual(
+                accelerator.deepspeed_config["zero_optimization"]["stage3_param_persistence_threshold"],
+                10 * hidden_size,
+            )
+            self.assertFalse(
+                accelerator.deepspeed_config["zero_optimization"]["stage3_gather_16bit_weights_on_model_save"]
+            )
+        accelerator.state.initialized = False
diff --git a/tests/test_examples.py b/tests/test_examples.py
index 8f277b932..6ad4a3ca8 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -32,7 +32,13 @@
 # Should mock `{script_name}.get_dataloaders` via:
 # @mock.patch("{script_name}.get_dataloaders", mocked_dataloaders)
 
-EXCLUDE_EXAMPLES = ["cross_validation.py", "multi_process_metrics.py", "memory.py", "fsdp_with_peak_mem_tracking.py"]
+EXCLUDE_EXAMPLES = [
+    "cross_validation.py",
+    "multi_process_metrics.py",
+    "memory.py",
+    "fsdp_with_peak_mem_tracking.py",
+    "deepspeed_with_config_support.py",
+]
 
 
 class ExampleDifferenceTests(unittest.TestCase):

```suggestion
                    "DeepSpeed is not installed => run `pip install deepspeed` or build it from source."
```
We usually just say pip in those error messages.
@@ -163,6 +172,25 @@ def __init__(
                 deepspeed_plugin, DeepSpeedPlugin
             ), "`deepspeed_plugin` must be a DeepSpeedPlugin object."
             os.environ["USE_DEEPSPEED"] = "true"  # use DeepSpeed if plugin is provided
+        if deepspeed_plugin:
+            if not is_deepspeed_available():
+                raise ImportError(
+                    "DeepSpeed is not installed => run `pip3 install deepspeed` or build it from source."
```suggestion
                        "Please run `pip install transformers`."
```
@@ -163,6 +172,25 @@ def __init__(
                 deepspeed_plugin, DeepSpeedPlugin
             ), "`deepspeed_plugin` must be a DeepSpeedPlugin object."
             os.environ["USE_DEEPSPEED"] = "true"  # use DeepSpeed if plugin is provided
+        if deepspeed_plugin:
+            if not is_deepspeed_available():
+                raise ImportError(
+                    "DeepSpeed is not installed => run `pip3 install deepspeed` or build it from source."
+                )
+            if compare_versions("deepspeed", "<", "0.6.4"):
+                raise ImportError("DeepSpeed version must be >= 0.6.4. Please update DeepSpeed.")
+            if os.environ.get("DEEPSPEED_ZERO3_INIT", "false") == "true" or deepspeed_plugin.zero3_init_flag:
+                if not is_transformers_available():
+                    raise Exception(
+                        "When `zero3_init_flag` is set, it requires Transformers to be installed. "
+                        "Please run `pip3 install transformers`."
Ultimately, we will want this object to live in Accelerate, not in Transformers. I don't know when is the best time to move it, just putting it as a general comment :-)
@@ -163,6 +172,25 @@ def __init__(
                 deepspeed_plugin, DeepSpeedPlugin
             ), "`deepspeed_plugin` must be a DeepSpeedPlugin object."
             os.environ["USE_DEEPSPEED"] = "true"  # use DeepSpeed if plugin is provided
+        if deepspeed_plugin:
+            if not is_deepspeed_available():
+                raise ImportError(
+                    "DeepSpeed is not installed => run `pip3 install deepspeed` or build it from source."
+                )
+            if compare_versions("deepspeed", "<", "0.6.4"):
+                raise ImportError("DeepSpeed version must be >= 0.6.4. Please update DeepSpeed.")
+            if os.environ.get("DEEPSPEED_ZERO3_INIT", "false") == "true" or deepspeed_plugin.zero3_init_flag:
+                if not is_transformers_available():
+                    raise Exception(
+                        "When `zero3_init_flag` is set, it requires Transformers to be installed. "
+                        "Please run `pip3 install transformers`."
+                    )
+                from transformers.deepspeed import HfDeepSpeedConfig
Put the comment above so everything fits on one line.
@@ -499,81 +527,116 @@ def _prepare_deepspeed(self, *args):
         deepspeed_plugin = self.state.deepspeed_plugin
         self.deepspeed_config = deepspeed_plugin.deepspeed_config
 
-        batch_sizes = [obj.batch_size for obj in args if hasattr(obj, "batch_size")]
-        if len(batch_sizes) == 0:
-            raise ValueError(
-                "You must specify a training or evaluation dataloader in `accelerate.prepare()` when using DeepSpeed."
-            )
-
-        batch_size_per_device = min(batch_sizes) if deepspeed_plugin.is_train_batch_min else max(batch_sizes)
-        if len(batch_sizes) > 1:
-            logger.info(
-                f"Since you passed both train and evaluation dataloader, `is_train_batch_min` (here \
-                {deepspeed_plugin.is_train_batch_min} will decide the `train_batch_size` ({batch_size_per_device})."
-            )
-
-        self.deepspeed_config["train_batch_size"] = (
-            batch_size_per_device * deepspeed_plugin.gradient_accumulation_steps * self.num_processes
-        )
-
         result = [
             self._prepare_one(obj, first_pass=True) if isinstance(obj, torch.utils.data.DataLoader) else obj
             for obj in args
         ]
 
+        if deepspeed_plugin.config_file == "none":
+            batch_sizes = [obj.batch_size for obj in args if hasattr(obj, "batch_size")]
+            if self.split_batches:
+                batch_sizes = [batch_size // self.num_processes for batch_size in batch_sizes]
+            if len(batch_sizes) == 0:
+                raise ValueError(
+                    "You must specify a training or evaluation dataloader in `accelerate.prepare()` when using DeepSpeed."
+                )
+
+            batch_size_per_device = min(batch_sizes) if deepspeed_plugin.is_train_batch_min else max(batch_sizes)
+            if len(batch_sizes) > 1:
+                logger.info(
+                    "Since you passed both train and evaluation dataloader, `is_train_batch_min` (here "
+                    f"{deepspeed_plugin.is_train_batch_min} will decide the `train_batch_size` ({batch_size_per_device})."
+                )
+
+            self.deepspeed_config["train_batch_size"] = (
+                batch_size_per_device * self.deepspeed_config["gradient_accumulation_steps"] * self.num_processes
+            )
+
         model = None
         optimizer = None
+        scheduler = None
         for obj in result:
             if isinstance(obj, torch.nn.Module):
                 model = obj
-            elif isinstance(obj, (torch.optim.Optimizer, dict)):
+            elif isinstance(obj, (torch.optim.Optimizer, DummyOptim)):
                 optimizer = obj
+            elif (isinstance(obj, (torch.optim.lr_scheduler._LRScheduler, DummyScheduler))) or (
+                type(obj).__name__ in deepspeed.runtime.lr_schedules.VALID_LR_SCHEDULES
+            ):
+                scheduler = obj
 
-        if deepspeed_plugin.auto_opt_mapping:
-            is_adam = isinstance(optimizer, torch.optim.Adam)
-            is_adamw = isinstance(optimizer, torch.optim.AdamW)
-            if (is_adam or is_adamw) and deepspeed_plugin.offload_optimizer_device == "cpu":
-                defaults = optimizer.defaults
-                params = []
-                for group in optimizer.param_groups:
-                    params.extend(group["params"])
-
-                optimizer = deepspeed.ops.adam.DeepSpeedCPUAdam(
-                    params,
-                    lr=defaults["lr"],
-                    bias_correction=True,
-                    betas=defaults["betas"],
-                    eps=defaults["eps"],
-                    weight_decay=defaults["weight_decay"],
-                    amsgrad=defaults["amsgrad"],
-                    adamw_mode=is_adamw,
-                )
+        if "optimizer" in self.deepspeed_config and not isinstance(optimizer, (DummyOptim)):
+            raise ValueError(
+                "You cannot specify an optimizer in the config file and in the code at the same time. "
+                "Please remove the optimizer from the config file or "
+                "create `accelerate.utils.DummyOptim` in the code."
+            )
+        elif "optimizer" not in self.deepspeed_config and isinstance(optimizer, (DummyOptim)):
+            raise ValueError("You cannot create a `DummyOptim` without specifying an optimizer in the config file.")
 
-        # useful when only eval_dataloader is given into `accelerator.prepare()`
-        if model is not None:
-            engine = DeepSpeedEngineWrapper(
-                args=None,
-                model=model,
-                optimizer=optimizer,
-                config_params=self.deepspeed_config,
-                dist_init_required=False,
+        if isinstance(optimizer, (torch.optim.Optimizer)):
+            self.deepspeed_config["zero_allow_untested_optimizer"] = True
+
+        if "scheduler" in self.deepspeed_config and not isinstance(scheduler, (DummyScheduler)):
+            raise ValueError(
+                "You cannot specify a scheduler in the config file and in the code at the same time. "
+                "Please remove the scheduler from the config file or "
+                "create `accelerate.utils.DummyScheduler` in the code."
+            )
+        elif "scheduler" not in self.deepspeed_config and isinstance(scheduler, (DummyScheduler)):
+            raise ValueError("You cannot create a `DummyScheduler` without specifying a scheduler in the config file.")
+
+        if isinstance(optimizer, (DummyOptim)) and not isinstance(scheduler, (DummyScheduler)):
+            raise ValueError(
+                "You can only specify `accelerate.utils.DummyScheduler` in the code when using "
+                "`accelerate.utils.DummyOptim`."
             )
+
+        if model is not None:
+            model_parameters = None
+            if isinstance(optimizer, (DummyOptim)):
+                model_parameters = optimizer.param_groups
+                kwargs = dict(model=model, model_parameters=model_parameters, config_params=self.deepspeed_config)
+            else:
+                kwargs = dict(model=model, optimizer=optimizer, config_params=self.deepspeed_config)
+                if type(scheduler).__name__ in deepspeed.runtime.lr_schedules.VALID_LR_SCHEDULES:
+                    kwargs["lr_scheduler"] = scheduler
+
+            engine, optimizer, _, lr_scheduler = deepspeed.initialize(**kwargs)
+            optimizer = DeepSpeedOptimizerWrapper(optimizer)
+            if lr_scheduler is None:
+                scheduler = AcceleratedScheduler(
+                    scheduler,
+                    optimizer,
+                    step_with_optimizer=self.step_scheduler_with_optimizer,
+                    split_batches=self.split_batches,
+                )
+            else:
+                scheduler = DeepSpeedSchedulerWrapper(lr_scheduler, optimizer)
+
             for i in range(len(result)):
                 if isinstance(result[i], torch.nn.Module):
                     result[i] = engine
-                elif isinstance(result[i], torch.optim.Optimizer):
-                    result[i] = DeepSpeedOptimizerWrapper(engine.optimizer, engine)
-            self.deepspeed_engine = engine  # pointing for deepspeed_engine.backward()
+                elif isinstance(result[i], (torch.optim.Optimizer, DummyOptim)):
+                    result[i] = optimizer
+                elif (isinstance(result[i], (torch.optim.lr_scheduler._LRScheduler, DummyScheduler))) or (
+                    type(result[i]).__name__ in deepspeed.runtime.lr_schedules.VALID_LR_SCHEDULES
+                ):
+                    result[i] = scheduler
+            self.deepspeed_engine_wrapped = DeepSpeedEngineWrapper(
+                engine
+            )  # pointing for deepspeed_engine_wrapped.backward()
Let's take this PR as an opportunity to remove those assers and use proper exceptions :-)
@@ -499,81 +527,116 @@ def _prepare_deepspeed(self, *args):
         deepspeed_plugin = self.state.deepspeed_plugin
         self.deepspeed_config = deepspeed_plugin.deepspeed_config
 
-        batch_sizes = [obj.batch_size for obj in args if hasattr(obj, "batch_size")]
-        if len(batch_sizes) == 0:
-            raise ValueError(
-                "You must specify a training or evaluation dataloader in `accelerate.prepare()` when using DeepSpeed."
-            )
-
-        batch_size_per_device = min(batch_sizes) if deepspeed_plugin.is_train_batch_min else max(batch_sizes)
-        if len(batch_sizes) > 1:
-            logger.info(
-                f"Since you passed both train and evaluation dataloader, `is_train_batch_min` (here \
-                {deepspeed_plugin.is_train_batch_min} will decide the `train_batch_size` ({batch_size_per_device})."
-            )
-
-        self.deepspeed_config["train_batch_size"] = (
-            batch_size_per_device * deepspeed_plugin.gradient_accumulation_steps * self.num_processes
-        )
-
         result = [
             self._prepare_one(obj, first_pass=True) if isinstance(obj, torch.utils.data.DataLoader) else obj
             for obj in args
         ]
 
+        if deepspeed_plugin.config_file == "none":
+            batch_sizes = [obj.batch_size for obj in args if hasattr(obj, "batch_size")]
+            if self.split_batches:
+                batch_sizes = [batch_size // self.num_processes for batch_size in batch_sizes]
+            if len(batch_sizes) == 0:
+                raise ValueError(
+                    "You must specify a training or evaluation dataloader in `accelerate.prepare()` when using DeepSpeed."
+                )
+
+            batch_size_per_device = min(batch_sizes) if deepspeed_plugin.is_train_batch_min else max(batch_sizes)
+            if len(batch_sizes) > 1:
+                logger.info(
+                    "Since you passed both train and evaluation dataloader, `is_train_batch_min` (here "
+                    f"{deepspeed_plugin.is_train_batch_min} will decide the `train_batch_size` ({batch_size_per_device})."
+                )
+
+            self.deepspeed_config["train_batch_size"] = (
+                batch_size_per_device * self.deepspeed_config["gradient_accumulation_steps"] * self.num_processes
+            )
+
         model = None
         optimizer = None
+        scheduler = None
         for obj in result:
             if isinstance(obj, torch.nn.Module):
                 model = obj
-            elif isinstance(obj, (torch.optim.Optimizer, dict)):
+            elif isinstance(obj, (torch.optim.Optimizer, DummyOptim)):
                 optimizer = obj
+            elif (isinstance(obj, (torch.optim.lr_scheduler._LRScheduler, DummyScheduler))) or (
+                type(obj).__name__ in deepspeed.runtime.lr_schedules.VALID_LR_SCHEDULES
+            ):
+                scheduler = obj
 
-        if deepspeed_plugin.auto_opt_mapping:
-            is_adam = isinstance(optimizer, torch.optim.Adam)
-            is_adamw = isinstance(optimizer, torch.optim.AdamW)
-            if (is_adam or is_adamw) and deepspeed_plugin.offload_optimizer_device == "cpu":
-                defaults = optimizer.defaults
-                params = []
-                for group in optimizer.param_groups:
-                    params.extend(group["params"])
-
-                optimizer = deepspeed.ops.adam.DeepSpeedCPUAdam(
-                    params,
-                    lr=defaults["lr"],
-                    bias_correction=True,
-                    betas=defaults["betas"],
-                    eps=defaults["eps"],
-                    weight_decay=defaults["weight_decay"],
-                    amsgrad=defaults["amsgrad"],
-                    adamw_mode=is_adamw,
-                )
+        if "optimizer" in self.deepspeed_config and not isinstance(optimizer, (DummyOptim)):
+            raise ValueError(
+                "You cannot specify an optimizer in the config file and in the code at the same time. "
+                "Please remove the optimizer from the config file or "
+                "create `accelerate.utils.DummyOptim` in the code."
+            )
+        elif "optimizer" not in self.deepspeed_config and isinstance(optimizer, (DummyOptim)):
+            raise ValueError("You cannot create a `DummyOptim` without specifying an optimizer in the config file.")
 
-        # useful when only eval_dataloader is given into `accelerator.prepare()`
-        if model is not None:
-            engine = DeepSpeedEngineWrapper(
-                args=None,
-                model=model,
-                optimizer=optimizer,
-                config_params=self.deepspeed_config,
-                dist_init_required=False,
+        if isinstance(optimizer, (torch.optim.Optimizer)):
+            self.deepspeed_config["zero_allow_untested_optimizer"] = True
+
+        if "scheduler" in self.deepspeed_config and not isinstance(scheduler, (DummyScheduler)):
+            raise ValueError(
+                "You cannot specify a scheduler in the config file and in the code at the same time. "
+                "Please remove the scheduler from the config file or "
+                "create `accelerate.utils.DummyScheduler` in the code."
+            )
+        elif "scheduler" not in self.deepspeed_config and isinstance(scheduler, (DummyScheduler)):
+            raise ValueError("You cannot create a `DummyScheduler` without specifying a scheduler in the config file.")
+
+        if isinstance(optimizer, (DummyOptim)) and not isinstance(scheduler, (DummyScheduler)):
+            raise ValueError(
+                "You can only specify `accelerate.utils.DummyScheduler` in the code when using "
+                "`accelerate.utils.DummyOptim`."
             )
+
+        if model is not None:
+            model_parameters = None
+            if isinstance(optimizer, (DummyOptim)):
+                model_parameters = optimizer.param_groups
+                kwargs = dict(model=model, model_parameters=model_parameters, config_params=self.deepspeed_config)
+            else:
+                kwargs = dict(model=model, optimizer=optimizer, config_params=self.deepspeed_config)
+                if type(scheduler).__name__ in deepspeed.runtime.lr_schedules.VALID_LR_SCHEDULES:
+                    kwargs["lr_scheduler"] = scheduler
+
+            engine, optimizer, _, lr_scheduler = deepspeed.initialize(**kwargs)
+            optimizer = DeepSpeedOptimizerWrapper(optimizer)
+            if lr_scheduler is None:
+                scheduler = AcceleratedScheduler(
+                    scheduler,
+                    optimizer,
+                    step_with_optimizer=self.step_scheduler_with_optimizer,
+                    split_batches=self.split_batches,
+                )
+            else:
+                scheduler = DeepSpeedSchedulerWrapper(lr_scheduler, optimizer)
+
             for i in range(len(result)):
                 if isinstance(result[i], torch.nn.Module):
                     result[i] = engine
-                elif isinstance(result[i], torch.optim.Optimizer):
-                    result[i] = DeepSpeedOptimizerWrapper(engine.optimizer, engine)
-            self.deepspeed_engine = engine  # pointing for deepspeed_engine.backward()
+                elif isinstance(result[i], (torch.optim.Optimizer, DummyOptim)):
+                    result[i] = optimizer
+                elif (isinstance(result[i], (torch.optim.lr_scheduler._LRScheduler, DummyScheduler))) or (
+                    type(result[i]).__name__ in deepspeed.runtime.lr_schedules.VALID_LR_SCHEDULES
+                ):
+                    result[i] = scheduler
+            self.deepspeed_engine_wrapped = DeepSpeedEngineWrapper(
+                engine
+            )  # pointing for deepspeed_engine_wrapped.backward()
             self._models.append(engine)
-            self._optimizers.append(engine.optimizer)
+            self._optimizers.append(optimizer)
+            self._schedulers.append(scheduler)
             assert (
                 len(self._models) == 1
             ), "You can't use same `Accelerator()` instance with 2 models when using DeepSpeed"
 
         if self.distributed_type == DistributedType.DEEPSPEED:
             assert hasattr(
-                self, "deepspeed_engine"
-            ), "You need to pass the model along the optimizer when using Deepspeed."
+                self, "deepspeed_engine_wrapped"
+            ), "You need to pass the model along with the optimizer when using Deepspeed."
```suggestion
                "Do you want to specify a json file to a DeepSpeed config? [yes/NO]: ",
```
@@ -77,24 +77,72 @@ def get_cluster_input():
             ), "DeepSpeed is not installed => run `pip3 install deepspeed` or build it from source"
 
         if distributed_type == DistributedType.DEEPSPEED:
-            deepspeed_config["zero_stage"] = _ask_field(
-                "What should be your DeepSpeed's ZeRO optimization stage (0, 1, 2, 3)? [2]: ",
-                lambda x: int(x),
-                default=2,
+            use_deepspeed_config = _ask_field(
+                "Do you want to specify DeepSpeed json config filepath? [yes/NO]: ",
```suggestion
                    "Please enter the path to the json DeepSpeed config file: ",
```
@@ -77,24 +77,72 @@ def get_cluster_input():
             ), "DeepSpeed is not installed => run `pip3 install deepspeed` or build it from source"
 
         if distributed_type == DistributedType.DEEPSPEED:
-            deepspeed_config["zero_stage"] = _ask_field(
-                "What should be your DeepSpeed's ZeRO optimization stage (0, 1, 2, 3)? [2]: ",
-                lambda x: int(x),
-                default=2,
+            use_deepspeed_config = _ask_field(
+                "Do you want to specify DeepSpeed json config filepath? [yes/NO]: ",
+                _convert_yes_no_to_bool,
+                default=False,
+                error_message="Please enter yes or no.",
             )
-
-            if deepspeed_config["zero_stage"] >= 2:
-                deepspeed_config["offload_optimizer_device"] = _ask_field(
-                    "Where to offload optimizer states? [NONE/cpu/nvme]: ",
+            if use_deepspeed_config:
+                deepspeed_config["deepspeed_config_file"] = _ask_field(
+                    "Please enter filepath of the DeepSpeed json config file: ",
Maybe put a comment to explain DeepSpeed will handle this?
@@ -648,13 +711,17 @@ def clip_grad_norm_(self, parameters, max_norm, norm_type=2):
                 if parameters == [p for p in model.parameters()]:
                     model.clip_grad_norm_(max_norm, norm_type)
                     return
+        elif self.distributed_type == DistributedType.DEEPSPEED:
+            return
Same here.
@@ -648,13 +711,17 @@ def clip_grad_norm_(self, parameters, max_norm, norm_type=2):
                 if parameters == [p for p in model.parameters()]:
                     model.clip_grad_norm_(max_norm, norm_type)
                     return
+        elif self.distributed_type == DistributedType.DEEPSPEED:
+            return
         self.unscale_gradients()
         torch.nn.utils.clip_grad_norm_(parameters, max_norm, norm_type=norm_type)
 
     def clip_grad_value_(self, parameters, clip_value):
         """
         Should be used in place of `torch.nn.utils.clip_grad_value_`.
         """
+        if self.distributed_type in [DistributedType.DEEPSPEED, DistributedType.FSDP]:
+            return
Not really sure it's worth adding all of those below. The DeepSpeed config is enough IMO.
@@ -57,6 +58,56 @@ def launch_command_parser(subparsers=None):
         action="store_true",
         help="Whether to use deepspeed.",
     )
+    parser.add_argument(
+        "--deepspeed_config_file",
+        default=None,
+        type=str,
+        help="DeepSpeed's config file (useful only when `use_deepspeed` flag is passed).",
+    )
+    parser.add_argument(
Not sure I understand the comments. Is it everything that happens in this step? If that's the case, can you add a comment stating it at the top?
@@ -12,58 +12,26 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ..optimizer import AcceleratedOptimizer
-from .imports import is_apex_available, is_deepspeed_available
-
-
-if is_deepspeed_available():
-    from deepspeed import DeepSpeedEngine
-
-if is_apex_available():
-    from apex import amp
-
-
-class DeepSpeedEngineWrapper(DeepSpeedEngine):
-    """
-    Wrapper over deepspeed.DeepSpeedEngine object
-    """
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+from accelerate.scheduler import AcceleratedScheduler
 
-        # overwriting micro_steps for user's gradient_accumulation
-        self.micro_steps = -1
+from ..optimizer import AcceleratedOptimizer
 
-    def step(self, lr_kwargs=None):
-        """DeepSpeedEngine.step() without `micro_steps` update & no profiling"""
-        if self.is_gradient_accumulation_boundary():  # it shouldn't matter whether we keep this line or not
-            if self.progressive_layer_drop:
-                self.progressive_layer_drop.update_state(self.global_steps)
 
-            self._take_model_step(lr_kwargs)
+class DeepSpeedEngineWrapper:
+    def __init__(self, engine):
+        self.engine = engine
 
     def backward(self, loss):
-        """DeepSpeedEngine.backward() with with no loss scaling; no profiling but with `micro_steps` update"""
-
-        if self.zero_optimization():
-            self.optimizer.is_gradient_accumulation_boundary = self.is_gradient_accumulation_boundary()
-            self.optimizer.backward(loss)
-        elif self.amp_enabled():
-            # AMP requires delaying unscale when inside gradient accumulation boundaries
-            # https://nvidia.github.io/apex/advanced.html#gradient-accumulation-across-iterations
-            delay_unscale = not self.is_gradient_accumulation_boundary()
-            with amp.scale_loss(loss, self.optimizer, delay_unscale=delay_unscale) as scaled_loss:
-                scaled_loss.backward()
-        elif self.fp16_enabled():
-            self.optimizer.backward(loss)
-        else:
-            loss.backward()
-
-        if self.enable_backward_allreduce:
-            self.allreduce_gradients()
+        # runs backpropagation and handles mixed precision
+        self.engine.backward(loss)
 
-        # this will ensure deepspeed gradient_accumulation matches user's accumulation
-        self.micro_steps += 1
+        # gradient accumulation check
+        # gradient clipping
+        # optimizer step
+        # zero grad
+        # checking overflow
+        # lr_scheduler step
+        self.engine.step()
```suggestion
        self.ds_config_dict = dict(zero2=config_zero2, zero3=config_zero3)
```
Fits in one line ;-)
@@ -0,0 +1,546 @@
+# 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 inspect
+import io
+import itertools
+import json
+import os
+import tempfile
+import unittest
+from copy import deepcopy
+from pathlib import Path
+
+import torch
+from torch.utils.data import DataLoader
+
+from accelerate.accelerator import Accelerator
+from accelerate.scheduler import AcceleratedScheduler
+from accelerate.state import AcceleratorState
+from accelerate.test_utils.testing import require_cuda, require_deepspeed
+from accelerate.test_utils.training import RegressionDataset
+from accelerate.utils.dataclasses import DeepSpeedPlugin
+from accelerate.utils.deepspeed import (
+    DeepSpeedEngineWrapper,
+    DeepSpeedOptimizerWrapper,
+    DeepSpeedSchedulerWrapper,
+    DummyOptim,
+    DummyScheduler,
+)
+from parameterized import parameterized
+from transformers import AutoModel, AutoModelForCausalLM, get_scheduler
+from transformers.deepspeed import HfDeepSpeedConfig
+from transformers.testing_utils import mockenv_context
+from transformers.trainer_utils import set_seed
+from transformers.utils import is_torch_bf16_available
+
+
+set_seed(42)
+
+T5_SMALL = "t5-small"
+T5_TINY = "patrickvonplaten/t5-tiny-random"
+GPT2_TINY = "sshleifer/tiny-gpt2"
+
+ZERO2 = "zero2"
+ZERO3 = "zero3"
+
+FP16 = "fp16"
+BF16 = "bf16"
+
+CUSTOM_OPTIMIZER = "custom_optimizer"
+CUSTOM_SCHEDULER = "custom_scheduler"
+DS_OPTIMIZER = "deepspeed_optimizer"
+DS_SCHEDULER = "deepspeed_scheduler"
+
+stages = [ZERO2, ZERO3]
+optims = [CUSTOM_OPTIMIZER, DS_OPTIMIZER]
+schedulers = [CUSTOM_SCHEDULER, DS_SCHEDULER]
+if is_torch_bf16_available():
+    dtypes = [FP16, BF16]
+else:
+    dtypes = [FP16]
+
+
+def parameterized_custom_name_func(func, param_num, param):
+    # customize the test name generator function as we want both params to appear in the sub-test
+    # name, as by default it shows only the first param
+    param_based_name = parameterized.to_safe_name("_".join(str(x) for x in param.args))
+    return f"{func.__name__}_{param_based_name}"
+
+
+# Cartesian-product of zero stages with models to test
+params = list(itertools.product(stages, dtypes))
+optim_scheduler_params = list(itertools.product(optims, schedulers))
+
+
+@require_deepspeed
+@require_cuda
+class DeepSpeedConfigIntegration(unittest.TestCase):
+    def setUp(self):
+        super().setUp()
+
+        self._test_file_path = inspect.getfile(self.__class__)
+        path = Path(self._test_file_path).resolve()
+        self.test_file_dir_str = str(path.parents[0])
+
+        self.ds_config_file = dict(
+            zero2=f"{self.test_file_dir_str}/ds_config_zero2.json",
+            zero3=f"{self.test_file_dir_str}/ds_config_zero3.json",
+        )
+
+        # use self.get_config_dict(stage) to use these to ensure the original is not modified
+        with io.open(self.ds_config_file[ZERO2], "r", encoding="utf-8") as f:
+            config_zero2 = json.load(f)
+        with io.open(self.ds_config_file[ZERO3], "r", encoding="utf-8") as f:
+            config_zero3 = json.load(f)
+            # The following setting slows things down, so don't enable it by default unless needed by a test.
+            # It's in the file as a demo for users since we want everything to work out of the box even if slower.
+            config_zero3["zero_optimization"]["stage3_gather_16bit_weights_on_model_save"] = False
+
+        self.ds_config_dict = dict(
+            zero2=config_zero2,
+            zero3=config_zero3,
+        )
I thought about this. The thing is, a weakref to `HFDeepSpeedConfig` is created in that file (transformers.deepspeed). This is important only when using ZeRO Stage-3 when we don't want to load Transformer models fully on CPU/GPU and we want to directly partition the model parameters across GPUs. This weakref `_hf_deepspeed_config_weak_ref` is used in transformers `modeling_utils.py` to check if deepspeed zero stage3 is enabled. If it is enabled DeepSpeed functionality of `zero.init` is used to directly partition model parameters across GPUs. It is used by `from_config` (when training from scratch) and `from_pretrained` (when finetuning) methods. 

Snippet in `modeling_utils.py`:
```python
if is_deepspeed_zero3_enabled():
    import deepspeed
    
    logger.info("Detected DeepSpeed ZeRO-3: activating zero.init() for this model")
    # this immediately partitions the model across all gpus, to avoid the overhead in time
    # and memory copying it on CPU or each GPU first
    with deepspeed.zero.Init(config_dict_or_path=deepspeed_config()):
        model = cls(config, **kwargs)
else:
    model = cls(config, **kwargs)
```
`is_deepspeed_zero3_enabled` from above snippet directly refers to the weakref in transformers.deepspeed
```python
def is_deepspeed_zero3_enabled():
    if _hf_deepspeed_config_weak_ref is not None and _hf_deepspeed_config_weak_ref() is not None:
        return _hf_deepspeed_config_weak_ref().is_zero3()
    else:
        return False
```

Due to above reasons I thought it would be good to let this be part of `transformers` repo as it is specifically used only in ZeRO Stage-3 for **efficiently loading models that are part of `transformers` repo.**

@@ -163,6 +172,25 @@ def __init__(
                 deepspeed_plugin, DeepSpeedPlugin
             ), "`deepspeed_plugin` must be a DeepSpeedPlugin object."
             os.environ["USE_DEEPSPEED"] = "true"  # use DeepSpeed if plugin is provided
+        if deepspeed_plugin:
+            if not is_deepspeed_available():
+                raise ImportError(
+                    "DeepSpeed is not installed => run `pip3 install deepspeed` or build it from source."
+                )
+            if compare_versions("deepspeed", "<", "0.6.4"):
+                raise ImportError("DeepSpeed version must be >= 0.6.4. Please update DeepSpeed.")
+            if os.environ.get("DEEPSPEED_ZERO3_INIT", "false") == "true" or deepspeed_plugin.zero3_init_flag:
+                if not is_transformers_available():
+                    raise Exception(
+                        "When `zero3_init_flag` is set, it requires Transformers to be installed. "
+                        "Please run `pip3 install transformers`."
+                    )
+                from transformers.deepspeed import HfDeepSpeedConfig
With the move, the weakref will disappear and we will rely on the `AcceleratorState` to know if zero3 is enabled inside Transformers. Again, not sure when is the right point to do the move (as it will make Accelerate a hard dep of Transformers) but want to flag this is the final destination :-)
@@ -163,6 +172,25 @@ def __init__(
                 deepspeed_plugin, DeepSpeedPlugin
             ), "`deepspeed_plugin` must be a DeepSpeedPlugin object."
             os.environ["USE_DEEPSPEED"] = "true"  # use DeepSpeed if plugin is provided
+        if deepspeed_plugin:
+            if not is_deepspeed_available():
+                raise ImportError(
+                    "DeepSpeed is not installed => run `pip3 install deepspeed` or build it from source."
+                )
+            if compare_versions("deepspeed", "<", "0.6.4"):
+                raise ImportError("DeepSpeed version must be >= 0.6.4. Please update DeepSpeed.")
+            if os.environ.get("DEEPSPEED_ZERO3_INIT", "false") == "true" or deepspeed_plugin.zero3_init_flag:
+                if not is_transformers_available():
+                    raise Exception(
+                        "When `zero3_init_flag` is set, it requires Transformers to be installed. "
+                        "Please run `pip3 install transformers`."
+                    )
+                from transformers.deepspeed import HfDeepSpeedConfig
I would focus the intro on Zero and offload which is the main feature DeepSpeed offers. Distributed training, mixed precision are done by Accelerate, not really specific to DeepSpeed ;-)
@@ -0,0 +1,542 @@
+<!--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.
+-->
+
+# DeepSpeed 
+
+Training advanced deep learning models is challenging. 
+Beyond model design, model scientists also need to set up the state-of-the-art training techniques such as 
+distributed training, mixed precision, gradient accumulation, and checkpointing. 
+Yet still, scientists may not achieve the desired system performance and convergence rate.
+Large model sizes are even more challenging: a large model easily runs out of memory with pure data parallelism and it is difficult to use model parallelism. 
+[DeepSpeed](https://www.deepspeed.ai/) addresses these challenges to accelerate model development and training.
A picture here or before to explain the differences from ZeRO stages would be great. Same, quickly explaining what the offload is doing would be great (here or above). Also, for me `ZeRO-Infinity` is the NVME offload (that's how they released it at least).

I'd focus on non-DeepSpeed marketing terms such a CPU offload and Disk offload (the disk should be an NVME for decent speed but it technically work on any Disk).
@@ -0,0 +1,542 @@
+<!--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.
+-->
+
+# DeepSpeed 
+
+Training advanced deep learning models is challenging. 
+Beyond model design, model scientists also need to set up the state-of-the-art training techniques such as 
+distributed training, mixed precision, gradient accumulation, and checkpointing. 
+Yet still, scientists may not achieve the desired system performance and convergence rate.
+Large model sizes are even more challenging: a large model easily runs out of memory with pure data parallelism and it is difficult to use model parallelism. 
+[DeepSpeed](https://www.deepspeed.ai/) addresses these challenges to accelerate model development and training.
+
+🤗 Accelerate integrates [DeepSpeed](https://github.com/microsoft/DeepSpeed) via 2 options:
+
+1. Integration of the DeepSpeed features via `deepspeed config file` specification in `accelerate config` . You just supply your custom config file or use our template. Most of
+   this document is focused on this feature. This supports all the core features of DeepSpeed and gives user a lot of flexibility. 
+   User may have to change few lines of code depending on the config.
+2. Integration via `deepspeed_plugin`.This supports subset of the DeepSpeed features and uses default options for the rest of the configurations. 
+   User need not change any code and is good for those who are fine with most of the default settings of DeepSpeed.
+
+## What is integrated?
+
+Training:
+
+1. DeepSpeed ZeRO training supports the full ZeRO stages 1, 2 and 3 with ZeRO-Infinity (CPU and NVME offload).
I didn't realize this meant via the `accelerate config` command, so I'd say it epxilcitly:
```suggestion
We will first look at easy to use integration via `accelerate config`. 
```
@@ -0,0 +1,542 @@
+<!--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.
+-->
+
+# DeepSpeed 
+
+Training advanced deep learning models is challenging. 
+Beyond model design, model scientists also need to set up the state-of-the-art training techniques such as 
+distributed training, mixed precision, gradient accumulation, and checkpointing. 
+Yet still, scientists may not achieve the desired system performance and convergence rate.
+Large model sizes are even more challenging: a large model easily runs out of memory with pure data parallelism and it is difficult to use model parallelism. 
+[DeepSpeed](https://www.deepspeed.ai/) addresses these challenges to accelerate model development and training.
+
+🤗 Accelerate integrates [DeepSpeed](https://github.com/microsoft/DeepSpeed) via 2 options:
+
+1. Integration of the DeepSpeed features via `deepspeed config file` specification in `accelerate config` . You just supply your custom config file or use our template. Most of
+   this document is focused on this feature. This supports all the core features of DeepSpeed and gives user a lot of flexibility. 
+   User may have to change few lines of code depending on the config.
+2. Integration via `deepspeed_plugin`.This supports subset of the DeepSpeed features and uses default options for the rest of the configurations. 
+   User need not change any code and is good for those who are fine with most of the default settings of DeepSpeed.
+
+## What is integrated?
+
+Training:
+
+1. DeepSpeed ZeRO training supports the full ZeRO stages 1, 2 and 3 with ZeRO-Infinity (CPU and NVME offload).
+
+Inference:
+
+1. DeepSpeed ZeRO Inference supports ZeRO stage 3 with ZeRO-Infinity. It uses the same ZeRO protocol as training, but
+   it doesn't use an optimizer and a lr scheduler and only stage 3 is relevant. For more details see:
+   [deepspeed-zero-inference](#deepspeed-zero-inference).
+
+
+## How it works?
+
+**Pre-Requisites**: Install DeepSpeed version >=0.6.5. Please refer to the [DeepSpeed Insallation details](https://github.com/microsoft/DeepSpeed#installation)
+for more information.
+
+We will first look at easy to use integration via `deepspeed_plugin`. 
```suggestion
and answer the questions asked. It will ask whether you want to use a config file for DeepSpeed to which you should answer no. Then answer the following questions to generate a basic DeepSpeed config.
```
@@ -0,0 +1,542 @@
+<!--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.
+-->
+
+# DeepSpeed 
+
+Training advanced deep learning models is challenging. 
+Beyond model design, model scientists also need to set up the state-of-the-art training techniques such as 
+distributed training, mixed precision, gradient accumulation, and checkpointing. 
+Yet still, scientists may not achieve the desired system performance and convergence rate.
+Large model sizes are even more challenging: a large model easily runs out of memory with pure data parallelism and it is difficult to use model parallelism. 
+[DeepSpeed](https://www.deepspeed.ai/) addresses these challenges to accelerate model development and training.
+
+🤗 Accelerate integrates [DeepSpeed](https://github.com/microsoft/DeepSpeed) via 2 options:
+
+1. Integration of the DeepSpeed features via `deepspeed config file` specification in `accelerate config` . You just supply your custom config file or use our template. Most of
+   this document is focused on this feature. This supports all the core features of DeepSpeed and gives user a lot of flexibility. 
+   User may have to change few lines of code depending on the config.
+2. Integration via `deepspeed_plugin`.This supports subset of the DeepSpeed features and uses default options for the rest of the configurations. 
+   User need not change any code and is good for those who are fine with most of the default settings of DeepSpeed.
+
+## What is integrated?
+
+Training:
+
+1. DeepSpeed ZeRO training supports the full ZeRO stages 1, 2 and 3 with ZeRO-Infinity (CPU and NVME offload).
+
+Inference:
+
+1. DeepSpeed ZeRO Inference supports ZeRO stage 3 with ZeRO-Infinity. It uses the same ZeRO protocol as training, but
+   it doesn't use an optimizer and a lr scheduler and only stage 3 is relevant. For more details see:
+   [deepspeed-zero-inference](#deepspeed-zero-inference).
+
+
+## How it works?
+
+**Pre-Requisites**: Install DeepSpeed version >=0.6.5. Please refer to the [DeepSpeed Insallation details](https://github.com/microsoft/DeepSpeed#installation)
+for more information.
+
+We will first look at easy to use integration via `deepspeed_plugin`. 
+Followed by more flexible and feature rich `deepspeed config file` integration. 
+
+### Accelerate DeepSpeed Plugin
+On your machine(s) just run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked. It will ask whether you want to use a config file for deepspeed to which you answer no and answer the following questions. 
It would be better to launch a basic example here.
@@ -0,0 +1,542 @@
+<!--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.
+-->
+
+# DeepSpeed 
+
+Training advanced deep learning models is challenging. 
+Beyond model design, model scientists also need to set up the state-of-the-art training techniques such as 
+distributed training, mixed precision, gradient accumulation, and checkpointing. 
+Yet still, scientists may not achieve the desired system performance and convergence rate.
+Large model sizes are even more challenging: a large model easily runs out of memory with pure data parallelism and it is difficult to use model parallelism. 
+[DeepSpeed](https://www.deepspeed.ai/) addresses these challenges to accelerate model development and training.
+
+🤗 Accelerate integrates [DeepSpeed](https://github.com/microsoft/DeepSpeed) via 2 options:
+
+1. Integration of the DeepSpeed features via `deepspeed config file` specification in `accelerate config` . You just supply your custom config file or use our template. Most of
+   this document is focused on this feature. This supports all the core features of DeepSpeed and gives user a lot of flexibility. 
+   User may have to change few lines of code depending on the config.
+2. Integration via `deepspeed_plugin`.This supports subset of the DeepSpeed features and uses default options for the rest of the configurations. 
+   User need not change any code and is good for those who are fine with most of the default settings of DeepSpeed.
+
+## What is integrated?
+
+Training:
+
+1. DeepSpeed ZeRO training supports the full ZeRO stages 1, 2 and 3 with ZeRO-Infinity (CPU and NVME offload).
+
+Inference:
+
+1. DeepSpeed ZeRO Inference supports ZeRO stage 3 with ZeRO-Infinity. It uses the same ZeRO protocol as training, but
+   it doesn't use an optimizer and a lr scheduler and only stage 3 is relevant. For more details see:
+   [deepspeed-zero-inference](#deepspeed-zero-inference).
+
+
+## How it works?
+
+**Pre-Requisites**: Install DeepSpeed version >=0.6.5. Please refer to the [DeepSpeed Insallation details](https://github.com/microsoft/DeepSpeed#installation)
+for more information.
+
+We will first look at easy to use integration via `deepspeed_plugin`. 
+Followed by more flexible and feature rich `deepspeed config file` integration. 
+
+### Accelerate DeepSpeed Plugin
+On your machine(s) just run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked. It will ask whether you want to use a config file for deepspeed to which you answer no and answer the following questions. 
+This will generate a config file that will be used automatically to properly set the
+default options when doing
+
+```bash
+accelerate launch my_script.py --args_to_my_script
+```
+
+For instance, here is how you would run the NLP example `examples/by_feature/deepspeed_with_config_support.py` (from the root of the repo) with DeepSpeed Plugin:
+
+**ZeRO Stage-2 DeepSpeed Plugin Example**
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+ gradient_accumulation_steps: 1
+ gradient_clipping: 1.0
+ offload_optimizer_device: none
+ offload_param_device: none
+ zero3_init_flag: true
+ zero_stage: 2
+distributed_type: DEEPSPEED
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 2
+use_cpu: false
+```
+
+```bash
+accelerate launch examples/by_feature/deepspeed_with_config_support.py \
Same here.
@@ -0,0 +1,542 @@
+<!--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.
+-->
+
+# DeepSpeed 
+
+Training advanced deep learning models is challenging. 
+Beyond model design, model scientists also need to set up the state-of-the-art training techniques such as 
+distributed training, mixed precision, gradient accumulation, and checkpointing. 
+Yet still, scientists may not achieve the desired system performance and convergence rate.
+Large model sizes are even more challenging: a large model easily runs out of memory with pure data parallelism and it is difficult to use model parallelism. 
+[DeepSpeed](https://www.deepspeed.ai/) addresses these challenges to accelerate model development and training.
+
+🤗 Accelerate integrates [DeepSpeed](https://github.com/microsoft/DeepSpeed) via 2 options:
+
+1. Integration of the DeepSpeed features via `deepspeed config file` specification in `accelerate config` . You just supply your custom config file or use our template. Most of
+   this document is focused on this feature. This supports all the core features of DeepSpeed and gives user a lot of flexibility. 
+   User may have to change few lines of code depending on the config.
+2. Integration via `deepspeed_plugin`.This supports subset of the DeepSpeed features and uses default options for the rest of the configurations. 
+   User need not change any code and is good for those who are fine with most of the default settings of DeepSpeed.
+
+## What is integrated?
+
+Training:
+
+1. DeepSpeed ZeRO training supports the full ZeRO stages 1, 2 and 3 with ZeRO-Infinity (CPU and NVME offload).
+
+Inference:
+
+1. DeepSpeed ZeRO Inference supports ZeRO stage 3 with ZeRO-Infinity. It uses the same ZeRO protocol as training, but
+   it doesn't use an optimizer and a lr scheduler and only stage 3 is relevant. For more details see:
+   [deepspeed-zero-inference](#deepspeed-zero-inference).
+
+
+## How it works?
+
+**Pre-Requisites**: Install DeepSpeed version >=0.6.5. Please refer to the [DeepSpeed Insallation details](https://github.com/microsoft/DeepSpeed#installation)
+for more information.
+
+We will first look at easy to use integration via `deepspeed_plugin`. 
+Followed by more flexible and feature rich `deepspeed config file` integration. 
+
+### Accelerate DeepSpeed Plugin
+On your machine(s) just run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked. It will ask whether you want to use a config file for deepspeed to which you answer no and answer the following questions. 
+This will generate a config file that will be used automatically to properly set the
+default options when doing
+
+```bash
+accelerate launch my_script.py --args_to_my_script
+```
+
+For instance, here is how you would run the NLP example `examples/by_feature/deepspeed_with_config_support.py` (from the root of the repo) with DeepSpeed Plugin:
+
+**ZeRO Stage-2 DeepSpeed Plugin Example**
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+ gradient_accumulation_steps: 1
+ gradient_clipping: 1.0
+ offload_optimizer_device: none
+ offload_param_device: none
+ zero3_init_flag: true
+ zero_stage: 2
+distributed_type: DEEPSPEED
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 2
+use_cpu: false
+```
+
+```bash
+accelerate launch examples/by_feature/deepspeed_with_config_support.py \
+--config_name "gpt2-large" \
+--tokenizer_name "gpt2-large" \
+--dataset_name "wikitext" \
+--dataset_config_name "wikitext-2-raw-v1" \
+--block_size 128 \
+--output_dir "./clm/clm_deepspeed_stage2_accelerate" \
+--learning_rate 5e-4 \
+--per_device_train_batch_size 24 \
+--per_device_eval_batch_size 24 \
+--num_train_epochs 3 \
+--with_tracking \
+--report_to "wandb"\
+```
+
+**ZeRO Stage-3 with CPU Offload DeepSpeed Plugin Example**
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+  gradient_accumulation_steps: 1
+  gradient_clipping: 1.0
+  offload_optimizer_device: cpu
+  offload_param_device: cpu
+  zero3_init_flag: true
+  zero3_save_16bit_model: true
+  zero_stage: 3
+distributed_type: DEEPSPEED
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 2
+use_cpu: false
+```
+
+```bash
+accelerate launch examples/by_feature/deepspeed_with_config_support.py \
Segway to the next section here: "To be able to tweak more options, you will need to use a DeepSpeed config file."
@@ -0,0 +1,542 @@
+<!--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.
+-->
+
+# DeepSpeed 
+
+Training advanced deep learning models is challenging. 
+Beyond model design, model scientists also need to set up the state-of-the-art training techniques such as 
+distributed training, mixed precision, gradient accumulation, and checkpointing. 
+Yet still, scientists may not achieve the desired system performance and convergence rate.
+Large model sizes are even more challenging: a large model easily runs out of memory with pure data parallelism and it is difficult to use model parallelism. 
+[DeepSpeed](https://www.deepspeed.ai/) addresses these challenges to accelerate model development and training.
+
+🤗 Accelerate integrates [DeepSpeed](https://github.com/microsoft/DeepSpeed) via 2 options:
+
+1. Integration of the DeepSpeed features via `deepspeed config file` specification in `accelerate config` . You just supply your custom config file or use our template. Most of
+   this document is focused on this feature. This supports all the core features of DeepSpeed and gives user a lot of flexibility. 
+   User may have to change few lines of code depending on the config.
+2. Integration via `deepspeed_plugin`.This supports subset of the DeepSpeed features and uses default options for the rest of the configurations. 
+   User need not change any code and is good for those who are fine with most of the default settings of DeepSpeed.
+
+## What is integrated?
+
+Training:
+
+1. DeepSpeed ZeRO training supports the full ZeRO stages 1, 2 and 3 with ZeRO-Infinity (CPU and NVME offload).
+
+Inference:
+
+1. DeepSpeed ZeRO Inference supports ZeRO stage 3 with ZeRO-Infinity. It uses the same ZeRO protocol as training, but
+   it doesn't use an optimizer and a lr scheduler and only stage 3 is relevant. For more details see:
+   [deepspeed-zero-inference](#deepspeed-zero-inference).
+
+
+## How it works?
+
+**Pre-Requisites**: Install DeepSpeed version >=0.6.5. Please refer to the [DeepSpeed Insallation details](https://github.com/microsoft/DeepSpeed#installation)
+for more information.
+
+We will first look at easy to use integration via `deepspeed_plugin`. 
+Followed by more flexible and feature rich `deepspeed config file` integration. 
+
+### Accelerate DeepSpeed Plugin
+On your machine(s) just run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked. It will ask whether you want to use a config file for deepspeed to which you answer no and answer the following questions. 
+This will generate a config file that will be used automatically to properly set the
+default options when doing
+
+```bash
+accelerate launch my_script.py --args_to_my_script
+```
+
+For instance, here is how you would run the NLP example `examples/by_feature/deepspeed_with_config_support.py` (from the root of the repo) with DeepSpeed Plugin:
+
+**ZeRO Stage-2 DeepSpeed Plugin Example**
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+ gradient_accumulation_steps: 1
+ gradient_clipping: 1.0
+ offload_optimizer_device: none
+ offload_param_device: none
+ zero3_init_flag: true
+ zero_stage: 2
+distributed_type: DEEPSPEED
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 2
+use_cpu: false
+```
+
+```bash
+accelerate launch examples/by_feature/deepspeed_with_config_support.py \
+--config_name "gpt2-large" \
+--tokenizer_name "gpt2-large" \
+--dataset_name "wikitext" \
+--dataset_config_name "wikitext-2-raw-v1" \
+--block_size 128 \
+--output_dir "./clm/clm_deepspeed_stage2_accelerate" \
+--learning_rate 5e-4 \
+--per_device_train_batch_size 24 \
+--per_device_eval_batch_size 24 \
+--num_train_epochs 3 \
+--with_tracking \
+--report_to "wandb"\
+```
+
+**ZeRO Stage-3 with CPU Offload DeepSpeed Plugin Example**
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+  gradient_accumulation_steps: 1
+  gradient_clipping: 1.0
+  offload_optimizer_device: cpu
+  offload_param_device: cpu
+  zero3_init_flag: true
+  zero3_save_16bit_model: true
+  zero_stage: 3
+distributed_type: DEEPSPEED
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 2
+use_cpu: false
+```
+
+```bash
+accelerate launch examples/by_feature/deepspeed_with_config_support.py \
+--config_name "gpt2-large" \
+--tokenizer_name "gpt2-large" \
+--dataset_name "wikitext" \
+--dataset_config_name "wikitext-2-raw-v1" \
+--block_size 128 \
+--output_dir "./clm/clm_deepspeed_stage3_offload_accelerate" \
+--learning_rate 5e-4 \
+--per_device_train_batch_size 32 \
+--per_device_eval_batch_size 32 \
+--num_train_epochs 3 \
+--with_tracking \
+--report_to "wandb"\
+```
+
+Currently, `Accelerate` supports following config through the CLI:
+
+```bash
+`zero_stage`: [0] Disabled, [1] optimizer state partitioning, [2] optimizer+gradient state partitioning and [3] optimizer+gradient+parameter partitioning
+`gradient_accumulation_steps`: Number of training steps to accumulate gradients before averaging and applying them.
+`gradient_clipping`: Enable gradient clipping with value.
+`offload_optimizer_device`: [none] Disable optimizer offloading, [cpu] offload optimizer to CPU, [nvme] offload optimizer to NVMe SSD. Only applicable with ZeRO >= Stage-2.
+`offload_param_device`: [none] Disable parameter offloading, [cpu] offload parameters to CPU, [nvme] offload parameters to NVMe SSD. Only applicable with ZeRO Stage-3.
+`zero3_init_flag`: Decides whether to enable `deepspeed.zero.Init` for constructing massive models. Only applicable with ZeRO Stage-3.
+`zero3_save_16bit_model`: Decides whether to save 16-bit model weights when using ZeRO Stage-3.
+`mixed_precision`: `no` for FP32 training, `fp16` for FP16 mixed-precision training and `bf16` for BF16 mixed-precision training. 
+```
+
```suggestion
   # Creates Dummy Optimizer if `optimizer` was specified in the config file else creates Adam Optimizer
```
@@ -0,0 +1,542 @@
+<!--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.
+-->
+
+# DeepSpeed 
+
+Training advanced deep learning models is challenging. 
+Beyond model design, model scientists also need to set up the state-of-the-art training techniques such as 
+distributed training, mixed precision, gradient accumulation, and checkpointing. 
+Yet still, scientists may not achieve the desired system performance and convergence rate.
+Large model sizes are even more challenging: a large model easily runs out of memory with pure data parallelism and it is difficult to use model parallelism. 
+[DeepSpeed](https://www.deepspeed.ai/) addresses these challenges to accelerate model development and training.
+
+🤗 Accelerate integrates [DeepSpeed](https://github.com/microsoft/DeepSpeed) via 2 options:
+
+1. Integration of the DeepSpeed features via `deepspeed config file` specification in `accelerate config` . You just supply your custom config file or use our template. Most of
+   this document is focused on this feature. This supports all the core features of DeepSpeed and gives user a lot of flexibility. 
+   User may have to change few lines of code depending on the config.
+2. Integration via `deepspeed_plugin`.This supports subset of the DeepSpeed features and uses default options for the rest of the configurations. 
+   User need not change any code and is good for those who are fine with most of the default settings of DeepSpeed.
+
+## What is integrated?
+
+Training:
+
+1. DeepSpeed ZeRO training supports the full ZeRO stages 1, 2 and 3 with ZeRO-Infinity (CPU and NVME offload).
+
+Inference:
+
+1. DeepSpeed ZeRO Inference supports ZeRO stage 3 with ZeRO-Infinity. It uses the same ZeRO protocol as training, but
+   it doesn't use an optimizer and a lr scheduler and only stage 3 is relevant. For more details see:
+   [deepspeed-zero-inference](#deepspeed-zero-inference).
+
+
+## How it works?
+
+**Pre-Requisites**: Install DeepSpeed version >=0.6.5. Please refer to the [DeepSpeed Insallation details](https://github.com/microsoft/DeepSpeed#installation)
+for more information.
+
+We will first look at easy to use integration via `deepspeed_plugin`. 
+Followed by more flexible and feature rich `deepspeed config file` integration. 
+
+### Accelerate DeepSpeed Plugin
+On your machine(s) just run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked. It will ask whether you want to use a config file for deepspeed to which you answer no and answer the following questions. 
+This will generate a config file that will be used automatically to properly set the
+default options when doing
+
+```bash
+accelerate launch my_script.py --args_to_my_script
+```
+
+For instance, here is how you would run the NLP example `examples/by_feature/deepspeed_with_config_support.py` (from the root of the repo) with DeepSpeed Plugin:
+
+**ZeRO Stage-2 DeepSpeed Plugin Example**
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+ gradient_accumulation_steps: 1
+ gradient_clipping: 1.0
+ offload_optimizer_device: none
+ offload_param_device: none
+ zero3_init_flag: true
+ zero_stage: 2
+distributed_type: DEEPSPEED
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 2
+use_cpu: false
+```
+
+```bash
+accelerate launch examples/by_feature/deepspeed_with_config_support.py \
+--config_name "gpt2-large" \
+--tokenizer_name "gpt2-large" \
+--dataset_name "wikitext" \
+--dataset_config_name "wikitext-2-raw-v1" \
+--block_size 128 \
+--output_dir "./clm/clm_deepspeed_stage2_accelerate" \
+--learning_rate 5e-4 \
+--per_device_train_batch_size 24 \
+--per_device_eval_batch_size 24 \
+--num_train_epochs 3 \
+--with_tracking \
+--report_to "wandb"\
+```
+
+**ZeRO Stage-3 with CPU Offload DeepSpeed Plugin Example**
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+  gradient_accumulation_steps: 1
+  gradient_clipping: 1.0
+  offload_optimizer_device: cpu
+  offload_param_device: cpu
+  zero3_init_flag: true
+  zero3_save_16bit_model: true
+  zero_stage: 3
+distributed_type: DEEPSPEED
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 2
+use_cpu: false
+```
+
+```bash
+accelerate launch examples/by_feature/deepspeed_with_config_support.py \
+--config_name "gpt2-large" \
+--tokenizer_name "gpt2-large" \
+--dataset_name "wikitext" \
+--dataset_config_name "wikitext-2-raw-v1" \
+--block_size 128 \
+--output_dir "./clm/clm_deepspeed_stage3_offload_accelerate" \
+--learning_rate 5e-4 \
+--per_device_train_batch_size 32 \
+--per_device_eval_batch_size 32 \
+--num_train_epochs 3 \
+--with_tracking \
+--report_to "wandb"\
+```
+
+Currently, `Accelerate` supports following config through the CLI:
+
+```bash
+`zero_stage`: [0] Disabled, [1] optimizer state partitioning, [2] optimizer+gradient state partitioning and [3] optimizer+gradient+parameter partitioning
+`gradient_accumulation_steps`: Number of training steps to accumulate gradients before averaging and applying them.
+`gradient_clipping`: Enable gradient clipping with value.
+`offload_optimizer_device`: [none] Disable optimizer offloading, [cpu] offload optimizer to CPU, [nvme] offload optimizer to NVMe SSD. Only applicable with ZeRO >= Stage-2.
+`offload_param_device`: [none] Disable parameter offloading, [cpu] offload parameters to CPU, [nvme] offload parameters to NVMe SSD. Only applicable with ZeRO Stage-3.
+`zero3_init_flag`: Decides whether to enable `deepspeed.zero.Init` for constructing massive models. Only applicable with ZeRO Stage-3.
+`zero3_save_16bit_model`: Decides whether to save 16-bit model weights when using ZeRO Stage-3.
+`mixed_precision`: `no` for FP32 training, `fp16` for FP16 mixed-precision training and `bf16` for BF16 mixed-precision training. 
+```
+
+### DeepSpeed Config File
+On your machine(s) just run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked. It will ask whether you want to use a config file for deepspeed to which you answer yes 
+and provide the path to the deepspeed config file. 
+This will generate a config file that will be used automatically to properly set the
+default options when doing
+
+```bash
+accelerate launch my_script.py --args_to_my_script
+```
+
+For instance, here is how you would run the NLP example `examples/by_feature/deepspeed_with_config_support.py` (from the root of the repo) with DeepSpeed Config File:
+
+**ZeRO Stage-2 DeepSpeed Config File Example**
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+ deepspeed_config_file: /home/ubuntu/accelerate/examples/configs/deepspeed_config_templates/zero_stage2_config.json
+ zero3_init_flag: true
+distributed_type: DEEPSPEED
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 2
+use_cpu: false
+```
+
+with the contents of `zero_stage2_config.json` being:
+```json
+{
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "optimizer": {
+        "type": "AdamW",
+        "params": {
+            "lr": "auto",
+            "weight_decay": "auto",
+            "torch_adam": true,
+            "adam_w_mode": true
+        }
+    },
+    "scheduler": {
+        "type": "WarmupDecayLR",
+        "params": {
+            "warmup_min_lr": "auto",
+            "warmup_max_lr": "auto",
+            "warmup_num_steps": "auto",
+            "total_num_steps": "auto"
+        }
+    },
+    "zero_optimization": {
+        "stage": 2,
+        "allgather_partitions": true,
+        "allgather_bucket_size": 2e8,
+        "overlap_comm": true,
+        "reduce_scatter": true,
+        "reduce_bucket_size": "auto",
+        "contiguous_gradients": true
+    },
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}
+```
+
+```bash
+accelerate launch examples/by_feature/deepspeed_with_config_support.py \
+--config_name "gpt2-large" \
+--tokenizer_name "gpt2-large" \
+--dataset_name "wikitext" \
+--dataset_config_name "wikitext-2-raw-v1" \
+--block_size 128 \
+--output_dir "./clm/clm_deepspeed_stage2_accelerate" \
+--learning_rate 5e-4 \
+--per_device_train_batch_size 24 \
+--per_device_eval_batch_size 24 \
+--num_train_epochs 3 \
+--with_tracking \
+--report_to "wandb"\
+```
+
+**ZeRO Stage-3 with CPU offload DeepSpeed Config File Example**
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+ deepspeed_config_file: /home/ubuntu/accelerate/examples/configs/deepspeed_config_templates/zero_stage3_offload_config.json
+ zero3_init_flag: true
+distributed_type: DEEPSPEED
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 2
+use_cpu: false
+```
+with the contents of `zero_stage3_offload_config.json` being:
+```json
+{
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "optimizer": {
+        "type": "AdamW",
+        "params": {
+            "lr": "auto",
+            "weight_decay": "auto"
+        }
+    },
+    "scheduler": {
+        "type": "WarmupDecayLR",
+        "params": {
+            "warmup_min_lr": "auto",
+            "warmup_max_lr": "auto",
+            "warmup_num_steps": "auto",
+            "total_num_steps": "auto"
+        }
+    },
+    "zero_optimization": {
+        "stage": 3,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "offload_param": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "reduce_bucket_size": "auto",
+        "stage3_prefetch_bucket_size": "auto",
+        "stage3_param_persistence_threshold": "auto",
+        "sub_group_size": 1e9,
+        "stage3_max_live_parameters": 1e9,
+        "stage3_max_reuse_distance": 1e9,
+        "stage3_gather_16bit_weights_on_model_save": "auto"
+    },
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}
+```
+
+```bash
+accelerate launch examples/by_feature/deepspeed_with_config_support.py \
+--config_name "gpt2-large" \
+--tokenizer_name "gpt2-large" \
+--dataset_name "wikitext" \
+--dataset_config_name "wikitext-2-raw-v1" \
+--block_size 128 \
+--output_dir "./clm/clm_deepspeed_stage3_offload_accelerate" \
+--learning_rate 5e-4 \
+--per_device_train_batch_size 32 \
+--per_device_eval_batch_size 32 \
+--num_train_epochs 3 \
+--with_tracking \
+--report_to "wandb"\
+```
+
+**Important code changes when using DeepSpeed Config File**
+
+1. DeepSpeed Optimizers and Schedulers. For more information on these, 
+see the [DeepSpeed Optimizers](https://deepspeed.readthedocs.io/en/latest/optimizers.html) and [DeepSpeed Schedulers](https://deepspeed.readthedocs.io/en/latest/schedulers.html) documentation.
+We will look at the changes needed in the code when using these.
+   
+   a. DS Optim + DS Scheduler: The case when both `optimizer` and `scheduler` keys present in the DeepSpeed config file.
+   In this situation, those will be used and user has to use `accelerate.utils.DummyOptim` and `accelerate.utils.DummyScheduler` to replace the PyTorch/Custom optimizers and schedulers in their code.
+   Below is the snippet from `examples/by_feature/deepspeed_with_config_support.py` showing this:
+   ```python
+   # Creates Dummy Optimizer if `optimizer` was spcified in the config file else creates Adam Optimizer
+   if "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config:
+       optimizer = torch.optim.AdamW(optimizer_grouped_parameters, lr=args.learning_rate)
+   else:
+       optimizer = DummyOptim(param_groups=optimizer_grouped_parameters)
+
+   # Creates Dummy Optimizer if `optimizer` was spcified in the config file else creates Adam Optimizer
This is a huge block, I'd document less auto option plus make sure we automatically set the ones we can!
@@ -0,0 +1,542 @@
+<!--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.
+-->
+
+# DeepSpeed 
+
+Training advanced deep learning models is challenging. 
+Beyond model design, model scientists also need to set up the state-of-the-art training techniques such as 
+distributed training, mixed precision, gradient accumulation, and checkpointing. 
+Yet still, scientists may not achieve the desired system performance and convergence rate.
+Large model sizes are even more challenging: a large model easily runs out of memory with pure data parallelism and it is difficult to use model parallelism. 
+[DeepSpeed](https://www.deepspeed.ai/) addresses these challenges to accelerate model development and training.
+
+🤗 Accelerate integrates [DeepSpeed](https://github.com/microsoft/DeepSpeed) via 2 options:
+
+1. Integration of the DeepSpeed features via `deepspeed config file` specification in `accelerate config` . You just supply your custom config file or use our template. Most of
+   this document is focused on this feature. This supports all the core features of DeepSpeed and gives user a lot of flexibility. 
+   User may have to change few lines of code depending on the config.
+2. Integration via `deepspeed_plugin`.This supports subset of the DeepSpeed features and uses default options for the rest of the configurations. 
+   User need not change any code and is good for those who are fine with most of the default settings of DeepSpeed.
+
+## What is integrated?
+
+Training:
+
+1. DeepSpeed ZeRO training supports the full ZeRO stages 1, 2 and 3 with ZeRO-Infinity (CPU and NVME offload).
+
+Inference:
+
+1. DeepSpeed ZeRO Inference supports ZeRO stage 3 with ZeRO-Infinity. It uses the same ZeRO protocol as training, but
+   it doesn't use an optimizer and a lr scheduler and only stage 3 is relevant. For more details see:
+   [deepspeed-zero-inference](#deepspeed-zero-inference).
+
+
+## How it works?
+
+**Pre-Requisites**: Install DeepSpeed version >=0.6.5. Please refer to the [DeepSpeed Insallation details](https://github.com/microsoft/DeepSpeed#installation)
+for more information.
+
+We will first look at easy to use integration via `deepspeed_plugin`. 
+Followed by more flexible and feature rich `deepspeed config file` integration. 
+
+### Accelerate DeepSpeed Plugin
+On your machine(s) just run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked. It will ask whether you want to use a config file for deepspeed to which you answer no and answer the following questions. 
+This will generate a config file that will be used automatically to properly set the
+default options when doing
+
+```bash
+accelerate launch my_script.py --args_to_my_script
+```
+
+For instance, here is how you would run the NLP example `examples/by_feature/deepspeed_with_config_support.py` (from the root of the repo) with DeepSpeed Plugin:
+
+**ZeRO Stage-2 DeepSpeed Plugin Example**
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+ gradient_accumulation_steps: 1
+ gradient_clipping: 1.0
+ offload_optimizer_device: none
+ offload_param_device: none
+ zero3_init_flag: true
+ zero_stage: 2
+distributed_type: DEEPSPEED
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 2
+use_cpu: false
+```
+
+```bash
+accelerate launch examples/by_feature/deepspeed_with_config_support.py \
+--config_name "gpt2-large" \
+--tokenizer_name "gpt2-large" \
+--dataset_name "wikitext" \
+--dataset_config_name "wikitext-2-raw-v1" \
+--block_size 128 \
+--output_dir "./clm/clm_deepspeed_stage2_accelerate" \
+--learning_rate 5e-4 \
+--per_device_train_batch_size 24 \
+--per_device_eval_batch_size 24 \
+--num_train_epochs 3 \
+--with_tracking \
+--report_to "wandb"\
+```
+
+**ZeRO Stage-3 with CPU Offload DeepSpeed Plugin Example**
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+  gradient_accumulation_steps: 1
+  gradient_clipping: 1.0
+  offload_optimizer_device: cpu
+  offload_param_device: cpu
+  zero3_init_flag: true
+  zero3_save_16bit_model: true
+  zero_stage: 3
+distributed_type: DEEPSPEED
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 2
+use_cpu: false
+```
+
+```bash
+accelerate launch examples/by_feature/deepspeed_with_config_support.py \
+--config_name "gpt2-large" \
+--tokenizer_name "gpt2-large" \
+--dataset_name "wikitext" \
+--dataset_config_name "wikitext-2-raw-v1" \
+--block_size 128 \
+--output_dir "./clm/clm_deepspeed_stage3_offload_accelerate" \
+--learning_rate 5e-4 \
+--per_device_train_batch_size 32 \
+--per_device_eval_batch_size 32 \
+--num_train_epochs 3 \
+--with_tracking \
+--report_to "wandb"\
+```
+
+Currently, `Accelerate` supports following config through the CLI:
+
+```bash
+`zero_stage`: [0] Disabled, [1] optimizer state partitioning, [2] optimizer+gradient state partitioning and [3] optimizer+gradient+parameter partitioning
+`gradient_accumulation_steps`: Number of training steps to accumulate gradients before averaging and applying them.
+`gradient_clipping`: Enable gradient clipping with value.
+`offload_optimizer_device`: [none] Disable optimizer offloading, [cpu] offload optimizer to CPU, [nvme] offload optimizer to NVMe SSD. Only applicable with ZeRO >= Stage-2.
+`offload_param_device`: [none] Disable parameter offloading, [cpu] offload parameters to CPU, [nvme] offload parameters to NVMe SSD. Only applicable with ZeRO Stage-3.
+`zero3_init_flag`: Decides whether to enable `deepspeed.zero.Init` for constructing massive models. Only applicable with ZeRO Stage-3.
+`zero3_save_16bit_model`: Decides whether to save 16-bit model weights when using ZeRO Stage-3.
+`mixed_precision`: `no` for FP32 training, `fp16` for FP16 mixed-precision training and `bf16` for BF16 mixed-precision training. 
+```
+
+### DeepSpeed Config File
+On your machine(s) just run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked. It will ask whether you want to use a config file for deepspeed to which you answer yes 
+and provide the path to the deepspeed config file. 
+This will generate a config file that will be used automatically to properly set the
+default options when doing
+
+```bash
+accelerate launch my_script.py --args_to_my_script
+```
+
+For instance, here is how you would run the NLP example `examples/by_feature/deepspeed_with_config_support.py` (from the root of the repo) with DeepSpeed Config File:
+
+**ZeRO Stage-2 DeepSpeed Config File Example**
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+ deepspeed_config_file: /home/ubuntu/accelerate/examples/configs/deepspeed_config_templates/zero_stage2_config.json
+ zero3_init_flag: true
+distributed_type: DEEPSPEED
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 2
+use_cpu: false
+```
+
+with the contents of `zero_stage2_config.json` being:
+```json
+{
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "optimizer": {
+        "type": "AdamW",
+        "params": {
+            "lr": "auto",
+            "weight_decay": "auto",
+            "torch_adam": true,
+            "adam_w_mode": true
+        }
+    },
+    "scheduler": {
+        "type": "WarmupDecayLR",
+        "params": {
+            "warmup_min_lr": "auto",
+            "warmup_max_lr": "auto",
+            "warmup_num_steps": "auto",
+            "total_num_steps": "auto"
+        }
+    },
+    "zero_optimization": {
+        "stage": 2,
+        "allgather_partitions": true,
+        "allgather_bucket_size": 2e8,
+        "overlap_comm": true,
+        "reduce_scatter": true,
+        "reduce_bucket_size": "auto",
+        "contiguous_gradients": true
+    },
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}
+```
+
+```bash
+accelerate launch examples/by_feature/deepspeed_with_config_support.py \
+--config_name "gpt2-large" \
+--tokenizer_name "gpt2-large" \
+--dataset_name "wikitext" \
+--dataset_config_name "wikitext-2-raw-v1" \
+--block_size 128 \
+--output_dir "./clm/clm_deepspeed_stage2_accelerate" \
+--learning_rate 5e-4 \
+--per_device_train_batch_size 24 \
+--per_device_eval_batch_size 24 \
+--num_train_epochs 3 \
+--with_tracking \
+--report_to "wandb"\
+```
+
+**ZeRO Stage-3 with CPU offload DeepSpeed Config File Example**
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+ deepspeed_config_file: /home/ubuntu/accelerate/examples/configs/deepspeed_config_templates/zero_stage3_offload_config.json
+ zero3_init_flag: true
+distributed_type: DEEPSPEED
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 2
+use_cpu: false
+```
+with the contents of `zero_stage3_offload_config.json` being:
+```json
+{
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "optimizer": {
+        "type": "AdamW",
+        "params": {
+            "lr": "auto",
+            "weight_decay": "auto"
+        }
+    },
+    "scheduler": {
+        "type": "WarmupDecayLR",
+        "params": {
+            "warmup_min_lr": "auto",
+            "warmup_max_lr": "auto",
+            "warmup_num_steps": "auto",
+            "total_num_steps": "auto"
+        }
+    },
+    "zero_optimization": {
+        "stage": 3,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "offload_param": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "reduce_bucket_size": "auto",
+        "stage3_prefetch_bucket_size": "auto",
+        "stage3_param_persistence_threshold": "auto",
+        "sub_group_size": 1e9,
+        "stage3_max_live_parameters": 1e9,
+        "stage3_max_reuse_distance": 1e9,
+        "stage3_gather_16bit_weights_on_model_save": "auto"
+    },
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}
+```
+
+```bash
+accelerate launch examples/by_feature/deepspeed_with_config_support.py \
+--config_name "gpt2-large" \
+--tokenizer_name "gpt2-large" \
+--dataset_name "wikitext" \
+--dataset_config_name "wikitext-2-raw-v1" \
+--block_size 128 \
+--output_dir "./clm/clm_deepspeed_stage3_offload_accelerate" \
+--learning_rate 5e-4 \
+--per_device_train_batch_size 32 \
+--per_device_eval_batch_size 32 \
+--num_train_epochs 3 \
+--with_tracking \
+--report_to "wandb"\
+```
+
+**Important code changes when using DeepSpeed Config File**
+
+1. DeepSpeed Optimizers and Schedulers. For more information on these, 
+see the [DeepSpeed Optimizers](https://deepspeed.readthedocs.io/en/latest/optimizers.html) and [DeepSpeed Schedulers](https://deepspeed.readthedocs.io/en/latest/schedulers.html) documentation.
+We will look at the changes needed in the code when using these.
+   
+   a. DS Optim + DS Scheduler: The case when both `optimizer` and `scheduler` keys present in the DeepSpeed config file.
+   In this situation, those will be used and user has to use `accelerate.utils.DummyOptim` and `accelerate.utils.DummyScheduler` to replace the PyTorch/Custom optimizers and schedulers in their code.
+   Below is the snippet from `examples/by_feature/deepspeed_with_config_support.py` showing this:
+   ```python
+   # Creates Dummy Optimizer if `optimizer` was spcified in the config file else creates Adam Optimizer
+   if "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config:
+       optimizer = torch.optim.AdamW(optimizer_grouped_parameters, lr=args.learning_rate)
+   else:
+       optimizer = DummyOptim(param_groups=optimizer_grouped_parameters)
+
+   # Creates Dummy Optimizer if `optimizer` was spcified in the config file else creates Adam Optimizer
+   if "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config:
+       lr_scheduler = get_scheduler(
+           name=args.lr_scheduler_type,
+           optimizer=optimizer,
+           num_warmup_steps=args.num_warmup_steps,
+           num_training_steps=args.max_train_steps,
+       )
+   else:
+       lr_scheduler = DummyScheduler()
+   ```
+   b. Custom Optim + Custom Scheduler: The case when both `optimizer` and `scheduler` keys are absent in the DeepSpeed config file.
+   In this situation, no code changes are needed from the user and this is the case when using integration via DeepSpeed Plugin.
+   In the above example we can see that the code reamins unchanged if the `optimizer` and `scheduler` keys are absent in the DeepSpeed config file.
+
+   c. Custom Optim + DS Scheduler: The case when only `scheduler` key is present in the DeepSpeed config file. 
+   In this situation, user has to use `accelerate.utils.DummyScheduler` to replace the PyTorch/Custom scheduler in their code. 
+
+   d. DS Optim + Custom Scheduler: The case when only `optimizer` key is present in the DeepSpeed config file. 
+   This will result in an error because one can only use DS Scheduler when using DS Optim.
+
+2. Notice the `auto` values in the above example DeepSpeed config files. 
+These have to specified by the user in the code by running `accelerator.state.deepspeed_plugin.deepspeed_config_process(**kwargs)`. 
+This allows user more flexibilty to fill these fields based on the code instead of hardcoding them. 
+Below is the snippet from `examples/by_feature/deepspeed_with_config_support.py` showing how to fill these: 
+```python
```suggestion
def checkpoint_model(checkpoint_folder, ckpt_id, model, epoch, last_global_step, **kwargs):
```
All-caps name are for constants, and here this is not a constant :-)
@@ -0,0 +1,753 @@
+#!/usr/bin/env python
+# coding=utf-8
+# Copyright 2022 The HuggingFace Inc. 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.
+"""
+Fine-tuning the library models for causal language modeling (GPT, GPT-2, CTRL, ...)
+on a text file or a dataset without using HuggingFace Trainer.
+
+Here is the full list of checkpoints on the hub that can be fine-tuned by this script:
+https://huggingface.co/models?filter=text-generation
+"""
+# You can also adapt this script on your own causal language modeling task. Pointers for this are left as comments.
+
+import argparse
+import json
+import logging
+import math
+import os
+import random
+from itertools import chain
+from pathlib import Path
+
+import torch
+from torch.utils.data import DataLoader
+
+import datasets
+import transformers
+from accelerate import Accelerator, DistributedType
+from accelerate.logging import get_logger
+from accelerate.utils import DummyOptim, DummyScheduler, set_seed
+from datasets import load_dataset
+from huggingface_hub import Repository
+from tqdm.auto import tqdm
+from transformers import (
+    CONFIG_MAPPING,
+    MODEL_MAPPING,
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    SchedulerType,
+    default_data_collator,
+    get_scheduler,
+)
+from transformers.utils import get_full_repo_name
+from transformers.utils.versions import require_version
+
+
+logger = get_logger(__name__)
+
+require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt")
+
+MODEL_CONFIG_CLASSES = list(MODEL_MAPPING.keys())
+MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Finetune a transformers model on a causal language modeling task")
+    parser.add_argument(
+        "--dataset_name",
+        type=str,
+        default=None,
+        help="The name of the dataset to use (via the datasets library).",
+    )
+    parser.add_argument(
+        "--dataset_config_name",
+        type=str,
+        default=None,
+        help="The configuration name of the dataset to use (via the datasets library).",
+    )
+    parser.add_argument(
+        "--train_file", type=str, default=None, help="A csv or a json file containing the training data."
+    )
+    parser.add_argument(
+        "--validation_file", type=str, default=None, help="A csv or a json file containing the validation data."
+    )
+    parser.add_argument(
+        "--validation_split_percentage",
+        default=5,
+        help="The percentage of the train set used as validation set in case there's no validation split",
+    )
+    parser.add_argument(
+        "--model_name_or_path",
+        type=str,
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+        required=False,
+    )
+    parser.add_argument(
+        "--config_name",
+        type=str,
+        default=None,
+        help="Pretrained config name or path if not the same as model_name",
+    )
+    parser.add_argument(
+        "--tokenizer_name",
+        type=str,
+        default=None,
+        help="Pretrained tokenizer name or path if not the same as model_name",
+    )
+    parser.add_argument(
+        "--use_slow_tokenizer",
+        action="store_true",
+        help="If passed, will use a slow tokenizer (not backed by the 🤗 Tokenizers library).",
+    )
+    parser.add_argument(
+        "--per_device_train_batch_size",
+        type=int,
+        default=8,
+        help="Batch size (per device) for the training dataloader.",
+    )
+    parser.add_argument(
+        "--per_device_eval_batch_size",
+        type=int,
+        default=8,
+        help="Batch size (per device) for the evaluation dataloader.",
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=5e-5,
+        help="Initial learning rate (after the potential warmup period) to use.",
+    )
+    parser.add_argument("--weight_decay", type=float, default=0.0, help="Weight decay to use.")
+    parser.add_argument("--num_train_epochs", type=int, default=3, help="Total number of training epochs to perform.")
+    parser.add_argument(
+        "--max_train_steps",
+        type=int,
+        default=None,
+        help="Total number of training steps to perform. If provided, overrides num_train_epochs.",
+    )
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help="Number of updates steps to accumulate before performing a backward/update pass.",
+    )
+    parser.add_argument(
+        "--lr_scheduler_type",
+        type=SchedulerType,
+        default="linear",
+        help="The scheduler type to use.",
+        choices=["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"],
+    )
+    parser.add_argument(
+        "--num_warmup_steps", type=int, default=0, help="Number of steps for the warmup in the lr scheduler."
+    )
+    parser.add_argument("--output_dir", type=str, default=None, help="Where to store the final model.")
+    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
+    parser.add_argument(
+        "--model_type",
+        type=str,
+        default=None,
+        help="Model type to use if training from scratch.",
+        choices=MODEL_TYPES,
+    )
+    parser.add_argument(
+        "--block_size",
+        type=int,
+        default=None,
+        help=(
+            "Optional input sequence length after tokenization. The training dataset will be truncated in block of"
+            " this size for training. Default to the model max input length for single sentence inputs (take into"
+            " account special tokens)."
+        ),
+    )
+    parser.add_argument(
+        "--preprocessing_num_workers",
+        type=int,
+        default=None,
+        help="The number of processes to use for the preprocessing.",
+    )
+    parser.add_argument(
+        "--overwrite_cache", type=bool, default=False, help="Overwrite the cached training and evaluation sets"
+    )
+    parser.add_argument(
+        "--no_keep_linebreaks", action="store_true", help="Do not keep line breaks when using TXT files."
+    )
+    parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")
+    parser.add_argument(
+        "--hub_model_id", type=str, help="The name of the repository to keep in sync with the local `output_dir`."
+    )
+    parser.add_argument("--hub_token", type=str, help="The token to use to push to the Model Hub.")
+    parser.add_argument(
+        "--checkpointing_steps",
+        type=str,
+        default=None,
+        help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.",
+    )
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        default=None,
+        help="If the training should continue from a checkpoint folder.",
+    )
+    # New Code #
+    # Whether to load the best model at the end of training
+    parser.add_argument(
+        "--load_best_model",
+        action="store_true",
+        help="Whether to load the best model at the end of training",
+    )
+    parser.add_argument(
+        "--with_tracking",
+        action="store_true",
+        help="Whether to enable experiment trackers for logging.",
+    )
+    parser.add_argument(
+        "--report_to",
+        type=str,
+        default="all",
+        help=(
+            'The integration to report the results and logs to. Supported platforms are `"tensorboard"`,'
+            ' `"wandb"` and `"comet_ml"`. Use `"all"` (default) to report to all integrations.'
+            "Only applicable when `--with_tracking` is passed."
+        ),
+    )
+    args = parser.parse_args()
+
+    # Sanity checks
+    if args.dataset_name is None and args.train_file is None and args.validation_file is None:
+        raise ValueError("Need either a dataset name or a training/validation file.")
+    else:
+        if args.train_file is not None:
+            extension = args.train_file.split(".")[-1]
+            assert extension in ["csv", "json", "txt"], "`train_file` should be a csv, json or txt file."
+        if args.validation_file is not None:
+            extension = args.validation_file.split(".")[-1]
+            assert extension in ["csv", "json", "txt"], "`validation_file` should be a csv, json or txt file."
+
+    if args.push_to_hub:
+        assert args.output_dir is not None, "Need an `output_dir` to create a repo when `--push_to_hub` is passed."
+
+    return args
+
+
+# New Code #
+def checkpoint_model(PATH, ckpt_id, model, epoch, last_global_step, **kwargs):
```suggestion
    success = model.save_checkpoint(checkpoint_folder, ckpt_id, checkpoint_state_dict)
    status_msg = f"checkpointing: checkpoint_folder={checkpoint_folder}, ckpt_id={ckpt_id}"
```
Variable name + f-strings only.
@@ -0,0 +1,753 @@
+#!/usr/bin/env python
+# coding=utf-8
+# Copyright 2022 The HuggingFace Inc. 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.
+"""
+Fine-tuning the library models for causal language modeling (GPT, GPT-2, CTRL, ...)
+on a text file or a dataset without using HuggingFace Trainer.
+
+Here is the full list of checkpoints on the hub that can be fine-tuned by this script:
+https://huggingface.co/models?filter=text-generation
+"""
+# You can also adapt this script on your own causal language modeling task. Pointers for this are left as comments.
+
+import argparse
+import json
+import logging
+import math
+import os
+import random
+from itertools import chain
+from pathlib import Path
+
+import torch
+from torch.utils.data import DataLoader
+
+import datasets
+import transformers
+from accelerate import Accelerator, DistributedType
+from accelerate.logging import get_logger
+from accelerate.utils import DummyOptim, DummyScheduler, set_seed
+from datasets import load_dataset
+from huggingface_hub import Repository
+from tqdm.auto import tqdm
+from transformers import (
+    CONFIG_MAPPING,
+    MODEL_MAPPING,
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    SchedulerType,
+    default_data_collator,
+    get_scheduler,
+)
+from transformers.utils import get_full_repo_name
+from transformers.utils.versions import require_version
+
+
+logger = get_logger(__name__)
+
+require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt")
+
+MODEL_CONFIG_CLASSES = list(MODEL_MAPPING.keys())
+MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Finetune a transformers model on a causal language modeling task")
+    parser.add_argument(
+        "--dataset_name",
+        type=str,
+        default=None,
+        help="The name of the dataset to use (via the datasets library).",
+    )
+    parser.add_argument(
+        "--dataset_config_name",
+        type=str,
+        default=None,
+        help="The configuration name of the dataset to use (via the datasets library).",
+    )
+    parser.add_argument(
+        "--train_file", type=str, default=None, help="A csv or a json file containing the training data."
+    )
+    parser.add_argument(
+        "--validation_file", type=str, default=None, help="A csv or a json file containing the validation data."
+    )
+    parser.add_argument(
+        "--validation_split_percentage",
+        default=5,
+        help="The percentage of the train set used as validation set in case there's no validation split",
+    )
+    parser.add_argument(
+        "--model_name_or_path",
+        type=str,
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+        required=False,
+    )
+    parser.add_argument(
+        "--config_name",
+        type=str,
+        default=None,
+        help="Pretrained config name or path if not the same as model_name",
+    )
+    parser.add_argument(
+        "--tokenizer_name",
+        type=str,
+        default=None,
+        help="Pretrained tokenizer name or path if not the same as model_name",
+    )
+    parser.add_argument(
+        "--use_slow_tokenizer",
+        action="store_true",
+        help="If passed, will use a slow tokenizer (not backed by the 🤗 Tokenizers library).",
+    )
+    parser.add_argument(
+        "--per_device_train_batch_size",
+        type=int,
+        default=8,
+        help="Batch size (per device) for the training dataloader.",
+    )
+    parser.add_argument(
+        "--per_device_eval_batch_size",
+        type=int,
+        default=8,
+        help="Batch size (per device) for the evaluation dataloader.",
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=5e-5,
+        help="Initial learning rate (after the potential warmup period) to use.",
+    )
+    parser.add_argument("--weight_decay", type=float, default=0.0, help="Weight decay to use.")
+    parser.add_argument("--num_train_epochs", type=int, default=3, help="Total number of training epochs to perform.")
+    parser.add_argument(
+        "--max_train_steps",
+        type=int,
+        default=None,
+        help="Total number of training steps to perform. If provided, overrides num_train_epochs.",
+    )
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help="Number of updates steps to accumulate before performing a backward/update pass.",
+    )
+    parser.add_argument(
+        "--lr_scheduler_type",
+        type=SchedulerType,
+        default="linear",
+        help="The scheduler type to use.",
+        choices=["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"],
+    )
+    parser.add_argument(
+        "--num_warmup_steps", type=int, default=0, help="Number of steps for the warmup in the lr scheduler."
+    )
+    parser.add_argument("--output_dir", type=str, default=None, help="Where to store the final model.")
+    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
+    parser.add_argument(
+        "--model_type",
+        type=str,
+        default=None,
+        help="Model type to use if training from scratch.",
+        choices=MODEL_TYPES,
+    )
+    parser.add_argument(
+        "--block_size",
+        type=int,
+        default=None,
+        help=(
+            "Optional input sequence length after tokenization. The training dataset will be truncated in block of"
+            " this size for training. Default to the model max input length for single sentence inputs (take into"
+            " account special tokens)."
+        ),
+    )
+    parser.add_argument(
+        "--preprocessing_num_workers",
+        type=int,
+        default=None,
+        help="The number of processes to use for the preprocessing.",
+    )
+    parser.add_argument(
+        "--overwrite_cache", type=bool, default=False, help="Overwrite the cached training and evaluation sets"
+    )
+    parser.add_argument(
+        "--no_keep_linebreaks", action="store_true", help="Do not keep line breaks when using TXT files."
+    )
+    parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")
+    parser.add_argument(
+        "--hub_model_id", type=str, help="The name of the repository to keep in sync with the local `output_dir`."
+    )
+    parser.add_argument("--hub_token", type=str, help="The token to use to push to the Model Hub.")
+    parser.add_argument(
+        "--checkpointing_steps",
+        type=str,
+        default=None,
+        help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.",
+    )
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        default=None,
+        help="If the training should continue from a checkpoint folder.",
+    )
+    # New Code #
+    # Whether to load the best model at the end of training
+    parser.add_argument(
+        "--load_best_model",
+        action="store_true",
+        help="Whether to load the best model at the end of training",
+    )
+    parser.add_argument(
+        "--with_tracking",
+        action="store_true",
+        help="Whether to enable experiment trackers for logging.",
+    )
+    parser.add_argument(
+        "--report_to",
+        type=str,
+        default="all",
+        help=(
+            'The integration to report the results and logs to. Supported platforms are `"tensorboard"`,'
+            ' `"wandb"` and `"comet_ml"`. Use `"all"` (default) to report to all integrations.'
+            "Only applicable when `--with_tracking` is passed."
+        ),
+    )
+    args = parser.parse_args()
+
+    # Sanity checks
+    if args.dataset_name is None and args.train_file is None and args.validation_file is None:
+        raise ValueError("Need either a dataset name or a training/validation file.")
+    else:
+        if args.train_file is not None:
+            extension = args.train_file.split(".")[-1]
+            assert extension in ["csv", "json", "txt"], "`train_file` should be a csv, json or txt file."
+        if args.validation_file is not None:
+            extension = args.validation_file.split(".")[-1]
+            assert extension in ["csv", "json", "txt"], "`validation_file` should be a csv, json or txt file."
+
+    if args.push_to_hub:
+        assert args.output_dir is not None, "Need an `output_dir` to create a repo when `--push_to_hub` is passed."
+
+    return args
+
+
+# New Code #
+def checkpoint_model(PATH, ckpt_id, model, epoch, last_global_step, **kwargs):
+    """Utility function for checkpointing model + optimizer dictionaries
+    The main purpose for this is to be able to resume training from that instant again
+    """
+    checkpoint_state_dict = {
+        "epoch": epoch,
+        "last_global_step": last_global_step,
+    }
+    # Add extra kwargs too
+    checkpoint_state_dict.update(kwargs)
+
+    success = model.save_checkpoint(PATH, ckpt_id, checkpoint_state_dict)
+    status_msg = "checkpointing: PATH={}, ckpt_id={}".format(PATH, ckpt_id)
As discussed offline this is a very big specific block. We should try to avoid this at all cost.
@@ -0,0 +1,753 @@
+#!/usr/bin/env python
+# coding=utf-8
+# Copyright 2022 The HuggingFace Inc. 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.
+"""
+Fine-tuning the library models for causal language modeling (GPT, GPT-2, CTRL, ...)
+on a text file or a dataset without using HuggingFace Trainer.
+
+Here is the full list of checkpoints on the hub that can be fine-tuned by this script:
+https://huggingface.co/models?filter=text-generation
+"""
+# You can also adapt this script on your own causal language modeling task. Pointers for this are left as comments.
+
+import argparse
+import json
+import logging
+import math
+import os
+import random
+from itertools import chain
+from pathlib import Path
+
+import torch
+from torch.utils.data import DataLoader
+
+import datasets
+import transformers
+from accelerate import Accelerator, DistributedType
+from accelerate.logging import get_logger
+from accelerate.utils import DummyOptim, DummyScheduler, set_seed
+from datasets import load_dataset
+from huggingface_hub import Repository
+from tqdm.auto import tqdm
+from transformers import (
+    CONFIG_MAPPING,
+    MODEL_MAPPING,
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    SchedulerType,
+    default_data_collator,
+    get_scheduler,
+)
+from transformers.utils import get_full_repo_name
+from transformers.utils.versions import require_version
+
+
+logger = get_logger(__name__)
+
+require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt")
+
+MODEL_CONFIG_CLASSES = list(MODEL_MAPPING.keys())
+MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Finetune a transformers model on a causal language modeling task")
+    parser.add_argument(
+        "--dataset_name",
+        type=str,
+        default=None,
+        help="The name of the dataset to use (via the datasets library).",
+    )
+    parser.add_argument(
+        "--dataset_config_name",
+        type=str,
+        default=None,
+        help="The configuration name of the dataset to use (via the datasets library).",
+    )
+    parser.add_argument(
+        "--train_file", type=str, default=None, help="A csv or a json file containing the training data."
+    )
+    parser.add_argument(
+        "--validation_file", type=str, default=None, help="A csv or a json file containing the validation data."
+    )
+    parser.add_argument(
+        "--validation_split_percentage",
+        default=5,
+        help="The percentage of the train set used as validation set in case there's no validation split",
+    )
+    parser.add_argument(
+        "--model_name_or_path",
+        type=str,
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+        required=False,
+    )
+    parser.add_argument(
+        "--config_name",
+        type=str,
+        default=None,
+        help="Pretrained config name or path if not the same as model_name",
+    )
+    parser.add_argument(
+        "--tokenizer_name",
+        type=str,
+        default=None,
+        help="Pretrained tokenizer name or path if not the same as model_name",
+    )
+    parser.add_argument(
+        "--use_slow_tokenizer",
+        action="store_true",
+        help="If passed, will use a slow tokenizer (not backed by the 🤗 Tokenizers library).",
+    )
+    parser.add_argument(
+        "--per_device_train_batch_size",
+        type=int,
+        default=8,
+        help="Batch size (per device) for the training dataloader.",
+    )
+    parser.add_argument(
+        "--per_device_eval_batch_size",
+        type=int,
+        default=8,
+        help="Batch size (per device) for the evaluation dataloader.",
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=5e-5,
+        help="Initial learning rate (after the potential warmup period) to use.",
+    )
+    parser.add_argument("--weight_decay", type=float, default=0.0, help="Weight decay to use.")
+    parser.add_argument("--num_train_epochs", type=int, default=3, help="Total number of training epochs to perform.")
+    parser.add_argument(
+        "--max_train_steps",
+        type=int,
+        default=None,
+        help="Total number of training steps to perform. If provided, overrides num_train_epochs.",
+    )
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help="Number of updates steps to accumulate before performing a backward/update pass.",
+    )
+    parser.add_argument(
+        "--lr_scheduler_type",
+        type=SchedulerType,
+        default="linear",
+        help="The scheduler type to use.",
+        choices=["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"],
+    )
+    parser.add_argument(
+        "--num_warmup_steps", type=int, default=0, help="Number of steps for the warmup in the lr scheduler."
+    )
+    parser.add_argument("--output_dir", type=str, default=None, help="Where to store the final model.")
+    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
+    parser.add_argument(
+        "--model_type",
+        type=str,
+        default=None,
+        help="Model type to use if training from scratch.",
+        choices=MODEL_TYPES,
+    )
+    parser.add_argument(
+        "--block_size",
+        type=int,
+        default=None,
+        help=(
+            "Optional input sequence length after tokenization. The training dataset will be truncated in block of"
+            " this size for training. Default to the model max input length for single sentence inputs (take into"
+            " account special tokens)."
+        ),
+    )
+    parser.add_argument(
+        "--preprocessing_num_workers",
+        type=int,
+        default=None,
+        help="The number of processes to use for the preprocessing.",
+    )
+    parser.add_argument(
+        "--overwrite_cache", type=bool, default=False, help="Overwrite the cached training and evaluation sets"
+    )
+    parser.add_argument(
+        "--no_keep_linebreaks", action="store_true", help="Do not keep line breaks when using TXT files."
+    )
+    parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")
+    parser.add_argument(
+        "--hub_model_id", type=str, help="The name of the repository to keep in sync with the local `output_dir`."
+    )
+    parser.add_argument("--hub_token", type=str, help="The token to use to push to the Model Hub.")
+    parser.add_argument(
+        "--checkpointing_steps",
+        type=str,
+        default=None,
+        help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.",
+    )
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        default=None,
+        help="If the training should continue from a checkpoint folder.",
+    )
+    # New Code #
+    # Whether to load the best model at the end of training
+    parser.add_argument(
+        "--load_best_model",
+        action="store_true",
+        help="Whether to load the best model at the end of training",
+    )
+    parser.add_argument(
+        "--with_tracking",
+        action="store_true",
+        help="Whether to enable experiment trackers for logging.",
+    )
+    parser.add_argument(
+        "--report_to",
+        type=str,
+        default="all",
+        help=(
+            'The integration to report the results and logs to. Supported platforms are `"tensorboard"`,'
+            ' `"wandb"` and `"comet_ml"`. Use `"all"` (default) to report to all integrations.'
+            "Only applicable when `--with_tracking` is passed."
+        ),
+    )
+    args = parser.parse_args()
+
+    # Sanity checks
+    if args.dataset_name is None and args.train_file is None and args.validation_file is None:
+        raise ValueError("Need either a dataset name or a training/validation file.")
+    else:
+        if args.train_file is not None:
+            extension = args.train_file.split(".")[-1]
+            assert extension in ["csv", "json", "txt"], "`train_file` should be a csv, json or txt file."
+        if args.validation_file is not None:
+            extension = args.validation_file.split(".")[-1]
+            assert extension in ["csv", "json", "txt"], "`validation_file` should be a csv, json or txt file."
+
+    if args.push_to_hub:
+        assert args.output_dir is not None, "Need an `output_dir` to create a repo when `--push_to_hub` is passed."
+
+    return args
+
+
+# New Code #
+def checkpoint_model(PATH, ckpt_id, model, epoch, last_global_step, **kwargs):
+    """Utility function for checkpointing model + optimizer dictionaries
+    The main purpose for this is to be able to resume training from that instant again
+    """
+    checkpoint_state_dict = {
+        "epoch": epoch,
+        "last_global_step": last_global_step,
+    }
+    # Add extra kwargs too
+    checkpoint_state_dict.update(kwargs)
+
+    success = model.save_checkpoint(PATH, ckpt_id, checkpoint_state_dict)
+    status_msg = "checkpointing: PATH={}, ckpt_id={}".format(PATH, ckpt_id)
+    if success:
+        logging.info(f"Success {status_msg}")
+    else:
+        logging.warning(f"Failure {status_msg}")
+    return
+
+
+# New Code #
+def load_training_checkpoint(model, load_dir, tag=None, **kwargs):
+    """Utility function for checkpointing model + optimizer dictionaries
+    The main purpose for this is to be able to resume training from that instant again
+    """
+    _, checkpoint_state_dict = model.load_checkpoint(load_dir, tag=tag, **kwargs)
+    epoch = checkpoint_state_dict["epoch"]
+    last_global_step = checkpoint_state_dict["last_global_step"]
+    del checkpoint_state_dict
+    return (epoch, last_global_step)
+
+
+# New Code #
+def evaluate(args, model, eval_dataloader, accelerator, eval_dataset):
+    model.eval()
+    losses = []
+    for step, batch in enumerate(eval_dataloader):
+        with torch.no_grad():
+            outputs = model(**batch)
+
+        loss = outputs.loss
+        losses.append(accelerator.gather(loss.repeat(args.per_device_eval_batch_size)))
+
+    losses = torch.cat(losses)
+    losses = losses[: len(eval_dataset)]
+    try:
+        eval_loss = torch.mean(losses)
+        perplexity = math.exp(eval_loss)
+    except OverflowError:
+        perplexity = float("inf")
+    return perplexity, eval_loss
+
+
+def main():
+    args = parse_args()
+
+    # Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
+    # If we're using tracking, we also need to initialize it here and it will by default pick up all supported trackers
+    # in the environment
+    accelerator = (
+        Accelerator(log_with=args.report_to, logging_dir=args.output_dir) if args.with_tracking else Accelerator()
+    )
+    # Make one log on every process with the configuration for debugging.
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO,
+    )
+    logger.info(accelerator.state, main_process_only=False)
+    if accelerator.is_local_main_process:
+        datasets.utils.logging.set_verbosity_warning()
+        transformers.utils.logging.set_verbosity_info()
+    else:
+        datasets.utils.logging.set_verbosity_error()
+        transformers.utils.logging.set_verbosity_error()
+
+    # If passed along, set the training seed now.
+    if args.seed is not None:
+        set_seed(args.seed)
+
+    # Handle the repository creation
+    if accelerator.is_main_process:
+        if args.push_to_hub:
+            if args.hub_model_id is None:
+                repo_name = get_full_repo_name(Path(args.output_dir).name, token=args.hub_token)
+            else:
+                repo_name = args.hub_model_id
+            repo = Repository(args.output_dir, clone_from=repo_name)
+
+            with open(os.path.join(args.output_dir, ".gitignore"), "w+") as gitignore:
+                if "step_*" not in gitignore:
+                    gitignore.write("step_*\n")
+                if "epoch_*" not in gitignore:
+                    gitignore.write("epoch_*\n")
+        elif args.output_dir is not None:
+            os.makedirs(args.output_dir, exist_ok=True)
+    accelerator.wait_for_everyone()
+
+    # Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)
+    # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
+    # (the dataset will be downloaded automatically from the datasets Hub).
+    #
+    # For CSV/JSON files, this script will use the column called 'text' or the first column if no column called
+    # 'text' is found. You can easily tweak this behavior (see below).
+    #
+    # In distributed training, the load_dataset function guarantee that only one local process can concurrently
+    # download the dataset.
+    if args.dataset_name is not None:
+        # Downloading and loading a dataset from the hub.
+        raw_datasets = load_dataset(args.dataset_name, args.dataset_config_name)
+        if "validation" not in raw_datasets.keys():
+            raw_datasets["validation"] = load_dataset(
+                args.dataset_name,
+                args.dataset_config_name,
+                split=f"train[:{args.validation_split_percentage}%]",
+            )
+            raw_datasets["train"] = load_dataset(
+                args.dataset_name,
+                args.dataset_config_name,
+                split=f"train[{args.validation_split_percentage}%:]",
+            )
+    else:
+        data_files = {}
+        dataset_args = {}
+        if args.train_file is not None:
+            data_files["train"] = args.train_file
+        if args.validation_file is not None:
+            data_files["validation"] = args.validation_file
+        extension = args.train_file.split(".")[-1]
+        if extension == "txt":
+            extension = "text"
+            dataset_args["keep_linebreaks"] = not args.no_keep_linebreaks
+        raw_datasets = load_dataset(extension, data_files=data_files, **dataset_args)
+        # If no validation data is there, validation_split_percentage will be used to divide the dataset.
+        if "validation" not in raw_datasets.keys():
+            raw_datasets["validation"] = load_dataset(
+                extension,
+                data_files=data_files,
+                split=f"train[:{args.validation_split_percentage}%]",
+                **dataset_args,
+            )
+            raw_datasets["train"] = load_dataset(
+                extension,
+                data_files=data_files,
+                split=f"train[{args.validation_split_percentage}%:]",
+                **dataset_args,
+            )
+
+    # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
+    # https://huggingface.co/docs/datasets/loading_datasets.html.
+
+    # Load pretrained model and tokenizer
+    #
+    # In distributed training, the .from_pretrained methods guarantee that only one local process can concurrently
+    # download model & vocab.
+    if args.config_name:
+        config = AutoConfig.from_pretrained(args.config_name)
+    elif args.model_name_or_path:
+        config = AutoConfig.from_pretrained(args.model_name_or_path)
+    else:
+        config = CONFIG_MAPPING[args.model_type]()
+        logger.warning("You are instantiating a new config instance from scratch.")
+
+    if args.tokenizer_name:
+        tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name, use_fast=not args.use_slow_tokenizer)
+    elif args.model_name_or_path:
+        tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, use_fast=not args.use_slow_tokenizer)
+    else:
+        raise ValueError(
+            "You are instantiating a new tokenizer from scratch. This is not supported by this script."
+            "You can do it from another script, save it, and load it from here, using --tokenizer_name."
+        )
+
+    if args.model_name_or_path:
+        model = AutoModelForCausalLM.from_pretrained(
+            args.model_name_or_path,
+            from_tf=bool(".ckpt" in args.model_name_or_path),
+            config=config,
+        )
+    else:
+        logger.info("Training new model from scratch")
+        model = AutoModelForCausalLM.from_config(config)
+
+    model.resize_token_embeddings(len(tokenizer))
+
+    # Preprocessing the datasets.
+    # First we tokenize all the texts.
+    column_names = raw_datasets["train"].column_names
+    text_column_name = "text" if "text" in column_names else column_names[0]
+
+    def tokenize_function(examples):
+        return tokenizer(examples[text_column_name])
+
+    with accelerator.main_process_first():
+        tokenized_datasets = raw_datasets.map(
+            tokenize_function,
+            batched=True,
+            num_proc=args.preprocessing_num_workers,
+            remove_columns=column_names,
+            load_from_cache_file=not args.overwrite_cache,
+            desc="Running tokenizer on dataset",
+        )
+
+    if args.block_size is None:
+        block_size = tokenizer.model_max_length
+        if block_size > 1024:
+            logger.warning(
+                f"The tokenizer picked seems to have a very large `model_max_length` ({tokenizer.model_max_length}). "
+                "Picking 1024 instead. You can change that default value by passing --block_size xxx."
+            )
+        block_size = 1024
+    else:
+        if args.block_size > tokenizer.model_max_length:
+            logger.warning(
+                f"The block_size passed ({args.block_size}) is larger than the maximum length for the model"
+                f"({tokenizer.model_max_length}). Using block_size={tokenizer.model_max_length}."
+            )
+        block_size = min(args.block_size, tokenizer.model_max_length)
+
+    # Main data processing function that will concatenate all texts from our dataset and generate chunks of block_size.
+    def group_texts(examples):
+        # Concatenate all texts.
+        concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}
+        total_length = len(concatenated_examples[list(examples.keys())[0]])
+        # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
+        # customize this part to your needs.
+        if total_length >= block_size:
+            total_length = (total_length // block_size) * block_size
+        # Split by chunks of max_len.
+        result = {
+            k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
+            for k, t in concatenated_examples.items()
+        }
+        result["labels"] = result["input_ids"].copy()
+        return result
+
+    # Note that with `batched=True`, this map processes 1,000 texts together, so group_texts throws away a remainder
+    # for each of those groups of 1,000 texts. You can adjust that batch_size here but a higher value might be slower
+    # to preprocess.
+    #
+    # To speed up this part, we use multiprocessing. See the documentation of the map method for more information:
+    # https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasets.Dataset.map
+
+    with accelerator.main_process_first():
+        lm_datasets = tokenized_datasets.map(
+            group_texts,
+            batched=True,
+            num_proc=args.preprocessing_num_workers,
+            load_from_cache_file=not args.overwrite_cache,
+            desc=f"Grouping texts in chunks of {block_size}",
+        )
+
+    train_dataset = lm_datasets["train"]
+    eval_dataset = lm_datasets["validation"]
+
+    # Log a few random samples from the training set:
+    for index in random.sample(range(len(train_dataset)), 3):
+        logger.info(f"Sample {index} of the training set: {train_dataset[index]}.")
+
+    # DataLoaders creation:
+    train_dataloader = DataLoader(
+        train_dataset, shuffle=True, collate_fn=default_data_collator, batch_size=args.per_device_train_batch_size
+    )
+    eval_dataloader = DataLoader(
+        eval_dataset, collate_fn=default_data_collator, batch_size=args.per_device_eval_batch_size
+    )
+
+    # Optimizer
+    # Split weights in two groups, one with weight decay and the other not.
+    no_decay = ["bias", "LayerNorm.weight"]
+    optimizer_grouped_parameters = [
+        {
+            "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
+            "weight_decay": args.weight_decay,
+        },
+        {
+            "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
+            "weight_decay": 0.0,
+        },
+    ]
+    # New Code #
+    # Creates Dummy Optimizer if `optimizer` was spcified in the config file else creates Adam Optimizer
+    if "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config:
+        optimizer = torch.optim.AdamW(optimizer_grouped_parameters, lr=args.learning_rate)
+    else:
+        optimizer = DummyOptim(param_groups=optimizer_grouped_parameters)
+
+    # On TPU, the tie weights in our model have been disconnected, so we need to restore the ties.
+    if accelerator.distributed_type == DistributedType.TPU:
+        model.tie_weights()
+
+    # Scheduler and math around the number of training steps.
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    if args.max_train_steps is None:
+        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+    else:
+        args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+    # New Code #
+    # Creates Dummy Scheduler if `scheduler` was spcified in the config file else creates `args.lr_scheduler_type` Scheduler
+    if "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config:
+        lr_scheduler = get_scheduler(
+            name=args.lr_scheduler_type,
+            optimizer=optimizer,
+            num_warmup_steps=args.num_warmup_steps,
+            num_training_steps=args.max_train_steps,
+        )
+    else:
+        lr_scheduler = DummyScheduler()
+
+    # New Code #
+    # Processing all the fields with `auto` values in the DeepSpeed Config file
+    if accelerator.state.deepspeed_plugin.config_file != "none":
Same here.
@@ -75,22 +51,65 @@ class DeepSpeedOptimizerWrapper(AcceleratedOptimizer):
             The optimizer to wrap.
     """
 
-    def __init__(self, optimizer, model: DeepSpeedEngineWrapper):
+    def __init__(self, optimizer):
         super().__init__(optimizer, device_placement=False, scaler=None)
 
-        self.model = model
-
     def zero_grad(self, set_to_none=None):
-        pass  # `model.step()` is doing that automatically. Therefore, it's implementation is not needed
+        pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, it's implementation is not needed
 
     def step(self):
-        """This will handle optimizer.step() & optimizer.zero_grad() with gradient_accumulation"""
-        self.model.step()
+        pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, it's implementation is not needed
 
     @property
-    def is_overflow(self):
+    def step_was_skipped(self):
         """Whether or not the optimizer step was done, or skipped because of gradient overflow."""
-        overflow = False
-        if hasattr(self.optimizer, "overflow"):
-            overflow = self.optimizer.overflow
-        return overflow
+        return self.optimizer.overflow
+
+
+class DeepSpeedSchedulerWrapper(AcceleratedScheduler):
+    """
+    Internal wrapper around a deepspeed scheduler.
+
+    Args:
+        scheduler (`torch.optim.lr_scheduler.LambdaLR`):
+            The scheduler to wrap.
+        optimizers (one or a list of `torch.optim.Optimizer`):
+    """
+
+    def __init__(self, scheduler, optimizers):
+        super().__init__(scheduler, optimizers)
+
+    def step(self):
+        pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, it's implementation is not needed
+
+
+class DummyOptim:
+    """
+    Dummy optimizer presents model parameters or param groups, this is primarily used to follow conventional training
+    loop when optimizer config is specified in the deepspeed config file.
+
+    Args:
+        params (list of `torch.Tensor`):
+            List of model parameters.
+        param_groups (list of dict):
+            List of dictionaries with parameters and their groups.
+    """
+
+    def __init__(self, params=None, param_groups=None):
+        self.param_groups = []
+        if params is None and param_groups is None:
+            raise ValueError("Either params or param_groups should be specified")
+        if param_groups is None:
+            self.param_groups.append({"params": params})
+        else:
+            self.param_groups = param_groups
+
+
+class DummyScheduler:
Done.
@@ -0,0 +1,542 @@
+<!--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.
+-->
+
+# DeepSpeed 
+
+Training advanced deep learning models is challenging. 
+Beyond model design, model scientists also need to set up the state-of-the-art training techniques such as 
+distributed training, mixed precision, gradient accumulation, and checkpointing. 
+Yet still, scientists may not achieve the desired system performance and convergence rate.
+Large model sizes are even more challenging: a large model easily runs out of memory with pure data parallelism and it is difficult to use model parallelism. 
+[DeepSpeed](https://www.deepspeed.ai/) addresses these challenges to accelerate model development and training.
+
+🤗 Accelerate integrates [DeepSpeed](https://github.com/microsoft/DeepSpeed) via 2 options:
+
+1. Integration of the DeepSpeed features via `deepspeed config file` specification in `accelerate config` . You just supply your custom config file or use our template. Most of
+   this document is focused on this feature. This supports all the core features of DeepSpeed and gives user a lot of flexibility. 
+   User may have to change few lines of code depending on the config.
+2. Integration via `deepspeed_plugin`.This supports subset of the DeepSpeed features and uses default options for the rest of the configurations. 
+   User need not change any code and is good for those who are fine with most of the default settings of DeepSpeed.
+
+## What is integrated?
+
+Training:
+
+1. DeepSpeed ZeRO training supports the full ZeRO stages 1, 2 and 3 with ZeRO-Infinity (CPU and NVME offload).
+
+Inference:
+
+1. DeepSpeed ZeRO Inference supports ZeRO stage 3 with ZeRO-Infinity. It uses the same ZeRO protocol as training, but
+   it doesn't use an optimizer and a lr scheduler and only stage 3 is relevant. For more details see:
+   [deepspeed-zero-inference](#deepspeed-zero-inference).
+
+
+## How it works?
+
+**Pre-Requisites**: Install DeepSpeed version >=0.6.5. Please refer to the [DeepSpeed Insallation details](https://github.com/microsoft/DeepSpeed#installation)
+for more information.
+
+We will first look at easy to use integration via `deepspeed_plugin`. 
+Followed by more flexible and feature rich `deepspeed config file` integration. 
+
+### Accelerate DeepSpeed Plugin
+On your machine(s) just run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked. It will ask whether you want to use a config file for deepspeed to which you answer no and answer the following questions. 
+This will generate a config file that will be used automatically to properly set the
+default options when doing
+
+```bash
+accelerate launch my_script.py --args_to_my_script
+```
+
+For instance, here is how you would run the NLP example `examples/by_feature/deepspeed_with_config_support.py` (from the root of the repo) with DeepSpeed Plugin:
+
+**ZeRO Stage-2 DeepSpeed Plugin Example**
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+ gradient_accumulation_steps: 1
+ gradient_clipping: 1.0
+ offload_optimizer_device: none
+ offload_param_device: none
+ zero3_init_flag: true
+ zero_stage: 2
+distributed_type: DEEPSPEED
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 2
+use_cpu: false
+```
+
+```bash
+accelerate launch examples/by_feature/deepspeed_with_config_support.py \
Done.
@@ -0,0 +1,542 @@
+<!--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.
+-->
+
+# DeepSpeed 
+
+Training advanced deep learning models is challenging. 
+Beyond model design, model scientists also need to set up the state-of-the-art training techniques such as 
+distributed training, mixed precision, gradient accumulation, and checkpointing. 
+Yet still, scientists may not achieve the desired system performance and convergence rate.
+Large model sizes are even more challenging: a large model easily runs out of memory with pure data parallelism and it is difficult to use model parallelism. 
+[DeepSpeed](https://www.deepspeed.ai/) addresses these challenges to accelerate model development and training.
+
+🤗 Accelerate integrates [DeepSpeed](https://github.com/microsoft/DeepSpeed) via 2 options:
+
+1. Integration of the DeepSpeed features via `deepspeed config file` specification in `accelerate config` . You just supply your custom config file or use our template. Most of
+   this document is focused on this feature. This supports all the core features of DeepSpeed and gives user a lot of flexibility. 
+   User may have to change few lines of code depending on the config.
+2. Integration via `deepspeed_plugin`.This supports subset of the DeepSpeed features and uses default options for the rest of the configurations. 
+   User need not change any code and is good for those who are fine with most of the default settings of DeepSpeed.
+
+## What is integrated?
+
+Training:
+
+1. DeepSpeed ZeRO training supports the full ZeRO stages 1, 2 and 3 with ZeRO-Infinity (CPU and NVME offload).
You should host the image in the `huggingface/documentation-images` dataset as we try to limit the size of the repo. Also those won't be copied when building the doc (the two images in this imgs folder are not used in the documentation, just the main README).
@@ -0,0 +1,508 @@
+<!--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.
+-->
+
+# DeepSpeed 
+
+[DeepSpeed](https://github.com/microsoft/DeepSpeed) implements everything described in the [ZeRO paper](https://arxiv.org/abs/1910.02054). Currently it provides full support for:
+
+1. Optimizer state partitioning (ZeRO stage 1)
+2. Gradient partitioning (ZeRO stage 2)
+3. Parameter partitioning (ZeRO stage 3)
+4. Custom mixed precision training handling
+5. A range of fast CUDA-extension-based optimizers
+6. ZeRO-Offload to CPU and Disk/NVMe
+
+ZeRO-Offload has its own dedicated paper: [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://arxiv.org/abs/2101.06840). And NVMe-support is described in the paper [ZeRO-Infinity: Breaking the GPU
+Memory Wall for Extreme Scale Deep Learning](https://arxiv.org/abs/2104.07857).
+
+DeepSpeed ZeRO-2 is primarily used only for training, as its features are of no use to inference.
+
+DeepSpeed ZeRO-3 can be used for inference as well, since it allows huge models to be loaded on multiple GPUs, which
+won't be possible on a single GPU.
+
+🤗 Accelerate integrates [DeepSpeed](https://github.com/microsoft/DeepSpeed) via 2 options:
+
+1. Integration of the DeepSpeed features via `deepspeed config file` specification in `accelerate config` . You just supply your custom config file or use our template. Most of
+   this document is focused on this feature. This supports all the core features of DeepSpeed and gives user a lot of flexibility. 
+   User may have to change few lines of code depending on the config.
+2. Integration via `deepspeed_plugin`.This supports subset of the DeepSpeed features and uses default options for the rest of the configurations. 
+   User need not change any code and is good for those who are fine with most of the default settings of DeepSpeed.
+
+## What is integrated?
+
+Training:
+
+1. DeepSpeed ZeRO training supports the full ZeRO stages 1, 2 and 3 as well as CPU/Disk offload of optimizer states, gradients and parameters. 
+Below is a short description of Data Parallelism using ZeRO - Zero Redundancy Optimizer along with diagram from this [blog post](https://www.microsoft.com/en-us/research/blog/zero-deepspeed-new-system-optimizations-enable-training-models-with-over-100-billion-parameters/)
+![ZeRO Data Parallelism](./imgs/parallelism-zero.png)
Done.
@@ -0,0 +1,508 @@
+<!--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.
+-->
+
+# DeepSpeed 
+
+[DeepSpeed](https://github.com/microsoft/DeepSpeed) implements everything described in the [ZeRO paper](https://arxiv.org/abs/1910.02054). Currently it provides full support for:
+
+1. Optimizer state partitioning (ZeRO stage 1)
+2. Gradient partitioning (ZeRO stage 2)
+3. Parameter partitioning (ZeRO stage 3)
+4. Custom mixed precision training handling
+5. A range of fast CUDA-extension-based optimizers
+6. ZeRO-Offload to CPU and Disk/NVMe
+
+ZeRO-Offload has its own dedicated paper: [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://arxiv.org/abs/2101.06840). And NVMe-support is described in the paper [ZeRO-Infinity: Breaking the GPU
+Memory Wall for Extreme Scale Deep Learning](https://arxiv.org/abs/2104.07857).
+
+DeepSpeed ZeRO-2 is primarily used only for training, as its features are of no use to inference.
+
+DeepSpeed ZeRO-3 can be used for inference as well, since it allows huge models to be loaded on multiple GPUs, which
+won't be possible on a single GPU.
+
+🤗 Accelerate integrates [DeepSpeed](https://github.com/microsoft/DeepSpeed) via 2 options:
+
+1. Integration of the DeepSpeed features via `deepspeed config file` specification in `accelerate config` . You just supply your custom config file or use our template. Most of
+   this document is focused on this feature. This supports all the core features of DeepSpeed and gives user a lot of flexibility. 
+   User may have to change few lines of code depending on the config.
+2. Integration via `deepspeed_plugin`.This supports subset of the DeepSpeed features and uses default options for the rest of the configurations. 
+   User need not change any code and is good for those who are fine with most of the default settings of DeepSpeed.
+
+## What is integrated?
+
+Training:
+
+1. DeepSpeed ZeRO training supports the full ZeRO stages 1, 2 and 3 as well as CPU/Disk offload of optimizer states, gradients and parameters. 
+Below is a short description of Data Parallelism using ZeRO - Zero Redundancy Optimizer along with diagram from this [blog post](https://www.microsoft.com/en-us/research/blog/zero-deepspeed-new-system-optimizations-enable-training-models-with-over-100-billion-parameters/)
+![ZeRO Data Parallelism](./imgs/parallelism-zero.png)
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/405
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 5020c6cc4..4613e1e8d 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -14,7 +14,7 @@ jobs:
     - name: Install Python dependencies
       run: pip install setuptools==59.5.0; pip install -e .[test,test_trackers]
     - name: Run Tests
-      run: make test
+      run: make test_cpu
       
   test_examples:
     runs-on: ubuntu-latest
@@ -27,4 +27,4 @@ jobs:
     - name: Install Python dependencies
       run: pip install setuptools==59.5.0; pip install -e .[test] tensorboard
     - name: Run Tests
-      run: make test_examples
\ No newline at end of file
+      run: make test_examples
diff --git a/Makefile b/Makefile
index db7893565..d2822153b 100644
--- a/Makefile
+++ b/Makefile
@@ -24,8 +24,12 @@ style:
 	python utils/style_doc.py src/accelerate docs/source --max_len 119
 	
 # Run tests for the library
-test:
-	python -m pytest -n auto --dist=loadfile -s -v ./tests/ --ignore=./tests/test_examples.py
+test_cpu:
+	python -m pytest -s -v ./tests/ --ignore=./tests/test_examples.py
+
+test_cuda:
+	python -m pytest -s -v ./tests/ --ignore=./tests/test_examples.py --ignore=./tests/test_scheduler.py --ignore=./tests/test_cpu.py
+	python -m pytest -s -v ./tests/test_cpu.py ./tests/test_scheduler.py
 
 test_examples:
-	python -m pytest -n auto --dist=loadfile -s -v ./tests/test_examples.py
+	python -m pytest -s -v ./tests/test_examples.py
diff --git a/tests/test_cpu.py b/tests/test_cpu.py
index 72a8b2e7d..63ee69a97 100644
--- a/tests/test_cpu.py
+++ b/tests/test_cpu.py
@@ -18,6 +18,6 @@
 from accelerate.test_utils import test_script
 
 
-class MultiTPUTester(unittest.TestCase):
+class MultiCPUTester(unittest.TestCase):
     def test_cpu(self):
         debug_launcher(test_script.main)

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/401
diff --git a/examples/by_feature/fsdp_with_peak_mem_tracking.py b/examples/by_feature/fsdp_with_peak_mem_tracking.py
index 8a289442d..8304bcefe 100644
--- a/examples/by_feature/fsdp_with_peak_mem_tracking.py
+++ b/examples/by_feature/fsdp_with_peak_mem_tracking.py
@@ -301,7 +301,7 @@ def collate_fn(examples):
                     {
                         "accuracy": eval_metric["accuracy"],
                         "f1": eval_metric["f1"],
-                        "train_loss": total_loss.item(),
+                        "train_loss": total_loss.item() / len(train_dataloader),
                     },
                     step=epoch,
                 )
diff --git a/examples/by_feature/tracking.py b/examples/by_feature/tracking.py
index 1ccc2c1b0..9538b21b1 100644
--- a/examples/by_feature/tracking.py
+++ b/examples/by_feature/tracking.py
@@ -222,7 +222,7 @@ def training_function(config, args):
                 {
                     "accuracy": eval_metric["accuracy"],
                     "f1": eval_metric["f1"],
-                    "train_loss": total_loss.item(),
+                    "train_loss": total_loss.item() / len(train_dataloader),
                     "epoch": epoch,
                 },
                 step=epoch,
diff --git a/examples/complete_cv_example.py b/examples/complete_cv_example.py
index fd830442d..bd2ed3b29 100644
--- a/examples/complete_cv_example.py
+++ b/examples/complete_cv_example.py
@@ -258,7 +258,12 @@ def training_function(config, args):
         accelerator.print(f"epoch {epoch}: {100 * eval_metric:.2f}")
         if args.with_tracking:
             accelerator.log(
-                {"accuracy": 100 * eval_metric, "total_loss": total_loss, "epoch": epoch}, step=overall_step
+                {
+                    "accuracy": 100 * eval_metric,
+                    "train_loss": total_loss.item() / len(train_dataloader),
+                    "epoch": epoch,
+                },
+                step=overall_step,
             )
         if checkpointing_steps == "epoch":
             output_dir = f"epoch_{epoch}"
diff --git a/examples/complete_nlp_example.py b/examples/complete_nlp_example.py
index abf85c8d9..9c8b89706 100644
--- a/examples/complete_nlp_example.py
+++ b/examples/complete_nlp_example.py
@@ -246,7 +246,7 @@ def collate_fn(examples):
                 {
                     "accuracy": eval_metric["accuracy"],
                     "f1": eval_metric["f1"],
-                    "train_loss": total_loss.item(),
+                    "train_loss": total_loss.item() / len(train_dataloader),
                     "epoch": epoch,
                 },
                 step=epoch,
diff --git a/tests/test_examples.py b/tests/test_examples.py
index af892f4eb..6ec2a143a 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -100,13 +100,13 @@ def test_cv_examples(self):
         cv_path = os.path.abspath(os.path.join("examples", "cv_example.py"))
         special_strings = [
             " " * 16 + "{\n\n",
-            " " * 18 + '"accuracy": eval_metric["accuracy"],\n\n',
-            " " * 18 + '"f1": eval_metric["f1"],\n\n',
-            " " * 18 + '"train_loss": total_loss.item(),\n\n',
-            " " * 18 + '"epoch": epoch,\n\n',
+            " " * 20 + '"accuracy": eval_metric["accuracy"],\n\n',
+            " " * 20 + '"f1": eval_metric["f1"],\n\n',
+            " " * 20 + '"train_loss": total_loss.item() / len(train_dataloader),\n\n',
+            " " * 20 + '"epoch": epoch,\n\n',
             " " * 16 + "},\n\n",
             " " * 16 + "step=epoch,\n",
-            " " * 8,
+            " " * 12,
         ]
         self.one_complete_example("complete_cv_example.py", True, cv_path, special_strings)
         self.one_complete_example("complete_cv_example.py", False, cv_path, special_strings)

_The documentation is not available anymore as the PR was closed or merged._
@pacman100 @sgugger In a multi-gpu environment, does the **total_loss** here represent the loss of all training data in the entire epoch or just the loss of training data on the main process? 

I tried to print the result of total_loss on each process, and it seems that accelerator.log() only records the total_loss on the main process. So, please, do I need to use accelerator.gather() to get the training loss on the whole epoch? I hope I can get your confirmation, many thanks!
Hello @benjpau , yes, you are correct, you would need to `gather` to get the total epoch loss. Will raise a PR to fix the examples to reflect the same.
https://api.github.com/repos/huggingface/accelerate/pulls/399
diff --git a/src/accelerate/utils/offload.py b/src/accelerate/utils/offload.py
index 660bc521a..a562e9283 100644
--- a/src/accelerate/utils/offload.py
+++ b/src/accelerate/utils/offload.py
@@ -21,44 +21,13 @@
 import torch
 
 
-def offload_state_dict(save_dir: Union[str, os.PathLike], state_dict: Dict[str, torch.Tensor]):
-    """
-    Offload a state dict in a given folder.
-
-    Args:
-        save_dir (`str` or `os.PathLike`): The directory in which to offload the state dict.
-        state_dict (`Dict[str, torch.Tensor]`): The dictionary of tensors to offload.
-    """
-    os.makedirs(save_dir, exist_ok=True)
-    index = {}
-    for name, parameter in state_dict.items():
-        tensor_file = os.path.join(save_dir, f"{name}.dat")
-        array = parameter.numpy()
-        index[name] = {"dtype": str(array.dtype), "shape": list(array.shape)}
-        if array.ndim == 0:
-            array = array[None]
-        file_array = np.memmap(tensor_file, dtype=array.dtype, mode="w+", shape=array.shape)
-        file_array[:] = array[:]
-        file_array.flush()
-
-    # Update index
-    index_file = os.path.join(save_dir, "index.json")
-    if os.path.isfile(index_file):
-        with open(index_file, "r", encoding="utf-8") as f:
-            current_index = json.load(f)
-    else:
-        current_index = {}
-    current_index.update(index)
-
-    with open(index_file, "w", encoding="utf-8") as f:
-        json.dump(current_index, f, indent=2)
-
-
 def offload_weight(weight, weight_name, offload_folder, index=None):
     array = weight.numpy()
     tensor_file = os.path.join(offload_folder, f"{weight_name}.dat")
     if index is not None:
         index[weight_name] = {"dtype": str(array.dtype), "shape": list(array.shape)}
+    if array.ndim == 0:
+        array = array[None]
     file_array = np.memmap(tensor_file, dtype=array.dtype, mode="w+", shape=array.shape)
     file_array[:] = array[:]
     file_array.flush()
@@ -82,6 +51,23 @@ def save_offload_index(index, offload_folder):
         json.dump(current_index, f, indent=2)
 
 
+def offload_state_dict(save_dir: Union[str, os.PathLike], state_dict: Dict[str, torch.Tensor]):
+    """
+    Offload a state dict in a given folder.
+
+    Args:
+        save_dir (`str` or `os.PathLike`): The directory in which to offload the state dict.
+        state_dict (`Dict[str, torch.Tensor]`): The dictionary of tensors to offload.
+    """
+    os.makedirs(save_dir, exist_ok=True)
+    index = {}
+    for name, parameter in state_dict.items():
+        index = offload_weight(parameter, name, save_dir, index=index)
+
+    # Update index
+    save_offload_index(index, save_dir)
+
+
 class PrefixedDataset(Mapping):
     """
     Will access keys in a given dataset by adding a prefix.

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/398
diff --git a/README.md b/README.md
index a3fc74cef..44b25778e 100644
--- a/README.md
+++ b/README.md
@@ -212,6 +212,7 @@ If you like the simplicity of 🤗 Accelerate but would prefer a higher-level ab
 
 * [Animus](https://github.com/Scitator/animus) is a minimalistic framework to run machine learning experiments. Animus highlights common "breakpoints" in ML experiments and provides a unified interface for them within [IExperiment](https://github.com/Scitator/animus/blob/main/animus/core.py#L76).
 * [Catalyst](https://github.com/catalyst-team/catalyst#getting-started) is a PyTorch framework for Deep Learning Research and Development. It focuses on reproducibility, rapid experimentation, and codebase reuse so you can create something new rather than write yet another train loop. Catalyst provides a [Runner](https://catalyst-team.github.io/catalyst/api/core.html#runner) to connect all parts of the experiment: hardware backend, data transformations, model train, and inference logic.
+* [fastai](https://github.com/fastai/fastai#installing) is a PyTorch framework for Deep Learning that simplifies training fast and accurate neural nets using modern best practices. fastai provides a [Learner](https://docs.fast.ai/learner.html#Learner) to handle the training, fine-tuning, and inference of deep learning algorithms.
 * [Kornia](https://kornia.readthedocs.io/en/latest/get-started/introduction.html) is a differentiable library that allows classical computer vision to be integrated into deep learning models. Kornia provides a [Trainer](https://kornia.readthedocs.io/en/latest/x.html#kornia.x.Trainer) with the specific purpose to train and fine-tune the supported deep learning algorithms within the library.
 * [pytorch-accelerated](https://github.com/Chris-hughes10/pytorch-accelerated) is a lightweight training library, with a streamlined feature set centred around a general-purpose [Trainer](https://pytorch-accelerated.readthedocs.io/en/latest/trainer.html), that places a huge emphasis on simplicity and transparency; enabling users to understand exactly what is going on under the hood, but without having to write and maintain the boilerplate themselves!
 

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/396
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index cb904a9b5..6dc1f847c 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -21,8 +21,6 @@
 
 import torch
 
-from packaging import version
-
 from .checkpointing import load_accelerator_state, load_custom_state, save_accelerator_state, save_custom_state
 from .data_loader import prepare_data_loader
 from .logging import get_logger
@@ -46,6 +44,7 @@
     gather,
     get_pretty_name,
     is_deepspeed_available,
+    is_torch_version,
     pad_across_processes,
     reduce,
     save,
@@ -173,7 +172,7 @@ def __init__(
             os.environ["USE_FSDP"] = "true"  # use FSDP if plugin is provided
 
         if os.environ.get("USE_FSDP", "false") == "true":
-            if version.parse(torch.__version__) < version.parse("1.12.0.dev20220418+cu113"):
+            if is_torch_version("<", "1.12.0.dev20220418+cu113"):
                 raise ValueError("FSDP requires PyTorch >= 1.12.0.dev20220418+cu113")
 
         # Kwargs handlers
@@ -212,7 +211,7 @@ def __init__(
         self.device_placement = device_placement
         self.split_batches = split_batches
         self.dispatch_batches = dispatch_batches
-        if dispatch_batches is True and version.parse(torch.__version__) < version.parse("1.8.0"):
+        if dispatch_batches is True and is_torch_version("<", "1.8.0"):
             raise ImportError(
                 "Using `DataLoaderDispatcher` requires PyTorch 1.8.0 minimum. You have {torch.__version__}."
             )
@@ -222,15 +221,15 @@ def __init__(
         self.scaler = None
         self.native_amp = False
         if self.state.mixed_precision == "fp16":
-            self.native_amp = version.parse(torch.__version__) >= version.parse("1.6")
-            if version.parse(torch.__version__) < version.parse("1.6"):
+            self.native_amp = is_torch_version(">=", "1.6")
+            if not self.native_amp:
                 raise ValueError("fp16 mixed precision requires PyTorch >= 1.6")
 
             kwargs = self.scaler_handler.to_kwargs() if self.scaler_handler is not None else {}
             self.scaler = torch.cuda.amp.GradScaler(**kwargs)
         elif self.state.mixed_precision == "bf16":
-            self.native_amp = version.parse(torch.__version__) >= version.parse("1.10")
-            if mixed_precision == "bf16" and version.parse(torch.__version__) < version.parse("1.10"):
+            self.native_amp = is_torch_version(">=", "1.10")
+            if mixed_precision == "bf16" and not self.native_amp:
                 raise ValueError("bf16 mixed precision requires PyTorch >= 1.10")
 
             kwargs = self.scaler_handler.to_kwargs() if self.scaler_handler is not None else {}
@@ -245,7 +244,7 @@ def __init__(
         # RNG Types
         self.rng_types = rng_types
         if self.rng_types is None:
-            self.rng_types = ["torch"] if version.parse(torch.__version__) <= version.parse("1.5.1") else ["generator"]
+            self.rng_types = ["torch"] if is_torch_version("<=", "1.5.1") else ["generator"]
 
     @property
     def distributed_type(self):
@@ -486,7 +485,7 @@ def prepare_model(self, model):
             kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {}
             model = torch.nn.parallel.DistributedDataParallel(model, **kwargs)
         if self.native_amp:
-            if self.mixed_precision == "fp16" and version.parse(torch.__version__) >= version.parse("1.10"):
+            if self.mixed_precision == "fp16" and is_torch_version(">=", "1.10"):
                 model.forward = torch.cuda.amp.autocast(dtype=torch.float16)(model.forward)
             elif self.mixed_precision == "bf16":
                 model.forward = torch.cuda.amp.autocast(dtype=torch.bfloat16)(model.forward)
@@ -924,7 +923,7 @@ def autocast(self):
         different will happen otherwise.
         """
         if self.native_amp:
-            if self.mixed_precision == "fp16" and version.parse(torch.__version__) >= version.parse("1.10"):
+            if self.mixed_precision == "fp16" and is_torch_version(">=", "1.10"):
                 autocast_context = torch.cuda.amp.autocast(dtype=torch.float16)
             elif self.mixed_precision == "bf16":
                 autocast_context = torch.cuda.amp.autocast(dtype=torch.bfloat16)
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index ff4522162..22c8e621c 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -33,8 +33,7 @@
     PrepareForLaunch,
     is_sagemaker_available,
 )
-from accelerate.utils.versions import torch_version
-from packaging import version
+from accelerate.utils.versions import is_torch_version
 
 
 def launch_command_parser(subparsers=None):
@@ -219,9 +218,9 @@ def simple_launcher(args):
 
 
 def multi_gpu_launcher(args):
-    if torch_version >= version.parse("1.10.0"):
+    if is_torch_version(">=", "1.10.0"):
         cmd = ["torchrun"]
-    elif torch_version >= version.parse("1.9.0"):
+    elif is_torch_version(">=", "1.9.0"):
         cmd = [sys.executable, "-m", "torch.distributed.run"]
     else:
         cmd = [sys.executable, "-m", "torch.distributed.launch", "--use_env"]
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index 50aeb57b6..79ba6b3f3 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -18,8 +18,6 @@
 import torch
 from torch.utils.data import BatchSampler, DataLoader, IterableDataset
 
-from packaging import version
-
 from .state import AcceleratorState, DistributedType, is_tpu_available
 from .utils import (
     RNGType,
@@ -29,6 +27,7 @@
     find_batch_size,
     get_data_structure,
     initialize_tensors,
+    is_torch_version,
     send_to_device,
     slice_tensors,
     synchronize_rng_states,
@@ -61,7 +60,7 @@
 }
 
 for v, additional_kwargs in _PYTORCH_DATALOADER_ADDITIONAL_KWARGS.items():
-    if version.parse(torch.__version__) >= version.parse(v):
+    if is_torch_version(">=", v):
         _PYTORCH_DATALOADER_KWARGS.update(additional_kwargs)
 
 
@@ -327,7 +326,7 @@ class DataLoaderDispatcher(DataLoader):
 
     def __init__(self, dataset, split_batches: bool = False, **kwargs):
         shuffle = False
-        if version.parse(torch.__version__) >= version.parse("1.11.0"):
+        if is_torch_version(">=", "1.11.0"):
             from torch.utils.data.datapipes.iter.combinatorics import ShufflerIterDataPipe
 
             # We need to save the shuffling state of the DataPipe
@@ -335,7 +334,7 @@ def __init__(self, dataset, split_batches: bool = False, **kwargs):
                 shuffle = dataset._shuffle_enabled
         super().__init__(dataset, **kwargs)
         self.split_batches = split_batches
-        if version.parse(torch.__version__) < version.parse("1.8.0"):
+        if is_torch_version("<", "1.8.0"):
             raise ImportError(
                 "Using `DataLoaderDispatcher` requires PyTorch 1.8.0 minimum. You have {torch.__version__}."
             )
@@ -487,7 +486,7 @@ def prepare_data_loader(
 
     </Tip>"""
     if dispatch_batches is None:
-        if version.parse(torch.__version__) < version.parse("1.8.0") or not put_on_device:
+        if is_torch_version("<", "1.8.0") or not put_on_device:
             dispatch_batches = False
         else:
             dispatch_batches = isinstance(dataloader.dataset, IterableDataset)
diff --git a/src/accelerate/launchers.py b/src/accelerate/launchers.py
index d3cd2e1a2..298db439b 100644
--- a/src/accelerate/launchers.py
+++ b/src/accelerate/launchers.py
@@ -19,10 +19,8 @@
 
 import torch
 
-from packaging import version
-
 from .state import AcceleratorState
-from .utils import PrecisionType, PrepareForLaunch, patch_environment
+from .utils import PrecisionType, PrepareForLaunch, is_torch_version, patch_environment
 
 
 def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, mixed_precision="no", use_port="29500"):
@@ -85,7 +83,7 @@ def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, mix
 
         if num_processes > 1:
             # Multi-GPU launch
-            if version.parse(torch.__version__) < version.parse("1.5.0"):
+            if is_torch_version("<", "1.5.0"):
                 raise ImportError(
                     "Using `notebook_launcher` for distributed training on GPUs require torch >= 1.5.0, got "
                     f"{torch.__version__}."
@@ -156,7 +154,7 @@ def debug_launcher(function, args=(), num_processes=2):
         num_processes (`int`, *optional*, defaults to 2):
             The number of processes to use for training.
     """
-    if version.parse(torch.__version__) < version.parse("1.5.0"):
+    if is_torch_version("<", "1.5.0"):
         raise ImportError(
             "Using `debug_launcher` for distributed training on GPUs require torch >= 1.5.0, got "
             f"{torch.__version__}."
diff --git a/src/accelerate/optimizer.py b/src/accelerate/optimizer.py
index dcf792592..3d5977968 100644
--- a/src/accelerate/optimizer.py
+++ b/src/accelerate/optimizer.py
@@ -17,10 +17,8 @@
 
 import torch
 
-from packaging import version
-
 from .state import AcceleratorState
-from .utils import DistributedType, honor_type, is_tpu_available
+from .utils import DistributedType, honor_type, is_torch_version, is_tpu_available
 
 
 if is_tpu_available():
@@ -103,7 +101,7 @@ def state_dict(self):
         return self.optimizer.state_dict()
 
     def zero_grad(self, set_to_none=None):
-        if version.parse(torch.__version__) < version.parse("1.7.0"):
+        if is_torch_version("<", "1.7.0"):
             if set_to_none is not None:
                 raise ValueError(
                     "`set_to_none` for Optimizer.zero_grad` was introduced in PyTorch 1.7.0 and can't be used for "
diff --git a/src/accelerate/test_utils/test_script.py b/src/accelerate/test_utils/test_script.py
index 48e7bcc49..a23179bd7 100644
--- a/src/accelerate/test_utils/test_script.py
+++ b/src/accelerate/test_utils/test_script.py
@@ -21,8 +21,7 @@
 from accelerate.data_loader import prepare_data_loader
 from accelerate.state import AcceleratorState
 from accelerate.test_utils import RegressionDataset, RegressionModel, are_the_same_tensors
-from accelerate.utils import DistributedType, gather, set_seed, synchronize_rng_states
-from packaging import version
+from accelerate.utils import DistributedType, gather, is_torch_version, set_seed, synchronize_rng_states
 
 
 def init_state_check():
@@ -40,7 +39,7 @@ def rng_sync_check():
     if state.distributed_type == DistributedType.MULTI_GPU:
         synchronize_rng_states(["cuda"])
         assert are_the_same_tensors(torch.cuda.get_rng_state()), "RNG states improperly synchronized on GPU."
-    if version.parse(torch.__version__) >= version.parse("1.6.0"):
+    if is_torch_version(">=", "1.6.0"):
         generator = torch.Generator()
         synchronize_rng_states(["generator"], generator=generator)
         assert are_the_same_tensors(generator.get_state()), "RNG states improperly synchronized in generator."
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
index 19d8b81d4..186f27bb7 100644
--- a/src/accelerate/utils/__init__.py
+++ b/src/accelerate/utils/__init__.py
@@ -73,6 +73,7 @@
     send_to_device,
     slice_tensors,
 )
+from .versions import compare_versions, is_torch_version
 
 
 if is_deepspeed_available():
diff --git a/src/accelerate/utils/constants.py b/src/accelerate/utils/constants.py
index c2c36a486..af8a6bf35 100644
--- a/src/accelerate/utils/constants.py
+++ b/src/accelerate/utils/constants.py
@@ -12,8 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import operator as op
+
+
 SCALER_NAME = "scaler.pt"
 MODEL_NAME = "pytorch_model"
 RNG_STATE_NAME = "random_states"
 OPTIMIZER_NAME = "optimizer"
 SCHEDULER_NAME = "scheduler"
+
+STR_OPERATION_TO_FUNC = {">": op.gt, ">=": op.ge, "==": op.eq, "!=": op.ne, "<=": op.le, "<": op.lt}
diff --git a/src/accelerate/utils/operations.py b/src/accelerate/utils/operations.py
index 3fc9d969c..954bb5dfd 100644
--- a/src/accelerate/utils/operations.py
+++ b/src/accelerate/utils/operations.py
@@ -23,11 +23,10 @@
 import torch
 from torch.distributed import ReduceOp
 
-from packaging import version
-
 from ..state import AcceleratorState
 from .dataclasses import DistributedType, TensorInformation
 from .imports import is_tpu_available
+from .versions import is_torch_version
 
 
 if is_tpu_available():
@@ -459,8 +458,7 @@ def _convert_to_fp32(tensor):
 
     def _is_fp16_bf16_tensor(tensor):
         return hasattr(tensor, "dtype") and (
-            tensor.dtype == torch.float16
-            or (version.parse(torch.__version__) >= version.parse("1.10") and tensor.dtype == torch.bfloat16)
+            tensor.dtype == torch.float16 or (is_torch_version(">=", "1.10") and tensor.dtype == torch.bfloat16)
         )
 
     return recursively_apply(_convert_to_fp32, tensor, test_type=_is_fp16_bf16_tensor)
diff --git a/src/accelerate/utils/versions.py b/src/accelerate/utils/versions.py
index 36dd7ac48..38674d4b3 100644
--- a/src/accelerate/utils/versions.py
+++ b/src/accelerate/utils/versions.py
@@ -13,8 +13,11 @@
 # limitations under the License.
 
 import sys
+from typing import Union
 
-from packaging.version import parse
+from packaging.version import Version, parse
+
+from .constants import STR_OPERATION_TO_FUNC
 
 
 if sys.version_info < (3, 8):
@@ -23,3 +26,36 @@
     import importlib.metadata as importlib_metadata
 
 torch_version = parse(importlib_metadata.version("torch"))
+
+
+def compare_versions(library_or_version: Union[str, Version], operation: str, requirement_version: str):
+    """
+    Compares a library version to some requirement using a given operation.
+
+    Args:
+        library_or_version (`str` or `packaging.version.Version`):
+            A library name or a version to check.
+        operation (`str`):
+            A string representation of an operator, such as `">"` or `"<="`.
+        requirement_version (`str`):
+            The version to compare the library version against
+    """
+    if operation not in STR_OPERATION_TO_FUNC.keys():
+        raise ValueError(f"`operation` must be one of {list(STR_OPERATION_TO_FUNC.keys())}, received {operation}")
+    operation = STR_OPERATION_TO_FUNC[operation]
+    if isinstance(library_or_version, str):
+        library_or_version = parse(importlib_metadata.version(library_or_version))
+    return operation(library_or_version, parse(requirement_version))
+
+
+def is_torch_version(operation: str, version: str):
+    """
+    Compares the current PyTorch version to a given reference with an operation.
+
+    Args:
+        operation (`str`):
+            A string representation of an operator, such as `">"` or `"<="`
+        version (`str`):
+            A string version of PyTorch
+    """
+    return compare_versions(torch_version, operation, version)

Not super fan of the name, wdyt of `torch_version_is("<", "1.12.0.dev20220418+cu113")` or `is_torch_version("<", "1.12.0.dev20220418+cu113")`?
@@ -173,7 +172,7 @@ def __init__(
             os.environ["USE_FSDP"] = "true"  # use FSDP if plugin is provided
 
         if os.environ.get("USE_FSDP", "false") == "true":
-            if version.parse(torch.__version__) < version.parse("1.12.0.dev20220418+cu113"):
+            if check_torch_version("<", "1.12.0.dev20220418+cu113"):
Meant to be `compare_version` here?
@@ -23,3 +26,34 @@
     import importlib.metadata as importlib_metadata
 
 torch_version = parse(importlib_metadata.version("torch"))
+
+
+def compare_versions(library_or_version: Union[str, Version], operation: str, requirement_version: str):
+    """Compares `library_or_version` to `requirement_version` with an `operation`
+
+    Args:
+        library_or_version (`str`, `packaging.version.Version]`):
+            A library name or Version to check
+        operation (`str`):
+            A string representation of an operator, such as ">" or "<="
+        requirement_version (`str`):
+            The version to compare the library version against
+    """
+    if operation not in STR_OPERATION_TO_FUNC.keys():
+        raise ValueError(f"`operation` must be one of {list(STR_OPERATION_TO_FUNC.keys())}, received {operation}")
+    operation = STR_OPERATION_TO_FUNC[operation]
+    if isinstance(library_or_version, str):
+        library_or_version = parse(importlib_metadata.version(library_or_version))
+    return operation(library_or_version, parse(requirement_version))
+
+
+def check_torch_version(operation: str, version: str):
+    """Compares the current PyTorch version to `version` with an `operation`
+
+    Args:
+        operation (`str`):
+            A string representation of an operator, such as ">" or "<="
+        version (`str`):
+            A string version of PyTorch
+    """
+    return check_torch_version(torch_version, operation, version)
Yup, nice catch! 
@@ -23,3 +26,34 @@
     import importlib.metadata as importlib_metadata
 
 torch_version = parse(importlib_metadata.version("torch"))
+
+
+def compare_versions(library_or_version: Union[str, Version], operation: str, requirement_version: str):
+    """Compares `library_or_version` to `requirement_version` with an `operation`
+
+    Args:
+        library_or_version (`str`, `packaging.version.Version]`):
+            A library name or Version to check
+        operation (`str`):
+            A string representation of an operator, such as ">" or "<="
+        requirement_version (`str`):
+            The version to compare the library version against
+    """
+    if operation not in STR_OPERATION_TO_FUNC.keys():
+        raise ValueError(f"`operation` must be one of {list(STR_OPERATION_TO_FUNC.keys())}, received {operation}")
+    operation = STR_OPERATION_TO_FUNC[operation]
+    if isinstance(library_or_version, str):
+        library_or_version = parse(importlib_metadata.version(library_or_version))
+    return operation(library_or_version, parse(requirement_version))
+
+
+def check_torch_version(operation: str, version: str):
+    """Compares the current PyTorch version to `version` with an `operation`
+
+    Args:
+        operation (`str`):
+            A string representation of an operator, such as ">" or "<="
+        version (`str`):
+            A string version of PyTorch
+    """
+    return check_torch_version(torch_version, operation, version)
`is_torch_version` sounds better, will go with that. 
@@ -173,7 +172,7 @@ def __init__(
             os.environ["USE_FSDP"] = "true"  # use FSDP if plugin is provided
 
         if os.environ.get("USE_FSDP", "false") == "true":
-            if version.parse(torch.__version__) < version.parse("1.12.0.dev20220418+cu113"):
+            if check_torch_version("<", "1.12.0.dev20220418+cu113"):
```suggestion
        library_or_version (`str` or `packaging.version.Version`):
```
@@ -23,3 +26,34 @@
     import importlib.metadata as importlib_metadata
 
 torch_version = parse(importlib_metadata.version("torch"))
+
+
+def compare_versions(library_or_version: Union[str, Version], operation: str, requirement_version: str):
+    """Compares `library_or_version` to `requirement_version` with an `operation`
+
+    Args:
+        library_or_version (`str`, `packaging.version.Version]`):
```suggestion
            A library name or a version to check.
```
@@ -23,3 +26,34 @@
     import importlib.metadata as importlib_metadata
 
 torch_version = parse(importlib_metadata.version("torch"))
+
+
+def compare_versions(library_or_version: Union[str, Version], operation: str, requirement_version: str):
+    """Compares `library_or_version` to `requirement_version` with an `operation`
+
+    Args:
+        library_or_version (`str`, `packaging.version.Version]`):
+            A library name or Version to check
```suggestion
    """
    Compares a library version to some requirement using a given operation.
```
I prefer to have complete English sentence describe the function :-)
@@ -23,3 +26,34 @@
     import importlib.metadata as importlib_metadata
 
 torch_version = parse(importlib_metadata.version("torch"))
+
+
+def compare_versions(library_or_version: Union[str, Version], operation: str, requirement_version: str):
+    """Compares `library_or_version` to `requirement_version` with an `operation`
```suggestion
            A string representation of an operator, such as `">"` or `"<="`.
```
@@ -23,3 +26,34 @@
     import importlib.metadata as importlib_metadata
 
 torch_version = parse(importlib_metadata.version("torch"))
+
+
+def compare_versions(library_or_version: Union[str, Version], operation: str, requirement_version: str):
+    """Compares `library_or_version` to `requirement_version` with an `operation`
+
+    Args:
+        library_or_version (`str`, `packaging.version.Version]`):
+            A library name or Version to check
+        operation (`str`):
+            A string representation of an operator, such as ">" or "<="
```suggestion
    """
    Compares the current PyTorch version to a given reference with an operation.
```
@@ -23,3 +26,34 @@
     import importlib.metadata as importlib_metadata
 
 torch_version = parse(importlib_metadata.version("torch"))
+
+
+def compare_versions(library_or_version: Union[str, Version], operation: str, requirement_version: str):
+    """Compares `library_or_version` to `requirement_version` with an `operation`
+
+    Args:
+        library_or_version (`str`, `packaging.version.Version]`):
+            A library name or Version to check
+        operation (`str`):
+            A string representation of an operator, such as ">" or "<="
+        requirement_version (`str`):
+            The version to compare the library version against
+    """
+    if operation not in STR_OPERATION_TO_FUNC.keys():
+        raise ValueError(f"`operation` must be one of {list(STR_OPERATION_TO_FUNC.keys())}, received {operation}")
+    operation = STR_OPERATION_TO_FUNC[operation]
+    if isinstance(library_or_version, str):
+        library_or_version = parse(importlib_metadata.version(library_or_version))
+    return operation(library_or_version, parse(requirement_version))
+
+
+def is_torch_version(operation: str, version: str):
+    """Compares the current PyTorch version to `version` with an `operation`
```suggestion
            A string representation of an operator, such as `">"` or `"<="`
```
@@ -23,3 +26,34 @@
     import importlib.metadata as importlib_metadata
 
 torch_version = parse(importlib_metadata.version("torch"))
+
+
+def compare_versions(library_or_version: Union[str, Version], operation: str, requirement_version: str):
+    """Compares `library_or_version` to `requirement_version` with an `operation`
+
+    Args:
+        library_or_version (`str`, `packaging.version.Version]`):
+            A library name or Version to check
+        operation (`str`):
+            A string representation of an operator, such as ">" or "<="
+        requirement_version (`str`):
+            The version to compare the library version against
+    """
+    if operation not in STR_OPERATION_TO_FUNC.keys():
+        raise ValueError(f"`operation` must be one of {list(STR_OPERATION_TO_FUNC.keys())}, received {operation}")
+    operation = STR_OPERATION_TO_FUNC[operation]
+    if isinstance(library_or_version, str):
+        library_or_version = parse(importlib_metadata.version(library_or_version))
+    return operation(library_or_version, parse(requirement_version))
+
+
+def is_torch_version(operation: str, version: str):
+    """Compares the current PyTorch version to `version` with an `operation`
+
+    Args:
+        operation (`str`):
+            A string representation of an operator, such as ">" or "<="
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/395
diff --git a/setup.py b/setup.py
index 7d463cf88..dcfe3fd8e 100644
--- a/setup.py
+++ b/setup.py
@@ -55,7 +55,7 @@
         ]
     },
     python_requires=">=3.6.0",
-    install_requires=["torch>=1.4.0", "pyyaml", "numpy>=1.17"],
+    install_requires=["numpy>=1.17", "packaging>=20.0", "pyyaml", "torch>=1.4.0"],
     extras_require=extras,
     classifiers=[
         "Development Status :: 5 - Production/Stable",

If we're alphabetically sorting, let's really do it ;-)
@@ -55,7 +55,7 @@
         ]
     },
     python_requires=">=3.6.0",
-    install_requires=["torch>=1.4.0", "pyyaml", "numpy>=1.17"],
+    install_requires=["numpy>=1.17", "pyyaml", "torch>=1.4.0", "packaging>=20.0"],
Fixed, heh
@@ -55,7 +55,7 @@
         ]
     },
     python_requires=">=3.6.0",
-    install_requires=["torch>=1.4.0", "pyyaml", "numpy>=1.17"],
+    install_requires=["numpy>=1.17", "pyyaml", "torch>=1.4.0", "packaging>=20.0"],
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/394
diff --git a/src/accelerate/hooks.py b/src/accelerate/hooks.py
index 0a5bca96d..dea10ba50 100644
--- a/src/accelerate/hooks.py
+++ b/src/accelerate/hooks.py
@@ -260,6 +260,24 @@ def detach_hook(self, module):
                     set_module_tensor_to_device(module, name, device, value=self.weights_map.get(name, None))
 
 
+def attach_execution_device_hook(module: torch.nn.Module, execution_device: Union[int, str, torch.device]):
+    """
+    Recursively attaches `AlignDevicesHook` to all submodules of a given model to make sure they have the right
+    execution device
+
+    Args:
+        module (`torch.nn.Module`):
+            The module where we want to attach the hooks.
+        execution_device (`int`, `str` or `torch.device`):
+            The device on which inputs and model weights should be placed before the forward pass.
+    """
+    if not hasattr(module, "_hf_hook") and len(module.state_dict()) > 0:
+        add_hook_to_module(module, AlignDevicesHook(execution_device))
+
+    for child in module.children():
+        attach_execution_device_hook(child, execution_device)
+
+
 def attach_align_device_hook(
     module: torch.nn.Module,
     execution_device: Optional[torch.device] = None,
@@ -383,6 +401,7 @@ def attach_align_device_hook_on_blocks(
             place_submodules=True,
         )
         add_hook_to_module(module, hook)
+        attach_execution_device_hook(module, execution_device[module_name])
     elif module_name in execution_device:
         attach_align_device_hook(
             module,
@@ -395,6 +414,7 @@ def attach_align_device_hook_on_blocks(
         if not hasattr(module, "_hf_hook"):
             hook = AlignDevicesHook(execution_device=execution_device[module_name], io_same_device=(module_name == ""))
             add_hook_to_module(module, hook)
+        attach_execution_device_hook(module, execution_device[module_name])
     elif module_name == "":
         hook = AlignDevicesHook(io_same_device=True)
         add_hook_to_module(module, hook)

_The documentation is not available anymore as the PR was closed or merged._
Not yet, will have a look before merging (though I don't expect much difference since there is not that many submodules).
https://api.github.com/repos/huggingface/accelerate/pulls/392
diff --git a/.github/workflows/build-docker-images.yml b/.github/workflows/build-docker-images.yml
new file mode 100644
index 000000000..6b7987e51
--- /dev/null
+++ b/.github/workflows/build-docker-images.yml
@@ -0,0 +1,53 @@
+name: Build Docker images (scheduled)
+
+on:
+  repository_dispatch:
+  schedule:
+    - cron: "0 1 * * *"
+
+concurrency:
+  group: docker-image-builds
+  cancel-in-progress: false
+
+jobs:
+  latest-cpu:
+    name: "Latest Accelerate CPU [dev]"
+    runs-on: ubuntu-latest
+    steps:
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v1
+      - name: Check out code
+        uses: actions/checkout@v2
+      - name: Login to DockerHub
+        uses: docker/login-action@v1
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_PASSWORD }}
+
+      - name: Build and Push CPU
+        uses: docker/build-push-action@v2
+        with:
+          context: ./docker/accelerate-cpu
+          push: true
+          tags: huggingface/accelerate-cpu
+
+  latest-cuda:
+    name: "Latest Accelerate GPU [dev]"
+    runs-on: ubuntu-latest
+    steps:
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v1
+      - name: Check out code
+        uses: actions/checkout@v2
+      - name: Login to DockerHub
+        uses: docker/login-action@v1
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_PASSWORD }}
+
+      - name: Build and Push GPU
+        uses: docker/build-push-action@v2
+        with:
+          context: ./docker/accelerate-gpu
+          push: true
+          tags: huggingface/accelerate-gpu
\ No newline at end of file
diff --git a/docker/accelerate-cpu/Dockerfile b/docker/accelerate-cpu/Dockerfile
index 06379e69b..c536fb70c 100644
--- a/docker/accelerate-cpu/Dockerfile
+++ b/docker/accelerate-cpu/Dockerfile
@@ -16,7 +16,6 @@ ENV VIRTUAL_ENV=/opt/venv
 RUN python3 -m venv ${VIRTUAL_ENV}
 # Make sure we use the virtualenv
 ENV PATH="${VIRTUAL_ENV}/bin:$PATH"
-
 WORKDIR /workspace
 # Install specific CPU torch wheel to save on space
 RUN python3 -m pip install --upgrade --no-cache-dir pip
@@ -28,6 +27,8 @@ RUN python3 -m pip install --no-cache-dir \
 # Stage 2
 FROM python:3.6-slim AS build-image
 COPY --from=compile-image /opt/venv /opt/venv
+RUN useradd -ms /bin/bash user
+USER user
 
 # Make sure we use the virtualenv
 ENV PATH="/opt/venv/bin:$PATH"
diff --git a/docker/accelerate-gpu/Dockerfile b/docker/accelerate-gpu/Dockerfile
index b371042e1..299f573d8 100644
--- a/docker/accelerate-gpu/Dockerfile
+++ b/docker/accelerate-gpu/Dockerfile
@@ -31,5 +31,8 @@ COPY --from=compile-image /opt/conda /opt/conda
 ENV PATH /opt/conda/bin:$PATH
 
 RUN echo "source activate accelerate" >> /.bashrc
+
+RUN useradd -ms /bin/bash user
+USER user
 # Activate the virtualenv
 CMD ["/bin/bash"]
\ No newline at end of file

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/391
diff --git a/.github/workflows/stale.yml b/.github/workflows/stale.yml
index 45181a092..c98ce0bb1 100644
--- a/.github/workflows/stale.yml
+++ b/.github/workflows/stale.yml
@@ -3,6 +3,7 @@ name: Stale Bot
 on:
   schedule:
     - cron: "0 15 * * *"
+  workflow_dispatch:
 
 jobs:
   close_stale_issues:
@@ -12,7 +13,7 @@ jobs:
     env:
       GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
     steps:
-    - uses: actions/@checkoutv2
+    - uses: actions/checkout@v2
     
     - name: Setup Python
       uses: actions/setup-python@v1

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/390
diff --git a/examples/cv_example.py b/examples/cv_example.py
index a80b42d13..483bddac3 100644
--- a/examples/cv_example.py
+++ b/examples/cv_example.py
@@ -73,7 +73,7 @@ def __getitem__(self, idx):
 
 def training_function(config, args):
     # Initialize accelerator
-    accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mix_precision)
+    accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
 
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
     lr = config["lr"]

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_388). All of your documentation changes will be reflected on that endpoint.
https://api.github.com/repos/huggingface/accelerate/pulls/388
diff --git a/.github/workflows/stale.yml b/.github/workflows/stale.yml
new file mode 100644
index 000000000..45181a092
--- /dev/null
+++ b/.github/workflows/stale.yml
@@ -0,0 +1,27 @@
+name: Stale Bot
+
+on:
+  schedule:
+    - cron: "0 15 * * *"
+
+jobs:
+  close_stale_issues:
+    name: Close Stale Issues
+    if: github.repository == 'huggingface/accelerate'
+    runs-on: ubuntu-latest
+    env:
+      GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+    steps:
+    - uses: actions/@checkoutv2
+    
+    - name: Setup Python
+      uses: actions/setup-python@v1
+      with:
+        python-version: 3.7
+    
+    - name: Install requirements
+      run: |
+        pip install PyGithub
+    - name: Close stale issues
+      run: |
+        python utils/stale.py
\ No newline at end of file
diff --git a/utils/stale.py b/utils/stale.py
new file mode 100644
index 000000000..1d8f9020b
--- /dev/null
+++ b/utils/stale.py
@@ -0,0 +1,66 @@
+# Copyright 2022 The HuggingFace Team, the AllenNLP library authors. 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.
+"""
+Script to close stale issue. Taken in part from the AllenNLP repository.
+https://github.com/allenai/allennlp.
+"""
+from datetime import datetime as dt
+import os
+
+from github import Github
+
+
+LABELS_TO_EXEMPT = [
+    "good first issue",
+    "feature request",
+    "wip",
+]
+
+
+def main():
+    g = Github(os.environ["GITHUB_TOKEN"])
+    repo = g.get_repo("huggingface/accelerate")
+    open_issues = repo.get_issues(state="open")
+
+    for issue in open_issues:
+        comments = sorted([comment for comment in issue.get_comments()], key=lambda i: i.created_at, reverse=True)
+        last_comment = comments[0] if len(comments) > 0 else None
+        current_time = dt.utcnow()
+        days_since_updated = (current_time - issue.updated_at).days
+        days_since_creation = (current_time - issue.created_at).days
+        if (
+            last_comment is not None and last_comment.user.login == "github-actions[bot]"
+            and days_since_updated > 7
+            and days_since_creation >= 30
+            and not any(label.name.lower() in LABELS_TO_EXEMPT for label in issue.get_labels())
+        ):
+            # Close issue since it has been 7 days of inactivity since bot mention.
+            issue.edit(state="closed")
+        elif (
+            days_since_updated > 23
+            and days_since_creation >= 30
+            and not any(label.name.lower() in LABELS_TO_EXEMPT for label in issue.get_labels())
+        ):
+            # Add stale comment
+            issue.create_comment(
+                "This issue has been automatically marked as stale because it has not had "
+                "recent activity. If you think this still needs to be addressed "
+                "please comment on this thread.\n\nPlease note that issues that do not follow the "
+                "[contributing guidelines](https://github.com/huggingface/accelerate/blob/main/CONTRIBUTING.md) "
+                "are likely to be ignored."
+            )
+
+
+if __name__ == "__main__":
+    main()

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/387
diff --git a/src/accelerate/scheduler.py b/src/accelerate/scheduler.py
index 8efce8dea..1fac5f012 100644
--- a/src/accelerate/scheduler.py
+++ b/src/accelerate/scheduler.py
@@ -61,7 +61,9 @@ def step(self, *args, **kwargs):
             # num_processes steps per training step
             num_processes = AcceleratorState().num_processes
             for _ in range(num_processes):
-                self.scheduler.step(*args, **kwargs)
+                # Special case when using OneCycle and `drop_last` was not used
+                if getattr(self.scheduler, "total_steps", 0) <= self.scheduler.last_epoch:
+                    self.scheduler.step(*args, **kwargs)
 
     # Passthroughs
     def get_last_lr(self):

Should we guard around the `last_epoch` attribute here as well somehow?
@@ -61,7 +61,9 @@ def step(self, *args, **kwargs):
             # num_processes steps per training step
             num_processes = AcceleratorState().num_processes
             for _ in range(num_processes):
-                self.scheduler.step(*args, **kwargs)
+                # Special case when using OneCycle and `drop_last` was not used
+                if getattr(self.scheduler, "total_steps", 0) <= self.scheduler.last_epoch:
`last_epoch` is part of the base scheduler object, so it's always present: https://github.com/pytorch/pytorch/blob/master/torch/optim/lr_scheduler.py#L24
@@ -61,7 +61,9 @@ def step(self, *args, **kwargs):
             # num_processes steps per training step
             num_processes = AcceleratorState().num_processes
             for _ in range(num_processes):
-                self.scheduler.step(*args, **kwargs)
+                # Special case when using OneCycle and `drop_last` was not used
+                if getattr(self.scheduler, "total_steps", 0) <= self.scheduler.last_epoch:
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/385
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
index f65b6b115..19d8b81d4 100644
--- a/src/accelerate/utils/__init__.py
+++ b/src/accelerate/utils/__init__.py
@@ -40,10 +40,18 @@
     get_max_memory,
     infer_auto_device_map,
     load_checkpoint_in_model,
+    load_offloaded_weights,
     named_module_tensors,
     set_module_tensor_to_device,
 )
-from .offload import OffloadedWeightsLoader, PrefixedDataset, extract_submodules_state_dict, offload_state_dict
+from .offload import (
+    OffloadedWeightsLoader,
+    PrefixedDataset,
+    extract_submodules_state_dict,
+    offload_state_dict,
+    offload_weight,
+    save_offload_index,
+)
 from .operations import (
     broadcast,
     broadcast_object_list,
diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
index 08c6f269d..0329f8786 100644
--- a/src/accelerate/utils/modeling.py
+++ b/src/accelerate/utils/modeling.py
@@ -25,6 +25,8 @@
 import torch
 import torch.nn as nn
 
+from .offload import offload_weight, save_offload_index
+
 
 WEIGHTS_INDEX_NAME = "pytorch_model.bin.index.json"
 
@@ -312,6 +314,18 @@ def clean_device_map(device_map: Dict[str, Union[int, str, torch.device]], modul
     return device_map
 
 
+def load_offloaded_weights(model, index, offload_folder):
+    if index is None or len(index) == 0:
+        # Nothing to do
+        return
+
+    for param_name, metadata in index.items():
+        tensor_file = os.path.join(offload_folder, f"{param_name}.dat")
+        shape = tuple(metadata["shape"])
+        weight = np.memmap(tensor_file, dtype=metadata["dtype"], mode="r", shape=shape)
+        set_module_tensor_to_device(model, param_name, "cpu", value=torch.tensor(weight))
+
+
 def infer_auto_device_map(
     model: nn.Module,
     max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None,
@@ -581,20 +595,10 @@ def load_checkpoint_in_model(
 
                 if param_device == "disk":
                     set_module_tensor_to_device(model, param_name, "meta")
-                    tensor_file = os.path.join(offload_folder, f"{param_name}.dat")
-                    array = param.numpy()
-                    offload_index[param_name] = {"dtype": str(array.dtype), "shape": list(array.shape)}
-                    file_array = np.memmap(tensor_file, dtype=array.dtype, mode="w+", shape=array.shape)
-                    file_array[:] = array[:]
-                    file_array.flush()
+                    offload_weight(param, param_name, offload_folder, index=offload_index)
                 elif param_device == "cpu" and offload_state_dict:
                     set_module_tensor_to_device(model, param_name, "meta")
-                    tensor_file = os.path.join(state_dict_folder, f"{param_name}.dat")
-                    array = param.numpy()
-                    state_dict_index[param_name] = {"dtype": str(array.dtype), "shape": list(array.shape)}
-                    file_array = np.memmap(tensor_file, dtype=array.dtype, mode="w+", shape=array.shape)
-                    file_array[:] = array[:]
-                    file_array.flush()
+                    offload_weight(param, param_name, state_dict_folder, index=state_dict_index)
                 else:
                     set_module_tensor_to_device(model, param_name, param_device, value=param)
 
@@ -602,23 +606,9 @@ def load_checkpoint_in_model(
         del checkpoint
         gc.collect()
 
-    if len(offload_index) > 0:
-        offload_index_file = os.path.join(offload_folder, "index.json")
-        if os.path.isfile(offload_index_file):
-            with open(offload_index_file, "r", encoding="utf-8") as f:
-                current_offload_index = json.load(f)
-        else:
-            current_offload_index = {}
-        current_offload_index.update(offload_index)
-
-        with open(offload_index_file, "w", encoding="utf-8") as f:
-            json.dump(current_offload_index, f, indent=2)
+    save_offload_index(offload_index, offload_folder)
 
     # Load back offloaded state dict on CPU
-    if offload_state_dict and len(state_dict_index) > 0:
-        for param_name, metadata in state_dict_index.items():
-            tensor_file = os.path.join(state_dict_folder, f"{param_name}.dat")
-            shape = tuple(metadata["shape"])
-            weight = np.memmap(tensor_file, dtype=metadata["dtype"], mode="r", shape=shape)
-            set_module_tensor_to_device(model, param_name, "cpu", value=torch.tensor(weight))
+    if offload_state_dict:
+        load_offloaded_weights(model, state_dict_index, state_dict_folder)
         shutil.rmtree(state_dict_folder)
diff --git a/src/accelerate/utils/offload.py b/src/accelerate/utils/offload.py
index 17f33d7c6..660bc521a 100644
--- a/src/accelerate/utils/offload.py
+++ b/src/accelerate/utils/offload.py
@@ -54,6 +54,34 @@ def offload_state_dict(save_dir: Union[str, os.PathLike], state_dict: Dict[str,
         json.dump(current_index, f, indent=2)
 
 
+def offload_weight(weight, weight_name, offload_folder, index=None):
+    array = weight.numpy()
+    tensor_file = os.path.join(offload_folder, f"{weight_name}.dat")
+    if index is not None:
+        index[weight_name] = {"dtype": str(array.dtype), "shape": list(array.shape)}
+    file_array = np.memmap(tensor_file, dtype=array.dtype, mode="w+", shape=array.shape)
+    file_array[:] = array[:]
+    file_array.flush()
+    return index
+
+
+def save_offload_index(index, offload_folder):
+    if index is None or len(index) == 0:
+        # Nothing to save
+        return
+
+    offload_index_file = os.path.join(offload_folder, "index.json")
+    if os.path.isfile(offload_index_file):
+        with open(offload_index_file, "r", encoding="utf-8") as f:
+            current_index = json.load(f)
+    else:
+        current_index = {}
+    current_index.update(index)
+
+    with open(offload_index_file, "w", encoding="utf-8") as f:
+        json.dump(current_index, f, indent=2)
+
+
 class PrefixedDataset(Mapping):
     """
     Will access keys in a given dataset by adding a prefix.

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/380
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index 3f10b83cf..b492f2f9c 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -12,8 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
 import importlib
+import sys
+
+
+# The package importlib_metadata is in a different place, depending on the Python version.
+if sys.version_info < (3, 8):
+    import importlib_metadata
+else:
+    import importlib.metadata as importlib_metadata
 
 
 try:
@@ -45,7 +52,15 @@ def is_tpu_available():
 
 
 def is_deepspeed_available():
-    return importlib.util.find_spec("deepspeed") is not None
+    package_exists = importlib.util.find_spec("deepspeed") is not None
+    # Check we're not importing a "deepspeed" directory somewhere but the actual library by trying to grab the version
+    # AND checking it has an author field in the metadata that is HuggingFace.
+    if package_exists:
+        try:
+            _ = importlib_metadata.metadata("deepspeed")
+            return True
+        except importlib_metadata.PackageNotFoundError:
+            return False
 
 
 def is_tensorflow_available():

```suggestion
```
If I'm not mistaken this is unnecessary. You can do `_deepspeed_metadata.get("version")` if it's needed?
@@ -45,7 +52,18 @@ def is_tpu_available():
 
 
 def is_deepspeed_available():
-    return importlib.util.find_spec("deepspeed") is not None
+    package_exists = importlib.util.find_spec("deepspeed") is not None
+    # Check we're not importing a "deepspeed" directory somewhere but the actual library by trying to grab the version
+    # AND checking it has an author field in the metadata that is HuggingFace.
+    if package_exists:
+        try:
+            _ = importlib_metadata.version("deepspeed")
`importlib_metadata.metadata` will also then raise the same error if it's not installed
@@ -45,7 +52,18 @@ def is_tpu_available():
 
 
 def is_deepspeed_available():
-    return importlib.util.find_spec("deepspeed") is not None
+    package_exists = importlib.util.find_spec("deepspeed") is not None
+    # Check we're not importing a "deepspeed" directory somewhere but the actual library by trying to grab the version
+    # AND checking it has an author field in the metadata that is HuggingFace.
+    if package_exists:
+        try:
+            _ = importlib_metadata.version("deepspeed")
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/379
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index f283dfbc5..ff4522162 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -324,7 +324,7 @@ def deepspeed_launcher(args):
     current_env["USE_DEEPSPEED"] = "true"
     current_env["DEEPSPEED_ZERO_STAGE"] = str(args.zero_stage)
     current_env["GRADIENT_ACCUMULATION_STEPS"] = str(args.gradient_accumulation_steps)
-    current_env["DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE"] = str(args.offload_optimizer_device)
+    current_env["DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE"] = str(args.offload_optimizer_device).lower()
 
     process = subprocess.Popen(cmd, env=current_env)
     process.wait()

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/378
diff --git a/docker/accelerate-cpu/Dockerfile b/docker/accelerate-cpu/Dockerfile
new file mode 100644
index 000000000..06379e69b
--- /dev/null
+++ b/docker/accelerate-cpu/Dockerfile
@@ -0,0 +1,34 @@
+# Builds CPU-only Docker image of PyTorch
+# Uses multi-staged approach to reduce size
+# Stage 1
+FROM python:3.6-slim as compile-image
+
+ARG DEBIAN_FRONTEND=noninteractive
+
+RUN apt update
+RUN apt-get install -y --no-install-recommends \
+    build-essential \
+    git \
+    gcc
+
+# Setup virtual environment for Docker
+ENV VIRTUAL_ENV=/opt/venv
+RUN python3 -m venv ${VIRTUAL_ENV}
+# Make sure we use the virtualenv
+ENV PATH="${VIRTUAL_ENV}/bin:$PATH"
+
+WORKDIR /workspace
+# Install specific CPU torch wheel to save on space
+RUN python3 -m pip install --upgrade --no-cache-dir pip
+RUN python3 -m pip install --no-cache-dir \
+    jupyter \
+    torch --extra-index-url https://download.pytorch.org/whl/cpu \
+    git+https://github.com/huggingface/accelerate#egg=accelerate[dev]
+    
+# Stage 2
+FROM python:3.6-slim AS build-image
+COPY --from=compile-image /opt/venv /opt/venv
+
+# Make sure we use the virtualenv
+ENV PATH="/opt/venv/bin:$PATH"
+CMD ["/bin/bash"]
\ No newline at end of file
diff --git a/docker/accelerate-gpu/Dockerfile b/docker/accelerate-gpu/Dockerfile
new file mode 100644
index 000000000..b371042e1
--- /dev/null
+++ b/docker/accelerate-gpu/Dockerfile
@@ -0,0 +1,35 @@
+# Builds GPU docker image of PyTorch
+# Uses multi-staged approach to reduce size
+# Stage 1
+# Use base conda image to reduce time
+FROM continuumio/miniconda3:latest AS compile-image
+# Specify py version
+ENV PYTHON_VERSION=3.6 
+# Install apt libs
+RUN apt-get update && \
+    apt-get install -y curl git wget && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists*
+
+# Create our conda env
+RUN conda create --name accelerate python=${PYTHON_VERSION} ipython jupyter pip
+# We don't install pytorch here yet since CUDA isn't available
+# instead we use the direct torch wheel
+ENV PATH /opt/conda/envs/accelerate/bin:$PATH
+# Activate our bash shell
+RUN chsh -s /bin/bash
+SHELL ["/bin/bash", "-c"]
+# Activate the conda env and install torch + accelerate
+RUN source activate accelerate && \
+    python3 -m pip install --no-cache-dir \
+    torch --extra-index-url https://download.pytorch.org/whl/cu113 \
+    git+https://github.com/huggingface/accelerate#egg=accelerate[dev]
+
+# Stage 2
+FROM nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04 AS build-image
+COPY --from=compile-image /opt/conda /opt/conda
+ENV PATH /opt/conda/bin:$PATH
+
+RUN echo "source activate accelerate" >> /.bashrc
+# Activate the virtualenv
+CMD ["/bin/bash"]
\ No newline at end of file

_The documentation is not available anymore as the PR was closed or merged._
I know nothing about Docker, so asking Lysandre for a review :-)
Great to have docker support! Helpful for production scenarios. Few comments:
1. The Python versions in CPU and GPU are different - 3.6 and 3.8, respectively. Would it be better to have same Python versions? Also, transformers will be discontinuing 3.6 support, so CPU version can be bumped up to match the GPU version?
2. Shouldn't the Dockerfile's be in the top-level directory to be able to copy accelerate codebase and install it? Maybe have `Dockerfile_cpu` and `Dockerfile_gpu` in top-level directory?
3. Getting below error while building GPU image :
```
W: GPG error: https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64  InRelease: The following signatures couldn't be verified because the public key is not available: NO_PUBKEY A4B469963BF863CC
E: The repository 'https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64  InRelease' is not signed.
```

> Great to have docker support! Helpful for production scenarios. Few comments:

These will also be used for having GPU and multi GPU tests, these mimic the transformers docker images for this exact purpose. (This is actually the true purpose behind this PR)

>     1. The Python versions in CPU and GPU are different - 3.6 and 3.8, respectively. Would it be better to have same Python versions? Also, transformers will be discontinuing 3.6 support, so CPU version can be bumped up to match the GPU version?

Yes they are, it's a partial limitation of the GPU image. I could go through the effort of having it use conda instead, but Sylvain and I discussed that 3.6 support will be dropped on the next release. Since it's so soon, it's easier to just have it be this way. 

>     2. Shouldn't the Dockerfile's be in the top-level directory to be able to copy accelerate codebase and install it? Maybe have `Dockerfile_cpu` and `Dockerfile_gpu` in top-level directory?

Nope, these assume the git repo is the top level directory and then copy it. So when building the dockerfiles it looks something like (assuming from a fresh accelerate clone and cd'd to it) `docker build . -f docker/accelerate-cpu/Dockerfile`. See the transformers repo for a similar behavior

>     3. Getting below error while building GPU image :
> 
> 
> ```
> W: GPG error: https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64  InRelease: The following signatures couldn't be verified because the public key is not available: NO_PUBKEY A4B469963BF863CC
> E: The repository 'https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64  InRelease' is not signed.
> ```

Unsure with this one. Do you have docker cuda properly setup? IIRC you might need to find the right keys to use as well. I was able to build it just fine yesterday.


Thinking on it more, having the python version be configurable in the cuda image would be nice. Will change this

Also would probably be better to do similar to this Dockerfile and clone the repo instead: https://github.com/huggingface/transformers/blob/main/docker/transformers-pytorch-gpu/Dockerfile#L11
https://api.github.com/repos/huggingface/accelerate/pulls/377
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index a827633c0..50aeb57b6 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -326,12 +326,21 @@ class DataLoaderDispatcher(DataLoader):
     """
 
     def __init__(self, dataset, split_batches: bool = False, **kwargs):
+        shuffle = False
+        if version.parse(torch.__version__) >= version.parse("1.11.0"):
+            from torch.utils.data.datapipes.iter.combinatorics import ShufflerIterDataPipe
+
+            # We need to save the shuffling state of the DataPipe
+            if isinstance(dataset, ShufflerIterDataPipe):
+                shuffle = dataset._shuffle_enabled
         super().__init__(dataset, **kwargs)
         self.split_batches = split_batches
         if version.parse(torch.__version__) < version.parse("1.8.0"):
             raise ImportError(
                 "Using `DataLoaderDispatcher` requires PyTorch 1.8.0 minimum. You have {torch.__version__}."
             )
+        if shuffle:
+            torch.utils.data.graph_settings.apply_shuffle_settings(dataset, shuffle=shuffle)
 
     def __iter__(self):
         state = AcceleratorState()

This will break compatibility with older versions of PyTorch as I believe datapipes were just introduced in PyTorch 1.11.
@@ -17,6 +17,7 @@
 
 import torch
 from torch.utils.data import BatchSampler, DataLoader, IterableDataset
+from torch.utils.data.datapipes.iter.combinatorics import ShufflerIterDataPipe
There should be a versioning test of PyTorch first, then this instance check. :-)
@@ -326,12 +327,15 @@ class DataLoaderDispatcher(DataLoader):
     """
 
     def __init__(self, dataset, split_batches: bool = False, **kwargs):
+        shuffle = dataset._shuffle_enabled if isinstance(dataset, ShufflerIterDataPipe) else False
```suggestion
        if version.parse(torch.__version__) >= version.parse("1.11.99"):
```
Just in case we get many micro releases ;-)
@@ -326,12 +326,21 @@ class DataLoaderDispatcher(DataLoader):
     """
 
     def __init__(self, dataset, split_batches: bool = False, **kwargs):
+        shuffle = False
+        if version.parse(torch.__version__) >= version.parse("1.11.0"):
Do you mean `"1.10.99"` ?
@@ -326,12 +326,21 @@ class DataLoaderDispatcher(DataLoader):
     """
 
     def __init__(self, dataset, split_batches: bool = False, **kwargs):
+        shuffle = False
+        if version.parse(torch.__version__) >= version.parse("1.11.0"):
Arf, review before coffee. I read > and thought for a minute it needed PyTorch 1.12, you can ignore :-)
@@ -326,12 +326,21 @@ class DataLoaderDispatcher(DataLoader):
     """
 
     def __init__(self, dataset, split_batches: bool = False, **kwargs):
+        shuffle = False
+        if version.parse(torch.__version__) >= version.parse("1.11.0"):
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/376
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 6fd24c321..cb904a9b5 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -881,14 +881,15 @@ def get_state_dict(self, model):
                 is_zero_3 = self.state.deepspeed_plugin.zero_stage == 3
 
         if is_zero_3:
-            state_dict = model._zero3_consolidated_fp16_state_dict()
+            state_dict = model._zero3_consolidated_16bit_state_dict()
         else:
             model = self.unwrap_model(model)
             state_dict = model.state_dict()
 
-        for k in state_dict:
-            if state_dict[k].dtype == torch.float16:
-                state_dict[k] = state_dict[k].float()
+        if state_dict is not None:
+            for k in state_dict:
+                if state_dict[k].dtype == torch.float16:
+                    state_dict[k] = state_dict[k].float()
 
         return state_dict
 

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/374
diff --git a/examples/by_feature/checkpointing.py b/examples/by_feature/checkpointing.py
index ee0716b6d..06edd8540 100644
--- a/examples/by_feature/checkpointing.py
+++ b/examples/by_feature/checkpointing.py
@@ -103,6 +103,13 @@ def collate_fn(examples):
     return train_dataloader, eval_dataloader
 
 
+# For testing only
+if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+    from accelerate.test_utils.training import mocked_dataloaders
+
+    get_dataloaders = mocked_dataloaders  # noqa: F811
+
+
 def training_function(config, args):
     # Initialize accelerator
     accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
diff --git a/examples/by_feature/fsdp_with_peak_mem_tracking.py b/examples/by_feature/fsdp_with_peak_mem_tracking.py
index 581c5fea0..8a289442d 100644
--- a/examples/by_feature/fsdp_with_peak_mem_tracking.py
+++ b/examples/by_feature/fsdp_with_peak_mem_tracking.py
@@ -74,6 +74,13 @@ def __exit__(self, *exc):
         # print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")
 
 
+# For testing only
+if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+    from accelerate.test_utils.training import mocked_dataloaders
+
+    get_dataloaders = mocked_dataloaders  # noqa: F811
+
+
 def training_function(config, args):
     # Initialize accelerator
     if args.with_tracking:
diff --git a/examples/by_feature/memory.py b/examples/by_feature/memory.py
index 6f3f3e0cb..140f14d75 100644
--- a/examples/by_feature/memory.py
+++ b/examples/by_feature/memory.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import argparse
+import os
 
 import torch
 from torch.utils.data import DataLoader
@@ -103,6 +104,13 @@ def collate_fn(examples):
     return train_dataloader, eval_dataloader
 
 
+# For testing only
+if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+    from accelerate.test_utils.training import mocked_dataloaders
+
+    get_dataloaders = mocked_dataloaders  # noqa: F811
+
+
 def training_function(config, args):
     # Initialize accelerator
     accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
diff --git a/examples/by_feature/multi_process_metrics.py b/examples/by_feature/multi_process_metrics.py
index db94d5773..6babf879a 100644
--- a/examples/by_feature/multi_process_metrics.py
+++ b/examples/by_feature/multi_process_metrics.py
@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import argparse
+import os
 
 import torch
 from torch.utils.data import DataLoader
@@ -104,6 +105,13 @@ def collate_fn(examples):
     return train_dataloader, eval_dataloader
 
 
+# For testing only
+if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+    from accelerate.test_utils.training import mocked_dataloaders
+
+    get_dataloaders = mocked_dataloaders  # noqa: F811
+
+
 def training_function(config, args):
     # Initialize accelerator
     accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
diff --git a/examples/by_feature/tracking.py b/examples/by_feature/tracking.py
index 47e15c0cf..1ccc2c1b0 100644
--- a/examples/by_feature/tracking.py
+++ b/examples/by_feature/tracking.py
@@ -103,6 +103,13 @@ def collate_fn(examples):
     return train_dataloader, eval_dataloader
 
 
+# For testing only
+if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+    from accelerate.test_utils.training import mocked_dataloaders
+
+    get_dataloaders = mocked_dataloaders  # noqa: F811
+
+
 def training_function(config, args):
     # Initialize Accelerator
 
diff --git a/src/accelerate/test_utils/training.py b/src/accelerate/test_utils/training.py
index 2ea86fbb1..c26587abe 100644
--- a/src/accelerate/test_utils/training.py
+++ b/src/accelerate/test_utils/training.py
@@ -14,6 +14,11 @@
 
 import numpy as np
 import torch
+from torch.utils.data import DataLoader
+
+from accelerate.utils.dataclasses import DistributedType
+from datasets import load_dataset
+from transformers import AutoTokenizer
 
 
 class RegressionDataset:
@@ -43,3 +48,40 @@ def forward(self, x=None):
             print(f"Model dtype: {self.a.dtype}, {self.b.dtype}. Input dtype: {x.dtype}")
             self.first_batch = False
         return x * self.a + self.b
+
+
+def mocked_dataloaders(accelerator, batch_size: int = 16):
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    data_files = {"train": "tests/test_samples/MRPC/train.csv", "validation": "tests/test_samples/MRPC/dev.csv"}
+    datasets = load_dataset("csv", data_files=data_files)
+    label_list = datasets["train"].unique("label")
+
+    label_to_id = {v: i for i, v in enumerate(label_list)}
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(
+            examples["sentence1"], examples["sentence2"], truncation=True, max_length=None, padding="max_length"
+        )
+        if "label" in examples:
+            outputs["labels"] = [label_to_id[l] for l in examples["label"]]
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["sentence1", "sentence2", "label"],
+    )
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=2)
+    eval_dataloader = DataLoader(tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=1)
+
+    return train_dataloader, eval_dataloader
diff --git a/tests/test_examples.py b/tests/test_examples.py
index b9759f8b5..af892f4eb 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -12,29 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import ast
 import os
-import sys
+import re
+import shutil
+import subprocess
 import tempfile
 import unittest
 from unittest import mock
 
-from torch.utils.data import DataLoader
-
-from accelerate import DistributedType
 from accelerate.test_utils.examples import compare_against_test
 from accelerate.test_utils.testing import TempDirTestCase, slow
-from datasets import load_dataset
-from transformers import AutoTokenizer
-
+from accelerate.utils import write_basic_config
 
-SRC_DIRS = [os.path.abspath(os.path.join("examples", "by_feature"))]
-sys.path.extend(SRC_DIRS)
-
-if SRC_DIRS is not None:
-    import checkpointing
-    import cross_validation
-    import multi_process_metrics
-    import tracking
 
 # DataLoaders built from `test_samples/MRPC` for quick testing
 # Should mock `{script_name}.get_dataloaders` via:
@@ -43,43 +33,6 @@
 EXCLUDE_EXAMPLES = ["cross_validation.py", "multi_process_metrics.py", "memory.py", "fsdp_with_peak_mem_tracking.py"]
 
 
-def mocked_dataloaders(accelerator, batch_size: int = 16):
-    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
-    data_files = {"train": "tests/test_samples/MRPC/train.csv", "validation": "tests/test_samples/MRPC/dev.csv"}
-    datasets = load_dataset("csv", data_files=data_files)
-    label_list = datasets["train"].unique("label")
-
-    label_to_id = {v: i for i, v in enumerate(label_list)}
-
-    def tokenize_function(examples):
-        # max_length=None => use the model max length (it's actually the default)
-        outputs = tokenizer(
-            examples["sentence1"], examples["sentence2"], truncation=True, max_length=None, padding="max_length"
-        )
-        if "label" in examples:
-            outputs["labels"] = [label_to_id[l] for l in examples["label"]]
-        return outputs
-
-    # Apply the method we just defined to all the examples in all the splits of the dataset
-    tokenized_datasets = datasets.map(
-        tokenize_function,
-        batched=True,
-        remove_columns=["sentence1", "sentence2", "label"],
-    )
-
-    def collate_fn(examples):
-        # On TPU it's best to pad everything to the same length or training will be very slow.
-        if accelerator.distributed_type == DistributedType.TPU:
-            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
-        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
-
-    # Instantiate dataloaders.
-    train_dataloader = DataLoader(tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=2)
-    eval_dataloader = DataLoader(tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=1)
-
-    return train_dataloader, eval_dataloader
-
-
 class ExampleDifferenceTests(unittest.TestCase):
     """
     This TestCase checks that all of the `complete_*` scripts contain all of the
@@ -159,88 +112,89 @@ def test_cv_examples(self):
         self.one_complete_example("complete_cv_example.py", False, cv_path, special_strings)
 
 
+@mock.patch.dict(os.environ, {"TESTING_MOCKED_DATALOADERS": "1"})
 class FeatureExamplesTests(TempDirTestCase):
     clear_on_setup = False
 
-    @mock.patch("checkpointing.get_dataloaders", mocked_dataloaders)
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        cls._tmpdir = tempfile.mkdtemp()
+        cls.configPath = os.path.join(cls._tmpdir, "default_config.yml")
+
+        write_basic_config(save_location=cls.configPath)
+        cls._launch_args = ["accelerate", "launch", "--config_file", cls.configPath]
+
+    @classmethod
+    def tearDownClass(cls):
+        super().tearDownClass()
+        shutil.rmtree(cls._tmpdir)
+
     def test_checkpointing_by_epoch(self):
         testargs = f"""
-        checkpointing.py
+        examples/by_feature/checkpointing.py
         --checkpointing_steps epoch
         --output_dir {self.tmpdir}
         """.split()
-        with mock.patch.object(sys, "argv", testargs):
-            checkpointing.main()
-            self.assertTrue(os.path.exists(os.path.join(self.tmpdir, "epoch_1")))
+        _ = subprocess.run(self._launch_args + testargs, stdout=subprocess.PIPE)
+        self.assertTrue(os.path.exists(os.path.join(self.tmpdir, "epoch_1")))
 
-    @mock.patch("checkpointing.get_dataloaders", mocked_dataloaders)
     def test_checkpointing_by_steps(self):
         testargs = f"""
-        checkpointing.py
+        examples/by_feature/checkpointing.py
         --checkpointing_steps 2
         --output_dir {self.tmpdir}
         """.split()
-        with mock.patch.object(sys, "argv", testargs):
-            checkpointing.main()
-            self.assertTrue(os.path.exists(os.path.join(self.tmpdir, "step_4")))
+        _ = subprocess.run(self._launch_args + testargs, stdout=subprocess.PIPE, env=os.environ)
+        self.assertTrue(os.path.exists(os.path.join(self.tmpdir, "step_4")))
 
-    @mock.patch("checkpointing.get_dataloaders", mocked_dataloaders)
     def test_load_states_by_epoch(self):
         testargs = f"""
-        checkpointing.py
+        examples/by_feature/checkpointing.py
         --resume_from_checkpoint {os.path.join(self.tmpdir, "epoch_1")}
         """.split()
-        dummy_results = {"accuracy": mock.ANY, "f1": mock.ANY}
-        with mock.patch("accelerate.Accelerator.print") as mocked_print:
-            with mock.patch.object(sys, "argv", testargs):
-                checkpointing.main()
-            with self.assertRaises(AssertionError):
-                mocked_print.assert_any_call("epoch 0:", dummy_results)
-            with self.assertRaises(AssertionError):
-                mocked_print.assert_any_call("epoch 1:", dummy_results)
-            mocked_print.assert_any_call("epoch 2:", dummy_results)
-
-    @mock.patch("checkpointing.get_dataloaders", mocked_dataloaders)
+        output = subprocess.run(
+            self._launch_args + testargs, universal_newlines=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE
+        ).stdout
+        self.assertNotIn("epoch 0:", output)
+        self.assertNotIn("epoch 1:", output)
+        self.assertIn("epoch 2:", output)
+
     def test_load_states_by_steps(self):
         testargs = f"""
-        checkpointing.py
+        examples/by_feature/checkpointing.py
         --resume_from_checkpoint {os.path.join(self.tmpdir, "step_4")}
         """.split()
-        dummy_results = {"accuracy": mock.ANY, "f1": mock.ANY}
-        with mock.patch("accelerate.Accelerator.print") as mocked_print:
-            with mock.patch.object(sys, "argv", testargs):
-                checkpointing.main()
-            with self.assertRaises(AssertionError):
-                mocked_print.assert_any_call("epoch 0:", dummy_results)
-            mocked_print.assert_any_call("epoch 1:", dummy_results)
-            mocked_print.assert_any_call("epoch 2:", dummy_results)
+        output = subprocess.run(
+            self._launch_args + testargs, universal_newlines=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE
+        ).stdout
+        self.assertNotIn("epoch 0:", output)
+        self.assertIn("epoch 1:", output)
+        self.assertIn("epoch 2:", output)
 
     @slow
     def test_cross_validation(self):
         testargs = """
-        cross_validation.py
+        examples/by_feature/cross_validation.py
         --num_folds 2
         """.split()
-        with mock.patch.object(sys, "argv", testargs):
-            with mock.patch("accelerate.Accelerator.print") as mocked_print:
-                cross_validation.main()
-                call = mocked_print.mock_calls[-1]
-                self.assertGreaterEqual(call.args[1]["accuracy"], 0.75)
+        with mock.patch.dict(os.environ, {"TESTING_MOCKED_DATALOADERS": "0"}):
+            output = subprocess.run(
+                self._launch_args + testargs, universal_newlines=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE
+            ).stdout
+            results = ast.literal_eval(re.findall("({.+})", output)[-1])
+            self.assertGreaterEqual(results["accuracy"], 0.75)
 
-    @mock.patch("multi_process_metrics.get_dataloaders", mocked_dataloaders)
     def test_multi_process_metrics(self):
-        testargs = ["multi_process_metrics.py"]
-        with mock.patch.object(sys, "argv", testargs):
-            multi_process_metrics.main()
+        testargs = ["examples/by_feature/multi_process_metrics.py"]
+        _ = subprocess.run(self._launch_args + testargs, stdout=subprocess.PIPE)
 
-    @mock.patch("tracking.get_dataloaders", mocked_dataloaders)
     def test_tracking(self):
         with tempfile.TemporaryDirectory() as tmpdir:
             testargs = f"""
-            tracking.py
+            examples/by_feature/tracking.py
             --with_tracking
             --logging_dir {tmpdir}
             """.split()
-            with mock.patch.object(sys, "argv", testargs):
-                tracking.main()
-                self.assertTrue(os.path.exists(os.path.join(tmpdir, "tracking")))
+            _ = subprocess.run(self._launch_args + testargs, stdout=subprocess.PIPE)
+            self.assertTrue(os.path.exists(os.path.join(tmpdir, "tracking")))

Since it's for testing only, maybe add it in the name of the env variable? `TESTING_MOCKED_DATALOADERS` for instance.
@@ -103,6 +103,13 @@ def collate_fn(examples):
     return train_dataloader, eval_dataloader
 
 
+# For testing only
+if os.environ.get("USE_MOCKED_DATALOADERS", None) == "1":
Same here.
@@ -74,6 +74,13 @@ def __exit__(self, *exc):
         # print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")
 
 
+# For testing only
+if os.environ.get("USE_MOCKED_DATALOADERS", None) == "1":
```suggestion
        super().setUpClass()
```
We're not in Python 2 ;-)
@@ -159,88 +112,89 @@ def test_cv_examples(self):
         self.one_complete_example("complete_cv_example.py", False, cv_path, special_strings)
 
 
+@mock.patch.dict(os.environ, {"USE_MOCKED_DATALOADERS": "1"})
 class FeatureExamplesTests(TempDirTestCase):
     clear_on_setup = False
 
-    @mock.patch("checkpointing.get_dataloaders", mocked_dataloaders)
+    @classmethod
+    def setUpClass(cls):
+        super(FeatureExamplesTests, cls).setUpClass()
```suggestion
        super().tearDownClass()
```
@@ -159,88 +112,89 @@ def test_cv_examples(self):
         self.one_complete_example("complete_cv_example.py", False, cv_path, special_strings)
 
 
+@mock.patch.dict(os.environ, {"USE_MOCKED_DATALOADERS": "1"})
 class FeatureExamplesTests(TempDirTestCase):
     clear_on_setup = False
 
-    @mock.patch("checkpointing.get_dataloaders", mocked_dataloaders)
+    @classmethod
+    def setUpClass(cls):
+        super(FeatureExamplesTests, cls).setUpClass()
+        cls._tmpdir = tempfile.mkdtemp()
+        cls.configPath = os.path.join(cls._tmpdir, "default_config.yml")
+
+        write_basic_config(save_location=cls.configPath)
+        cls._launch_args = ["accelerate", "launch", "--config_file", cls.configPath]
+
+    @classmethod
+    def tearDownClass(cls):
+        super(FeatureExamplesTests, cls).tearDownClass()
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/373
diff --git a/docs/source/internal.mdx b/docs/source/internal.mdx
index 94803a62f..9b496b30e 100644
--- a/docs/source/internal.mdx
+++ b/docs/source/internal.mdx
@@ -67,3 +67,5 @@ The main work on your PyTorch `DataLoader` is done by the following function:
 [[autodoc]] utils.synchronize_rng_states
 
 [[autodoc]] utils.wait_for_everyone
+
+[[autodoc]] utils.write_basic_config
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
index 58b33c1ad..f65b6b115 100644
--- a/src/accelerate/utils/__init__.py
+++ b/src/accelerate/utils/__init__.py
@@ -72,5 +72,12 @@
 
 from .launch import PrepareForLaunch
 from .memory import find_executable_batch_size
-from .other import extract_model_from_parallel, get_pretty_name, patch_environment, save, wait_for_everyone
+from .other import (
+    extract_model_from_parallel,
+    get_pretty_name,
+    patch_environment,
+    save,
+    wait_for_everyone,
+    write_basic_config,
+)
 from .random import set_seed, synchronize_rng_state, synchronize_rng_states
diff --git a/src/accelerate/utils/other.py b/src/accelerate/utils/other.py
index 35c96a21a..0d4f5b7e2 100644
--- a/src/accelerate/utils/other.py
+++ b/src/accelerate/utils/other.py
@@ -14,9 +14,12 @@
 
 import os
 from contextlib import contextmanager
+from pathlib import Path
 
 import torch
 
+from ..commands.config.cluster import ClusterConfig
+from ..commands.config.config_args import default_json_config_file
 from ..state import AcceleratorState
 from .dataclasses import DistributedType
 from .imports import is_deepspeed_available, is_tpu_available
@@ -109,3 +112,45 @@ def get_pretty_name(obj):
     if hasattr(obj, "__name__"):
         return obj.__name__
     return str(obj)
+
+
+def write_basic_config(mixed_precision="no", save_location: str = default_json_config_file):
+    """
+    Creates and saves a basic cluster config to be used on a local machine with potentially multiple GPUs. Will also
+    set CPU if it is a CPU-only machine.
+
+    Args:
+        mixed_precision (`str`, *optional*, defaults to "no"):
+            Mixed Precision to use. Should be one of "no", "fp16", or "bf16"
+        save_location (`str`, *optional*, defaults to `default_json_config_file`):
+            Optional custom save location. Should be passed to `--config_file` when using `accelerate launch`. Default
+            location is inside the huggingface cache folder (`~/.cache/huggingface`) but can be overriden by setting
+            the `HF_HOME` environmental variable, followed by `accelerate/default_config.yaml`.
+    """
+    path = Path(save_location)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    if path.exists():
+        print(
+            f"Configuration already exists at {save_location}, will not override. Run `accelerate config` manually or pass a different `save_location`."
+        )
+        return
+    mixed_precision = mixed_precision.lower()
+    if mixed_precision not in ["no", "fp16", "bf16"]:
+        raise ValueError(f"`mixed_precision` should be one of 'no', 'fp16', or 'bf16'. Received {mixed_precision}")
+    config = {"compute_environment": "LOCAL_MACHINE", "mixed_precision": mixed_precision}
+    if torch.cuda.is_available():
+        num_gpus = torch.cuda.device_count()
+        config["num_processes"] = num_gpus
+        config["use_cpu"] = False
+        if num_gpus > 1:
+            config["distributed_type"] = "MULTI_GPU"
+        else:
+            config["distributed_type"] = "NO"
+    else:
+        num_gpus = 0
+        config["use_cpu"] = True
+        config["num_processes"] = 1
+        config["distributed_type"] = "NO"
+    if not path.exists():
+        config = ClusterConfig(**config)
+        config.to_json_file(path)

This utility should take an optional path to store the config (you can pass any config file to `accelerate launch` with the arg `--config_file`.
@@ -109,3 +112,41 @@ def get_pretty_name(obj):
     if hasattr(obj, "__name__"):
         return obj.__name__
     return str(obj)
+
+
+def write_basic_config(mixed_precision="no"):
+    """
+    Creates and saves a basic cluster config to be used on a local machine with potentially multiple GPUs. Will also
+    set CPU if it is a CPU-only machine.
+
+    Args:
+        mixed_precision (`str`, *optional*, defaults to "no"):
+            Mixed Precision to use. Should be one of "no", "fp16", or "bf16"
+    """
+    path = Path(default_json_config_file)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    if path.exists():
+        print(
+            "User configuration already setup, will not override existing configuration. Run `accelerate config` manually."
+        )
+        return
This is the default already, so I don't think this part is useful.
@@ -109,3 +112,45 @@ def get_pretty_name(obj):
     if hasattr(obj, "__name__"):
         return obj.__name__
     return str(obj)
+
+
+def write_basic_config(mixed_precision="no", save_location: str = default_json_config_file):
+    """
+    Creates and saves a basic cluster config to be used on a local machine with potentially multiple GPUs. Will also
+    set CPU if it is a CPU-only machine.
+
+    Args:
+        mixed_precision (`str`, *optional*, defaults to "no"):
+            Mixed Precision to use. Should be one of "no", "fp16", or "bf16"
+        save_location (`str`, *optional*):
+            Optional custom save location. Should be passed to `--config_file` when using `accelerate launch`.
+    """
+    if save_location is None:
+        save_location = default_json_config_file
The default should be indicated there.
@@ -109,3 +112,45 @@ def get_pretty_name(obj):
     if hasattr(obj, "__name__"):
         return obj.__name__
     return str(obj)
+
+
+def write_basic_config(mixed_precision="no", save_location: str = default_json_config_file):
+    """
+    Creates and saves a basic cluster config to be used on a local machine with potentially multiple GPUs. Will also
+    set CPU if it is a CPU-only machine.
+
+    Args:
+        mixed_precision (`str`, *optional*, defaults to "no"):
+            Mixed Precision to use. Should be one of "no", "fp16", or "bf16"
+        save_location (`str`, *optional*):
+            Optional custom save location. Should be passed to `--config_file` when using `accelerate launch`.
This won't be the default location for everyone, so I'd just say defaults to the default config path and expand a bit in the docstring that it's inside the default cache of hugging face, and may add which env variable control it.
@@ -122,11 +122,9 @@ def write_basic_config(mixed_precision="no", save_location: str = default_json_c
     Args:
         mixed_precision (`str`, *optional*, defaults to "no"):
             Mixed Precision to use. Should be one of "no", "fp16", or "bf16"
-        save_location (`str`, *optional*):
+        save_location (`str`, *optional*, defaults to "~/.cache/huggingface/accelerate/default_config.yaml"):
```suggestion
            location is inside the huggingface cache folder (`~/.cache/huggingface` but can be overriden by setting the 
            `HF_HOME` environmental variable) followed by `accelerate/default_config.yaml`.
```
@@ -122,8 +122,10 @@ def write_basic_config(mixed_precision="no", save_location: str = default_json_c
     Args:
         mixed_precision (`str`, *optional*, defaults to "no"):
             Mixed Precision to use. Should be one of "no", "fp16", or "bf16"
-        save_location (`str`, *optional*, defaults to "~/.cache/huggingface/accelerate/default_config.yaml"):
-            Optional custom save location. Should be passed to `--config_file` when using `accelerate launch`.
+        save_location (`str`, *optional*, defaults to `default_json_config_file`):
+            Optional custom save location. Should be passed to `--config_file` when using `accelerate launch`. Default
+            location is inside the huggingface cache folder as default_config.yaml, can be overriden by setting the 
+            `HF_HOME` environmental variable.
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/371
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 6fd24c321..b9780330a 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -881,7 +881,7 @@ def get_state_dict(self, model):
                 is_zero_3 = self.state.deepspeed_plugin.zero_stage == 3
 
         if is_zero_3:
-            state_dict = model._zero3_consolidated_fp16_state_dict()
+            state_dict = model._zero3_consolidated_16bit_state_dict()
         else:
             model = self.unwrap_model(model)
             state_dict = model.state_dict()

_The documentation is not available anymore as the PR was closed or merged._
> Thanks for fixing! Should we have a minimal version of DeepSpeed? I don't think this was always supported.

Hello, yes, we should have minimal version of DeepSpeed. I will club it with all the other changes that I am planning to have to improve the DeeSpeed integration.
https://api.github.com/repos/huggingface/accelerate/pulls/370
diff --git a/src/accelerate/checkpointing.py b/src/accelerate/checkpointing.py
index b0e2c04a2..b043d5d20 100644
--- a/src/accelerate/checkpointing.py
+++ b/src/accelerate/checkpointing.py
@@ -102,7 +102,7 @@ def save_accelerator_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())
+        states["xm_seed"] = 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}")

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/365
diff --git a/src/accelerate/launchers.py b/src/accelerate/launchers.py
index cfd8929d9..d3cd2e1a2 100644
--- a/src/accelerate/launchers.py
+++ b/src/accelerate/launchers.py
@@ -80,7 +80,7 @@ def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, mix
     else:
         if num_processes is None:
             raise ValueError(
-                "You have to specify the number of GPUs you would like to use, add `num_process=...` to your call."
+                "You have to specify the number of GPUs you would like to use, add `num_processes=...` to your call."
             )
 
         if num_processes > 1:

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/363
diff --git a/examples/by_feature/fsdp_with_peak_mem_tracking.py b/examples/by_feature/fsdp_with_peak_mem_tracking.py
index a8546c3c8..581c5fea0 100644
--- a/examples/by_feature/fsdp_with_peak_mem_tracking.py
+++ b/examples/by_feature/fsdp_with_peak_mem_tracking.py
@@ -103,11 +103,12 @@ def training_function(config, args):
 
     # We need to initialize the trackers we use, and also store our configuration
     if args.with_tracking:
-        run = os.path.split(__file__)[-1].split(".")[0]
-        if args.logging_dir:
-            run = os.path.join(args.logging_dir, run)
-            accelerator.print(run)
-        accelerator.init_trackers(run, config)
+        if accelerator.is_main_process:
+            run = os.path.split(__file__)[-1].split(".")[0]
+            if args.logging_dir:
+                run = os.path.join(args.logging_dir, run)
+                accelerator.print(run)
+            accelerator.init_trackers(run, config)
 
     tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
     datasets = load_dataset("glue", "mrpc")
@@ -293,7 +294,7 @@ def collate_fn(examples):
                     {
                         "accuracy": eval_metric["accuracy"],
                         "f1": eval_metric["f1"],
-                        "train_loss": total_loss,
+                        "train_loss": total_loss.item(),
                     },
                     step=epoch,
                 )
diff --git a/examples/by_feature/tracking.py b/examples/by_feature/tracking.py
index 19f8e9dea..47e15c0cf 100644
--- a/examples/by_feature/tracking.py
+++ b/examples/by_feature/tracking.py
@@ -161,10 +161,11 @@ def training_function(config, args):
     # New Code #
     # We need to initalize the trackers we use. Overall configurations can also be stored
     if args.with_tracking:
-        run = os.path.split(__file__)[-1].split(".")[0]
-        if args.logging_dir:
-            run = os.path.join(args.logging_dir, run)
-        accelerator.init_trackers(run, config)
+        if accelerator.is_main_process:
+            run = os.path.split(__file__)[-1].split(".")[0]
+            if args.logging_dir:
+                run = os.path.join(args.logging_dir, run)
+            accelerator.init_trackers(run, config)
 
     # Now we train the model
     for epoch in range(num_epochs):
@@ -208,15 +209,16 @@ def training_function(config, args):
 
         # New Code #
         # To actually log, we call `Accelerator.log`
-        # The values passed can be of `str`, `int`, or `float`
+        # The values passed can be of `str`, `int`, `float` or `dict` of `str` to `float`/`int`
         if args.with_tracking:
             accelerator.log(
                 {
                     "accuracy": eval_metric["accuracy"],
                     "f1": eval_metric["f1"],
-                    "train_loss": total_loss,
+                    "train_loss": total_loss.item(),
                     "epoch": epoch,
-                }
+                },
+                step=epoch,
             )
 
     # New Code #
diff --git a/examples/complete_cv_example.py b/examples/complete_cv_example.py
index 79e9289b8..fd830442d 100644
--- a/examples/complete_cv_example.py
+++ b/examples/complete_cv_example.py
@@ -104,10 +104,11 @@ def training_function(config, args):
 
     # We need to initialize the trackers we use, and also store our configuration
     if args.with_tracking:
-        run = os.path.split(__file__)[-1].split(".")[0]
-        if args.logging_dir:
-            run = os.path.join(args.logging_dir, run)
-        accelerator.init_trackers(run, config)
+        if accelerator.is_main_process:
+            run = os.path.split(__file__)[-1].split(".")[0]
+            if args.logging_dir:
+                run = os.path.join(args.logging_dir, run)
+            accelerator.init_trackers(run, config)
 
     # Grab all the image filenames
     file_names = [os.path.join(args.data_dir, fname) for fname in os.listdir(args.data_dir) if fname.endswith(".jpg")]
diff --git a/examples/complete_nlp_example.py b/examples/complete_nlp_example.py
index a231566dc..abf85c8d9 100644
--- a/examples/complete_nlp_example.py
+++ b/examples/complete_nlp_example.py
@@ -81,10 +81,11 @@ def training_function(config, args):
 
     # We need to initialize the trackers we use, and also store our configuration
     if args.with_tracking:
-        run = os.path.split(__file__)[-1].split(".")[0]
-        if args.logging_dir:
-            run = os.path.join(args.logging_dir, run)
-        accelerator.init_trackers(run, config)
+        if accelerator.is_main_process:
+            run = os.path.split(__file__)[-1].split(".")[0]
+            if args.logging_dir:
+                run = os.path.join(args.logging_dir, run)
+            accelerator.init_trackers(run, config)
 
     tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
     datasets = load_dataset("glue", "mrpc")
@@ -245,9 +246,10 @@ def collate_fn(examples):
                 {
                     "accuracy": eval_metric["accuracy"],
                     "f1": eval_metric["f1"],
-                    "train_loss": total_loss,
+                    "train_loss": total_loss.item(),
                     "epoch": epoch,
-                }
+                },
+                step=epoch,
             )
 
         if checkpointing_steps == "epoch":
diff --git a/src/accelerate/tracking.py b/src/accelerate/tracking.py
index 218ddf2df..3efff77c9 100644
--- a/src/accelerate/tracking.py
+++ b/src/accelerate/tracking.py
@@ -136,8 +136,9 @@ def log(self, values: dict, step: Optional[int] = None):
         Logs `values` to the current run.
 
         Args:
-            values (Dictionary `str` to `str`, `float`, or `int`):
-                Values to be logged as key-value pairs. The values need to have type `str`, `float`, or `int`.
+            values (Dictionary `str` to `str`, `float`, `int` or `dict` of `str` to `float`/`int`):
+                Values to be logged as key-value pairs. The values need to have type `str`, `float`, `int` or `dict` of
+                `str` to `float`/`int`.
             step (`int`, *optional*):
                 The run step. If included, the log will be affiliated with this step.
         """
@@ -146,6 +147,8 @@ def log(self, values: dict, step: Optional[int] = None):
                 self.writer.add_scalar(k, v, global_step=step)
             elif isinstance(v, str):
                 self.writer.add_text(k, v, global_step=step)
+            elif isinstance(v, dict):
+                self.writer.add_scalars(k, v, global_step=step)
         self.writer.flush()
         logger.info("Successfully logged to TensorBoard")
 
@@ -170,7 +173,7 @@ class WandBTracker(GeneralTracker):
 
     def __init__(self, run_name: str):
         self.run_name = run_name
-        self.run = wandb.init(self.run_name)
+        self.run = wandb.init(project=self.run_name)
         logger.info(f"Initialized WandB project {self.run_name}")
         logger.info(
             "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
@@ -193,8 +196,9 @@ def log(self, values: dict, step: Optional[int] = None):
         Logs `values` to the current run.
 
         Args:
-            values (Dictionary `str` to `str`, `float`, or `int`):
-                Values to be logged as key-value pairs. The values need to have type `str`, `float`, or `int`.
+            values (Dictionary `str` to `str`, `float`, `int` or `dict` of `str` to `float`/`int`):
+                Values to be logged as key-value pairs. The values need to have type `str`, `float`, `int` or `dict` of
+                `str` to `float`/`int`.
             step (`int`, *optional*):
                 The run step. If included, the log will be affiliated with this step.
         """
@@ -247,14 +251,21 @@ def log(self, values: dict, step: Optional[int] = None):
         Logs `values` to the current run.
 
         Args:
-            values (Dictionary `str` to `str`, `float`, or `int`):
-                Values to be logged as key-value pairs. The values need to have type `str`, `float`, or `int`.
+            values (Dictionary `str` to `str`, `float`, `int` or `dict` of `str` to `float`/`int`):
+                Values to be logged as key-value pairs. The values need to have type `str`, `float`, `int` or `dict` of
+                `str` to `float`/`int`.
             step (`int`, *optional*):
                 The run step. If included, the log will be affiliated with this step.
         """
         if step is not None:
             self.writer.set_step(step)
-        self.writer.log_others(values)
+        for k, v in values.items():
+            if isinstance(v, (int, float)):
+                self.writer.log_metric(k, v, step=step)
+            elif isinstance(v, str):
+                self.writer.log_other(k, v)
+            elif isinstance(v, dict):
+                self.writer.log_metrics(v, step=step)
         logger.info("Successfully logged to CometML")
 
     def finish(self):
diff --git a/tests/test_examples.py b/tests/test_examples.py
index 84b6adaef..b9759f8b5 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -149,9 +149,10 @@ def test_cv_examples(self):
             " " * 16 + "{\n\n",
             " " * 18 + '"accuracy": eval_metric["accuracy"],\n\n',
             " " * 18 + '"f1": eval_metric["f1"],\n\n',
-            " " * 18 + '"train_loss": total_loss,\n\n',
+            " " * 18 + '"train_loss": total_loss.item(),\n\n',
             " " * 18 + '"epoch": epoch,\n\n',
-            " " * 16 + "}\n",
+            " " * 16 + "},\n\n",
+            " " * 16 + "step=epoch,\n",
             " " * 8,
         ]
         self.one_complete_example("complete_cv_example.py", True, cv_path, special_strings)
diff --git a/tests/test_tracking.py b/tests/test_tracking.py
index e9a2170ee..f02db8d21 100644
--- a/tests/test_tracking.py
+++ b/tests/test_tracking.py
@@ -203,6 +203,9 @@ def get_value_from_key(log_list, key: str, is_param: bool = False):
             if "log_other" in j.keys():
                 if j["log_other"]["key"] == key:
                     return j["log_other"]["val"]
+            if "metric" in j.keys():
+                if j["metric"]["metricName"] == key:
+                    return j["metric"]["metricValue"]
 
     def test_init_trackers(self):
         with tempfile.TemporaryDirectory() as d:

Do we really have either a Tensor or something else? Can't it always be `.item()` here?
@@ -293,7 +294,7 @@ def collate_fn(examples):
                     {
                         "accuracy": eval_metric["accuracy"],
                         "f1": eval_metric["f1"],
-                        "train_loss": total_loss,
+                        "train_loss": total_loss.item() if type(total_loss) == torch.Tensor else total_loss,
Same comment
@@ -208,15 +209,16 @@ def training_function(config, args):
 
         # New Code #
         # To actually log, we call `Accelerator.log`
-        # The values passed can be of `str`, `int`, or `float`
+        # The values passed can be of `str`, `int`, `float` or `dict` of `str` to `float`/`int`
         if args.with_tracking:
             accelerator.log(
                 {
                     "accuracy": eval_metric["accuracy"],
                     "f1": eval_metric["f1"],
-                    "train_loss": total_loss,
+                    "train_loss": total_loss.item() if type(total_loss) == torch.Tensor else total_loss,
Same here.
@@ -245,9 +246,10 @@ def collate_fn(examples):
                 {
                     "accuracy": eval_metric["accuracy"],
                     "f1": eval_metric["f1"],
-                    "train_loss": total_loss,
+                    "train_loss": total_loss.item() if type(total_loss) == torch.Tensor else total_loss,
_The documentation is not available anymore as the PR was closed or merged._
Oh, also make sure to edit the title of the PR ;-)
https://api.github.com/repos/huggingface/accelerate/pulls/361
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index e3ce3fd62..f283dfbc5 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -33,6 +33,8 @@
     PrepareForLaunch,
     is_sagemaker_available,
 )
+from accelerate.utils.versions import torch_version
+from packaging import version
 
 
 def launch_command_parser(subparsers=None):
@@ -131,6 +133,12 @@ def launch_command_parser(subparsers=None):
         action="store_true",
         help="Skip prepending the training script with 'python' - just execute it directly. Useful when the script is not a Python script.",
     )
+    parser.add_argument(
+        "--num_cpu_threads_per_process",
+        type=int,
+        default=1,
+        help="The number of CPU threads per process. Can be tuned for optimal performance.",
+    )
     parser.add_argument(
         "--aws_access_key_id",
         type=str,
@@ -211,7 +219,12 @@ def simple_launcher(args):
 
 
 def multi_gpu_launcher(args):
-    cmd = [sys.executable, "-m", "torch.distributed.launch", "--use_env"]
+    if torch_version >= version.parse("1.10.0"):
+        cmd = ["torchrun"]
+    elif torch_version >= version.parse("1.9.0"):
+        cmd = [sys.executable, "-m", "torch.distributed.run"]
+    else:
+        cmd = [sys.executable, "-m", "torch.distributed.launch", "--use_env"]
     if args.num_machines > 1:
         cmd.extend(
             [
@@ -259,6 +272,7 @@ def multi_gpu_launcher(args):
         current_env["FSDP_OFFLOAD_PARAMS"] = str(args.offload_params).lower()
         current_env["FSDP_MIN_NUM_PARAMS"] = str(args.min_num_params)
         current_env["FSDP_SHARDING_STRATEGY"] = str(args.sharding_strategy)
+    current_env["OMP_NUM_THREADS"] = str(args.num_cpu_threads_per_process)
     process = subprocess.Popen(cmd, env=current_env)
     process.wait()
     if process.returncode != 0:
diff --git a/src/accelerate/utils/versions.py b/src/accelerate/utils/versions.py
new file mode 100644
index 000000000..36dd7ac48
--- /dev/null
+++ b/src/accelerate/utils/versions.py
@@ -0,0 +1,25 @@
+# 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 sys
+
+from packaging.version import parse
+
+
+if sys.version_info < (3, 8):
+    import importlib_metadata
+else:
+    import importlib.metadata as importlib_metadata
+
+torch_version = parse(importlib_metadata.version("torch"))

No need to make it private since we reuse it.
@@ -0,0 +1,25 @@
+# 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 sys
+
+from packaging.version import parse
+
+
+if sys.version_info < (3, 8):
+    import importlib_metadata
+else:
+    import importlib.metadata as importlib_metadata
+
+_torch_version = parse(importlib_metadata.version("torch"))
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/360
diff --git a/src/accelerate/__init__.py b/src/accelerate/__init__.py
index b74c683b4..cb38abce6 100644
--- a/src/accelerate/__init__.py
+++ b/src/accelerate/__init__.py
@@ -10,6 +10,7 @@
     DeepSpeedPlugin,
     DistributedDataParallelKwargs,
     DistributedType,
+    FullyShardedDataParallelPlugin,
     GradScalerKwargs,
     InitProcessGroupKwargs,
     find_executable_batch_size,
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index ba34587e2..6fd24c321 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -163,15 +163,18 @@ def __init__(
             assert isinstance(
                 deepspeed_plugin, DeepSpeedPlugin
             ), "`deepspeed_plugin` must be a DeepSpeedPlugin object."
+            os.environ["USE_DEEPSPEED"] = "true"  # use DeepSpeed if plugin is provided
+
+        if fsdp_plugin is None:  # init from env variables
+            fsdp_plugin = FullyShardedDataParallelPlugin() if os.environ.get("USE_FSDP", "false") == "true" else None
+        else:
+            if not isinstance(fsdp_plugin, FullyShardedDataParallelPlugin):
+                raise TypeError("`fsdp_plugin` must be a FullyShardedDataParallelPlugin object.")
+            os.environ["USE_FSDP"] = "true"  # use FSDP if plugin is provided
 
         if os.environ.get("USE_FSDP", "false") == "true":
             if version.parse(torch.__version__) < version.parse("1.12.0.dev20220418+cu113"):
                 raise ValueError("FSDP requires PyTorch >= 1.12.0.dev20220418+cu113")
-            if fsdp_plugin is None:  # init from env variables
-                fsdp_plugin = FullyShardedDataParallelPlugin()
-            else:
-                if not isinstance(fsdp_plugin, FullyShardedDataParallelPlugin):
-                    raise TypeError("`fsdp_plugin` must be a FullyShardedDataParallelPlugin object.")
 
         # Kwargs handlers
         self.ddp_handler = None

_The documentation is not available anymore as the PR was closed or merged._
Thanks a lot!! @sgugger @pacman100 
https://api.github.com/repos/huggingface/accelerate/pulls/356
diff --git a/examples/by_feature/fsdp_with_peak_mem_tracking.py b/examples/by_feature/fsdp_with_peak_mem_tracking.py
index e39d3ede8..a8546c3c8 100644
--- a/examples/by_feature/fsdp_with_peak_mem_tracking.py
+++ b/examples/by_feature/fsdp_with_peak_mem_tracking.py
@@ -263,6 +263,7 @@ def collate_fn(examples):
         # context manager to track the peak memory usage during the evaluation
         with TorchTracemalloc() as tracemalloc:
             model.eval()
+            samples_seen = 0
             for step, batch in enumerate(eval_dataloader):
                 # We could avoid this line since we set the accelerator with `device_placement=True`.
                 batch.to(accelerator.device)
@@ -270,7 +271,15 @@ def collate_fn(examples):
                     outputs = model(**batch)
                 predictions = outputs.logits.argmax(dim=-1)
                 # It is slightly faster to call this once, than multiple times
-                predictions, references = accelerator.gather((predictions, batch["labels"]))
+                predictions, references = accelerator.gather(
+                    (predictions, batch["labels"])
+                )  # If we are in a multiprocess environment, the last batch has duplicates
+                if accelerator.num_processes > 1:
+                    if step == len(eval_dataloader) - 1:
+                        predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
+                        references = references[: len(eval_dataloader.dataset) - samples_seen]
+                    else:
+                        samples_seen += references.shape[0]
                 metric.add_batch(
                     predictions=predictions,
                     references=references,
diff --git a/examples/by_feature/multi_process_metrics.py b/examples/by_feature/multi_process_metrics.py
index 1b5bf1d58..db94d5773 100644
--- a/examples/by_feature/multi_process_metrics.py
+++ b/examples/by_feature/multi_process_metrics.py
@@ -177,21 +177,13 @@ def training_function(config, args):
             # First we check if it's a distributed system
             if accelerator.num_processes > 1:
                 # Then see if we're on the last batch of our eval dataloader
-                if step == len(eval_dataloader):
+                if step == len(eval_dataloader) - 1:
                     # Last batch needs to be truncated on distributed systems as it contains additional samples
                     predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
                     references = references[: len(eval_dataloader.dataset) - samples_seen]
                 else:
                     # Otherwise we add the number of samples seen
                     samples_seen += references.shape[0]
-            predictions, references = accelerator.gather((predictions, batch["labels"]))
-            # If we are in a multiprocess environment, the last batch has duplicates
-            if accelerator.num_processes > 1:
-                if step == len(eval_dataloader):
-                    predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
-                    references = references[: len(eval_dataloader.dataset) - samples_seen]
-                else:
-                    samples_seen += references.shape[0]
             metric.add_batch(
                 predictions=predictions,
                 references=references,
diff --git a/examples/complete_cv_example.py b/examples/complete_cv_example.py
index adb46d957..79e9289b8 100644
--- a/examples/complete_cv_example.py
+++ b/examples/complete_cv_example.py
@@ -242,7 +242,7 @@ def training_function(config, args):
             predictions = outputs.argmax(dim=-1)
             predictions, references = accelerator.gather((predictions, batch["label"]))
             if accelerator.num_processes > 1:
-                if step == len(eval_dataloader):
+                if step == len(eval_dataloader) - 1:
                     predictions = predictions[: len(eval_dataloader) - samples_seen]
                     references = references[: len(eval_dataloader) - samples_seen]
                 else:
diff --git a/examples/complete_nlp_example.py b/examples/complete_nlp_example.py
index d4c8ddf0d..a231566dc 100644
--- a/examples/complete_nlp_example.py
+++ b/examples/complete_nlp_example.py
@@ -227,7 +227,7 @@ def collate_fn(examples):
                 (predictions, batch["labels"])
             )  # If we are in a multiprocess environment, the last batch has duplicates
             if accelerator.num_processes > 1:
-                if step == len(eval_dataloader):
+                if step == len(eval_dataloader) - 1:
                     predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
                     references = references[: len(eval_dataloader.dataset) - samples_seen]
                 else:

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/355
diff --git a/examples/by_feature/multi_process_metrics.py b/examples/by_feature/multi_process_metrics.py
index abba1a746..1b5bf1d58 100644
--- a/examples/by_feature/multi_process_metrics.py
+++ b/examples/by_feature/multi_process_metrics.py
@@ -184,6 +184,14 @@ def training_function(config, args):
                 else:
                     # Otherwise we add the number of samples seen
                     samples_seen += references.shape[0]
+            predictions, references = accelerator.gather((predictions, batch["labels"]))
+            # If we are in a multiprocess environment, the last batch has duplicates
+            if accelerator.num_processes > 1:
+                if step == len(eval_dataloader):
+                    predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
+                    references = references[: len(eval_dataloader.dataset) - samples_seen]
+                else:
+                    samples_seen += references.shape[0]
             metric.add_batch(
                 predictions=predictions,
                 references=references,
diff --git a/examples/complete_cv_example.py b/examples/complete_cv_example.py
index 80a083a8d..adb46d957 100644
--- a/examples/complete_cv_example.py
+++ b/examples/complete_cv_example.py
@@ -232,7 +232,7 @@ def training_function(config, args):
                     accelerator.save_state(output_dir)
         model.eval()
         accurate = 0
-        num_elems = 0
+        samples_seen = 0
         for step, batch in enumerate(eval_dataloader):
             # We could avoid this line since we set the accelerator with `device_placement=True`.
             batch = {k: v.to(accelerator.device) for k, v in batch.items()}
@@ -240,11 +240,19 @@ def training_function(config, args):
             with torch.no_grad():
                 outputs = model(inputs)
             predictions = outputs.argmax(dim=-1)
-            accurate_preds = accelerator.gather(predictions) == accelerator.gather(batch["label"])
-            num_elems += accurate_preds.shape[0]
+            predictions, references = accelerator.gather((predictions, batch["label"]))
+            if accelerator.num_processes > 1:
+                if step == len(eval_dataloader):
+                    predictions = predictions[: len(eval_dataloader) - samples_seen]
+                    references = references[: len(eval_dataloader) - samples_seen]
+                else:
+                    samples_seen += references.shape[0]
+            else:
+                samples_seen += references.shape[0]
+            accurate_preds = predictions == references
             accurate += accurate_preds.long().sum()
 
-        eval_metric = accurate.item() / num_elems
+        eval_metric = accurate.item() / samples_seen
         # Use accelerator.print to print only on the main process.
         accelerator.print(f"epoch {epoch}: {100 * eval_metric:.2f}")
         if args.with_tracking:
diff --git a/examples/complete_nlp_example.py b/examples/complete_nlp_example.py
index 1910a4f83..d4c8ddf0d 100644
--- a/examples/complete_nlp_example.py
+++ b/examples/complete_nlp_example.py
@@ -215,6 +215,7 @@ def collate_fn(examples):
                     accelerator.save_state(output_dir)
 
         model.eval()
+        samples_seen = 0
         for step, batch in enumerate(eval_dataloader):
             # We could avoid this line since we set the accelerator with `device_placement=True`.
             batch.to(accelerator.device)
@@ -222,7 +223,15 @@ def collate_fn(examples):
                 outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1)
             # It is slightly faster to call this once, than multiple times
-            predictions, references = accelerator.gather((predictions, batch["labels"]))
+            predictions, references = accelerator.gather(
+                (predictions, batch["labels"])
+            )  # If we are in a multiprocess environment, the last batch has duplicates
+            if accelerator.num_processes > 1:
+                if step == len(eval_dataloader):
+                    predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
+                    references = references[: len(eval_dataloader.dataset) - samples_seen]
+                else:
+                    samples_seen += references.shape[0]
             metric.add_batch(
                 predictions=predictions,
                 references=references,

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/352
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index 8ff69e048..01a55e646 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -133,17 +133,24 @@ def get_cluster_input():
         main_training_function = "main"
 
     if distributed_type in [DistributedType.MULTI_CPU, DistributedType.MULTI_GPU, DistributedType.TPU]:
-        machine_type = str(distributed_type).split(".")[0].replace("MULTI_", "")
+        machine_type = str(distributed_type).split(".")[1].replace("MULTI_", "")
         if machine_type == "TPU":
             machine_type += " cores"
         else:
-            machine_type += "'s"
+            machine_type += "(s)"
         num_processes = _ask_field(
             f"How many {machine_type} should be used for distributed training? [1]:",
             lambda x: int(x),
             default=1,
             error_message="Please enter an integer.",
         )
+    elif distributed_type in [DistributedType.FSDP, DistributedType.DEEPSPEED]:
+        num_processes = _ask_field(
+            "How many GPU(s) should be used for distributed training? [1]:",
+            lambda x: int(x),
+            default=1,
+            error_message="Please enter an integer.",
+        )
     else:
         num_processes = 1
 

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/347
diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml
index d35024d66..130659196 100644
--- a/docs/source/_toctree.yml
+++ b/docs/source/_toctree.yml
@@ -7,6 +7,8 @@
     title: Installation
   title: Get started
 - sections:
+  - local: big_modeling
+    title: Handling big models
   - local: sagemaker
     title: Amazon SageMaker
   title: Guides
diff --git a/docs/source/big_modeling.mdx b/docs/source/big_modeling.mdx
new file mode 100644
index 000000000..4ca316870
--- /dev/null
+++ b/docs/source/big_modeling.mdx
@@ -0,0 +1,232 @@
+<!--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.
+-->
+
+# Handling big models
+
+When loading a pretrained model in PyTorch, the usual workflow looks like this:
+
+```py
+import torch
+
+my_model = ModelClass(...)
+state_dict = torch.load(checkpoint_file)
+my_model.load_state_dict(state_dict)
+```
+
+In plain English, those steps are:
+1. Create the model with randomly initialized weights
+2. Load the model weights (in a dictionary usually called a state dict) from the disk
+3. Load those weights inside the model
+
+While this works very well for regularly sized models, this workflow has some clear limitation when we deal with a huge model: in step 1, we load a full version of the model in RAM, and spend some time randomly initializing the weights (which will be discarded in step 3). In step 2, we load another full version of the model in RAM, with the pretrained weights. If you're loading a model with 6 billions parameters, this needs you will need 24GB of RAM for each copy of the model, so 48GB in total (half of it to load the model in FP16).
+
+<Tip warning={true}>
+
+This API is quite new and still in its experimental stage. While we strive to provide a stable API, it's possible some small parts of the public API will change in the future.
+
+</Tip>
+
+## Instantiating an empty model
+
+The first tool Accelerate introduces to help with big models is a context manager [`init_empty_weights`] that helps you initialize a model without using any RAM, so that step 1 can be done on models of any size. Here is how it works:
+
+```py
+from accelerate import init_empty_weights
+
+with init_empty_weights():
+    my_model = ModelClass(...)
+```
+
+For instance:
+
+```py
+with init_empty_weights():
+    model = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])
+```
+
+initializes an empty model with a bit more than 100B parameters. Behind the scenes, this relies on the meta device introduced in PyTorch 1.9. During the initialization under the context manager, each time a parameter is created, it is instantly moved on that device.
+
+<Tip warning={true}>
+
+You can't move a model initialized like this on CPU or another device directly, since it doesn't have any data. It's also very likely that a forward pass with that empty model will fail, as not all operations are supported on the meta device.
+
+</Tip>
+
+## Sharded checkpoints
+
+It's possible your model is so big that even a single copy won't fit in RAM. That doesn't mean it can't be loaded: if you have one or several GPUs, this is more memory available to store your model. In this case, it's better if your checkpoint is split in several smaller files that we call checkpoint shards.
+
+Accelerate will handle sharded checkpoints as long as you follow the following format: your checkpoint should be in a folder, with several files containing the partial state dicts, and there should be an index in the JSON format that contains a dictionary mapping parameter names to the file containing their weights. For instance we could have a folder containing:
+
+```bash
+first_state_dict.bin
+index.json
+second_state_dict.bin
+```
+
+with index.json being the following file:
+
+```
+{
+  "linear1.weight": "first_state_dict.bin",
+  "linear1.bias": "first_state_dict.bin",
+  "linear2.weight": "second_state_dict.bin",
+  "linear2.bias": "second_state_dict.bin"
+}
+```
+
+and `first_state_dict.bin` containing the weights for `"linear1.weight"` and `"linear1.bias"`, `second_state_dict.bin` the ones for `"linear2.weight"` and `"linear2.bias"`
+
+## Loading weights
+
+The second tool Accelerate introduces is a function [`load_checkpoint_and_dispatch`], that will allow you to load a checkpoint inside your empty model. This supports full checkpoints (a single file containing the whole state dict) as well as sharded checkpoints. It will also automatically dispatch those weights across the devices you have available (GPUs, CPU RAM), so if you are loading a sharded checkpoint, the maximum RAM usage will be the size of the biggest shard.
+
+Here is how we can use this to load the [GPT-J-6B](https://huggingface.co/EleutherAI/gpt-j-6B) model. You clone the sharded version of this model with:
+
+```bash
+git clone https://huggingface.co/sgugger/sharded-gpt-j-6B
+cd sharded-gpt-j-6B
+git-lfs install
+git pull
+```
+
+then we can initialize the model with
+
+```py
+from accelerate import init_empty_weights
+from transformers import AutoConfig, AutoModelForCausalLM
+
+checkpoint = "EleutherAI/gpt-j-6B"
+config = AutoConfig.from_pretrained(checkpoint)
+
+with init_empty_weights():
+    model = AutoModelForCausalLM.from_config(config)
+```
+
+and load the checkpoint we just downloaded with:
+
+```py
+from accelerate import load_checkpoint_and_dispatch
+
+model = load_checkpoint_and_dispatch(
+    model, "sharded-gpt-j-6B", device_map="auto", no_split_module_classes=["GPTJBlock"]
+)
+```
+
+By passing `device_map="auto"`, we tell Accelerate to determine automatically where to put each layer of the model depending on the available resources:
+- first we use the maximum space available on the GPU(s)
+- if we still need space, we store the remaining weights on the CPU
+- if there is not enough RAM, we store the remaining weights on the hard drive as memory-mapped tensors
+
+`no_split_module_classes=["GPTJBlock"]` indicates that the modules that are `GPTJBlock` should not be split on different devices. You should set here all blocks that include a residual connection of some kind.
+
+You can see the `device_map` that Accelerate picked by accessing the `hf_device_map` attribute of your model:
+
+```py
+model.hf_device_map
+```
+
+```python out
+{'transformer.wte': 0,
+ 'transformer.drop': 0,
+ 'transformer.h.0': 0,
+ 'transformer.h.1': 0,
+ 'transformer.h.2': 0,
+ 'transformer.h.3': 0,
+ 'transformer.h.4': 0,
+ 'transformer.h.5': 0,
+ 'transformer.h.6': 0,
+ 'transformer.h.7': 0,
+ 'transformer.h.8': 0,
+ 'transformer.h.9': 0,
+ 'transformer.h.10': 0,
+ 'transformer.h.11': 0,
+ 'transformer.h.12': 0,
+ 'transformer.h.13': 0,
+ 'transformer.h.14': 0,
+ 'transformer.h.15': 0,
+ 'transformer.h.16': 0,
+ 'transformer.h.17': 0,
+ 'transformer.h.18': 0,
+ 'transformer.h.19': 0,
+ 'transformer.h.20': 0,
+ 'transformer.h.21': 0,
+ 'transformer.h.22': 0,
+ 'transformer.h.23': 0,
+ 'transformer.h.24': 1,
+ 'transformer.h.25': 1,
+ 'transformer.h.26': 1,
+ 'transformer.h.27': 1,
+ 'transformer.ln_f': 1,
+ 'lm_head': 1}
+ ```
+
+You can also design your `device_map` yourself, if you prefer to explicitly decide where each layer should be. In this case, the command above becomes:
+
+```py
+model = load_checkpoint_and_dispatch(model, "sharded-gpt-j-6B", device_map=my_device_map)
+```
+
+## Run the model
+
+Now that we have done this, our model lies across several devices, and maybe the hard drive. But it can still be used as a regular PyTorch model:
+
+```py
+from transformers import AutoTokenizer
+
+tokenizer = AutoTokenizer.from_pretrained(checkpoint)
+inputs = tokenizer("Hello, my name is", return_tensors="pt")
+inputs = inputs.to(0)
+output = model.generate(inputs["input_ids"])
+tokenizer.decode(output[0].tolist())
+```
+
+Behind the scenes, Accelerate added hooks to the model, so that:
+- at each layer, the inputs are put on the right device (so even if your model is spread across several GPUs, it works)
+- for the weights offloaded on the CPU, they are put on a GPU just before the forward pass, and cleaned up just after
+- for the weights offloaded on the hard drive, they are loaded in RAM then put on a GPU just before the forward pass, and cleaned up just after
+
+This way, you model can run for inference even if it doesn't fit on one of the GPUs or the CPU RAM!
+
+<Tip warning={true}>
+
+This only supports inference of your model, not training. Most of the computation happens behind `torch.no_grad()` context managers to avoid spending some GPU memory with intermediate activations.
+
+</Tip>
+
+## Limits and further development
+
+We are aware of the current limitations in the API:
+
+- While this could theoretically work just one CPU with potential disk offload, you need at least one GPU to run this API. This will be fixed in further development.
+- [`infer_auto_device_map`] (or `device_map="auto"` in [`load_checkpoint_and_dispatch`]) tries to maximize GPU and CPU RAM it sees available when you execute it. While PyTorch is very good at managing GPU RAM efficiently (and giving it back when not needed), it's not entirely true with Python and CPU RAM. Therefore, an automatically computed device map might be too intense on the CPU. Move a few modules to the disk device if you get crashes due to lack of RAM.
+- [`infer_auto_device_map`] (or `device_map="auto"` in [`load_checkpoint_and_dispatch`]) attributes devices sequentially (to avoid moving things back and forth) so if your first layer is bigger than the size of the GPU you have, it will end up with everything on the CPU/Disk.
+- [`load_checkpoint_and_dispatch`] and [`load_checkpoint_in_model`] do not perform any check on the correctness of your state dict compared to your model at the moment (this will be fixed in a future version), so you may get some weird errors if trying to load a checkpoint with mismatched or missing keys.
+- The model parallelism used when your model is split on several GPUs is naive and not optimized, meaning that only one GPU works at a given time and the other sits idle.
+- When weights are offloaded on the CPU/hard drive, there is no pre-fetching (yet, we will work on this for future versions) which means the weights are put on the GPU when they are needed and not before.
+- Hard-drive offloading might be very slow if the hardware you run on does not have fast communication between disk and CPU (like NVMes).
+
+## API doc
+
+[[autodoc]] cpu_offload
+
+[[autodoc]] disk_offload
+
+[[autodoc]] dispatch_model
+
+[[autodoc]] infer_auto_device_map
+
+[[autodoc]] init_empty_weights
+
+[[autodoc]] load_checkpoint_and_dispatch
+
+[[autodoc]] load_checkpoint_in_model
diff --git a/src/accelerate/__init__.py b/src/accelerate/__init__.py
index b74c683b4..2f7ac4af2 100644
--- a/src/accelerate/__init__.py
+++ b/src/accelerate/__init__.py
@@ -5,6 +5,7 @@
 __version__ = "0.8.0.dev0"
 
 from .accelerator import Accelerator
+from .big_modeling import cpu_offload, disk_offload, dispatch_model, init_empty_weights, load_checkpoint_and_dispatch
 from .launchers import debug_launcher, notebook_launcher
 from .utils import (
     DeepSpeedPlugin,
@@ -13,5 +14,7 @@
     GradScalerKwargs,
     InitProcessGroupKwargs,
     find_executable_batch_size,
+    infer_auto_device_map,
+    load_checkpoint_in_model,
     synchronize_rng_states,
 )
diff --git a/src/accelerate/big_modeling.py b/src/accelerate/big_modeling.py
new file mode 100644
index 000000000..9db015427
--- /dev/null
+++ b/src/accelerate/big_modeling.py
@@ -0,0 +1,285 @@
+# 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
+from contextlib import contextmanager
+from typing import Dict, List, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from .hooks import AlignDevicesHook, add_hook_to_module, attach_align_device_hook, attach_align_device_hook_on_blocks
+from .utils import (
+    OffloadedWeightsLoader,
+    check_device_map,
+    extract_submodules_state_dict,
+    infer_auto_device_map,
+    load_checkpoint_in_model,
+    offload_state_dict,
+)
+
+
+@contextmanager
+def init_empty_weights(include_buffers: bool = False):
+    """
+    A context manager under which models are initialized with all parameters on the meta device, therefore creating an
+    empty model. Useful when just initializing the model would blow the available RAM.
+
+    Args:
+        include_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to also put all buffers on the meta device while initializing.
+
+    Example:
+
+    ```pyton
+    import torch.nn as nn
+    from accelerate import init_empty_weights
+
+    # Initialize a model with 100 billions parameters in no time and without using any RAM.
+    with init_empty_weights():
+        tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])
+    ```
+
+    <Tip warning={true}>
+
+    Any model created under this context manager has no weights. As such you can't do something like
+    `model.to(some_device)` with it. To load weights inside your empty model, see [`load_checkpoint_and_dispatch`].
+
+    </Tip>
+    """
+    old_register_parameter = nn.Module.register_parameter
+    if include_buffers:
+        old_register_buffer = nn.Module.register_buffer
+
+    def register_empty_parameter(module, name, param):
+        old_register_parameter(module, name, param)
+        if param is not None:
+            module._parameters[name] = nn.Parameter(module._parameters[name].to(torch.device("meta")))
+
+    def register_empty_buffer(module, name, buffer):
+        old_register_buffer(module, name, buffer)
+        if buffer is not None:
+            module._buffers[name] = module._buffers[name].to(torch.device("meta"))
+
+    try:
+        nn.Module.register_parameter = register_empty_parameter
+        if include_buffers:
+            nn.Module.register_buffer = register_empty_buffer
+        yield
+    finally:
+        nn.Module.register_parameter = old_register_parameter
+        if include_buffers:
+            nn.Module.register_buffer = old_register_buffer
+
+
+def cpu_offload(
+    model: nn.Module,
+    execution_device: Optional[torch.device] = None,
+    offload_buffers: bool = False,
+    state_dict: Optional[Dict[str, torch.Tensor]] = None,
+):
+    """
+    Activates full CPU offload for a model. As a result, all parameters of the model will be offloaded and only one
+    copy of the state dict of the model will be kept. During the forward pass, parameters will be extracted from that
+    state dict and put on the execution device passed as they are needed, then offloaded again.
+
+    Args:
+        model (`torch.nn.Module`):
+            The model to offload.
+        execution_device (`torch.device`, *optional*):
+            The device on which the forward pass of the model will be executed (should be a GPU). Will default to the
+            model first parameter device.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+        state_dict (`Dict[str, torch.Tensor]`, *optional*):
+            The state dict of the model that will be kept on CPU.
+    """
+    if execution_device is None:
+        execution_device = next(iter(model.parameters())).device
+    if state_dict is None:
+        state_dict = {n: p.to("cpu") for n, p in model.state_dict().items()}
+    attach_align_device_hook(
+        model, execution_device=execution_device, offload=True, offload_buffers=offload_buffers, weights_map=state_dict
+    )
+    add_hook_to_module(model, AlignDevicesHook(io_same_device=True))
+    return model
+
+
+def disk_offload(
+    model: nn.Module,
+    offload_dir: Union[str, os.PathLike],
+    execution_device: Optional[torch.device] = None,
+    offload_buffers: bool = False,
+):
+    """
+    Activates full disk offload for a model. As a result, all parameters of the model will be offloaded as
+    memory-mapped array in a given folder. During the forward pass, parameters will be accessed from that folder and
+    put on the execution device passed as they are needed, then offloaded again.
+
+    Args:
+        model (`torch.nn.Module`): The model to offload.
+        offload_dir (`str` or `os.PathLike`):
+            The folder in which to offload the model weights (or where the model weights are already offloaded).
+        execution_device (`torch.device`, *optional*):
+            The device on which the forward pass of the model will be executed (should be a GPU). Will default to the
+            model's first parameter device.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+    """
+    if not os.path.isdir(offload_dir) or not os.path.isfile(os.path.join(offload_dir, "index.json")):
+        offload_state_dict(offload_dir, model.state_dict())
+    if execution_device is None:
+        execution_device = next(iter(model.parameters())).device
+    weights_map = OffloadedWeightsLoader(save_folder=offload_dir)
+    attach_align_device_hook(
+        model,
+        execution_device=execution_device,
+        offload=True,
+        offload_buffers=offload_buffers,
+        weights_map=weights_map,
+    )
+    add_hook_to_module(model, AlignDevicesHook(io_same_device=True))
+    return model
+
+
+def dispatch_model(
+    model: nn.Module,
+    device_map: Dict[str, Union[str, int, torch.device]],
+    main_device: Optional[torch.device] = None,
+    state_dict: Optional[Dict[str, torch.Tensor]] = None,
+    offload_dir: Union[str, os.PathLike] = None,
+    offload_buffers: bool = False,
+):
+    """
+    Dispatches a model according to a given device map. Layers of the model might be spread across GPUs, offloaded on
+    the CPU or even the disk.
+
+    Args:
+        model (`torch.nn.Module`):
+            The model to dispatch.
+        device_map (`Dict[str, Union[str, int, torch.device]]`):
+            A dictionary mapping module names in the models `state_dict` to the device they should go to. Note that
+            `"disk"` is accepted even if it's not a proper value for `torch.device`.
+        main_device (`str`, `int` or `torch.device`, *optional*):
+            The main execution device. Will default to the first device in the `device_map` different from `"cpu"` or
+            `"disk"`.
+        state_dict (`Dict[str, torch.Tensor]`, *optional*):
+            The state dict of the part of the model that will be kept on CPU.
+        offload_dir (`str` or `os.PathLike`):
+            The folder in which to offload the model weights (or where the model weights are already offloaded).
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+    """
+    # Error early if the device map is incomplete.
+    check_device_map(model, device_map)
+
+    if main_device is None:
+        main_device = [d for d in device_map.values() if d not in ["cpu", "disk"]][0]
+
+    cpu_modules = [name for name, device in device_map.items() if device == "cpu"]
+    if state_dict is None and len(cpu_modules) > 0:
+        state_dict = extract_submodules_state_dict(model.state_dict(), cpu_modules)
+
+    disk_modules = [name for name, device in device_map.items() if device == "disk"]
+    if offload_dir is None and len(disk_modules) > 0:
+        raise ValueError(
+            "We need an `offload_dir` to dispatch this model according to this `device_map`, the following submodules "
+            f"need to be offloaded: {', '.join(disk_modules)}."
+        )
+    if len(disk_modules) > 0 and (
+        not os.path.isdir(offload_dir) or not os.path.isfile(os.path.join(offload_dir, "index.json"))
+    ):
+        disk_state_dict = extract_submodules_state_dict(model.state_dict(), disk_modules)
+        offload_state_dict(offload_dir, disk_state_dict)
+
+    execution_device = {
+        name: main_device if device in ["cpu", "disk"] else device for name, device in device_map.items()
+    }
+    offload = {name: device in ["cpu", "disk"] for name, device in device_map.items()}
+    save_folder = offload_dir if len(disk_modules) > 0 else None
+    if state_dict is not None or save_folder is not None:
+        weights_map = OffloadedWeightsLoader(state_dict=state_dict, save_folder=save_folder)
+    else:
+        weights_map = None
+
+    attach_align_device_hook_on_blocks(
+        model,
+        execution_device=execution_device,
+        offload=offload,
+        offload_buffers=offload_buffers,
+        weights_map=weights_map,
+    )
+    model.hf_device_map = device_map
+    return model
+
+
+def load_checkpoint_and_dispatch(
+    model: nn.Module,
+    checkpoint: Union[str, os.PathLike],
+    device_map: Optional[Union[str, Dict[str, Union[int, str, torch.device]]]] = None,
+    max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None,
+    no_split_module_classes: Optional[List[str]] = None,
+    offload_folder: Optional[Union[str, os.PathLike]] = None,
+    offload_buffers: bool = False,
+    dtype: Optional[Union[str, torch.dtype]] = None,
+    offload_state_dict: bool = False,
+):
+    """
+    Loads a (potentially sharded) checkpoint inside a model, potentially sending weights to a given device as they are
+    loaded and adds the various hooks that will make this model run properly (even if split across devices).
+
+    Args:
+        model (`torch.nn.Module`): The model in which we want to load a checkpoint.
+        checkpoint (`str` or `os.PathLike`):
+            The folder checkpoint to load. It can be:
+            - a path to a file containing a whole model state dict
+            - a path to a `.json` file containing the index to a sharded checkpoint
+            - a path to a folder containing a unique `.index.json` file and the shards of a checkpoint.
+        device_map (`Dict[str, Union[int, str, torch.device]]`, *optional*):
+            A map that specifies where each submodule should go. It doesn't need to be refined to each parameter/buffer
+            name, once a given module name is inside, every submodule of it will be sent to the same device.
+
+            To have Accelerate compute the most optimized `device_map` automatically, set `device_map="auto"`.
+        max_memory (`Dict`, *optional*):
+            A dictionary device identifier to maximum memory. Will default to the maximum memory available for each GPU
+            and the available CPU RAM if unset.
+        no_split_module_classes (`List[str]`, *optional*):
+            A list of layer class names that should never be split across device (for instance any layer that has a
+            residual connection).
+        offload_folder (`str` or `os.PathLike`, *optional*):
+            If the `device_map` contains any value `"disk"`, the folder where we will offload weights.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            In the layers that are offloaded on the CPU or the hard drive, whether or not to offload the buffers as
+            well as the parameters.
+        dtype (`str` or `torch.dtype`, *optional*):
+            If provided, the weights will be converted to that type when loaded.
+        offload_state_dict (`bool`, *optional*, defaults to `False`):
+            If `True`, will temporarily offload the CPU state dict on the hard drive to avoig getting out of CPU RAM if
+            the weight of the CPU state dict + the biggest shard does not fit.
+    """
+    if device_map == "auto":
+        device_map = infer_auto_device_map(
+            model, max_memory=max_memory, no_split_module_classes=no_split_module_classes, dtype=dtype
+        )
+    load_checkpoint_in_model(
+        model,
+        checkpoint,
+        device_map=device_map,
+        offload_folder=offload_folder,
+        dtype=dtype,
+        offload_state_dict=offload_state_dict,
+    )
+    if device_map is None:
+        return model
+    return dispatch_model(model, device_map=device_map, offload_dir=offload_folder, offload_buffers=offload_buffers)
diff --git a/src/accelerate/hooks.py b/src/accelerate/hooks.py
new file mode 100644
index 000000000..0a5bca96d
--- /dev/null
+++ b/src/accelerate/hooks.py
@@ -0,0 +1,411 @@
+# 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 functools
+from typing import Dict, Mapping, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from .utils import PrefixedDataset, find_device, named_module_tensors, send_to_device, set_module_tensor_to_device
+
+
+class ModelHook:
+    """
+    A hook that contains callbacks to be executed just before and after the forward method of a model. The difference
+    with PyTorch existing hooks is that they get passed along the kwargs.
+
+    Class attribute:
+    - **no_grad** (`bool`, *optional*, defaults to `False`) -- Whether or not to execute the actual forward pass under
+      the `torch.no_grad()` context manager.
+    """
+
+    no_grad = False
+
+    def init_hook(self, module):
+        """
+        To be executed when the hook is attached to the module.
+
+        Args:
+            module (`torch.nn.Module`): The module attached to this hook.
+        """
+        return module
+
+    def pre_forward(self, module, *args, **kwargs):
+        """
+        To be executed just before the forward method of the model.
+
+        Args:
+            module (`torch.nn.Module`): The module whose forward pass will be executed just after this event.
+            args (`Tuple[Any]`): The positional arguments passed to the module.
+            kwargs (`Dict[Str, Any]`): The keyword arguments passed to the module.
+
+        Returns:
+            `Tuple[Tuple[Any], Dict[Str, Any]]`: A tuple with the treated `args` and `kwargs`.
+        """
+        return args, kwargs
+
+    def post_forward(self, module, output):
+        """
+        To be executed just after the forward method of the model.
+
+        Args:
+            module (`torch.nn.Module`): The module whose forward pass been executed just before this event.
+            output (`Any`): The output of the module.
+
+        Returns:
+            `Any`: The processed `output`.
+        """
+        return output
+
+    def detach_hook(self, module):
+        """
+        To be executed when the hook is deached from a module.
+
+        Args:
+            module (`torch.nn.Module`): The module detached from this hook.
+        """
+        return module
+
+
+class SequentialHook(ModelHook):
+    """
+    A hook that can contain several hooks and iterates through them at each event.
+    """
+
+    def __init__(self, *hooks):
+        self.hooks = hooks
+
+    def init_hook(self, module):
+        for hook in self.hooks:
+            module = hook.init_hook(module)
+        return module
+
+    def pre_forward(self, module, *args, **kwargs):
+        for hook in self.hooks:
+            args, kwargs = hook.pre_forward(module, *args, **kwargs)
+        return args, kwargs
+
+    def post_forward(self, module, output):
+        for hook in self.hooks:
+            output = hook.post_forward(module, output)
+        return output
+
+    def detach_hook(self, module):
+        for hook in self.hooks:
+            module = hook.detach_hook(module)
+        return module
+
+
+def add_hook_to_module(module: nn.Module, hook: ModelHook):
+    """
+    Adds a hook to a given module. This will rewrite the `forward` method of the module to include the hook, to remove
+    this behavior and restore the original `forward` method, use `remove_hook_from_module`.
+
+    <Tip warning={true}>
+
+    If the module already contains a hook, this will replace it with the new hook passed. To chain two hooks together,
+    use the `SequentialHook` class.
+
+    </Tip>
+
+    Args:
+        module (`torch.nn.Module`): The module to attach a hook to.
+        hook (`ModelHook`): The hook to attach.
+
+    Returns:
+        `torch.nn.Module`: The same module, with the hook attached (the module is modified in place, so the result can
+        be discarded).
+    """
+    if hasattr(module, "_hf_hook") and hasattr(module, "_old_forward"):
+        # If we already put some hook on this module, we replace it with the new one.
+        old_forward = module._old_forward
+    else:
+        old_forward = module.forward
+        module._old_forward = old_forward
+
+    module = hook.init_hook(module)
+    module._hf_hook = hook
+
+    @functools.wraps(old_forward)
+    def new_forward(*args, **kwargs):
+        args, kwargs = module._hf_hook.pre_forward(module, *args, **kwargs)
+        if module._hf_hook.no_grad:
+            with torch.no_grad():
+                output = old_forward(*args, **kwargs)
+        else:
+            output = old_forward(*args, **kwargs)
+        return module._hf_hook.post_forward(module, output)
+
+    module.forward = new_forward
+    return module
+
+
+def remove_hook_from_module(module: nn.Module):
+    """
+    Removes any hook attached to a module via `add_hook_to_module`.
+
+    Args:
+        module (`torch.nn.Module`): The module to attach a hook to.
+
+    Returns:
+        `torch.nn.Module`: The same module, with the hook detached (the module is modified in place, so the result can
+        be discarded).
+    """
+    if hasattr(module, "_hf_hook"):
+        module._hf_hook.detach_hook(module)
+        delattr(module, "_hf_hook")
+
+    if hasattr(module, "_old_forward"):
+        module.forward = module._old_forward
+        delattr(module, "_old_forward")
+
+    return module
+
+
+class AlignDevicesHook(ModelHook):
+    """
+    A generic `ModelHook` that ensures inputs and model weights are on the same device for the forward pass of the
+    associated module, potentially offloading the weights after the forward pass.
+
+    Args:
+        execution_device (`torch.device`, *optional*):
+            The device on which inputs and model weights should be placed before the forward pass.
+        offload (`bool`, *optional*, defauts to `False`):
+            Whether or not the weights should be offloaded after the forward pass.
+        io_same_device (`bool`, *optional*, defaults to `False`):
+            Whether or not the output should be placed on the same device as the input was.
+        weights_map (`Mapping[str, torch.Tensor]`, *optional*):
+            When the model weights are offloaded, a (potentially lazy) map from param names to the tensor values.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to include the associated module's buffers when offloading.
+        place_submodules (`bool`, *optional*, defaults to `False`):
+            Whether to place the submodules on `execution_device` during the `init_hook` event.
+    """
+
+    def __init__(
+        self,
+        execution_device: Optional[Union[int, str, torch.device]] = None,
+        offload: bool = False,
+        io_same_device: bool = False,
+        weights_map: Optional[Mapping] = None,
+        offload_buffers: bool = False,
+        place_submodules: bool = False,
+    ):
+        self.execution_device = execution_device
+        self.offload = offload
+        self.io_same_device = io_same_device
+        self.weights_map = weights_map
+        self.offload_buffers = offload_buffers
+        self.place_submodules = place_submodules
+
+        # Will contain the input device when `io_same_device=True`.
+        self.input_device = None
+        self.param_original_devices = {}
+        self.buffer_original_devices = {}
+
+    def init_hook(self, module):
+        if not self.offload and self.execution_device is not None:
+            for name, _ in named_module_tensors(module, recurse=self.place_submodules):
+                set_module_tensor_to_device(module, name, self.execution_device)
+        elif self.offload:
+            self.original_devices = {name: param.device for name, param in named_module_tensors(module)}
+            if self.weights_map is None:
+                self.weights_map = {
+                    name: param.to("cpu")
+                    for name, param in named_module_tensors(module, include_buffers=self.offload_buffers)
+                }
+
+            for name, _ in named_module_tensors(module, include_buffers=self.offload_buffers):
+                set_module_tensor_to_device(module, name, "meta")
+            if not self.offload_buffers and self.execution_device is not None:
+                for name, _ in module.named_buffers(recurse=False):
+                    set_module_tensor_to_device(module, name, self.execution_device)
+        return module
+
+    def pre_forward(self, module, *args, **kwargs):
+        if self.io_same_device:
+            self.input_device = find_device([args, kwargs])
+        if self.offload:
+            for name, _ in named_module_tensors(module, include_buffers=self.offload_buffers):
+                set_module_tensor_to_device(module, name, self.execution_device, value=self.weights_map[name])
+
+        return send_to_device(args, self.execution_device), send_to_device(kwargs, self.execution_device)
+
+    def post_forward(self, module, output):
+        if self.offload:
+            for name, _ in named_module_tensors(module, include_buffers=self.offload_buffers):
+                set_module_tensor_to_device(module, name, "meta")
+
+        if self.io_same_device and self.input_device is not None:
+            output = send_to_device(output, self.input_device)
+
+        return output
+
+    def detach_hook(self, module):
+        if self.offload:
+            for name, device in self.original_devices.items():
+                if device != torch.device("meta"):
+                    set_module_tensor_to_device(module, name, device, value=self.weights_map.get(name, None))
+
+
+def attach_align_device_hook(
+    module: torch.nn.Module,
+    execution_device: Optional[torch.device] = None,
+    offload: bool = False,
+    weights_map: Optional[Mapping] = None,
+    offload_buffers: bool = False,
+    module_name: str = "",
+):
+    """
+    Recursively attaches `AlignDevicesHook` to all submodules of a given model that have direct parameters and/or
+    buffers.
+
+    Args:
+        module (`torch.nn.Module`):
+            The module where we want to attach the hooks.
+        execution_device (`torch.device`, *optional*):
+            The device on which inputs and model weights should be placed before the forward pass.
+        offload (`bool`, *optional*, defauts to `False`):
+            Whether or not the weights should be offloaded after the forward pass.
+        weights_map (`Mapping[str, torch.Tensor]`, *optional*):
+            When the model weights are offloaded, a (potentially lazy) map from param names to the tensor values.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to include the associated module's buffers when offloading.
+        module_name (`str`, *optional*, defaults to `""`):
+            The name of the module.
+    """
+    # Attach the hook on this module if it has any direct tensor.
+    directs = named_module_tensors(module)
+    if len(list(directs)) > 0:
+        if weights_map is not None:
+            prefix = f"{module_name}." if len(module_name) > 0 else ""
+            prefixed_weights_map = PrefixedDataset(weights_map, prefix)
+        else:
+            prefixed_weights_map = None
+        hook = AlignDevicesHook(
+            execution_device=execution_device,
+            offload=offload,
+            weights_map=prefixed_weights_map,
+            offload_buffers=offload_buffers,
+        )
+        add_hook_to_module(module, hook)
+
+    # Recurse on all children of the module.
+    for child_name, child in module.named_children():
+        child_name = f"{module_name}.{child_name}" if len(module_name) > 0 else child_name
+        attach_align_device_hook(
+            child,
+            execution_device=execution_device,
+            offload=offload,
+            weights_map=weights_map,
+            offload_buffers=offload_buffers,
+            module_name=child_name,
+        )
+
+
+def remove_hook_from_submodules(module: nn.Module):
+    """
+    Recursively removes all hooks attached on the submodules of a given model.
+
+    Args:
+        module (`torch.nn.Module`): The module on which to remove all hooks.
+    """
+    remove_hook_from_module(module)
+    for child in module.children():
+        remove_hook_from_submodules(child)
+
+
+def attach_align_device_hook_on_blocks(
+    module: nn.Module,
+    execution_device: Optional[Union[torch.device, Dict[str, torch.device]]] = None,
+    offload: Union[bool, Dict[str, bool]] = False,
+    weights_map: Mapping = None,
+    offload_buffers: bool = False,
+    module_name: str = "",
+):
+    """
+    Attaches `AlignDevicesHook` to all blocks of a given model as needed.
+
+    Args:
+        module (`torch.nn.Module`):
+            The module where we want to attach the hooks.
+        execution_device (`torch.device` or `Dict[str, torch.device]`, *optional*):
+            The device on which inputs and model weights should be placed before the forward pass. It can be one device
+            for the whole module, or a dictionary mapping module name to device.
+        offload (`bool`, *optional*, defauts to `False`):
+            Whether or not the weights should be offloaded after the forward pass. It can be one boolean for the whole
+            module, or a dictionary mapping module name to boolean.
+        weights_map (`Mapping[str, torch.Tensor]`, *optional*):
+            When the model weights are offloaded, a (potentially lazy) map from param names to the tensor values.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to include the associated module's buffers when offloading.
+        module_name (`str`, *optional*, defaults to `""`):
+            The name of the module.
+    """
+    # If one device and one offload, we've got one hook.
+    if not isinstance(execution_device, Mapping) and not isinstance(offload, dict):
+        if not offload:
+            hook = AlignDevicesHook(execution_device=execution_device, io_same_device=True, place_submodules=True)
+            add_hook_to_module(module, hook)
+        else:
+            attach_align_device_hook(
+                module,
+                execution_device=execution_device,
+                offload=True,
+                weights_map=weights_map,
+                offload_buffers=offload_buffers,
+                module_name=module_name,
+            )
+        return
+
+    if not isinstance(execution_device, Mapping):
+        execution_device = {key: offload for key in offload.keys()}
+    if not isinstance(offload, Mapping):
+        offload = {key: offload for key in execution_device.keys()}
+
+    if module_name in execution_device and not offload[module_name]:
+        hook = AlignDevicesHook(
+            execution_device=execution_device[module_name],
+            offload_buffers=offload_buffers,
+            io_same_device=(module_name == ""),
+            place_submodules=True,
+        )
+        add_hook_to_module(module, hook)
+    elif module_name in execution_device:
+        attach_align_device_hook(
+            module,
+            execution_device=execution_device[module_name],
+            offload=True,
+            weights_map=weights_map,
+            offload_buffers=offload_buffers,
+            module_name=module_name,
+        )
+        if not hasattr(module, "_hf_hook"):
+            hook = AlignDevicesHook(execution_device=execution_device[module_name], io_same_device=(module_name == ""))
+            add_hook_to_module(module, hook)
+    elif module_name == "":
+        hook = AlignDevicesHook(io_same_device=True)
+        add_hook_to_module(module, hook)
+
+    for child_name, child in module.named_children():
+        child_name = f"{module_name}.{child_name}" if len(module_name) > 0 else child_name
+        attach_align_device_hook_on_blocks(
+            child,
+            execution_device=execution_device,
+            offload=offload,
+            weights_map=weights_map,
+            offload_buffers=offload_buffers,
+            module_name=child_name,
+        )
diff --git a/src/accelerate/test_utils/__init__.py b/src/accelerate/test_utils/__init__.py
index 03aef45b3..46723ebfa 100644
--- a/src/accelerate/test_utils/__init__.py
+++ b/src/accelerate/test_utils/__init__.py
@@ -2,5 +2,5 @@
 # There's no way to ignore "F401 '...' imported but unused" warnings in this
 # module, but to preserve other warnings. So, don't check this module at all.
 
-from .testing import are_the_same_tensors, execute_subprocess_async, require_cuda, require_multi_gpu, require_tpu
+from .testing import are_the_same_tensors, execute_subprocess_async, require_cuda, require_multi_gpu, require_tpu, slow
 from .training import RegressionDataset, RegressionModel
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
index a43c7d261..58b33c1ad 100644
--- a/src/accelerate/utils/__init__.py
+++ b/src/accelerate/utils/__init__.py
@@ -30,6 +30,20 @@
     is_tpu_available,
     is_wandb_available,
 )
+from .modeling import (
+    check_device_map,
+    compute_module_sizes,
+    convert_file_size_to_int,
+    dtype_byte_size,
+    find_tied_parameters,
+    get_max_layer_size,
+    get_max_memory,
+    infer_auto_device_map,
+    load_checkpoint_in_model,
+    named_module_tensors,
+    set_module_tensor_to_device,
+)
+from .offload import OffloadedWeightsLoader, PrefixedDataset, extract_submodules_state_dict, offload_state_dict
 from .operations import (
     broadcast,
     broadcast_object_list,
@@ -37,6 +51,7 @@
     convert_outputs_to_fp32,
     convert_to_fp32,
     find_batch_size,
+    find_device,
     gather,
     gather_object,
     get_data_structure,
diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
new file mode 100644
index 000000000..08c6f269d
--- /dev/null
+++ b/src/accelerate/utils/modeling.py
@@ -0,0 +1,624 @@
+# 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 gc
+import json
+import os
+import re
+import shutil
+import tempfile
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+WEIGHTS_INDEX_NAME = "pytorch_model.bin.index.json"
+
+
+def convert_file_size_to_int(size: Union[int, str]):
+    """
+    Converts a size expressed as a string with digits an unit (like `"5MB"`) to an integer (in bytes).
+
+    Args:
+        size (`int` or `str`): The size to convert. Will be directly returned if an `int`.
+
+    Example:
+
+    ```py
+    >>> convert_file_size_to_int("1MiB")
+    1048576
+    ```
+    """
+    if isinstance(size, int):
+        return size
+    if size.upper().endswith("GIB"):
+        return int(size[:-3]) * (2**30)
+    if size.upper().endswith("MIB"):
+        return int(size[:-3]) * (2**20)
+    if size.upper().endswith("KIB"):
+        return int(size[:-3]) * (2**10)
+    if size.upper().endswith("GB"):
+        int_size = int(size[:-2]) * (10**9)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("MB"):
+        int_size = int(size[:-2]) * (10**6)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("KB"):
+        int_size = int(size[:-2]) * (10**3)
+        return int_size // 8 if size.endswith("b") else int_size
+    raise ValueError("`size` is not in a valid format. Use an integer followed by the unit, e.g., '5GB'.")
+
+
+def dtype_byte_size(dtype: torch.dtype):
+    """
+    Returns the size (in bytes) occupied by one parameter of type `dtype`.
+
+    Example:
+
+    ```py
+    >>> dtype_byte_size(torch.float32)
+    4
+    ```
+    """
+    if dtype == torch.bool:
+        return 1 / 8
+    bit_search = re.search("[^\d](\d+)$", str(dtype))
+    if bit_search is None:
+        raise ValueError(f"`dtype` is not a valid dtype: {dtype}.")
+    bit_size = int(bit_search.groups()[0])
+    return bit_size // 8
+
+
+def set_module_tensor_to_device(
+    module: nn.Module, tensor_name: str, device: Union[int, str, torch.device], value: Optional[torch.Tensor] = None
+):
+    """
+    A helper function to set a given tensor (parameter of buffer) of a module on a specific device (note that doing
+    `param.to(device)` creates a new tensor not linked to the parameter, which is why we need this function).
+
+    Args:
+        module (`torch.nn.Module`): The module in which the tensor we want to move lives.
+        param_name (`str`): The full name of the parameter/buffer.
+        device (`int`, `str` or `torch.device`): The device on which to set the tensor.
+        value (`torch.Tensor`, *optional*): The value of the tensor (useful when going from the meta device to any
+            other device).
+    """
+    # Recurse if needed
+    if "." in tensor_name:
+        splits = tensor_name.split(".")
+        for split in splits[:-1]:
+            new_module = getattr(module, split)
+            if new_module is None:
+                raise ValueError(f"{module} has no attribute {split}.")
+            module = new_module
+        tensor_name = splits[-1]
+
+    if tensor_name not in module._parameters and tensor_name not in module._buffers:
+        raise ValueError(f"{module} does not have a parameter or a buffer named {tensor_name}.")
+    is_buffer = tensor_name in module._buffers
+    old_value = getattr(module, tensor_name)
+
+    if old_value.device == torch.device("meta") and device not in ["meta", torch.device("meta")] and value is None:
+        raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {device}.")
+
+    with torch.no_grad():
+        if value is None:
+            new_value = old_value.to(device)
+        elif isinstance(value, torch.Tensor):
+            new_value = value.to(device)
+        else:
+            new_value = torch.tensor(value, device=device)
+    if is_buffer:
+        module._buffers[tensor_name] = new_value
+    else:
+        new_value = nn.Parameter(new_value, requires_grad=old_value.requires_grad)
+        module._parameters[tensor_name] = new_value
+
+
+def named_module_tensors(module: nn.Module, include_buffers: bool = True, recurse: bool = False):
+    """
+    A helper function that gathers all the tensors (parameters + buffers) of a given module. If `include_buffers=True`
+    it's the same as doing `module.named_parameters(recurse=recurse) + module.named_buffers(recurse=recurse)`.
+
+    Args:
+        module (`torch.nn.Module`): The module we want the tensors or.
+        include_buffer (`bool`, *optional*, defaults to `True`): Whether or not to include the buffers in the result.
+        recurse (`bool`, *optional`, defaults to `False`):
+            Whether or not to go look in every submodule or just return the direct parameters and buffers.
+    """
+    for named_parameter in module.named_parameters(recurse=recurse):
+        yield named_parameter
+
+    if include_buffers:
+        for named_buffer in module.named_buffers(recurse=recurse):
+            yield named_buffer
+
+
+def find_tied_parameters(model: nn.Module, **kwargs):
+    """
+    Find the tied parameters in a given model.
+
+    Args:
+        model (`torch.nn.Module`): The model to inspect.
+
+    <Tip warning={true}>
+
+    The signature accepts keyword arguments, but they are for the recursive part of this function and you should ignore
+    them.
+
+    </Tip>
+
+    Example:
+
+
+    ```py
+    >>> from collections import OrderedDict
+    >>> import torch.nn as nn
+
+    >>> model = nn.Sequential(OrderedDict([("linear1", nn.Linear(4, 4)), ("linear2", nn.Linear(4, 4))]))
+    >>> model.linear2.weight = test_model.linear1.weight
+    >>> find_tied_parameters(test_model)
+    {'linear1.weight': 'linear2.weight'}
+    ```
+
+    Returns:
+        Dict[str, str]: A dictionary mapping tied parameter names to the name of the parameter they are tied to.
+    """
+    # Initialize result and named_parameters before recursing.
+    named_parameters = kwargs.get("named_parameters", None)
+    prefix = kwargs.get("prefix", "")
+    result = kwargs.get("result", {})
+
+    if named_parameters is None:
+        named_parameters = {n: p for n, p in model.named_parameters()}
+    else:
+        # A tied parameter will not be in the full `named_parameters` seen above but will be in the `named_parameters`
+        # of the submodule it belongs to. So while recursing we track the names that are not in the initial
+        # `named_parameters`.
+        for name, parameter in model.named_parameters():
+            full_name = name if prefix == "" else f"{prefix}.{name}"
+            if full_name not in named_parameters:
+                # When we find one, it has to be one of the existing parameters.
+                for new_name, new_param in named_parameters.items():
+                    if new_param is parameter:
+                        result[new_name] = full_name
+
+    # Once we have treated direct parameters, we move to the child modules.
+    for name, child in model.named_children():
+        child_name = name if prefix == "" else f"{prefix}.{name}"
+        find_tied_parameters(child, named_parameters=named_parameters, prefix=child_name, result=result)
+
+    return result
+
+
+def compute_module_sizes(model: nn.Module, dtype: Optional[Union[str, torch.device]] = None):
+    """
+    Compute the size of each submodule of a given model.
+    """
+    if isinstance(dtype, str):
+        # We accept "torch.float16" or just "float16"
+        dtype = dtype.replace("torch.", "")
+        dtype = getattr(torch, dtype)
+    if dtype is not None:
+        dtype_size = dtype_byte_size(dtype)
+    module_sizes = defaultdict(int)
+    for name, tensor in named_module_tensors(model, recurse=True):
+        if dtype is None:
+            size = tensor.numel() * dtype_byte_size(tensor.dtype)
+        else:
+            size = tensor.numel() * min(dtype_size, dtype_byte_size(tensor.dtype))
+        name_parts = name.split(".")
+        for idx in range(len(name_parts) + 1):
+            module_sizes[".".join(name_parts[:idx])] += size
+
+    return module_sizes
+
+
+def get_max_layer_size(
+    modules: List[Tuple[str, torch.nn.Module]], module_sizes: Dict[str, int], no_split_module_classes: List[str]
+):
+    """
+    Utility function that will scan a list of named modules and return the maximum size used by one full layer. The
+    definition of a layer being:
+    - a module with no direct children (just parameters and buffers)
+    - a module whose class name is in the list `no_split_module_classes`
+
+    Args:
+        modules (`List[Tuple[str, torch.nn.Module]]`):
+            The list of named modules where we want to determine the maximum layer size.
+        module_sizes (`Dict[str, int]`):
+            A dictionary mapping each layer name to its size (as generated by `compute_module_sizes`).
+        no_split_module_classes (`List[str]`):
+            A list of class names for layers we don't want to be split.
+
+    Returns:
+        `Tuple[int, List[str]]`: The maximum size of a layer with the list of layer names realizing that maximum size.
+    """
+    max_size = 0
+    layer_names = []
+    modules_to_treat = modules.copy()
+    while len(modules_to_treat) > 0:
+        module_name, module = modules_to_treat.pop(0)
+        modules_children = list(module.named_children())
+        if len(modules_children) == 0 or module.__class__.__name__ in no_split_module_classes:
+            # No splitting this one so we compare to the max_size
+            size = module_sizes[module_name]
+            if size > max_size:
+                max_size = size
+                layer_names = [module_name]
+            elif size == max_size:
+                layer_names.append(module_name)
+        else:
+            modules_to_treat = [(f"{module_name}.{n}", v) for n, v in modules_children] + modules_to_treat
+    return max_size, layer_names
+
+
+def get_max_memory(max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None):
+    """
+    Get the maximum memory available if nothing is passed, converts string to int otherwise.
+    """
+    import psutil
+
+    if max_memory is None:
+        if not torch.cuda.is_available():
+            max_memory = {}
+        else:
+            # Make sure CUDA is initialized on each GPU to have the right memory info.
+            for i in range(torch.cuda.device_count()):
+                _ = torch.tensor([0], device=i)
+            max_memory = {i: torch.cuda.mem_get_info(i)[0] for i in range(torch.cuda.device_count())}
+        max_memory["cpu"] = psutil.virtual_memory().available
+        return max_memory
+
+    for key in max_memory:
+        if isinstance(max_memory[key], str):
+            max_memory[key] = convert_file_size_to_int(max_memory[key])
+    return max_memory
+
+
+def clean_device_map(device_map: Dict[str, Union[int, str, torch.device]], module_name: str = ""):
+    """
+    Cleans a device_map by grouping all submodules that go on the same device together.
+    """
+    # Get the value of the current module and if there is only one split across several keys, regroup it.
+    prefix = "" if module_name == "" else f"{module_name}."
+    values = [v for k, v in device_map.items() if k.startswith(prefix)]
+    if len(set(values)) == 1 and len(values) > 1:
+        for k in [k for k in device_map if k.startswith(prefix)]:
+            del device_map[k]
+        device_map[module_name] = values[0]
+
+    # Recurse over the children
+    children_modules = [k for k in device_map.keys() if k.startswith(module_name) and len(k) > len(module_name)]
+    idx = len(module_name.split(".")) + 1 if len(module_name) > 0 else 1
+    children_modules = set(".".join(k.split(".")[:idx]) for k in children_modules)
+    for child in children_modules:
+        clean_device_map(device_map, module_name=child)
+
+    return device_map
+
+
+def infer_auto_device_map(
+    model: nn.Module,
+    max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None,
+    no_split_module_classes: Optional[List[str]] = None,
+    dtype: Optional[Union[str, torch.dtype]] = None,
+):
+    """
+    Compute a device map for a given model giving priority to GPUs, then offload on CPU and finally offload to disk,
+    such that:
+    - we don't exceed the memory available of any of the GPU.
+    - if offload to the CPU is needed, there is always room left on GPU 0 to put back the layer offloaded on CPU that
+      has the largest size.
+    - if offload to the CPU is needed,we don't exceed the RAM available on the CPU.
+    - if offload to the disk is needed, there is always room left on the CPU to put back the layer offloaded on disk
+      that has the largest size.
+
+    <Tip>
+
+    All computation is done analyzing sizes and dtypes of the model parameters. As a result, the model can be on the
+    meta device (as it would if initialized within the `init_empty_weights` context manager).
+
+    </Tip>
+
+    Args:
+        model (`torch.nn.Module`): The model to analyze.
+        max_memory (`Dict`, *optional*):
+            A dictionary device identifier to maximum memory. Will default to the maximum memory available if unset.
+        no_split_module_classes (`List[str]`, *optional*):
+            A list of layer class names that should never be split across device (for instance any layer that has a
+            residual connection).
+        dtype (`str` or `torch.dtype`, *optional*):
+            If provided, the weights will be converted to that type when loaded.
+    """
+    # Get default / clean up max_memory
+    max_memory = get_max_memory(max_memory)
+    if no_split_module_classes is None:
+        no_split_module_classes = []
+    elif not isinstance(no_split_module_classes, (list, tuple)):
+        no_split_module_classes = [no_split_module_classes]
+
+    devices = list(max_memory.keys())
+    gpus = [device for device in devices if device != "cpu"]
+    if "disk" not in devices:
+        devices.append("disk")
+
+    # Devices that need to keep space for a potential offloaded layer.
+    main_devices = [gpus[0], "cpu"] if len(gpus) > 0 else ["cpu"]
+
+    module_sizes = compute_module_sizes(model, dtype=dtype)
+    tied_parameters = find_tied_parameters(model)
+
+    device_map = {}
+    current_device = 0
+    current_memory_used = 0
+
+    # Direct submodules and parameters
+    modules_to_treat = list(model.named_parameters(recurse=False)) + list(model.named_children())
+    # Initialize maximum largest layer, to know which space to keep in memory
+    max_layer_size, max_layer_names = get_max_layer_size(modules_to_treat, module_sizes, no_split_module_classes)
+
+    # Ready ? This is going to be a bit messy.
+    while len(modules_to_treat) > 0:
+        name, module = modules_to_treat.pop(0)
+        # Max size in the remaining layers may have changed since we took one, so we maybe update it.
+        max_layer_names = [n for n in max_layer_names if not n.startswith(name)]
+        if len(max_layer_names) == 0:
+            max_layer_size, max_layer_names = get_max_layer_size(
+                [(n, m) for n, m in modules_to_treat if isinstance(m, torch.nn.Module)],
+                module_sizes,
+                no_split_module_classes,
+            )
+        # Assess size needed
+        module_size = module_sizes[name]
+        tied_params = [v for k, v in tied_parameters.items() if name in k]
+        # We ignore parameters that are tied when they're tied to > 1 one
+        tied_param = tied_params[0] if len(tied_params) == 1 else None
+
+        device = devices[current_device]
+        current_max_size = max_memory[device] if device != "disk" else None
+        # Reduce max size available by the largest layer.
+        if devices[current_device] in main_devices:
+            current_max_size = current_max_size - max_layer_size
+        # Case 1 -> We're too big!
+        if current_max_size is not None and current_memory_used + module_size > current_max_size:
+            # Split or not split?
+            modules_children = list(module.named_children())
+            if len(modules_children) == 0 or module.__class__.__name__ in no_split_module_classes:
+                # -> no split, we go to the next device
+                current_device += 1
+                modules_to_treat = [(name, module)] + modules_to_treat
+                current_memory_used = 0
+            else:
+                # -> split, we replace the module studied by its children + parameters
+                modules_children = list(module.named_parameters(recurse=False)) + modules_children
+                modules_to_treat = [(f"{name}.{n}", v) for n, v in modules_children] + modules_to_treat
+                # Update the max layer size.
+                max_layer_size, max_layer_names = get_max_layer_size(
+                    [(n, m) for n, m in modules_to_treat if isinstance(m, torch.nn.Module)],
+                    module_sizes,
+                    no_split_module_classes,
+                )
+
+        # Case 2, it fits! We're not entirely out of the wood though, because we may have some tied parameters.
+        elif tied_param is not None:
+            # Determine the sized occupied by this module + the module containing the tied parameter
+            tied_module_size = module_size
+            tied_module_index = [i for i, (n, _) in enumerate(modules_to_treat) if n in tied_param][0]
+            tied_module_name, tied_module = modules_to_treat[tied_module_index]
+            tied_module_size += module_sizes[tied_module_name] - module_sizes[tied_param]
+            if current_max_size is not None and current_memory_used + tied_module_size > current_max_size:
+                # Split or not split?
+                tied_module_children = list(tied_module.named_children())
+                if len(tied_module_children) == 0 or tied_module.__class__.__name__ in no_split_module_classes:
+                    # If the tied module is not split, we go to the next device
+                    current_device += 1
+                    modules_to_treat = [(name, module)] + modules_to_treat
+                    current_memory_used = 0
+                else:
+                    # Otherwise, we replace the tied module by its children.
+                    tied_module_children = list(tied_module.named_parameters(recurse=False)) + tied_module_children
+                    tied_module_children = [(f"{tied_module_name}.{n}", v) for n, v in tied_module_children]
+                    modules_to_treat = (
+                        [(name, module)]
+                        + modules_to_treat[:tied_module_index]
+                        + tied_module_children
+                        + modules_to_treat[tied_module_index + 1 :]
+                    )
+                    # Update the max layer size.
+                    max_layer_size, max_layer_names = get_max_layer_size(
+                        [(n, m) for n, m in modules_to_treat if isinstance(m, torch.nn.Module)],
+                        module_sizes,
+                        no_split_module_classes,
+                    )
+            else:
+                # We really really fit!
+                current_memory_used += tied_module_size
+                device_map[name] = devices[current_device]
+                modules_to_treat.pop(tied_module_index)
+                device_map[tied_module_name] = devices[current_device]
+        else:
+            current_memory_used += module_size
+            device_map[name] = devices[current_device]
+
+    return clean_device_map(device_map)
+
+
+def check_device_map(model: nn.Module, device_map: Dict[str, Union[int, str, torch.device]]):
+    """
+    Checks a device map covers everything in a given model.
+
+    Args:
+        model (`torch.nn.Module`): The model to check the device map against.
+        device_map (`Dict[str, Union[int, str, torch.device]]`): The device map to check.
+    """
+    all_model_tensors = [name for name, _ in model.state_dict().items()]
+    for module_name in device_map.keys():
+        all_model_tensors = [name for name in all_model_tensors if not name.startswith(module_name)]
+    if len(all_model_tensors) > 0:
+        non_covered_params = ", ".join(all_model_tensors)
+        raise ValueError(
+            f"The device_map provided does not give any device for the following parameters: {non_covered_params}"
+        )
+
+
+def load_checkpoint_in_model(
+    model: nn.Module,
+    checkpoint: Union[str, os.PathLike],
+    device_map: Optional[Dict[str, Union[int, str, torch.device]]] = None,
+    offload_folder: Optional[Union[str, os.PathLike]] = None,
+    dtype: Optional[Union[str, torch.dtype]] = None,
+    offload_state_dict: bool = False,
+):
+    """
+    Loads a (potentially sharded) checkpoint inside a model, potentially sending weights to a given device as they are
+    loaded.
+
+    <Tip warning={true}>
+
+    Once loaded across devices, you still need to call [`dispatch_model`] on your model to make it able to run. To
+    group the checkpoint loading and dispatch in one single call, use [`load_checkpoint_and_dispatch`].
+
+    </Tip>
+
+    Args:
+        model (`torch.nn.Module`): The model in which we want to load a checkpoint.
+        checkpoint (`str` or `os.PathLike`):
+            The folder checkpoint to load. It can be:
+            - a path to a file containing a whole model state dict
+            - a path to a `.json` file containing the index to a sharded checkpoint
+            - a path to a folder containing a unique `.index.json` file and the shards of a checkpoint.
+        device_map (`Dict[str, Union[int, str, torch.device]]`, *optional*):
+            A map that specifies where each submodule should go. It doesn't need to be refined to each parameter/buffer
+            name, once a given module name is inside, every submodule of it will be sent to the same device.
+        offload_folder (`str` or `os.PathLike`, *optional*):
+            If the `device_map` contains any value `"disk"`, the folder where we will offload weights.
+        dtype (`str` or `torch.dtype`, *optional*):
+            If provided, the weights will be converted to that type when loaded.
+        offload_state_dict (`bool`, *optional*, defaults to `False`):
+            If `True`, will temporarily offload the CPU state dict on the hard drive to avoig getting out of CPU RAM if
+            the weight of the CPU state dict + the biggest shard does not fit.
+    """
+    if offload_folder is None and device_map is not None and "disk" in device_map.values():
+        raise ValueError(
+            "At least one of the model submodule will be offloaded to disk, please pass along an `offload_folder`."
+        )
+    elif offload_folder is not None and device_map is not None and "disk" in device_map.values():
+        os.makedirs(offload_folder, exist_ok=True)
+
+    if isinstance(dtype, str):
+        # We accept "torch.float16" or just "float16"
+        dtype = dtype.replace("torch.", "")
+        dtype = getattr(torch, dtype)
+
+    checkpoint_files = None
+    index_filename = None
+    if os.path.isfile(checkpoint):
+        if str(checkpoint).endswith(".json"):
+            index_filename = checkpoint
+        else:
+            checkpoint_files = [checkpoint]
+    elif os.path.isdir(checkpoint):
+        potential_index = [f for f in os.listdir(checkpoint) if f.endswith(".index.json")]
+        if len(potential_index) == 0:
+            raise ValueError(f"{checkpoint} is not a folder containing a `.index.json` file.")
+        elif len(potential_index) == 1:
+            index_filename = os.path.join(checkpoint, potential_index[0])
+        else:
+            raise ValueError(f"{checkpoint} containing mote than one `.index.json` file, delete the irrelevant ones.")
+    else:
+        raise ValueError(
+            "`checkpoint` should be the path to a file containing a whole state dict, or the index of a sharded "
+            f"checkpoint, or a folder containing a sharded checkpoint, but got {checkpoint}."
+        )
+
+    if index_filename is not None:
+        checkpoint_folder = os.path.split(index_filename)[0]
+        with open(index_filename, "r") as f:
+            index = json.loads(f.read())
+
+        if "weight_map" in index:
+            index = index["weight_map"]
+        checkpoint_files = sorted(list(set(index.values())))
+        checkpoint_files = [os.path.join(checkpoint_folder, f) for f in checkpoint_files]
+
+    # Logic for missing/unexepected keys goes here.
+
+    offload_index = {}
+    if offload_state_dict:
+        state_dict_folder = tempfile.mkdtemp()
+        state_dict_index = {}
+
+    for checkpoint_file in checkpoint_files:
+        checkpoint = torch.load(checkpoint_file)
+        if device_map is None:
+            model.load_state_dict(checkpoint, strict=False)
+        else:
+            for param_name, param in checkpoint.items():
+                module_name = param_name
+                if dtype is not None:
+                    param = param.to(dtype)
+                while len(module_name) > 0 and module_name not in device_map:
+                    module_name = ".".join(module_name.split(".")[:-1])
+                if module_name == "" and "" not in device_map:
+                    # TODO: group all errors and raise at the end.
+                    raise ValueError(f"{param_name} doesn't have any device set.")
+                param_device = device_map[module_name]
+
+                if param_device == "disk":
+                    set_module_tensor_to_device(model, param_name, "meta")
+                    tensor_file = os.path.join(offload_folder, f"{param_name}.dat")
+                    array = param.numpy()
+                    offload_index[param_name] = {"dtype": str(array.dtype), "shape": list(array.shape)}
+                    file_array = np.memmap(tensor_file, dtype=array.dtype, mode="w+", shape=array.shape)
+                    file_array[:] = array[:]
+                    file_array.flush()
+                elif param_device == "cpu" and offload_state_dict:
+                    set_module_tensor_to_device(model, param_name, "meta")
+                    tensor_file = os.path.join(state_dict_folder, f"{param_name}.dat")
+                    array = param.numpy()
+                    state_dict_index[param_name] = {"dtype": str(array.dtype), "shape": list(array.shape)}
+                    file_array = np.memmap(tensor_file, dtype=array.dtype, mode="w+", shape=array.shape)
+                    file_array[:] = array[:]
+                    file_array.flush()
+                else:
+                    set_module_tensor_to_device(model, param_name, param_device, value=param)
+
+        # Force Python to clean up.
+        del checkpoint
+        gc.collect()
+
+    if len(offload_index) > 0:
+        offload_index_file = os.path.join(offload_folder, "index.json")
+        if os.path.isfile(offload_index_file):
+            with open(offload_index_file, "r", encoding="utf-8") as f:
+                current_offload_index = json.load(f)
+        else:
+            current_offload_index = {}
+        current_offload_index.update(offload_index)
+
+        with open(offload_index_file, "w", encoding="utf-8") as f:
+            json.dump(current_offload_index, f, indent=2)
+
+    # Load back offloaded state dict on CPU
+    if offload_state_dict and len(state_dict_index) > 0:
+        for param_name, metadata in state_dict_index.items():
+            tensor_file = os.path.join(state_dict_folder, f"{param_name}.dat")
+            shape = tuple(metadata["shape"])
+            weight = np.memmap(tensor_file, dtype=metadata["dtype"], mode="r", shape=shape)
+            set_module_tensor_to_device(model, param_name, "cpu", value=torch.tensor(weight))
+        shutil.rmtree(state_dict_folder)
diff --git a/src/accelerate/utils/offload.py b/src/accelerate/utils/offload.py
new file mode 100644
index 000000000..17f33d7c6
--- /dev/null
+++ b/src/accelerate/utils/offload.py
@@ -0,0 +1,143 @@
+# 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 json
+import os
+from collections.abc import Mapping
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+import torch
+
+
+def offload_state_dict(save_dir: Union[str, os.PathLike], state_dict: Dict[str, torch.Tensor]):
+    """
+    Offload a state dict in a given folder.
+
+    Args:
+        save_dir (`str` or `os.PathLike`): The directory in which to offload the state dict.
+        state_dict (`Dict[str, torch.Tensor]`): The dictionary of tensors to offload.
+    """
+    os.makedirs(save_dir, exist_ok=True)
+    index = {}
+    for name, parameter in state_dict.items():
+        tensor_file = os.path.join(save_dir, f"{name}.dat")
+        array = parameter.numpy()
+        index[name] = {"dtype": str(array.dtype), "shape": list(array.shape)}
+        if array.ndim == 0:
+            array = array[None]
+        file_array = np.memmap(tensor_file, dtype=array.dtype, mode="w+", shape=array.shape)
+        file_array[:] = array[:]
+        file_array.flush()
+
+    # Update index
+    index_file = os.path.join(save_dir, "index.json")
+    if os.path.isfile(index_file):
+        with open(index_file, "r", encoding="utf-8") as f:
+            current_index = json.load(f)
+    else:
+        current_index = {}
+    current_index.update(index)
+
+    with open(index_file, "w", encoding="utf-8") as f:
+        json.dump(current_index, f, indent=2)
+
+
+class PrefixedDataset(Mapping):
+    """
+    Will access keys in a given dataset by adding a prefix.
+
+    Args:
+        dataset (`Mapping`): Any map with string keys.
+        prefix (`str`): A prefix to add when trying to access any element in the underlying dataset.
+    """
+
+    def __init__(self, dataset: Mapping, prefix: str):
+        self.dataset = dataset
+        self.prefix = prefix
+
+    def __getitem__(self, key):
+        return self.dataset[f"{self.prefix}{key}"]
+
+    def __iter__(self):
+        return iter([key for key in self.dataset if key.startswith(self.prefix)])
+
+    def __len__(self):
+        return len(self.dataset)
+
+
+class OffloadedWeightsLoader(Mapping):
+    """
+    A collection that loads weights stored in a given state dict or memory-mapped on disk.
+
+    Args:
+        state_dict (`Dict[str, torch.Tensor]`, *optional*):
+            A dictionary parameter name to tensor.
+        save_folder (`str` or `os.PathLike`, *optional*):
+            The directory in which the weights are stored (by `offload_state_dict` for instance).
+        index (`Dict`, *optional*):
+            A dictionary from weight name to their information (`dtype` and `shape`). Will default to the index saved
+            in `save_folder`.
+    """
+
+    def __init__(
+        self,
+        state_dict: Dict[str, torch.Tensor] = None,
+        save_folder: Optional[Union[str, os.PathLike]] = None,
+        index: Mapping = None,
+    ):
+        if state_dict is None and save_folder is None:
+            raise ValueError("Need either a `state_dict` or a `save_folder` containing offloaded weights.")
+
+        self.state_dict = {} if state_dict is None else state_dict
+        self.save_folder = save_folder
+        if index is None and save_folder is not None:
+            with open(os.path.join(save_folder, "index.json")) as f:
+                index = json.load(f)
+        self.index = {} if index is None else index
+        self.all_keys = list(self.state_dict.keys())
+        self.all_keys.extend([key for key in self.index if key not in self.all_keys])
+
+    def __getitem__(self, key: str):
+        # State dict gets priority
+        if key in self.state_dict:
+            return self.state_dict[key]
+        weight_info = self.index[key]
+        weight_file = os.path.join(self.save_folder, f"{key}.dat")
+        shape = tuple(weight_info["shape"])
+        if shape == ():
+            weight = np.memmap(weight_file, dtype=weight_info["dtype"], shape=(1,), mode="r")[0]
+        else:
+            weight = np.memmap(weight_file, dtype=weight_info["dtype"], shape=shape, mode="r")
+        return torch.tensor(weight)
+
+    def __iter__(self):
+        return iter(self.all_keys)
+
+    def __len__(self):
+        return len(self.all_keys)
+
+
+def extract_submodules_state_dict(state_dict: Dict[str, torch.Tensor], submodule_names: List[str]):
+    """
+    Extract the sub state-dict corresponding to a list of given submodules.
+
+    Args:
+        state_dict (`Dict[str, torch.Tensor]`): The state dict to extract from.
+        submodule_names (`List[str]`): The list of submodule names we want to extract.
+    """
+    result = {}
+    for module_name in submodule_names:
+        result.update({key: param for key, param in state_dict.items() if key.startswith(module_name)})
+    return result
diff --git a/src/accelerate/utils/operations.py b/src/accelerate/utils/operations.py
index b61ae069b..3fc9d969c 100644
--- a/src/accelerate/utils/operations.py
+++ b/src/accelerate/utils/operations.py
@@ -490,3 +490,24 @@ def __call__(self, *args, **kwargs):
 
 
 convert_outputs_to_fp32 = ConvertOutputsToFp32
+
+
+def find_device(data):
+    """
+    Finds the device on which a nested dict/list/tuple of tensors lies (assuming they are all on the same device).
+
+    Args:
+        (nested list/tuple/dictionary of `torch.Tensor`): The data we want to know the device of.
+    """
+    if isinstance(data, Mapping):
+        for obj in data.values():
+            device = find_device(obj)
+            if device is not None:
+                return device
+    elif isinstance(data, (tuple, list)):
+        for obj in data:
+            device = find_device(obj)
+            if device is not None:
+                return device
+    elif isinstance(data, torch.Tensor):
+        return data.device
diff --git a/tests/test_big_modeling.py b/tests/test_big_modeling.py
new file mode 100644
index 000000000..be2238794
--- /dev/null
+++ b/tests/test_big_modeling.py
@@ -0,0 +1,276 @@
+# 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 unittest
+from tempfile import TemporaryDirectory
+
+import torch
+import torch.nn as nn
+
+from accelerate.big_modeling import (
+    cpu_offload,
+    disk_offload,
+    dispatch_model,
+    init_empty_weights,
+    load_checkpoint_and_dispatch,
+)
+from accelerate.hooks import remove_hook_from_submodules
+from accelerate.test_utils import require_cuda, require_multi_gpu, slow
+from accelerate.utils import offload_state_dict
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+class ModelForTest(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear1 = nn.Linear(3, 4)
+        self.batchnorm = nn.BatchNorm1d(4)
+        self.linear2 = nn.Linear(4, 5)
+
+    def forward(self, x):
+        return self.linear2(self.batchnorm(self.linear1(x)))
+
+
+class BiggerModelForTest(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear1 = nn.Linear(3, 4)
+        self.linear2 = nn.Linear(4, 5)
+        self.batchnorm = nn.BatchNorm1d(5)
+        self.linear3 = nn.Linear(5, 6)
+        self.linear4 = nn.Linear(6, 5)
+
+    def forward(self, x):
+        return self.linear4(self.linear3(self.batchnorm(self.linear2(self.linear1(x)))))
+
+
+class BigModelingTester(unittest.TestCase):
+    def test_init_empty_weights(self):
+        # base use
+        with init_empty_weights():
+            module = nn.Linear(4, 5)
+        self.assertEqual(module.weight.device, torch.device("meta"))
+
+        # base use with buffers, they are not touched
+        with init_empty_weights():
+            module = nn.BatchNorm1d(4)
+        self.assertEqual(module.weight.device, torch.device("meta"))
+        self.assertEqual(module.running_mean.device, torch.device("cpu"))
+
+        # Use with include_buffers=True
+        with init_empty_weights(include_buffers=True):
+            module = nn.BatchNorm1d(4)
+        self.assertEqual(module.weight.device, torch.device("meta"))
+        self.assertEqual(module.running_mean.device, torch.device("meta"))
+
+        # Double check we didn't break PyTorch
+        module = nn.BatchNorm1d(4)
+        self.assertEqual(module.weight.device, torch.device("cpu"))
+        self.assertEqual(module.running_mean.device, torch.device("cpu"))
+
+    def test_init_empty_weights_very_large_model(self):
+        # This is a 100 billion parameters model.
+        with init_empty_weights():
+            _ = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])
+
+    def test_cpu_offload(self):
+        model = ModelForTest()
+        x = torch.randn(2, 3)
+        expected = model(x)
+
+        device = torch.device(0 if torch.cuda.is_available() else "cpu")
+
+        cpu_offload(model, execution_device=device)
+        output = model(x)
+        self.assertTrue(torch.allclose(expected, output.cpu()))
+
+        # Clean up for next test.
+        remove_hook_from_submodules(model)
+
+        cpu_offload(model, execution_device=device, offload_buffers=True)
+        output = model(x)
+        self.assertTrue(torch.allclose(expected, output.cpu()))
+
+    @slow
+    @require_cuda
+    def test_cpu_offload_gpt2(self):
+        tokenizer = AutoTokenizer.from_pretrained("gpt2")
+        inputs = tokenizer("Hello world! My name is", return_tensors="pt").to(0)
+
+        gpt2 = AutoModelForCausalLM.from_pretrained("gpt2")
+        cpu_offload(gpt2, execution_device=0)
+        outputs = gpt2.generate(inputs["input_ids"])
+        self.assertEqual(
+            tokenizer.decode(outputs[0].tolist()),
+            "Hello world! My name is Kiyoshi, and I'm a student at the University of Tokyo",
+        )
+
+    def test_disk_offload(self):
+        model = ModelForTest()
+        x = torch.randn(2, 3)
+        expected = model(x)
+
+        device = torch.device(0 if torch.cuda.is_available() else "cpu")
+
+        with TemporaryDirectory() as tmp_dir:
+            disk_offload(model, tmp_dir, execution_device=device)
+            output = model(x)
+            self.assertTrue(torch.allclose(expected, output.cpu()))
+
+            # Clean up for next test.
+            remove_hook_from_submodules(model)
+
+        with TemporaryDirectory() as tmp_dir:
+            disk_offload(model, tmp_dir, execution_device=device, offload_buffers=True)
+            output = model(x)
+            self.assertTrue(torch.allclose(expected, output.cpu()))
+
+    @slow
+    @require_cuda
+    def test_disk_offload_gpt2(self):
+        tokenizer = AutoTokenizer.from_pretrained("gpt2")
+        inputs = tokenizer("Hello world! My name is", return_tensors="pt").to(0)
+
+        gpt2 = AutoModelForCausalLM.from_pretrained("gpt2")
+        with TemporaryDirectory() as tmp_dir:
+            disk_offload(gpt2, tmp_dir, execution_device=0)
+            outputs = gpt2.generate(inputs["input_ids"])
+            self.assertEqual(
+                tokenizer.decode(outputs[0].tolist()),
+                "Hello world! My name is Kiyoshi, and I'm a student at the University of Tokyo",
+            )
+
+    @require_cuda
+    def test_dispatch_model(self):
+        model = ModelForTest()
+        device_map = {"linear1": "disk", "batchnorm": "cpu", "linear2": 0}
+
+        x = torch.randn(2, 3)
+        expected = model(x)
+
+        with TemporaryDirectory() as tmp_dir:
+            dispatch_model(model, device_map, offload_dir=tmp_dir)
+            output = model(x)
+            self.assertTrue(torch.allclose(expected, output.cpu(), atol=1e-5))
+
+    @require_multi_gpu
+    def test_dispatch_model_multi_gpu(self):
+        model = BiggerModelForTest()
+        device_map = {"linear1": "cpu", "linear2": "disk", "batchnorm": "cpu", "linear3": 0, "linear4": 1}
+
+        x = torch.randn(2, 3)
+        expected = model(x)
+
+        with TemporaryDirectory() as tmp_dir:
+            dispatch_model(model, device_map, offload_dir=tmp_dir)
+            output = model(x)
+            self.assertTrue(torch.allclose(expected, output.cpu(), atol=1e-5))
+
+    @slow
+    @require_multi_gpu
+    def test_dispatch_model_gpt2_on_two_gpus(self):
+        tokenizer = AutoTokenizer.from_pretrained("gpt2")
+        inputs = tokenizer("Hello world! My name is", return_tensors="pt").to(0)
+
+        gpt2 = AutoModelForCausalLM.from_pretrained("gpt2")
+        # Dispatch on GPUs 0 and 1
+        device_map = {
+            "transformer.wte": 0,
+            "transformer.wpe": 0,
+            "transformer.ln_f": 1,
+            "lm_head": 1,
+        }
+        for i in range(12):
+            device_map[f"transformer.h.{i}"] = 0 if i <= 5 else 1
+
+        gpt2 = dispatch_model(gpt2, device_map)
+        outputs = gpt2.generate(inputs["input_ids"])
+        self.assertEqual(
+            tokenizer.decode(outputs[0].tolist()),
+            "Hello world! My name is Kiyoshi, and I'm a student at the University of Tokyo",
+        )
+
+        # Dispatch with a bit of CPU offload
+        gpt2 = AutoModelForCausalLM.from_pretrained("gpt2")
+        for i in range(4):
+            device_map[f"transformer.h.{i}"] = "cpu"
+        gpt2 = dispatch_model(gpt2, device_map)
+        outputs = gpt2.generate(inputs["input_ids"])
+        self.assertEqual(
+            tokenizer.decode(outputs[0].tolist()),
+            "Hello world! My name is Kiyoshi, and I'm a student at the University of Tokyo",
+        )
+        # Dispatch with a bit of CPU and disk offload
+        gpt2 = AutoModelForCausalLM.from_pretrained("gpt2")
+        for i in range(2):
+            device_map[f"transformer.h.{i}"] = "disk"
+
+        with TemporaryDirectory() as tmp_dir:
+            state_dict = {
+                k: p for k, p in gpt2.state_dict().items() if "transformer.h.0" in k or "transformer.h.1" in k
+            }
+            offload_state_dict(tmp_dir, state_dict)
+            gpt2 = dispatch_model(gpt2, device_map, offload_dir=tmp_dir)
+            outputs = gpt2.generate(inputs["input_ids"])
+            self.assertEqual(
+                tokenizer.decode(outputs[0].tolist()),
+                "Hello world! My name is Kiyoshi, and I'm a student at the University of Tokyo",
+            )
+
+    @require_cuda
+    def test_load_checkpoint_and_dispatch(self):
+        model = ModelForTest()
+        device_map = {"linear1": "cpu", "batchnorm": "cpu", "linear2": 0}
+
+        x = torch.randn(2, 3)
+        expected = model(x)
+
+        with TemporaryDirectory() as tmp_dir:
+            checkpoint = os.path.join(tmp_dir, "pt_model.bin")
+            torch.save(model.state_dict(), checkpoint)
+
+            new_model = ModelForTest()
+            new_model = load_checkpoint_and_dispatch(new_model, checkpoint, device_map=device_map)
+
+        # CPU-offloaded weights are on the meta device while waiting for the forward pass.
+        self.assertEqual(new_model.linear1.weight.device, torch.device("meta"))
+        self.assertEqual(new_model.linear2.weight.device, torch.device(0))
+
+        output = new_model(x)
+        self.assertTrue(torch.allclose(expected, output.cpu(), atol=1e-5))
+
+    @require_multi_gpu
+    def test_load_checkpoint_and_dispatch_multi_gpu(self):
+        model = BiggerModelForTest()
+        device_map = {"linear1": "cpu", "linear2": "cpu", "batchnorm": 0, "linear3": 0, "linear4": 1}
+
+        x = torch.randn(2, 3)
+        expected = model(x)
+
+        with TemporaryDirectory() as tmp_dir:
+            checkpoint = os.path.join(tmp_dir, "pt_model.bin")
+            torch.save(model.state_dict(), checkpoint)
+
+            new_model = BiggerModelForTest()
+            new_model = load_checkpoint_and_dispatch(new_model, checkpoint, device_map=device_map)
+
+        # CPU-offloaded weights are on the meta device while waiting for the forward pass.
+        self.assertEqual(new_model.linear1.weight.device, torch.device("meta"))
+        self.assertEqual(new_model.linear2.weight.device, torch.device("meta"))
+        self.assertEqual(new_model.linear3.weight.device, torch.device(0))
+        self.assertEqual(new_model.linear4.weight.device, torch.device(1))
+
+        output = new_model(x)
+        self.assertTrue(torch.allclose(expected, output.cpu(), atol=1e-5))
diff --git a/tests/test_hooks.py b/tests/test_hooks.py
new file mode 100644
index 000000000..51c434ab1
--- /dev/null
+++ b/tests/test_hooks.py
@@ -0,0 +1,330 @@
+# 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 inspect
+import unittest
+
+import torch
+import torch.nn as nn
+
+from accelerate.hooks import (
+    AlignDevicesHook,
+    ModelHook,
+    SequentialHook,
+    add_hook_to_module,
+    attach_align_device_hook,
+    remove_hook_from_module,
+    remove_hook_from_submodules,
+)
+from accelerate.test_utils import require_multi_gpu
+
+
+class ModelForTest(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear1 = nn.Linear(3, 4)
+        self.batchnorm = nn.BatchNorm1d(4)
+        self.linear2 = nn.Linear(4, 5)
+
+    def forward(self, x):
+        return self.linear2(self.batchnorm(self.linear1(x)))
+
+
+class PreForwardHook(ModelHook):
+    def pre_forward(self, module, *args, **kwargs):
+        return (args[0] + 1,) + args[1:], kwargs
+
+
+class PostForwardHook(ModelHook):
+    def post_forward(self, module, output):
+        return output + 1
+
+
+class HooksModelTester(unittest.TestCase):
+    def test_add_and_remove_hooks(self):
+        test_model = ModelForTest()
+        test_hook = ModelHook()
+
+        add_hook_to_module(test_model, test_hook)
+        self.assertEqual(test_model._hf_hook, test_hook)
+        self.assertTrue(hasattr(test_model, "_old_forward"))
+
+        # Check adding the hook did not change the name or the signature
+        self.assertEqual(test_model.forward.__name__, "forward")
+        self.assertListEqual(list(inspect.signature(test_model.forward).parameters), ["x"])
+
+        remove_hook_from_module(test_model)
+        self.assertFalse(hasattr(test_model, "_hf_hook"))
+        self.assertFalse(hasattr(test_model, "_old_forward"))
+
+    def test_pre_forward_hook_is_executed(self):
+        test_model = ModelForTest()
+        x = torch.randn(2, 3)
+        expected = test_model(x + 1)
+        expected2 = test_model(x + 2)
+
+        test_hook = PreForwardHook()
+        add_hook_to_module(test_model, test_hook)
+        output1 = test_model(x)
+        self.assertTrue(torch.allclose(output1, expected))
+
+        # Attaching a hook to a model when it already has one replaces, does not chain
+        test_hook = PreForwardHook()
+        add_hook_to_module(test_model, test_hook)
+        output1 = test_model(x)
+        self.assertTrue(torch.allclose(output1, expected))
+
+        # You need to use the sequential hook to chain two or more hooks
+        test_hook = SequentialHook(PreForwardHook(), PreForwardHook())
+        add_hook_to_module(test_model, test_hook)
+
+        output2 = test_model(x)
+        assert torch.allclose(output2, expected2)
+
+    def test_post_forward_hook_is_executed(self):
+        test_model = ModelForTest()
+        x = torch.randn(2, 3)
+        output = test_model(x)
+
+        test_hook = PostForwardHook()
+        add_hook_to_module(test_model, test_hook)
+        output1 = test_model(x)
+        self.assertTrue(torch.allclose(output1, output + 1))
+
+        # Attaching a hook to a model when it already has one replaces, does not chain
+        test_hook = PostForwardHook()
+        add_hook_to_module(test_model, test_hook)
+        output1 = test_model(x)
+        self.assertTrue(torch.allclose(output1, output + 1))
+
+        # You need to use the sequential hook to chain two or more hooks
+        test_hook = SequentialHook(PostForwardHook(), PostForwardHook())
+        add_hook_to_module(test_model, test_hook)
+
+        output2 = test_model(x)
+        assert torch.allclose(output2, output + 2)
+
+    def test_no_grad_in_hook(self):
+        test_model = ModelForTest()
+        x = torch.randn(2, 3)
+        output = test_model(x)
+
+        test_hook = PostForwardHook()
+        add_hook_to_module(test_model, test_hook)
+        output1 = test_model(x)
+        self.assertTrue(torch.allclose(output1, output + 1))
+        self.assertTrue(output1.requires_grad)
+
+        test_hook.no_grad = True
+        output1 = test_model(x)
+        self.assertFalse(output1.requires_grad)
+
+    @require_multi_gpu
+    def test_align_devices_as_model_parallelism(self):
+        model = ModelForTest()
+        # Everything is on CPU
+        self.assertEqual(model.linear1.weight.device, torch.device("cpu"))
+        self.assertEqual(model.batchnorm.weight.device, torch.device("cpu"))
+        self.assertEqual(model.linear2.weight.device, torch.device("cpu"))
+
+        # This will move each submodule on different devices
+        add_hook_to_module(model.linear1, AlignDevicesHook(execution_device=0))
+        add_hook_to_module(model.batchnorm, AlignDevicesHook(execution_device=0))
+        add_hook_to_module(model.linear2, AlignDevicesHook(execution_device=1))
+
+        self.assertEqual(model.linear1.weight.device, torch.device(0))
+        self.assertEqual(model.batchnorm.weight.device, torch.device(0))
+        self.assertEqual(model.batchnorm.running_mean.device, torch.device(0))
+        self.assertEqual(model.linear2.weight.device, torch.device(1))
+
+        # We can still make a forward pass. The input does not need to be on any particular device
+        x = torch.randn(2, 3)
+        output = model(x)
+        self.assertEqual(output.device, torch.device(1))
+
+        # We can add a general hook to put back output on same device as input.
+        add_hook_to_module(model, AlignDevicesHook(io_same_device=True))
+        x = torch.randn(2, 3).to(0)
+        output = model(x)
+        self.assertEqual(output.device, torch.device(0))
+
+    def test_align_devices_as_cpu_offload(self):
+        model = ModelForTest()
+
+        # Everything is on CPU
+        self.assertEqual(model.linear1.weight.device, torch.device("cpu"))
+        self.assertEqual(model.batchnorm.weight.device, torch.device("cpu"))
+        self.assertEqual(model.linear2.weight.device, torch.device("cpu"))
+
+        # This will move each submodule on different devices
+        hook_kwargs = {"execution_device": 0 if torch.cuda.is_available() else "cpu", "offload": True}
+
+        add_hook_to_module(model.linear1, AlignDevicesHook(**hook_kwargs))
+        add_hook_to_module(model.batchnorm, AlignDevicesHook(**hook_kwargs))
+        add_hook_to_module(model.linear2, AlignDevicesHook(**hook_kwargs))
+
+        # Parameters have been offloaded, so on the meta device
+        self.assertEqual(model.linear1.weight.device, torch.device("meta"))
+        self.assertEqual(model.batchnorm.weight.device, torch.device("meta"))
+        self.assertEqual(model.linear2.weight.device, torch.device("meta"))
+        # Buffers are not included in the offload by default, so are on the execution device
+        device = torch.device(hook_kwargs["execution_device"])
+        self.assertEqual(model.batchnorm.running_mean.device, device)
+
+        x = torch.randn(2, 3)
+        output = model(x)
+        self.assertEqual(output.device, device)
+
+        # Removing hooks loads back the weights in the model.
+        remove_hook_from_module(model.linear1)
+        remove_hook_from_module(model.batchnorm)
+        remove_hook_from_module(model.linear2)
+        self.assertEqual(model.linear1.weight.device, torch.device("cpu"))
+        self.assertEqual(model.batchnorm.weight.device, torch.device("cpu"))
+        self.assertEqual(model.linear2.weight.device, torch.device("cpu"))
+
+        # Now test with buffers included in the offload
+        hook_kwargs = {
+            "execution_device": 0 if torch.cuda.is_available() else "cpu",
+            "offload": True,
+            "offload_buffers": True,
+        }
+
+        add_hook_to_module(model.linear1, AlignDevicesHook(**hook_kwargs))
+        add_hook_to_module(model.batchnorm, AlignDevicesHook(**hook_kwargs))
+        add_hook_to_module(model.linear2, AlignDevicesHook(**hook_kwargs))
+
+        # Parameters have been offloaded, so on the meta device, buffers included
+        self.assertEqual(model.linear1.weight.device, torch.device("meta"))
+        self.assertEqual(model.batchnorm.weight.device, torch.device("meta"))
+        self.assertEqual(model.linear2.weight.device, torch.device("meta"))
+        self.assertEqual(model.batchnorm.running_mean.device, torch.device("meta"))
+
+        x = torch.randn(2, 3)
+        output = model(x)
+        self.assertEqual(output.device, device)
+
+        # Removing hooks loads back the weights in the model.
+        remove_hook_from_module(model.linear1)
+        remove_hook_from_module(model.batchnorm)
+        remove_hook_from_module(model.linear2)
+        self.assertEqual(model.linear1.weight.device, torch.device("cpu"))
+        self.assertEqual(model.batchnorm.weight.device, torch.device("cpu"))
+        self.assertEqual(model.linear2.weight.device, torch.device("cpu"))
+
+    def test_attach_align_device_hook_as_cpu_offload(self):
+        model = ModelForTest()
+
+        # Everything is on CPU
+        self.assertEqual(model.linear1.weight.device, torch.device("cpu"))
+        self.assertEqual(model.batchnorm.weight.device, torch.device("cpu"))
+        self.assertEqual(model.linear2.weight.device, torch.device("cpu"))
+
+        # This will move each submodule on different devices
+        execution_device = 0 if torch.cuda.is_available() else "cpu"
+        attach_align_device_hook(model, execution_device=execution_device, offload=True)
+
+        # Parameters have been offloaded, so on the meta device
+        self.assertEqual(model.linear1.weight.device, torch.device("meta"))
+        self.assertEqual(model.batchnorm.weight.device, torch.device("meta"))
+        self.assertEqual(model.linear2.weight.device, torch.device("meta"))
+        # Buffers are not included in the offload by default, so are on the execution device
+        device = torch.device(execution_device)
+        self.assertEqual(model.batchnorm.running_mean.device, device)
+
+        x = torch.randn(2, 3)
+        output = model(x)
+        self.assertEqual(output.device, device)
+
+        # Removing hooks loads back the weights in the model.
+        remove_hook_from_submodules(model)
+        self.assertEqual(model.linear1.weight.device, torch.device("cpu"))
+        self.assertEqual(model.batchnorm.weight.device, torch.device("cpu"))
+        self.assertEqual(model.linear2.weight.device, torch.device("cpu"))
+
+        # Now test with buffers included in the offload
+        attach_align_device_hook(model, execution_device=execution_device, offload=True, offload_buffers=True)
+
+        # Parameters have been offloaded, so on the meta device, buffers included
+        self.assertEqual(model.linear1.weight.device, torch.device("meta"))
+        self.assertEqual(model.batchnorm.weight.device, torch.device("meta"))
+        self.assertEqual(model.linear2.weight.device, torch.device("meta"))
+        self.assertEqual(model.batchnorm.running_mean.device, torch.device("meta"))
+
+        x = torch.randn(2, 3)
+        output = model(x)
+        self.assertEqual(output.device, device)
+
+        # Removing hooks loads back the weights in the model.
+        remove_hook_from_submodules(model)
+        self.assertEqual(model.linear1.weight.device, torch.device("cpu"))
+        self.assertEqual(model.batchnorm.weight.device, torch.device("cpu"))
+        self.assertEqual(model.linear2.weight.device, torch.device("cpu"))
+
+    def test_attach_align_device_hook_as_cpu_offload_with_weight_map(self):
+        model = ModelForTest()
+
+        # Everything is on CPU
+        self.assertEqual(model.linear1.weight.device, torch.device("cpu"))
+        self.assertEqual(model.batchnorm.weight.device, torch.device("cpu"))
+        self.assertEqual(model.linear2.weight.device, torch.device("cpu"))
+
+        # This will move each submodule on different devices
+        execution_device = 0 if torch.cuda.is_available() else "cpu"
+        attach_align_device_hook(
+            model, execution_device=execution_device, offload=True, weights_map=model.state_dict()
+        )
+
+        # Parameters have been offloaded, so on the meta device
+        self.assertEqual(model.linear1.weight.device, torch.device("meta"))
+        self.assertEqual(model.batchnorm.weight.device, torch.device("meta"))
+        self.assertEqual(model.linear2.weight.device, torch.device("meta"))
+        # Buffers are not included in the offload by default, so are on the execution device
+        device = torch.device(execution_device)
+        self.assertEqual(model.batchnorm.running_mean.device, device)
+
+        x = torch.randn(2, 3)
+        output = model(x)
+        self.assertEqual(output.device, device)
+
+        # Removing hooks loads back the weights in the model.
+        remove_hook_from_submodules(model)
+        self.assertEqual(model.linear1.weight.device, torch.device("cpu"))
+        self.assertEqual(model.batchnorm.weight.device, torch.device("cpu"))
+        self.assertEqual(model.linear2.weight.device, torch.device("cpu"))
+
+        # Now test with buffers included in the offload
+        attach_align_device_hook(
+            model,
+            execution_device=execution_device,
+            offload=True,
+            weights_map=model.state_dict(),
+            offload_buffers=True,
+        )
+
+        # Parameters have been offloaded, so on the meta device, buffers included
+        self.assertEqual(model.linear1.weight.device, torch.device("meta"))
+        self.assertEqual(model.batchnorm.weight.device, torch.device("meta"))
+        self.assertEqual(model.linear2.weight.device, torch.device("meta"))
+        self.assertEqual(model.batchnorm.running_mean.device, torch.device("meta"))
+
+        x = torch.randn(2, 3)
+        output = model(x)
+        self.assertEqual(output.device, device)
+
+        # Removing hooks loads back the weights in the model.
+        remove_hook_from_submodules(model)
+        self.assertEqual(model.linear1.weight.device, torch.device("cpu"))
+        self.assertEqual(model.batchnorm.weight.device, torch.device("cpu"))
+        self.assertEqual(model.linear2.weight.device, torch.device("cpu"))
diff --git a/tests/test_modeling_utils.py b/tests/test_modeling_utils.py
new file mode 100644
index 000000000..15e164218
--- /dev/null
+++ b/tests/test_modeling_utils.py
@@ -0,0 +1,360 @@
+# 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 json
+import os
+import tempfile
+import unittest
+
+import torch
+import torch.nn as nn
+
+from accelerate.test_utils import require_cuda, require_multi_gpu
+from accelerate.utils.modeling import (
+    check_device_map,
+    clean_device_map,
+    compute_module_sizes,
+    find_tied_parameters,
+    infer_auto_device_map,
+    load_checkpoint_in_model,
+    named_module_tensors,
+    set_module_tensor_to_device,
+)
+
+
+class ModelForTest(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear1 = nn.Linear(3, 4)
+        self.batchnorm = nn.BatchNorm1d(4)
+        self.linear2 = nn.Linear(4, 5)
+
+    def forward(self, x):
+        return self.linear2(self.batchnorm(self.linear1(x)))
+
+
+class ModelingUtilsTester(unittest.TestCase):
+    def check_set_module_tensor_for_device(self, model, device1, device2):
+        self.assertEqual(model.linear1.weight.device, torch.device(device1))
+
+        with self.subTest("Access by submodule and direct name for a parameter"):
+            set_module_tensor_to_device(model.linear1, "weight", device2)
+            self.assertEqual(model.linear1.weight.device, torch.device(device2))
+
+            if torch.device(device2) == torch.device("meta"):
+                with self.assertRaises(ValueError):
+                    # We need a `value` to set the weight back on device1
+                    set_module_tensor_to_device(model.linear1, "weight", device1)
+
+                set_module_tensor_to_device(model.linear1, "weight", device1, value=torch.randn(4, 3))
+            else:
+                set_module_tensor_to_device(model.linear1, "weight", device1)
+            self.assertEqual(model.linear1.weight.device, torch.device(device1))
+
+        with self.subTest("Access by module and full name for a parameter"):
+            set_module_tensor_to_device(model, "linear1.weight", device2)
+            self.assertEqual(model.linear1.weight.device, torch.device(device2))
+
+            if torch.device(device2) == torch.device("meta"):
+                with self.assertRaises(ValueError):
+                    # We need a `value` to set the weight back on device1
+                    set_module_tensor_to_device(model, "linear1.weight", device1)
+                set_module_tensor_to_device(model, "linear1.weight", device1, value=torch.randn(4, 3))
+            else:
+                set_module_tensor_to_device(model, "linear1.weight", device1)
+            self.assertEqual(model.linear1.weight.device, torch.device(device1))
+
+        self.assertEqual(model.batchnorm.running_mean.device, torch.device(device1))
+
+        with self.subTest("Access by submodule and direct name for a buffer"):
+            set_module_tensor_to_device(model.batchnorm, "running_mean", device2)
+            self.assertEqual(model.batchnorm.running_mean.device, torch.device(device2))
+
+            if torch.device(device2) == torch.device("meta"):
+                with self.assertRaises(ValueError):
+                    # We need a `value` to set the weight back on device1
+                    set_module_tensor_to_device(model.batchnorm, "running_mean", device1)
+                set_module_tensor_to_device(model.batchnorm, "running_mean", device1, value=torch.randn(4))
+            else:
+                set_module_tensor_to_device(model.batchnorm, "running_mean", device1)
+            self.assertEqual(model.batchnorm.running_mean.device, torch.device(device1))
+
+        with self.subTest("Access by module and full name for a parameter"):
+            set_module_tensor_to_device(model, "batchnorm.running_mean", device2)
+            self.assertEqual(model.batchnorm.running_mean.device, torch.device(device2))
+
+            if torch.device(device2) == torch.device("meta"):
+                with self.assertRaises(ValueError):
+                    # We need a `value` to set the weight back on CPU
+                    set_module_tensor_to_device(model, "batchnorm.running_mean", device1)
+
+                set_module_tensor_to_device(model, "batchnorm.running_mean", device1, value=torch.randn(4))
+            else:
+                set_module_tensor_to_device(model, "batchnorm.running_mean", device1)
+            self.assertEqual(model.batchnorm.running_mean.device, torch.device(device1))
+
+    def test_set_module_tensor_to_meta_and_cpu(self):
+        model = ModelForTest()
+        self.check_set_module_tensor_for_device(model, "cpu", "meta")
+
+    @require_cuda
+    def test_set_module_tensor_to_cpu_and_gpu(self):
+        model = ModelForTest()
+        self.check_set_module_tensor_for_device(model, "cpu", 0)
+
+    @require_cuda
+    def test_set_module_tensor_to_meta_and_gpu(self):
+        model = ModelForTest().to(0)
+        self.check_set_module_tensor_for_device(model, 0, "meta")
+
+    @require_multi_gpu
+    def test_set_module_tensor_between_gpus(self):
+        model = ModelForTest().to(0)
+        self.check_set_module_tensor_for_device(model, 0, 1)
+
+    def test_named_tensors(self):
+        model = nn.BatchNorm1d(4)
+        named_tensors = named_module_tensors(model)
+        self.assertListEqual(
+            [name for name, _ in named_tensors],
+            ["weight", "bias", "running_mean", "running_var", "num_batches_tracked"],
+        )
+
+        named_tensors = named_module_tensors(model, include_buffers=False)
+        self.assertListEqual([name for name, _ in named_tensors], ["weight", "bias"])
+
+        model = ModelForTest()
+        named_tensors = named_module_tensors(model)
+        self.assertListEqual([name for name, _ in named_tensors], [])
+
+        named_tensors = named_module_tensors(model, recurse=True)
+        self.assertListEqual(
+            [name for name, _ in named_tensors],
+            [
+                "linear1.weight",
+                "linear1.bias",
+                "batchnorm.weight",
+                "batchnorm.bias",
+                "linear2.weight",
+                "linear2.bias",
+                "batchnorm.running_mean",
+                "batchnorm.running_var",
+                "batchnorm.num_batches_tracked",
+            ],
+        )
+
+        named_tensors = named_module_tensors(model, include_buffers=False, recurse=True)
+        self.assertListEqual(
+            [name for name, _ in named_tensors],
+            ["linear1.weight", "linear1.bias", "batchnorm.weight", "batchnorm.bias", "linear2.weight", "linear2.bias"],
+        )
+
+    def test_find_tied_parameters(self):
+        model = ModelForTest()
+        self.assertDictEqual(find_tied_parameters(model), {})
+        model.linear2.weight = model.linear1.weight
+        self.assertDictEqual(find_tied_parameters(model), {"linear1.weight": "linear2.weight"})
+
+    def test_compute_module_sizes(self):
+        model = ModelForTest()
+        expected_sizes = {"": 236, "linear1": 64, "linear1.weight": 48, "linear1.bias": 16}
+        expected_sizes.update({"linear2": 100, "linear2.weight": 80, "linear2.bias": 20})
+        expected_sizes.update({"batchnorm": 72, "batchnorm.weight": 16, "batchnorm.bias": 16})
+        expected_sizes.update(
+            {"batchnorm.running_mean": 16, "batchnorm.running_var": 16, "batchnorm.num_batches_tracked": 8}
+        )
+
+        module_sizes = compute_module_sizes(model)
+        self.assertDictEqual(module_sizes, expected_sizes)
+
+        model.half()
+        expected_sizes = {k: s // 2 for k, s in expected_sizes.items()}
+        # This one is not converted to half.
+        expected_sizes["batchnorm.num_batches_tracked"] = 8
+        # This impacts batchnorm and total
+        expected_sizes["batchnorm"] += 4
+        expected_sizes[""] += 4
+
+        module_sizes = compute_module_sizes(model)
+        self.assertDictEqual(module_sizes, expected_sizes)
+
+    def test_check_device_map(self):
+        model = ModelForTest()
+        check_device_map(model, {"": 0})
+        with self.assertRaises(ValueError):
+            check_device_map(model, {"linear1": 0, "linear2": 1})
+
+        check_device_map(model, {"linear1": 0, "linear2": 1, "batchnorm": 1})
+
+    def shard_test_model(self, model, tmp_dir):
+        module_index = {
+            "linear1": "checkpoint_part1.bin",
+            "batchnorm": "checkpoint_part2.bin",
+            "linear2": "checkpoint_part3.bin",
+        }
+        index = {}
+        for name, _ in model.state_dict().items():
+            module = name.split(".")[0]
+            index[name] = module_index[module]
+
+        with open(os.path.join(tmp_dir, "weight_map.index.json"), "w") as f:
+            json.dump(index, f)
+
+        for module, fname in module_index.items():
+            state_dict = {k: v for k, v in model.state_dict().items() if k.startswith(module)}
+            full_fname = os.path.join(tmp_dir, fname)
+            torch.save(state_dict, full_fname)
+
+    def test_load_checkpoint_in_model(self):
+        # Check with whole checkpoint
+        model = ModelForTest()
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            fname = os.path.join(tmp_dir, "pt_model.bin")
+            torch.save(model.state_dict(), fname)
+            load_checkpoint_in_model(model, fname)
+
+        # Check with sharded index
+        model = ModelForTest()
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            self.shard_test_model(model, tmp_dir)
+            index_file = os.path.join(tmp_dir, "weight_map.index.json")
+            load_checkpoint_in_model(model, index_file)
+
+        # Check with sharded checkpoint
+        model = ModelForTest()
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            self.shard_test_model(model, tmp_dir)
+            load_checkpoint_in_model(model, tmp_dir)
+
+    @require_cuda
+    def test_load_checkpoint_in_model_one_gpu(self):
+        device_map = {"linear1": 0, "batchnorm": "cpu", "linear2": "cpu"}
+
+        # Check with whole checkpoint
+        model = ModelForTest()
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            fname = os.path.join(tmp_dir, "pt_model.bin")
+            torch.save(model.state_dict(), fname)
+            load_checkpoint_in_model(model, fname, device_map=device_map)
+        self.assertEqual(model.linear1.weight.device, torch.device(0))
+        self.assertEqual(model.batchnorm.weight.device, torch.device("cpu"))
+        self.assertEqual(model.linear2.weight.device, torch.device("cpu"))
+
+        # Check with sharded index
+        model = ModelForTest()
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            self.shard_test_model(model, tmp_dir)
+            index_file = os.path.join(tmp_dir, "weight_map.index.json")
+            load_checkpoint_in_model(model, index_file, device_map=device_map)
+
+        self.assertEqual(model.linear1.weight.device, torch.device(0))
+        self.assertEqual(model.batchnorm.weight.device, torch.device("cpu"))
+        self.assertEqual(model.linear2.weight.device, torch.device("cpu"))
+
+        # Check with sharded checkpoint folder
+        model = ModelForTest()
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            self.shard_test_model(model, tmp_dir)
+            load_checkpoint_in_model(model, tmp_dir, device_map=device_map)
+
+        self.assertEqual(model.linear1.weight.device, torch.device(0))
+        self.assertEqual(model.batchnorm.weight.device, torch.device("cpu"))
+        self.assertEqual(model.linear2.weight.device, torch.device("cpu"))
+
+    @require_multi_gpu
+    def test_load_checkpoint_in_model_two_gpu(self):
+        device_map = {"linear1": 0, "batchnorm": "cpu", "linear2": 1}
+
+        # Check with whole checkpoint
+        model = ModelForTest()
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            fname = os.path.join(tmp_dir, "pt_model.bin")
+            torch.save(model.state_dict(), fname)
+            load_checkpoint_in_model(model, fname, device_map=device_map)
+        self.assertEqual(model.linear1.weight.device, torch.device(0))
+        self.assertEqual(model.batchnorm.weight.device, torch.device("cpu"))
+        self.assertEqual(model.linear2.weight.device, torch.device(1))
+
+        # Check with sharded index
+        model = ModelForTest()
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            self.shard_test_model(model, tmp_dir)
+            index_file = os.path.join(tmp_dir, "weight_map.index.json")
+            load_checkpoint_in_model(model, index_file, device_map=device_map)
+
+        self.assertEqual(model.linear1.weight.device, torch.device(0))
+        self.assertEqual(model.batchnorm.weight.device, torch.device("cpu"))
+        self.assertEqual(model.linear2.weight.device, torch.device(1))
+
+        # Check with sharded checkpoint
+        model = ModelForTest()
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            self.shard_test_model(model, tmp_dir)
+            load_checkpoint_in_model(model, tmp_dir, device_map=device_map)
+
+        self.assertEqual(model.linear1.weight.device, torch.device(0))
+        self.assertEqual(model.batchnorm.weight.device, torch.device("cpu"))
+        self.assertEqual(model.linear2.weight.device, torch.device(1))
+
+    def test_clean_device_map(self):
+        # Regroup everything if all is on the same device
+        self.assertDictEqual(clean_device_map({"a": 0, "b": 0, "c": 0}), {"": 0})
+        # Regroups children of level 1 on the same device
+        self.assertDictEqual(
+            clean_device_map({"a.x": 0, "a.y": 0, "b.x": 1, "b.y": 1, "c": 1}), {"a": 0, "b": 1, "c": 1}
+        )
+        # Regroups children of level 2 on the same device
+        self.assertDictEqual(
+            clean_device_map({"a.x": 0, "a.y": 0, "b.x.0": 1, "b.x.1": 1, "b.y.0": 2, "b.y.1": 2, "c": 2}),
+            {"a": 0, "b.x": 1, "b.y": 2, "c": 2},
+        )
+
+    def test_infer_auto_device_map(self):
+        model = ModelForTest()
+        # model has size 236: linear1 64, batchnorm 72, linear2 100
+
+        device_map = infer_auto_device_map(model, max_memory={0: 200, 1: 200})
+        # only linear1 fits on device 0 as we keep memory available for the maximum layer in case of offload
+        self.assertDictEqual(device_map, {"linear1": 0, "batchnorm": 1, "linear2": 1})
+
+        device_map = infer_auto_device_map(model, max_memory={0: 200, 1: 172, 2: 200})
+        # On device 1, we don't care about keeping size available for the max layer, so even if there is just the
+        # size available for batchnorm + linear2, they fit here.
+        self.assertDictEqual(device_map, {"linear1": 0, "batchnorm": 1, "linear2": 1})
+
+        model.linear1.weight = model.linear2.weight
+        device_map = infer_auto_device_map(model, max_memory={0: 200, 1: 200})
+        # By tying weights, the whole model fits on device 0
+        self.assertDictEqual(device_map, {"": 0})
+
+        # When splitting a bigger model, the split is done at the layer level
+        model = nn.Sequential(ModelForTest(), ModelForTest(), ModelForTest())
+        device_map = infer_auto_device_map(model, max_memory={0: 500, 1: 500})
+        self.assertDictEqual(device_map, {"0": 0, "1.linear1": 0, "1.batchnorm": 0, "1.linear2": 1, "2": 1})
+
+        # With no_split_module_classes, it's done at that module level
+        model = nn.Sequential(ModelForTest(), ModelForTest(), ModelForTest())
+        device_map = infer_auto_device_map(
+            model, max_memory={0: 500, 1: 500}, no_split_module_classes=["ModelForTest"]
+        )
+        self.assertDictEqual(device_map, {"0": 0, "1": 1, "2": 1})
+
+        # Now if we have weights tied inside submodules, tied weights are on the same device.
+        model = nn.Sequential(ModelForTest(), ModelForTest(), ModelForTest())
+        layer0 = getattr(model, "0")
+        layer2 = getattr(model, "2")
+        layer0.linear2.weight = layer2.linear2.weight
+        device_map = infer_auto_device_map(model, max_memory={0: 400, 1: 500})
+        expected = {"0": 0, "2.linear2": 0, "1": 1, "2.linear1": 1, "2.batchnorm": 1}
+        self.assertDictEqual(device_map, expected)
diff --git a/tests/test_offload.py b/tests/test_offload.py
new file mode 100644
index 000000000..515d5094d
--- /dev/null
+++ b/tests/test_offload.py
@@ -0,0 +1,87 @@
+# 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 unittest
+from tempfile import TemporaryDirectory
+
+import torch
+import torch.nn as nn
+
+from accelerate.utils import OffloadedWeightsLoader, offload_state_dict
+
+
+class ModelForTest(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear1 = nn.Linear(3, 4)
+        self.batchnorm = nn.BatchNorm1d(4)
+        self.linear2 = nn.Linear(4, 5)
+
+    def forward(self, x):
+        return self.linear2(self.batchnorm(self.linear1(x)))
+
+
+class OffloadTester(unittest.TestCase):
+    def test_offload_state_dict(self):
+        from tempfile import TemporaryDirectory
+
+        model = ModelForTest()
+        with TemporaryDirectory() as tmp_dir:
+            offload_state_dict(tmp_dir, model.state_dict())
+            index_file = os.path.join(tmp_dir, "index.json")
+            self.assertTrue(os.path.isfile(index_file))
+            # TODO: add tests on what is inside the index
+
+            for key in ["linear1.weight", "linear1.bias", "linear2.weight", "linear2.bias"]:
+                weight_file = os.path.join(tmp_dir, f"{key}.dat")
+                self.assertTrue(os.path.isfile(weight_file))
+                # TODO: add tests on the fact weights are properly loaded
+
+    def test_offload_weights_loader(self):
+        model = ModelForTest()
+        state_dict = model.state_dict()
+        cpu_part = {k: v for k, v in state_dict.items() if "linear2" not in k}
+        disk_part = {k: v for k, v in state_dict.items() if "linear2" in k}
+
+        with TemporaryDirectory() as tmp_dir:
+            offload_state_dict(tmp_dir, disk_part)
+            weight_map = OffloadedWeightsLoader(state_dict=cpu_part, save_folder=tmp_dir)
+
+            # Every key is there with the right value
+            self.assertEqual(sorted(weight_map), sorted(state_dict.keys()))
+            for key, param in state_dict.items():
+                self.assertTrue(torch.allclose(param, weight_map[key]))
+
+        cpu_part = {k: v for k, v in state_dict.items() if "weight" in k}
+        disk_part = {k: v for k, v in state_dict.items() if "weight" not in k}
+
+        with TemporaryDirectory() as tmp_dir:
+            offload_state_dict(tmp_dir, disk_part)
+            weight_map = OffloadedWeightsLoader(state_dict=cpu_part, save_folder=tmp_dir)
+
+            # Every key is there with the right value
+            self.assertEqual(sorted(weight_map), sorted(state_dict.keys()))
+            for key, param in state_dict.items():
+                self.assertTrue(torch.allclose(param, weight_map[key]))
+
+        with TemporaryDirectory() as tmp_dir:
+            offload_state_dict(tmp_dir, state_dict)
+            # Duplicates are removed
+            weight_map = OffloadedWeightsLoader(state_dict=cpu_part, save_folder=tmp_dir)
+
+            # Every key is there with the right value
+            self.assertEqual(sorted(weight_map), sorted(state_dict.keys()))
+            for key, param in state_dict.items():
+                self.assertTrue(torch.allclose(param, weight_map[key]))
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 0beba353a..f2414dfe5 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -20,7 +20,7 @@
 import torch
 
 from accelerate.test_utils.training import RegressionModel
-from accelerate.utils import convert_outputs_to_fp32, patch_environment, send_to_device
+from accelerate.utils import convert_outputs_to_fp32, find_device, patch_environment, send_to_device
 
 
 TestNamedTuple = namedtuple("TestNamedTuple", "a b c")
@@ -78,3 +78,8 @@ def test_convert_to_32_lets_model_pickle(self):
         model = RegressionModel()
         model.forward = convert_outputs_to_fp32(model.forward)
         _ = pickle.dumps(model)
+
+    def test_find_device(self):
+        self.assertEqual(find_device([1, "a", torch.tensor([1, 2, 3])]), torch.device("cpu"))
+        self.assertEqual(find_device({"a": 1, "b": torch.tensor([1, 2, 3])}), torch.device("cpu"))
+        self.assertIsNone(find_device([1, "a"]))

Nice! Think that's a good use of a context manager mechanism - it's quite intuitive that inside the context init uses no RAM.


@@ -0,0 +1,203 @@
+# 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
+from contextlib import contextmanager
+from typing import Dict, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from .hooks import AlignDevicesHook, add_hook_to_module, attach_align_device_hook, attach_align_device_hook_on_blocks
+from .utils import OffloadedWeightsLoader, extract_submodules_state_dict, offload_state_dict
+
+
+@contextmanager
+def init_empty_weights(include_buffers: bool = False):
```suggestion
            model's first parameter device.
```
@@ -0,0 +1,203 @@
+# 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
+from contextlib import contextmanager
+from typing import Dict, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from .hooks import AlignDevicesHook, add_hook_to_module, attach_align_device_hook, attach_align_device_hook_on_blocks
+from .utils import OffloadedWeightsLoader, extract_submodules_state_dict, offload_state_dict
+
+
+@contextmanager
+def init_empty_weights(include_buffers: bool = False):
+    """
+    A context manager under which models are initialized with all parameters on the meta device, therefore creating an
+    empty model. Useful when just initializing the model would blow the available RAM.
+
+    Args:
+        include_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to also put all bufeers on the meta device while initializing.
+
+    Example:
+
+    ```pyton
+    import torch.nn as nn
+    from accelerate import init_empty_weights
+
+    # Initialize a model with 100 billions parameters in no time and without using any RAM.
+    with init_empty_weights():
+        tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])
+    ```
+    """
+    old_register_parameter = nn.Module.register_parameter
+    if include_buffers:
+        old_register_buffer = nn.Module.register_buffer
+
+    def register_empty_parameter(module, name, param):
+        old_register_parameter(module, name, param)
+        if param is not None:
+            module._parameters[name] = nn.Parameter(module._parameters[name].to(torch.device("meta")))
+
+    def register_empty_buffer(module, name, buffer):
+        old_register_buffer(module, name, buffer)
+        if buffer is not None:
+            module._buffers[name] = module._buffers[name].to(torch.device("meta"))
+
+    try:
+        nn.Module.register_parameter = register_empty_parameter
+        if include_buffers:
+            nn.Module.register_buffer = register_empty_buffer
+        yield
+    finally:
+        nn.Module.register_parameter = old_register_parameter
+        if include_buffers:
+            nn.Module.register_buffer = old_register_buffer
+
+
+def cpu_offload(
+    model: nn.Module,
+    execution_device: Optional[torch.device] = None,
+    offload_buffers: bool = False,
+    state_dict: Optional[Dict[str, torch.Tensor]] = None,
+):
+    """
+    Activates full CPU offload for a model. As a result, all parameters of the model will be offloaded and only one
+    copy of the state dict of the model will be kept. During the forward pass, parameters will be extracted from that
+    state dict and put on the execution device passed as they are needed, then offloaded again.
+
+    Args:
+        model (`torch.nn.Module`): The model to offload.
+        execution_device (`torch.device`, *optional*):
+            The device on which the forward pass of the model will be executed (should be a GPU). Will default to the
+            model first parameter device.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+        state_dict (`Dict[str, torch.Tensor]`, *optional*):
+            The state dict of the model that will be kept on CPU.
+    """
+    if execution_device is None:
+        execution_device = next(iter(model.parameter())).device
+    if state_dict is None:
+        state_dict = {n: p.to("cpu") for n, p in model.state_dict().items()}
+    attach_align_device_hook(
+        model, execution_device=execution_device, offload=True, offload_buffers=offload_buffers, weights_map=state_dict
+    )
+    add_hook_to_module(model, AlignDevicesHook(io_same_device=True))
+    return model
+
+
+def disk_offload(
+    model: nn.Module,
+    offload_dir: Union[str, os.PathLike],
+    execution_device: Optional[torch.device] = None,
+    offload_buffers: bool = False,
+):
+    """
+    Activates full disk offload for a model. As a result, all parameters of the model will be offloaded as
+    memory-mapped array in a given folder. During the forward pass, parameters will be accessed from that folder and
+    put on the execution device passed as they are needed, then offloaded again.
+
+    Args:
+        model (`torch.nn.Module`): The model to offload.
+        offload_dir (`str` or `os.PathLike`):
+            The folder in which to offload the model weights (or where the model weights are already offloaded).
+        execution_device (`torch.device`, *optional*):
+            The device on which the forward pass of the model will be executed (should be a GPU). Will default to the
+            model first parameter device.
Could it make sense to have models automatically create such a map for themselves? Think this would have to be model-specific and to be in Transformers no?
@@ -0,0 +1,204 @@
+# 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 unittest
+from tempfile import TemporaryDirectory
+
+import torch
+import torch.nn as nn
+
+from accelerate.big_modeling import cpu_offload, disk_offload, dispatch_model, init_empty_weights
+from accelerate.hooks import remove_hook_from_submodules
+from accelerate.test_utils import require_cuda, require_multi_gpu, slow
+from accelerate.utils import offload_state_dict
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+class ModelForTest(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear1 = nn.Linear(3, 4)
+        self.batchnorm = nn.BatchNorm1d(4)
+        self.linear2 = nn.Linear(4, 5)
+
+    def forward(self, x):
+        return self.linear2(self.batchnorm(self.linear1(x)))
+
+
+class BiggerModelForTest(nn.Sequential):
+    def __init__(self):
+        super().__init__()
+        self.linear1 = nn.Linear(3, 4)
+        self.linear2 = nn.Linear(4, 5)
+        self.batchnorm = nn.BatchNorm1d(5)
+        self.linear3 = nn.Linear(5, 6)
+        self.linear4 = nn.Linear(6, 5)
+
+    def forward(self, x):
+        return self.linear4(self.linear3(self.batchnorm(self.linear2(self.linear1(x)))))
+
+
+class BigModelingTester(unittest.TestCase):
+    def test_init_empty_weights(self):
+        # base use
+        with init_empty_weights():
+            module = nn.Linear(4, 5)
+        self.assertEqual(module.weight.device, torch.device("meta"))
+
+        # base use with buffers, they are not touched
+        with init_empty_weights():
+            module = nn.BatchNorm1d(4)
+        self.assertEqual(module.weight.device, torch.device("meta"))
+        self.assertEqual(module.running_mean.device, torch.device("cpu"))
+
+        # Use with include_buffers=True
+        with init_empty_weights(include_buffers=True):
+            module = nn.BatchNorm1d(4)
+        self.assertEqual(module.weight.device, torch.device("meta"))
+        self.assertEqual(module.running_mean.device, torch.device("meta"))
+
+        # Double check we didn't break PyTorch
+        module = nn.BatchNorm1d(4)
+        self.assertEqual(module.weight.device, torch.device("cpu"))
+        self.assertEqual(module.running_mean.device, torch.device("cpu"))
+
+    def test_init_empty_weights_very_large_model(self):
+        # This is a 100 billion parameters model.
+        with init_empty_weights():
+            _ = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])
+
+    def test_cpu_offload(self):
+        model = ModelForTest()
+        x = torch.randn(2, 3)
+        expected = model(x)
+
+        device = torch.device(0 if torch.cuda.is_available() else "cpu")
+
+        cpu_offload(model, execution_device=device)
+        output = model(x)
+        self.assertTrue(torch.allclose(expected, output.cpu()))
+
+        # Clean up for next test.
+        remove_hook_from_submodules(model)
+
+        cpu_offload(model, execution_device=device, offload_buffers=True)
+        output = model(x)
+        self.assertTrue(torch.allclose(expected, output.cpu()))
+
+    @slow
+    @require_cuda
+    def test_cpu_offload_gpt2(self):
+        tokenizer = AutoTokenizer.from_pretrained("gpt2")
+        inputs = tokenizer("Hello world! My name is", return_tensors="pt").to(0)
+
+        gpt2 = AutoModelForCausalLM.from_pretrained("gpt2")
+        cpu_offload(gpt2, execution_device=0)
+        _ = gpt2.generate(inputs["input_ids"])
+
+    def test_disk_offload(self):
+        model = ModelForTest()
+        x = torch.randn(2, 3)
+        expected = model(x)
+
+        device = torch.device(0 if torch.cuda.is_available() else "cpu")
+
+        with TemporaryDirectory() as tmp_dir:
+            disk_offload(model, tmp_dir, execution_device=device)
+            output = model(x)
+            self.assertTrue(torch.allclose(expected, output.cpu()))
+
+            # Clean up for next test.
+            remove_hook_from_submodules(model)
+
+        with TemporaryDirectory() as tmp_dir:
+            disk_offload(model, tmp_dir, execution_device=device, offload_buffers=True)
+            output = model(x)
+            self.assertTrue(torch.allclose(expected, output.cpu()))
+
+    @slow
+    @require_cuda
+    def test_disk_offload_gpt2(self):
+        tokenizer = AutoTokenizer.from_pretrained("gpt2")
+        inputs = tokenizer("Hello world! My name is", return_tensors="pt").to(0)
+
+        gpt2 = AutoModelForCausalLM.from_pretrained("gpt2")
+        with TemporaryDirectory() as tmp_dir:
+            disk_offload(gpt2, tmp_dir, execution_device=0)
+            _ = gpt2.generate(inputs["input_ids"])
+
+    @require_cuda
+    def test_dispatch_model(self):
+        model = ModelForTest()
+        device_map = {"linear1": "disk", "batchnorm": "cpu", "linear2": 0}
+
+        x = torch.randn(2, 3)
+        expected = model(x)
+
+        with TemporaryDirectory() as tmp_dir:
+            dispatch_model(model, device_map, offload_dir=tmp_dir)
+            output = model(x)
+            self.assertTrue(torch.allclose(expected, output.cpu(), atol=1e-5))
+
+    @require_multi_gpu
+    def test_dispatch_model_multi_gpu(self):
+        model = BiggerModelForTest()
+        device_map = {"linear1": "cpu", "linear2": "disk", "batchnorm": "cpu", "linear3": 0, "linear4": 1}
Maybe worth adding a quick test here in the beginning that the parameter names of the model all have to be present in the `device_map()` so that people would get a quick / good error message?
@@ -0,0 +1,203 @@
+# 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
+from contextlib import contextmanager
+from typing import Dict, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from .hooks import AlignDevicesHook, add_hook_to_module, attach_align_device_hook, attach_align_device_hook_on_blocks
+from .utils import OffloadedWeightsLoader, extract_submodules_state_dict, offload_state_dict
+
+
+@contextmanager
+def init_empty_weights(include_buffers: bool = False):
+    """
+    A context manager under which models are initialized with all parameters on the meta device, therefore creating an
+    empty model. Useful when just initializing the model would blow the available RAM.
+
+    Args:
+        include_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to also put all bufeers on the meta device while initializing.
+
+    Example:
+
+    ```pyton
+    import torch.nn as nn
+    from accelerate import init_empty_weights
+
+    # Initialize a model with 100 billions parameters in no time and without using any RAM.
+    with init_empty_weights():
+        tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])
+    ```
+    """
+    old_register_parameter = nn.Module.register_parameter
+    if include_buffers:
+        old_register_buffer = nn.Module.register_buffer
+
+    def register_empty_parameter(module, name, param):
+        old_register_parameter(module, name, param)
+        if param is not None:
+            module._parameters[name] = nn.Parameter(module._parameters[name].to(torch.device("meta")))
+
+    def register_empty_buffer(module, name, buffer):
+        old_register_buffer(module, name, buffer)
+        if buffer is not None:
+            module._buffers[name] = module._buffers[name].to(torch.device("meta"))
+
+    try:
+        nn.Module.register_parameter = register_empty_parameter
+        if include_buffers:
+            nn.Module.register_buffer = register_empty_buffer
+        yield
+    finally:
+        nn.Module.register_parameter = old_register_parameter
+        if include_buffers:
+            nn.Module.register_buffer = old_register_buffer
+
+
+def cpu_offload(
+    model: nn.Module,
+    execution_device: Optional[torch.device] = None,
+    offload_buffers: bool = False,
+    state_dict: Optional[Dict[str, torch.Tensor]] = None,
+):
+    """
+    Activates full CPU offload for a model. As a result, all parameters of the model will be offloaded and only one
+    copy of the state dict of the model will be kept. During the forward pass, parameters will be extracted from that
+    state dict and put on the execution device passed as they are needed, then offloaded again.
+
+    Args:
+        model (`torch.nn.Module`): The model to offload.
+        execution_device (`torch.device`, *optional*):
+            The device on which the forward pass of the model will be executed (should be a GPU). Will default to the
+            model first parameter device.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+        state_dict (`Dict[str, torch.Tensor]`, *optional*):
+            The state dict of the model that will be kept on CPU.
+    """
+    if execution_device is None:
+        execution_device = next(iter(model.parameter())).device
+    if state_dict is None:
+        state_dict = {n: p.to("cpu") for n, p in model.state_dict().items()}
+    attach_align_device_hook(
+        model, execution_device=execution_device, offload=True, offload_buffers=offload_buffers, weights_map=state_dict
+    )
+    add_hook_to_module(model, AlignDevicesHook(io_same_device=True))
+    return model
+
+
+def disk_offload(
+    model: nn.Module,
+    offload_dir: Union[str, os.PathLike],
+    execution_device: Optional[torch.device] = None,
+    offload_buffers: bool = False,
+):
+    """
+    Activates full disk offload for a model. As a result, all parameters of the model will be offloaded as
+    memory-mapped array in a given folder. During the forward pass, parameters will be accessed from that folder and
+    put on the execution device passed as they are needed, then offloaded again.
+
+    Args:
+        model (`torch.nn.Module`): The model to offload.
+        offload_dir (`str` or `os.PathLike`):
+            The folder in which to offload the model weights (or where the model weights are already offloaded).
+        execution_device (`torch.device`, *optional*):
+            The device on which the forward pass of the model will be executed (should be a GPU). Will default to the
+            model first parameter device.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+    """
+    if not os.path.isdir(offload_dir) or not os.path.isfile(os.path.join(offload_dir, "index.json")):
+        offload_state_dict(offload_dir, model.state_dict())
+    if execution_device is None:
+        execution_device = next(iter(model.parameter())).device
+    weights_map = OffloadedWeightsLoader(save_folder=offload_dir)
+    attach_align_device_hook(
+        model,
+        execution_device=execution_device,
+        offload=True,
+        offload_buffers=offload_buffers,
+        weights_map=weights_map,
+    )
+    add_hook_to_module(model, AlignDevicesHook(io_same_device=True))
+    return model
+
+
+def dispatch_model(
+    model: nn.Module,
+    device_map: Dict[str, Union[str, int, torch.device]],
+    main_device: Optional[torch.device] = None,
+    state_dict: Optional[Dict[str, torch.Tensor]] = None,
+    offload_dir: Union[str, os.PathLike] = None,
+    offload_buffers: bool = False,
+):
+    """
+    Dispatches a model according to a given device map. Layers of the model might be spread across GPUs, offloaded on
+    the CPU or even the disk.
+
+    Args:
+        model (`torch.nn.Module`): The model to dispatch.
+        device_map (`Dict[str, Union[str, int, torch.device]]`):
+            A dictionary mapping module names to the device they should go to. Note that `"disk"` is accepted even if
+            it's not a proper value for `torch.device`.
+        main_device (`str`, `int` or `torch.device`, *optional*):
+            The main execution device. Will default to the first device in the `device_map` different from `"cpu"` or
+            `"disk"`.
+        state_dict (`Dict[str, torch.Tensor]`, *optional*):
+            The state dict of the part of the model that will be kept on CPU.
+        offload_dir (`str` or `os.PathLike`):
+            The folder in which to offload the model weights (or where the model weights are already offloaded).
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+    """
Think a type hint for `max_memory` that it should be a dict could be useful here
@@ -0,0 +1,454 @@
+# 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 gc
+import json
+import os
+import re
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+WEIGHTS_INDEX_NAME = "pytorch_model.bin.index.json"
+
+
+def convert_file_size_to_int(size: Union[int, str]):
+    """
+    Converts a size expressed as a string with digits an unit (like `"5MB"`) to an integer (in bytes).
+
+    Args:
+        size (`int` or `str`): The size to convert. Will be directly returned if an `int`.
+
+    Example:
+
+    ```py
+    >>> convert_file_size_to_int("1MiB")
+    1048576
+    ```
+    """
+    if isinstance(size, int):
+        return size
+    if size.upper().endswith("GIB"):
+        return int(size[:-3]) * (2**30)
+    if size.upper().endswith("MIB"):
+        return int(size[:-3]) * (2**20)
+    if size.upper().endswith("KIB"):
+        return int(size[:-3]) * (2**10)
+    if size.upper().endswith("GB"):
+        int_size = int(size[:-2]) * (10**9)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("MB"):
+        int_size = int(size[:-2]) * (10**6)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("KB"):
+        int_size = int(size[:-2]) * (10**3)
+        return int_size // 8 if size.endswith("b") else int_size
+    raise ValueError("`size` is not in a valid format. Use an integer followed by the unit, e.g., '5GB'.")
+
+
+def dtype_byte_size(dtype):
+    """
+    Returns the size (in bytes) occupied by one parameter of type `dtype`.
+
+    Example:
+
+    ```py
+    >>> dtype_byte_size(torch.float32)
+    4
+    ```
+    """
+    if dtype == torch.bool:
+        return 1 / 8
+    bit_search = re.search("[^\d](\d+)$", str(dtype))
+    if bit_search is None:
+        raise ValueError(f"`dtype` is not a valid dtype: {dtype}.")
+    bit_size = int(bit_search.groups()[0])
+    return bit_size // 8
+
+
+def set_module_tensor_to_device(
+    module: nn.Module, tensor_name: str, device: Union[int, str, torch.device], value: Optional[torch.Tensor] = None
+):
+    """
+    A helper function to set a given tensor (parameter of buffer) of a module on a specific device (note that doing
+    `param.to(device)` creates a new tensor not linked to the parameter, which is why we need this function).
+
+    Args:
+        module (`torch.nn.Module`): The module in which the tensor we want to move lives.
+        param_name (`str`): The full name of the parameter/buffer.
+        device (`int`, `str` or `torch.device`): The device on which to set the tensor.
+        value (`torch.Tensor`, *optional*): The value of the tensor (useful when going from the meta device to any
+            other device).
+    """
+    # Recurse if needed
+    if "." in tensor_name:
+        splits = tensor_name.split(".")
+        for split in splits[:-1]:
+            new_module = getattr(module, split)
+            if new_module is None:
+                raise ValueError(f"{module} has no attribute {split}.")
+            module = new_module
+        tensor_name = splits[-1]
+
+    if tensor_name not in module._parameters and tensor_name not in module._buffers:
+        raise ValueError(f"{module} does not have a parameter or a buffer named {tensor_name}.")
+    is_buffer = tensor_name in module._buffers
+    old_value = getattr(module, tensor_name)
+
+    if old_value.device == torch.device("meta") and device != torch.device("meta") and value is None:
+        raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {device}.")
+
+    with torch.no_grad():
+        if value is None:
+            new_value = old_value.to(device)
+        elif isinstance(value, torch.Tensor):
+            new_value = value.to(device)
+        else:
+            new_value = torch.tensor(value, device=device)
+    if is_buffer:
+        module._buffers[tensor_name] = new_value
+    else:
+        new_value = nn.Parameter(new_value, requires_grad=old_value.requires_grad)
+        module._parameters[tensor_name] = new_value
+
+
+def named_module_tensors(module: nn.Module, include_buffers: bool = True, recurse: bool = False):
+    """
+    A helper function that gathers all the tensors (parameters + buffers) of a given module. If `include_buffers=True`
+    it's the same as doing `module.named_parameters(recurse=recurse) + module.named_buffers(recurse=recurse)`.
+
+    Args:
+        module (`torch.nn.Module`): The module we want the tensors or.
+        include_buffer (`bool`, *optional*, defaults to `True`): Whether or not to include the buffers in the result.
+        recurse (`bool`, *optional`, defaults to `False`):
+            Whether or not to go look in every submodule or just return the direct parameters and buffers.
+    """
+    for named_parameter in module.named_parameters(recurse=recurse):
+        yield named_parameter
+
+    if include_buffers:
+        for named_buffer in module.named_buffers(recurse=recurse):
+            yield named_buffer
+
+
+def find_tied_parameters(model: nn.Module, named_parameters=None, prefix="", result=None):
+    """
+    Find the tied parameters in a given model.
+
+    Args:
+        model (`torch.nn.Module`): The model to inspect.
+
+    <Tip warning={true}>
+
+    The signature has more arguments, but they are for the recursive part of this function and you should ignore them.
+
+    </Tip>
+
+    Example:
+
+
+    ```py
+    >>> from collections import OrderedDict
+    >>> import torch.nn as nn
+
+    >>> model = nn.Sequential(OrderedDict([("linear1", nn.Linear(4, 4)), ("linear2", nn.Linear(4, 4))]))
+    >>> model.linear2.weight = test_model.linear1.weight
+    >>> find_tied_parameters(test_model)
+    {'linear1.weight': 'linear2.weight'}
+    ```
+
+    Returns:
+        Dict[str, str]: A dictionary mapping tied parameter names to the name of the parameter they are tied to.
+    """
+    # Initialize result and named_parameters before recursing.
+    if result is None:
+        result = {}
+    if named_parameters is None:
+        named_parameters = {n: p for n, p in model.named_parameters()}
+    else:
+        # A tied parameter will not be in the full `named_parameters` seen above but will be in the `named_parameters`
+        # of the submodule it belongs to. So while recursing we track the names that are not in the initial
+        # `named_parameters`.
+        for name, parameter in model.named_parameters():
+            full_name = name if prefix == "" else f"{prefix}.{name}"
+            if full_name not in named_parameters:
+                # When we find one, it has to be one of the existing parameters.
+                for new_name, new_param in named_parameters.items():
+                    if new_param is parameter:
+                        result[new_name] = full_name
+
+    # Once we have treated direct parameters, we move to the child modules.
+    for name, child in model.named_children():
+        child_name = name if prefix == "" else f"{prefix}.{name}"
+        find_tied_parameters(child, named_parameters, prefix=child_name, result=result)
+
+    return result
+
+
+def compute_module_sizes(model: nn.Module):
+    """
+    Compute the size of each submodule of a given model.
+    """
+    module_sizes = defaultdict(int)
+    for name, tensor in named_module_tensors(model, recurse=True):
+        size = tensor.numel() * dtype_byte_size(tensor.dtype)
+        name_parts = name.split(".")
+        for idx in range(len(name_parts) + 1):
+            module_sizes[".".join(name_parts[:idx])] += size
+
+    return module_sizes
+
+
+def get_max_layer_size(
+    modules: List[Tuple[str, torch.nn.Module]], module_sizes: Dict[str, int], no_split_module_classes: List[str]
+):
+    """
+    Utility function that will scan a list of named modules and return the maximum size used by one full layer. The
+    definition of a layer being:
+    - a module with no direct children (just parameters and buffers)
+    - a module whose class name is in the list `no_split_module_classes`
+
+    Args:
+        modules (`List[Tuple[str, torch.nn.Module]]`):
+            The list of named modules where we want to determine the maximum layer size.
+        module_sizes (`Dict[str, int]`):
+            A dictionary mapping each layer name to its size (as generated by `compute_module_sizes`).
+        no_split_module_classes (`List[str]`):
+            A list of class names for layers we don't want to be split.
+
+    Returns:
+        `Tuple[int, List[str]]`: The maximum size of a layer with the list of layer names realizing that maximum size.
+    """
+    max_size = 0
+    layer_names = []
+    modules_to_treat = modules.copy()
+    while len(modules_to_treat) > 0:
+        module_name, module = modules_to_treat.pop(0)
+        modules_children = list(module.named_children())
+        if len(modules_children) == 0 or module.__class__.__name__ in no_split_module_classes:
+            # No splitting this one so we compare to the max_size
+            size = module_sizes[module_name]
+            if size > max_size:
+                max_size = size
+                layer_names = [module_name]
+            elif size == max_size:
+                layer_names.append(module_name)
+        else:
+            modules_to_treat = [(f"{module_name}.{n}", v) for n, v in modules_children] + modules_to_treat
+    return max_size, layer_names
+
+
+def get_max_memory(max_memory=None, buffer="1500MB"):
+    """
+    Get the maximum memory available if nothing is passed, converts string to int otherwise.
+    """
+    import psutil
+
+    if max_memory is None:
+        if not torch.cuda.is_available():
+            max_memory = {}
+        else:
+            # Make sure CUDA is initialized on each GPU to have the right memory info.
+            for i in range(torch.cuda.device_count()):
+                _ = torch.tensor([0], device=i)
+            max_memory = {i: mem - convert_file_size_to_int(buffer) for i, mem in enumerate(torch.cuda.mem_get_info())}
+        max_memory["cpu"] = psutil.virtual_memory().available
+        return max_memory
+
+    for key in max_memory:
+        if isinstance(max_memory[key], str):
+            max_memory[key] = convert_file_size_to_int(max_memory[key])
+    return max_memory
+
+
+def infer_auto_device_map(model, max_memory=None, no_split_module_classes=None):
This is what the `infer_auto_device_map` function does.
@@ -0,0 +1,204 @@
+# 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 unittest
+from tempfile import TemporaryDirectory
+
+import torch
+import torch.nn as nn
+
+from accelerate.big_modeling import cpu_offload, disk_offload, dispatch_model, init_empty_weights
+from accelerate.hooks import remove_hook_from_submodules
+from accelerate.test_utils import require_cuda, require_multi_gpu, slow
+from accelerate.utils import offload_state_dict
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+class ModelForTest(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear1 = nn.Linear(3, 4)
+        self.batchnorm = nn.BatchNorm1d(4)
+        self.linear2 = nn.Linear(4, 5)
+
+    def forward(self, x):
+        return self.linear2(self.batchnorm(self.linear1(x)))
+
+
+class BiggerModelForTest(nn.Sequential):
+    def __init__(self):
+        super().__init__()
+        self.linear1 = nn.Linear(3, 4)
+        self.linear2 = nn.Linear(4, 5)
+        self.batchnorm = nn.BatchNorm1d(5)
+        self.linear3 = nn.Linear(5, 6)
+        self.linear4 = nn.Linear(6, 5)
+
+    def forward(self, x):
+        return self.linear4(self.linear3(self.batchnorm(self.linear2(self.linear1(x)))))
+
+
+class BigModelingTester(unittest.TestCase):
+    def test_init_empty_weights(self):
+        # base use
+        with init_empty_weights():
+            module = nn.Linear(4, 5)
+        self.assertEqual(module.weight.device, torch.device("meta"))
+
+        # base use with buffers, they are not touched
+        with init_empty_weights():
+            module = nn.BatchNorm1d(4)
+        self.assertEqual(module.weight.device, torch.device("meta"))
+        self.assertEqual(module.running_mean.device, torch.device("cpu"))
+
+        # Use with include_buffers=True
+        with init_empty_weights(include_buffers=True):
+            module = nn.BatchNorm1d(4)
+        self.assertEqual(module.weight.device, torch.device("meta"))
+        self.assertEqual(module.running_mean.device, torch.device("meta"))
+
+        # Double check we didn't break PyTorch
+        module = nn.BatchNorm1d(4)
+        self.assertEqual(module.weight.device, torch.device("cpu"))
+        self.assertEqual(module.running_mean.device, torch.device("cpu"))
+
+    def test_init_empty_weights_very_large_model(self):
+        # This is a 100 billion parameters model.
+        with init_empty_weights():
+            _ = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])
+
+    def test_cpu_offload(self):
+        model = ModelForTest()
+        x = torch.randn(2, 3)
+        expected = model(x)
+
+        device = torch.device(0 if torch.cuda.is_available() else "cpu")
+
+        cpu_offload(model, execution_device=device)
+        output = model(x)
+        self.assertTrue(torch.allclose(expected, output.cpu()))
+
+        # Clean up for next test.
+        remove_hook_from_submodules(model)
+
+        cpu_offload(model, execution_device=device, offload_buffers=True)
+        output = model(x)
+        self.assertTrue(torch.allclose(expected, output.cpu()))
+
+    @slow
+    @require_cuda
+    def test_cpu_offload_gpt2(self):
+        tokenizer = AutoTokenizer.from_pretrained("gpt2")
+        inputs = tokenizer("Hello world! My name is", return_tensors="pt").to(0)
+
+        gpt2 = AutoModelForCausalLM.from_pretrained("gpt2")
+        cpu_offload(gpt2, execution_device=0)
+        _ = gpt2.generate(inputs["input_ids"])
+
+    def test_disk_offload(self):
+        model = ModelForTest()
+        x = torch.randn(2, 3)
+        expected = model(x)
+
+        device = torch.device(0 if torch.cuda.is_available() else "cpu")
+
+        with TemporaryDirectory() as tmp_dir:
+            disk_offload(model, tmp_dir, execution_device=device)
+            output = model(x)
+            self.assertTrue(torch.allclose(expected, output.cpu()))
+
+            # Clean up for next test.
+            remove_hook_from_submodules(model)
+
+        with TemporaryDirectory() as tmp_dir:
+            disk_offload(model, tmp_dir, execution_device=device, offload_buffers=True)
+            output = model(x)
+            self.assertTrue(torch.allclose(expected, output.cpu()))
+
+    @slow
+    @require_cuda
+    def test_disk_offload_gpt2(self):
+        tokenizer = AutoTokenizer.from_pretrained("gpt2")
+        inputs = tokenizer("Hello world! My name is", return_tensors="pt").to(0)
+
+        gpt2 = AutoModelForCausalLM.from_pretrained("gpt2")
+        with TemporaryDirectory() as tmp_dir:
+            disk_offload(gpt2, tmp_dir, execution_device=0)
+            _ = gpt2.generate(inputs["input_ids"])
+
+    @require_cuda
+    def test_dispatch_model(self):
+        model = ModelForTest()
+        device_map = {"linear1": "disk", "batchnorm": "cpu", "linear2": 0}
+
+        x = torch.randn(2, 3)
+        expected = model(x)
+
+        with TemporaryDirectory() as tmp_dir:
+            dispatch_model(model, device_map, offload_dir=tmp_dir)
+            output = model(x)
+            self.assertTrue(torch.allclose(expected, output.cpu(), atol=1e-5))
+
+    @require_multi_gpu
+    def test_dispatch_model_multi_gpu(self):
+        model = BiggerModelForTest()
+        device_map = {"linear1": "cpu", "linear2": "disk", "batchnorm": "cpu", "linear3": 0, "linear4": 1}
Yes, definitely needed!
@@ -0,0 +1,203 @@
+# 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
+from contextlib import contextmanager
+from typing import Dict, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from .hooks import AlignDevicesHook, add_hook_to_module, attach_align_device_hook, attach_align_device_hook_on_blocks
+from .utils import OffloadedWeightsLoader, extract_submodules_state_dict, offload_state_dict
+
+
+@contextmanager
+def init_empty_weights(include_buffers: bool = False):
+    """
+    A context manager under which models are initialized with all parameters on the meta device, therefore creating an
+    empty model. Useful when just initializing the model would blow the available RAM.
+
+    Args:
+        include_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to also put all bufeers on the meta device while initializing.
+
+    Example:
+
+    ```pyton
+    import torch.nn as nn
+    from accelerate import init_empty_weights
+
+    # Initialize a model with 100 billions parameters in no time and without using any RAM.
+    with init_empty_weights():
+        tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])
+    ```
+    """
+    old_register_parameter = nn.Module.register_parameter
+    if include_buffers:
+        old_register_buffer = nn.Module.register_buffer
+
+    def register_empty_parameter(module, name, param):
+        old_register_parameter(module, name, param)
+        if param is not None:
+            module._parameters[name] = nn.Parameter(module._parameters[name].to(torch.device("meta")))
+
+    def register_empty_buffer(module, name, buffer):
+        old_register_buffer(module, name, buffer)
+        if buffer is not None:
+            module._buffers[name] = module._buffers[name].to(torch.device("meta"))
+
+    try:
+        nn.Module.register_parameter = register_empty_parameter
+        if include_buffers:
+            nn.Module.register_buffer = register_empty_buffer
+        yield
+    finally:
+        nn.Module.register_parameter = old_register_parameter
+        if include_buffers:
+            nn.Module.register_buffer = old_register_buffer
+
+
+def cpu_offload(
+    model: nn.Module,
+    execution_device: Optional[torch.device] = None,
+    offload_buffers: bool = False,
+    state_dict: Optional[Dict[str, torch.Tensor]] = None,
+):
+    """
+    Activates full CPU offload for a model. As a result, all parameters of the model will be offloaded and only one
+    copy of the state dict of the model will be kept. During the forward pass, parameters will be extracted from that
+    state dict and put on the execution device passed as they are needed, then offloaded again.
+
+    Args:
+        model (`torch.nn.Module`): The model to offload.
+        execution_device (`torch.device`, *optional*):
+            The device on which the forward pass of the model will be executed (should be a GPU). Will default to the
+            model first parameter device.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+        state_dict (`Dict[str, torch.Tensor]`, *optional*):
+            The state dict of the model that will be kept on CPU.
+    """
+    if execution_device is None:
+        execution_device = next(iter(model.parameter())).device
+    if state_dict is None:
+        state_dict = {n: p.to("cpu") for n, p in model.state_dict().items()}
+    attach_align_device_hook(
+        model, execution_device=execution_device, offload=True, offload_buffers=offload_buffers, weights_map=state_dict
+    )
+    add_hook_to_module(model, AlignDevicesHook(io_same_device=True))
+    return model
+
+
+def disk_offload(
+    model: nn.Module,
+    offload_dir: Union[str, os.PathLike],
+    execution_device: Optional[torch.device] = None,
+    offload_buffers: bool = False,
+):
+    """
+    Activates full disk offload for a model. As a result, all parameters of the model will be offloaded as
+    memory-mapped array in a given folder. During the forward pass, parameters will be accessed from that folder and
+    put on the execution device passed as they are needed, then offloaded again.
+
+    Args:
+        model (`torch.nn.Module`): The model to offload.
+        offload_dir (`str` or `os.PathLike`):
+            The folder in which to offload the model weights (or where the model weights are already offloaded).
+        execution_device (`torch.device`, *optional*):
+            The device on which the forward pass of the model will be executed (should be a GPU). Will default to the
+            model first parameter device.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+    """
+    if not os.path.isdir(offload_dir) or not os.path.isfile(os.path.join(offload_dir, "index.json")):
+        offload_state_dict(offload_dir, model.state_dict())
+    if execution_device is None:
+        execution_device = next(iter(model.parameter())).device
+    weights_map = OffloadedWeightsLoader(save_folder=offload_dir)
+    attach_align_device_hook(
+        model,
+        execution_device=execution_device,
+        offload=True,
+        offload_buffers=offload_buffers,
+        weights_map=weights_map,
+    )
+    add_hook_to_module(model, AlignDevicesHook(io_same_device=True))
+    return model
+
+
+def dispatch_model(
+    model: nn.Module,
+    device_map: Dict[str, Union[str, int, torch.device]],
+    main_device: Optional[torch.device] = None,
+    state_dict: Optional[Dict[str, torch.Tensor]] = None,
+    offload_dir: Union[str, os.PathLike] = None,
+    offload_buffers: bool = False,
+):
+    """
+    Dispatches a model according to a given device map. Layers of the model might be spread across GPUs, offloaded on
+    the CPU or even the disk.
+
+    Args:
+        model (`torch.nn.Module`): The model to dispatch.
+        device_map (`Dict[str, Union[str, int, torch.device]]`):
+            A dictionary mapping module names to the device they should go to. Note that `"disk"` is accepted even if
+            it's not a proper value for `torch.device`.
+        main_device (`str`, `int` or `torch.device`, *optional*):
+            The main execution device. Will default to the first device in the `device_map` different from `"cpu"` or
+            `"disk"`.
+        state_dict (`Dict[str, torch.Tensor]`, *optional*):
+            The state dict of the part of the model that will be kept on CPU.
+        offload_dir (`str` or `os.PathLike`):
+            The folder in which to offload the model weights (or where the model weights are already offloaded).
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+    """
```suggestion
            Whether or not to also put all buffers on the meta device while initializing.
```
@@ -0,0 +1,206 @@
+# 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
+from contextlib import contextmanager
+from typing import Dict, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from .hooks import AlignDevicesHook, add_hook_to_module, attach_align_device_hook, attach_align_device_hook_on_blocks
+from .utils import OffloadedWeightsLoader, check_device_map, extract_submodules_state_dict, offload_state_dict
+
+
+@contextmanager
+def init_empty_weights(include_buffers: bool = False):
+    """
+    A context manager under which models are initialized with all parameters on the meta device, therefore creating an
+    empty model. Useful when just initializing the model would blow the available RAM.
+
+    Args:
+        include_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to also put all bufeers on the meta device while initializing.
```suggestion
        model (`torch.nn.Module`): 
            The model to offload.
```
@@ -0,0 +1,206 @@
+# 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
+from contextlib import contextmanager
+from typing import Dict, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from .hooks import AlignDevicesHook, add_hook_to_module, attach_align_device_hook, attach_align_device_hook_on_blocks
+from .utils import OffloadedWeightsLoader, check_device_map, extract_submodules_state_dict, offload_state_dict
+
+
+@contextmanager
+def init_empty_weights(include_buffers: bool = False):
+    """
+    A context manager under which models are initialized with all parameters on the meta device, therefore creating an
+    empty model. Useful when just initializing the model would blow the available RAM.
+
+    Args:
+        include_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to also put all bufeers on the meta device while initializing.
+
+    Example:
+
+    ```pyton
+    import torch.nn as nn
+    from accelerate import init_empty_weights
+
+    # Initialize a model with 100 billions parameters in no time and without using any RAM.
+    with init_empty_weights():
+        tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])
+    ```
+    """
+    old_register_parameter = nn.Module.register_parameter
+    if include_buffers:
+        old_register_buffer = nn.Module.register_buffer
+
+    def register_empty_parameter(module, name, param):
+        old_register_parameter(module, name, param)
+        if param is not None:
+            module._parameters[name] = nn.Parameter(module._parameters[name].to(torch.device("meta")))
+
+    def register_empty_buffer(module, name, buffer):
+        old_register_buffer(module, name, buffer)
+        if buffer is not None:
+            module._buffers[name] = module._buffers[name].to(torch.device("meta"))
+
+    try:
+        nn.Module.register_parameter = register_empty_parameter
+        if include_buffers:
+            nn.Module.register_buffer = register_empty_buffer
+        yield
+    finally:
+        nn.Module.register_parameter = old_register_parameter
+        if include_buffers:
+            nn.Module.register_buffer = old_register_buffer
+
+
+def cpu_offload(
+    model: nn.Module,
+    execution_device: Optional[torch.device] = None,
+    offload_buffers: bool = False,
+    state_dict: Optional[Dict[str, torch.Tensor]] = None,
+):
+    """
+    Activates full CPU offload for a model. As a result, all parameters of the model will be offloaded and only one
+    copy of the state dict of the model will be kept. During the forward pass, parameters will be extracted from that
+    state dict and put on the execution device passed as they are needed, then offloaded again.
+
+    Args:
+        model (`torch.nn.Module`): The model to offload.
```suggestion
        model (`torch.nn.Module`): 
            The model to dispatch.
```
@@ -0,0 +1,206 @@
+# 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
+from contextlib import contextmanager
+from typing import Dict, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from .hooks import AlignDevicesHook, add_hook_to_module, attach_align_device_hook, attach_align_device_hook_on_blocks
+from .utils import OffloadedWeightsLoader, check_device_map, extract_submodules_state_dict, offload_state_dict
+
+
+@contextmanager
+def init_empty_weights(include_buffers: bool = False):
+    """
+    A context manager under which models are initialized with all parameters on the meta device, therefore creating an
+    empty model. Useful when just initializing the model would blow the available RAM.
+
+    Args:
+        include_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to also put all bufeers on the meta device while initializing.
+
+    Example:
+
+    ```pyton
+    import torch.nn as nn
+    from accelerate import init_empty_weights
+
+    # Initialize a model with 100 billions parameters in no time and without using any RAM.
+    with init_empty_weights():
+        tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])
+    ```
+    """
+    old_register_parameter = nn.Module.register_parameter
+    if include_buffers:
+        old_register_buffer = nn.Module.register_buffer
+
+    def register_empty_parameter(module, name, param):
+        old_register_parameter(module, name, param)
+        if param is not None:
+            module._parameters[name] = nn.Parameter(module._parameters[name].to(torch.device("meta")))
+
+    def register_empty_buffer(module, name, buffer):
+        old_register_buffer(module, name, buffer)
+        if buffer is not None:
+            module._buffers[name] = module._buffers[name].to(torch.device("meta"))
+
+    try:
+        nn.Module.register_parameter = register_empty_parameter
+        if include_buffers:
+            nn.Module.register_buffer = register_empty_buffer
+        yield
+    finally:
+        nn.Module.register_parameter = old_register_parameter
+        if include_buffers:
+            nn.Module.register_buffer = old_register_buffer
+
+
+def cpu_offload(
+    model: nn.Module,
+    execution_device: Optional[torch.device] = None,
+    offload_buffers: bool = False,
+    state_dict: Optional[Dict[str, torch.Tensor]] = None,
+):
+    """
+    Activates full CPU offload for a model. As a result, all parameters of the model will be offloaded and only one
+    copy of the state dict of the model will be kept. During the forward pass, parameters will be extracted from that
+    state dict and put on the execution device passed as they are needed, then offloaded again.
+
+    Args:
+        model (`torch.nn.Module`): The model to offload.
+        execution_device (`torch.device`, *optional*):
+            The device on which the forward pass of the model will be executed (should be a GPU). Will default to the
+            model first parameter device.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+        state_dict (`Dict[str, torch.Tensor]`, *optional*):
+            The state dict of the model that will be kept on CPU.
+    """
+    if execution_device is None:
+        execution_device = next(iter(model.parameter())).device
+    if state_dict is None:
+        state_dict = {n: p.to("cpu") for n, p in model.state_dict().items()}
+    attach_align_device_hook(
+        model, execution_device=execution_device, offload=True, offload_buffers=offload_buffers, weights_map=state_dict
+    )
+    add_hook_to_module(model, AlignDevicesHook(io_same_device=True))
+    return model
+
+
+def disk_offload(
+    model: nn.Module,
+    offload_dir: Union[str, os.PathLike],
+    execution_device: Optional[torch.device] = None,
+    offload_buffers: bool = False,
+):
+    """
+    Activates full disk offload for a model. As a result, all parameters of the model will be offloaded as
+    memory-mapped array in a given folder. During the forward pass, parameters will be accessed from that folder and
+    put on the execution device passed as they are needed, then offloaded again.
+
+    Args:
+        model (`torch.nn.Module`): The model to offload.
+        offload_dir (`str` or `os.PathLike`):
+            The folder in which to offload the model weights (or where the model weights are already offloaded).
+        execution_device (`torch.device`, *optional*):
+            The device on which the forward pass of the model will be executed (should be a GPU). Will default to the
+            model's first parameter device.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+    """
+    if not os.path.isdir(offload_dir) or not os.path.isfile(os.path.join(offload_dir, "index.json")):
+        offload_state_dict(offload_dir, model.state_dict())
+    if execution_device is None:
+        execution_device = next(iter(model.parameter())).device
+    weights_map = OffloadedWeightsLoader(save_folder=offload_dir)
+    attach_align_device_hook(
+        model,
+        execution_device=execution_device,
+        offload=True,
+        offload_buffers=offload_buffers,
+        weights_map=weights_map,
+    )
+    add_hook_to_module(model, AlignDevicesHook(io_same_device=True))
+    return model
+
+
+def dispatch_model(
+    model: nn.Module,
+    device_map: Dict[str, Union[str, int, torch.device]],
+    main_device: Optional[torch.device] = None,
+    state_dict: Optional[Dict[str, torch.Tensor]] = None,
+    offload_dir: Union[str, os.PathLike] = None,
+    offload_buffers: bool = False,
+):
+    """
+    Dispatches a model according to a given device map. Layers of the model might be spread across GPUs, offloaded on
+    the CPU or even the disk.
+
+    Args:
+        model (`torch.nn.Module`): The model to dispatch.
I don't wonder if there's a clever way to check when someone tries this with a model named `nn.Sequential`, and give them hints as to where those are named. Because at first as a newbie I'd think "maybe that's not supported"? When in reality I'd need to check `net.state_dict()`.

One potential solution would be changing the phrasing some:

```suggestion
        device_map (`Dict[str, Union[str, int, torch.device]]`):
            A dictionary mapping module names in the models `state_dict` to the device they should go to. Note that `"disk"` is accepted even if
            it's not a proper value for `torch.device`.
```

@@ -0,0 +1,206 @@
+# 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
+from contextlib import contextmanager
+from typing import Dict, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from .hooks import AlignDevicesHook, add_hook_to_module, attach_align_device_hook, attach_align_device_hook_on_blocks
+from .utils import OffloadedWeightsLoader, check_device_map, extract_submodules_state_dict, offload_state_dict
+
+
+@contextmanager
+def init_empty_weights(include_buffers: bool = False):
+    """
+    A context manager under which models are initialized with all parameters on the meta device, therefore creating an
+    empty model. Useful when just initializing the model would blow the available RAM.
+
+    Args:
+        include_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to also put all bufeers on the meta device while initializing.
+
+    Example:
+
+    ```pyton
+    import torch.nn as nn
+    from accelerate import init_empty_weights
+
+    # Initialize a model with 100 billions parameters in no time and without using any RAM.
+    with init_empty_weights():
+        tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])
+    ```
+    """
+    old_register_parameter = nn.Module.register_parameter
+    if include_buffers:
+        old_register_buffer = nn.Module.register_buffer
+
+    def register_empty_parameter(module, name, param):
+        old_register_parameter(module, name, param)
+        if param is not None:
+            module._parameters[name] = nn.Parameter(module._parameters[name].to(torch.device("meta")))
+
+    def register_empty_buffer(module, name, buffer):
+        old_register_buffer(module, name, buffer)
+        if buffer is not None:
+            module._buffers[name] = module._buffers[name].to(torch.device("meta"))
+
+    try:
+        nn.Module.register_parameter = register_empty_parameter
+        if include_buffers:
+            nn.Module.register_buffer = register_empty_buffer
+        yield
+    finally:
+        nn.Module.register_parameter = old_register_parameter
+        if include_buffers:
+            nn.Module.register_buffer = old_register_buffer
+
+
+def cpu_offload(
+    model: nn.Module,
+    execution_device: Optional[torch.device] = None,
+    offload_buffers: bool = False,
+    state_dict: Optional[Dict[str, torch.Tensor]] = None,
+):
+    """
+    Activates full CPU offload for a model. As a result, all parameters of the model will be offloaded and only one
+    copy of the state dict of the model will be kept. During the forward pass, parameters will be extracted from that
+    state dict and put on the execution device passed as they are needed, then offloaded again.
+
+    Args:
+        model (`torch.nn.Module`): The model to offload.
+        execution_device (`torch.device`, *optional*):
+            The device on which the forward pass of the model will be executed (should be a GPU). Will default to the
+            model first parameter device.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+        state_dict (`Dict[str, torch.Tensor]`, *optional*):
+            The state dict of the model that will be kept on CPU.
+    """
+    if execution_device is None:
+        execution_device = next(iter(model.parameter())).device
+    if state_dict is None:
+        state_dict = {n: p.to("cpu") for n, p in model.state_dict().items()}
+    attach_align_device_hook(
+        model, execution_device=execution_device, offload=True, offload_buffers=offload_buffers, weights_map=state_dict
+    )
+    add_hook_to_module(model, AlignDevicesHook(io_same_device=True))
+    return model
+
+
+def disk_offload(
+    model: nn.Module,
+    offload_dir: Union[str, os.PathLike],
+    execution_device: Optional[torch.device] = None,
+    offload_buffers: bool = False,
+):
+    """
+    Activates full disk offload for a model. As a result, all parameters of the model will be offloaded as
+    memory-mapped array in a given folder. During the forward pass, parameters will be accessed from that folder and
+    put on the execution device passed as they are needed, then offloaded again.
+
+    Args:
+        model (`torch.nn.Module`): The model to offload.
+        offload_dir (`str` or `os.PathLike`):
+            The folder in which to offload the model weights (or where the model weights are already offloaded).
+        execution_device (`torch.device`, *optional*):
+            The device on which the forward pass of the model will be executed (should be a GPU). Will default to the
+            model's first parameter device.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+    """
+    if not os.path.isdir(offload_dir) or not os.path.isfile(os.path.join(offload_dir, "index.json")):
+        offload_state_dict(offload_dir, model.state_dict())
+    if execution_device is None:
+        execution_device = next(iter(model.parameter())).device
+    weights_map = OffloadedWeightsLoader(save_folder=offload_dir)
+    attach_align_device_hook(
+        model,
+        execution_device=execution_device,
+        offload=True,
+        offload_buffers=offload_buffers,
+        weights_map=weights_map,
+    )
+    add_hook_to_module(model, AlignDevicesHook(io_same_device=True))
+    return model
+
+
+def dispatch_model(
+    model: nn.Module,
+    device_map: Dict[str, Union[str, int, torch.device]],
+    main_device: Optional[torch.device] = None,
+    state_dict: Optional[Dict[str, torch.Tensor]] = None,
+    offload_dir: Union[str, os.PathLike] = None,
+    offload_buffers: bool = False,
+):
+    """
+    Dispatches a model according to a given device map. Layers of the model might be spread across GPUs, offloaded on
+    the CPU or even the disk.
+
+    Args:
+        model (`torch.nn.Module`): The model to dispatch.
+        device_map (`Dict[str, Union[str, int, torch.device]]`):
+            A dictionary mapping module names to the device they should go to. Note that `"disk"` is accepted even if
+            it's not a proper value for `torch.device`.
While there is a warning in the docstring, I'm not wondering if perhaps add an extra warning here too? Or do we think that'd be too much feedback for the user
@@ -0,0 +1,411 @@
+# 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 functools
+from typing import Dict, Mapping, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from .utils import PrefixedDataset, find_device, named_module_tensors, send_to_device, set_module_tensor_to_device
+
+
+class ModelHook:
+    """
+    A hook that contains callbacks to be executed just before and after the forward method of a model. The difference
+    with PyTorch existing hooks is that they get passed along the kwargs.
+
+    Class attribute:
+    - **no_grad** (`bool`, *optional*, defaults to `False`) -- Whether or not to execute the actual forward pass under
+      the `torch.no_grad()` context manager.
+    """
+
+    no_grad = False
+
+    def init_hook(self, module):
+        """
+        To be executed when the hook is attached to the module.
+
+        Args:
+            module (`torch.nn.Module`): The module attached to this hook.
+        """
+        return module
+
+    def pre_forward(self, module, *args, **kwargs):
+        """
+        To be executed just before the forward method of the model.
+
+        Args:
+            module (`torch.nn.Module`): The module whose forward pass will be executed just after this event.
+            args (`Tuple[Any]`): The positional arguments passed to the module.
+            kwargs (`Dict[Str, Any]`): The keyword arguments passed to the module.
+
+        Returns:
+            `Tuple[Tuple[Any], Dict[Str, Any]]`: A tuple with the treated `args` and `kwargs`.
+        """
+        return args, kwargs
+
+    def post_forward(self, module, output):
+        """
+        To be executed just after the forward method of the model.
+
+        Args:
+            module (`torch.nn.Module`): The module whose forward pass been executed just before this event.
+            output (`Any`): The output of the module.
+
+        Returns:
+            `Any`: The processed `output`.
+        """
+        return output
+
+    def detach_hook(self, module):
+        """
+        To be executed when the hook is deached from a module.
+
+        Args:
+            module (`torch.nn.Module`): The module detached from this hook.
+        """
+        return module
+
+
+class SequentialHook(ModelHook):
+    """
+    A hook that can contain several hooks and iterates through them at each event.
+    """
+
+    def __init__(self, *hooks):
+        self.hooks = hooks
+
+    def init_hook(self, module):
+        for hook in self.hooks:
+            module = hook.init_hook(module)
+        return module
+
+    def pre_forward(self, module, *args, **kwargs):
+        for hook in self.hooks:
+            args, kwargs = hook.pre_forward(module, *args, **kwargs)
+        return args, kwargs
+
+    def post_forward(self, module, output):
+        for hook in self.hooks:
+            output = hook.post_forward(module, output)
+        return output
+
+    def detach_hook(self, module):
+        for hook in self.hooks:
+            module = hook.detach_hook(module)
+        return module
+
+
+def add_hook_to_module(module: nn.Module, hook: ModelHook):
+    """
+    Adds a hook to a given module. This will rewrite the `forward` method of the module to include the hook, to remove
+    this behavior and restore the original `forward` method, use `remove_hook_from_module`.
+
+    <Tip warning={true}>
+
+    If the module already contains a hook, this will replace it with the new hook passed. To chain two hooks together,
+    use the `SequentialHook` class.
+
+    </Tip>
+
+    Args:
+        module (`torch.nn.Module`): The module to attach a hook to.
+        hook (`ModelHook`): The hook to attach.
+
+    Returns:
+        `torch.nn.Module`: The same module, with the hook attached (the module is modified in place, so the result can
+        be discarded).
+    """
+    if hasattr(module, "_hf_hook") and hasattr(module, "_old_forward"):
+        # If we already put some hook on this module, we replace it with the new one.
+        old_forward = module._old_forward
We should just keep those as kwargs then I think, since the user should never need to worry about them
@@ -0,0 +1,489 @@
+# 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 gc
+import json
+import os
+import re
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+WEIGHTS_INDEX_NAME = "pytorch_model.bin.index.json"
+
+
+def convert_file_size_to_int(size: Union[int, str]):
+    """
+    Converts a size expressed as a string with digits an unit (like `"5MB"`) to an integer (in bytes).
+
+    Args:
+        size (`int` or `str`): The size to convert. Will be directly returned if an `int`.
+
+    Example:
+
+    ```py
+    >>> convert_file_size_to_int("1MiB")
+    1048576
+    ```
+    """
+    if isinstance(size, int):
+        return size
+    if size.upper().endswith("GIB"):
+        return int(size[:-3]) * (2**30)
+    if size.upper().endswith("MIB"):
+        return int(size[:-3]) * (2**20)
+    if size.upper().endswith("KIB"):
+        return int(size[:-3]) * (2**10)
+    if size.upper().endswith("GB"):
+        int_size = int(size[:-2]) * (10**9)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("MB"):
+        int_size = int(size[:-2]) * (10**6)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("KB"):
+        int_size = int(size[:-2]) * (10**3)
+        return int_size // 8 if size.endswith("b") else int_size
+    raise ValueError("`size` is not in a valid format. Use an integer followed by the unit, e.g., '5GB'.")
+
+
+def dtype_byte_size(dtype):
+    """
+    Returns the size (in bytes) occupied by one parameter of type `dtype`.
+
+    Example:
+
+    ```py
+    >>> dtype_byte_size(torch.float32)
+    4
+    ```
+    """
+    if dtype == torch.bool:
+        return 1 / 8
+    bit_search = re.search("[^\d](\d+)$", str(dtype))
+    if bit_search is None:
+        raise ValueError(f"`dtype` is not a valid dtype: {dtype}.")
+    bit_size = int(bit_search.groups()[0])
+    return bit_size // 8
+
+
+def set_module_tensor_to_device(
+    module: nn.Module, tensor_name: str, device: Union[int, str, torch.device], value: Optional[torch.Tensor] = None
+):
+    """
+    A helper function to set a given tensor (parameter of buffer) of a module on a specific device (note that doing
+    `param.to(device)` creates a new tensor not linked to the parameter, which is why we need this function).
+
+    Args:
+        module (`torch.nn.Module`): The module in which the tensor we want to move lives.
+        param_name (`str`): The full name of the parameter/buffer.
+        device (`int`, `str` or `torch.device`): The device on which to set the tensor.
+        value (`torch.Tensor`, *optional*): The value of the tensor (useful when going from the meta device to any
+            other device).
+    """
+    # Recurse if needed
+    if "." in tensor_name:
+        splits = tensor_name.split(".")
+        for split in splits[:-1]:
+            new_module = getattr(module, split)
+            if new_module is None:
+                raise ValueError(f"{module} has no attribute {split}.")
+            module = new_module
+        tensor_name = splits[-1]
+
+    if tensor_name not in module._parameters and tensor_name not in module._buffers:
+        raise ValueError(f"{module} does not have a parameter or a buffer named {tensor_name}.")
+    is_buffer = tensor_name in module._buffers
+    old_value = getattr(module, tensor_name)
+
+    if old_value.device == torch.device("meta") and device != torch.device("meta") and value is None:
+        raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {device}.")
+
+    with torch.no_grad():
+        if value is None:
+            new_value = old_value.to(device)
+        elif isinstance(value, torch.Tensor):
+            new_value = value.to(device)
+        else:
+            new_value = torch.tensor(value, device=device)
+    if is_buffer:
+        module._buffers[tensor_name] = new_value
+    else:
+        new_value = nn.Parameter(new_value, requires_grad=old_value.requires_grad)
+        module._parameters[tensor_name] = new_value
+
+
+def named_module_tensors(module: nn.Module, include_buffers: bool = True, recurse: bool = False):
+    """
+    A helper function that gathers all the tensors (parameters + buffers) of a given module. If `include_buffers=True`
+    it's the same as doing `module.named_parameters(recurse=recurse) + module.named_buffers(recurse=recurse)`.
+
+    Args:
+        module (`torch.nn.Module`): The module we want the tensors or.
+        include_buffer (`bool`, *optional*, defaults to `True`): Whether or not to include the buffers in the result.
+        recurse (`bool`, *optional`, defaults to `False`):
+            Whether or not to go look in every submodule or just return the direct parameters and buffers.
+    """
+    for named_parameter in module.named_parameters(recurse=recurse):
+        yield named_parameter
+
+    if include_buffers:
+        for named_buffer in module.named_buffers(recurse=recurse):
+            yield named_buffer
+
+
+def find_tied_parameters(model: nn.Module, named_parameters=None, prefix="", result=None):
+    """
+    Find the tied parameters in a given model.
+
+    Args:
+        model (`torch.nn.Module`): The model to inspect.
+
+    <Tip warning={true}>
+
+    The signature has more arguments, but they are for the recursive part of this function and you should ignore them.
+
+    </Tip>
Just to make sure I understand correctly, we just check to make sure it *runs*, we don't worry about if the outputs make sense?

Should we just to make sure the params didn't get messed with at all, etc?
@@ -0,0 +1,204 @@
+# 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 unittest
+from tempfile import TemporaryDirectory
+
+import torch
+import torch.nn as nn
+
+from accelerate.big_modeling import cpu_offload, disk_offload, dispatch_model, init_empty_weights
+from accelerate.hooks import remove_hook_from_submodules
+from accelerate.test_utils import require_cuda, require_multi_gpu, slow
+from accelerate.utils import offload_state_dict
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+class ModelForTest(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear1 = nn.Linear(3, 4)
+        self.batchnorm = nn.BatchNorm1d(4)
+        self.linear2 = nn.Linear(4, 5)
+
+    def forward(self, x):
+        return self.linear2(self.batchnorm(self.linear1(x)))
+
+
+class BiggerModelForTest(nn.Sequential):
+    def __init__(self):
+        super().__init__()
+        self.linear1 = nn.Linear(3, 4)
+        self.linear2 = nn.Linear(4, 5)
+        self.batchnorm = nn.BatchNorm1d(5)
+        self.linear3 = nn.Linear(5, 6)
+        self.linear4 = nn.Linear(6, 5)
+
+    def forward(self, x):
+        return self.linear4(self.linear3(self.batchnorm(self.linear2(self.linear1(x)))))
+
+
+class BigModelingTester(unittest.TestCase):
+    def test_init_empty_weights(self):
+        # base use
+        with init_empty_weights():
+            module = nn.Linear(4, 5)
+        self.assertEqual(module.weight.device, torch.device("meta"))
+
+        # base use with buffers, they are not touched
+        with init_empty_weights():
+            module = nn.BatchNorm1d(4)
+        self.assertEqual(module.weight.device, torch.device("meta"))
+        self.assertEqual(module.running_mean.device, torch.device("cpu"))
+
+        # Use with include_buffers=True
+        with init_empty_weights(include_buffers=True):
+            module = nn.BatchNorm1d(4)
+        self.assertEqual(module.weight.device, torch.device("meta"))
+        self.assertEqual(module.running_mean.device, torch.device("meta"))
+
+        # Double check we didn't break PyTorch
+        module = nn.BatchNorm1d(4)
+        self.assertEqual(module.weight.device, torch.device("cpu"))
+        self.assertEqual(module.running_mean.device, torch.device("cpu"))
+
+    def test_init_empty_weights_very_large_model(self):
+        # This is a 100 billion parameters model.
+        with init_empty_weights():
+            _ = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])
+
+    def test_cpu_offload(self):
+        model = ModelForTest()
+        x = torch.randn(2, 3)
+        expected = model(x)
+
+        device = torch.device(0 if torch.cuda.is_available() else "cpu")
+
+        cpu_offload(model, execution_device=device)
+        output = model(x)
+        self.assertTrue(torch.allclose(expected, output.cpu()))
+
+        # Clean up for next test.
+        remove_hook_from_submodules(model)
+
+        cpu_offload(model, execution_device=device, offload_buffers=True)
+        output = model(x)
+        self.assertTrue(torch.allclose(expected, output.cpu()))
+
+    @slow
+    @require_cuda
+    def test_cpu_offload_gpt2(self):
+        tokenizer = AutoTokenizer.from_pretrained("gpt2")
+        inputs = tokenizer("Hello world! My name is", return_tensors="pt").to(0)
+
+        gpt2 = AutoModelForCausalLM.from_pretrained("gpt2")
Too much feedback I believe. This is also very internal, I don't think the average user will use those hooks.
@@ -0,0 +1,411 @@
+# 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 functools
+from typing import Dict, Mapping, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from .utils import PrefixedDataset, find_device, named_module_tensors, send_to_device, set_module_tensor_to_device
+
+
+class ModelHook:
+    """
+    A hook that contains callbacks to be executed just before and after the forward method of a model. The difference
+    with PyTorch existing hooks is that they get passed along the kwargs.
+
+    Class attribute:
+    - **no_grad** (`bool`, *optional*, defaults to `False`) -- Whether or not to execute the actual forward pass under
+      the `torch.no_grad()` context manager.
+    """
+
+    no_grad = False
+
+    def init_hook(self, module):
+        """
+        To be executed when the hook is attached to the module.
+
+        Args:
+            module (`torch.nn.Module`): The module attached to this hook.
+        """
+        return module
+
+    def pre_forward(self, module, *args, **kwargs):
+        """
+        To be executed just before the forward method of the model.
+
+        Args:
+            module (`torch.nn.Module`): The module whose forward pass will be executed just after this event.
+            args (`Tuple[Any]`): The positional arguments passed to the module.
+            kwargs (`Dict[Str, Any]`): The keyword arguments passed to the module.
+
+        Returns:
+            `Tuple[Tuple[Any], Dict[Str, Any]]`: A tuple with the treated `args` and `kwargs`.
+        """
+        return args, kwargs
+
+    def post_forward(self, module, output):
+        """
+        To be executed just after the forward method of the model.
+
+        Args:
+            module (`torch.nn.Module`): The module whose forward pass been executed just before this event.
+            output (`Any`): The output of the module.
+
+        Returns:
+            `Any`: The processed `output`.
+        """
+        return output
+
+    def detach_hook(self, module):
+        """
+        To be executed when the hook is deached from a module.
+
+        Args:
+            module (`torch.nn.Module`): The module detached from this hook.
+        """
+        return module
+
+
+class SequentialHook(ModelHook):
+    """
+    A hook that can contain several hooks and iterates through them at each event.
+    """
+
+    def __init__(self, *hooks):
+        self.hooks = hooks
+
+    def init_hook(self, module):
+        for hook in self.hooks:
+            module = hook.init_hook(module)
+        return module
+
+    def pre_forward(self, module, *args, **kwargs):
+        for hook in self.hooks:
+            args, kwargs = hook.pre_forward(module, *args, **kwargs)
+        return args, kwargs
+
+    def post_forward(self, module, output):
+        for hook in self.hooks:
+            output = hook.post_forward(module, output)
+        return output
+
+    def detach_hook(self, module):
+        for hook in self.hooks:
+            module = hook.detach_hook(module)
+        return module
+
+
+def add_hook_to_module(module: nn.Module, hook: ModelHook):
+    """
+    Adds a hook to a given module. This will rewrite the `forward` method of the module to include the hook, to remove
+    this behavior and restore the original `forward` method, use `remove_hook_from_module`.
+
+    <Tip warning={true}>
+
+    If the module already contains a hook, this will replace it with the new hook passed. To chain two hooks together,
+    use the `SequentialHook` class.
+
+    </Tip>
+
+    Args:
+        module (`torch.nn.Module`): The module to attach a hook to.
+        hook (`ModelHook`): The hook to attach.
+
+    Returns:
+        `torch.nn.Module`: The same module, with the hook attached (the module is modified in place, so the result can
+        be discarded).
+    """
+    if hasattr(module, "_hf_hook") and hasattr(module, "_old_forward"):
+        # If we already put some hook on this module, we replace it with the new one.
+        old_forward = module._old_forward
That would be nice to add in that test indeed!
@@ -0,0 +1,204 @@
+# 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 unittest
+from tempfile import TemporaryDirectory
+
+import torch
+import torch.nn as nn
+
+from accelerate.big_modeling import cpu_offload, disk_offload, dispatch_model, init_empty_weights
+from accelerate.hooks import remove_hook_from_submodules
+from accelerate.test_utils import require_cuda, require_multi_gpu, slow
+from accelerate.utils import offload_state_dict
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+class ModelForTest(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear1 = nn.Linear(3, 4)
+        self.batchnorm = nn.BatchNorm1d(4)
+        self.linear2 = nn.Linear(4, 5)
+
+    def forward(self, x):
+        return self.linear2(self.batchnorm(self.linear1(x)))
+
+
+class BiggerModelForTest(nn.Sequential):
+    def __init__(self):
+        super().__init__()
+        self.linear1 = nn.Linear(3, 4)
+        self.linear2 = nn.Linear(4, 5)
+        self.batchnorm = nn.BatchNorm1d(5)
+        self.linear3 = nn.Linear(5, 6)
+        self.linear4 = nn.Linear(6, 5)
+
+    def forward(self, x):
+        return self.linear4(self.linear3(self.batchnorm(self.linear2(self.linear1(x)))))
+
+
+class BigModelingTester(unittest.TestCase):
+    def test_init_empty_weights(self):
+        # base use
+        with init_empty_weights():
+            module = nn.Linear(4, 5)
+        self.assertEqual(module.weight.device, torch.device("meta"))
+
+        # base use with buffers, they are not touched
+        with init_empty_weights():
+            module = nn.BatchNorm1d(4)
+        self.assertEqual(module.weight.device, torch.device("meta"))
+        self.assertEqual(module.running_mean.device, torch.device("cpu"))
+
+        # Use with include_buffers=True
+        with init_empty_weights(include_buffers=True):
+            module = nn.BatchNorm1d(4)
+        self.assertEqual(module.weight.device, torch.device("meta"))
+        self.assertEqual(module.running_mean.device, torch.device("meta"))
+
+        # Double check we didn't break PyTorch
+        module = nn.BatchNorm1d(4)
+        self.assertEqual(module.weight.device, torch.device("cpu"))
+        self.assertEqual(module.running_mean.device, torch.device("cpu"))
+
+    def test_init_empty_weights_very_large_model(self):
+        # This is a 100 billion parameters model.
+        with init_empty_weights():
+            _ = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])
+
+    def test_cpu_offload(self):
+        model = ModelForTest()
+        x = torch.randn(2, 3)
+        expected = model(x)
+
+        device = torch.device(0 if torch.cuda.is_available() else "cpu")
+
+        cpu_offload(model, execution_device=device)
+        output = model(x)
+        self.assertTrue(torch.allclose(expected, output.cpu()))
+
+        # Clean up for next test.
+        remove_hook_from_submodules(model)
+
+        cpu_offload(model, execution_device=device, offload_buffers=True)
+        output = model(x)
+        self.assertTrue(torch.allclose(expected, output.cpu()))
+
+    @slow
+    @require_cuda
+    def test_cpu_offload_gpt2(self):
+        tokenizer = AutoTokenizer.from_pretrained("gpt2")
+        inputs = tokenizer("Hello world! My name is", return_tensors="pt").to(0)
+
+        gpt2 = AutoModelForCausalLM.from_pretrained("gpt2")
Is there a situation where we wouldn't want to include buffers?
@@ -0,0 +1,260 @@
+# 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
+from contextlib import contextmanager
+from typing import Dict, List, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from .hooks import AlignDevicesHook, add_hook_to_module, attach_align_device_hook, attach_align_device_hook_on_blocks
+from .utils import (
+    OffloadedWeightsLoader,
+    check_device_map,
+    extract_submodules_state_dict,
+    infer_auto_device_map,
+    load_checkpoint_in_model,
+    offload_state_dict,
+)
+
+
+@contextmanager
+def init_empty_weights(include_buffers: bool = False):
+    """
+    A context manager under which models are initialized with all parameters on the meta device, therefore creating an
+    empty model. Useful when just initializing the model would blow the available RAM.
+
+    Args:
+        include_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to also put all buffers on the meta device while initializing.
```suggestion
        # Will contain the input device when `io_same_device=True`.
```
@@ -0,0 +1,411 @@
+# 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 functools
+from typing import Dict, Mapping, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from .utils import PrefixedDataset, find_device, named_module_tensors, send_to_device, set_module_tensor_to_device
+
+
+class ModelHook:
+    """
+    A hook that contains callbacks to be executed just before and after the forward method of a model. The difference
+    with PyTorch existing hooks is that they get passed along the kwargs.
+
+    Class attribute:
+    - **no_grad** (`bool`, *optional*, defaults to `False`) -- Whether or not to execute the actual forward pass under
+      the `torch.no_grad()` context manager.
+    """
+
+    no_grad = False
+
+    def init_hook(self, module):
+        """
+        To be executed when the hook is attached to the module.
+
+        Args:
+            module (`torch.nn.Module`): The module attached to this hook.
+        """
+        return module
+
+    def pre_forward(self, module, *args, **kwargs):
+        """
+        To be executed just before the forward method of the model.
+
+        Args:
+            module (`torch.nn.Module`): The module whose forward pass will be executed just after this event.
+            args (`Tuple[Any]`): The positional arguments passed to the module.
+            kwargs (`Dict[Str, Any]`): The keyword arguments passed to the module.
+
+        Returns:
+            `Tuple[Tuple[Any], Dict[Str, Any]]`: A tuple with the treated `args` and `kwargs`.
+        """
+        return args, kwargs
+
+    def post_forward(self, module, output):
+        """
+        To be executed just after the forward method of the model.
+
+        Args:
+            module (`torch.nn.Module`): The module whose forward pass been executed just before this event.
+            output (`Any`): The output of the module.
+
+        Returns:
+            `Any`: The processed `output`.
+        """
+        return output
+
+    def detach_hook(self, module):
+        """
+        To be executed when the hook is deached from a module.
+
+        Args:
+            module (`torch.nn.Module`): The module detached from this hook.
+        """
+        return module
+
+
+class SequentialHook(ModelHook):
+    """
+    A hook that can contain several hooks and iterates through them at each event.
+    """
+
+    def __init__(self, *hooks):
+        self.hooks = hooks
+
+    def init_hook(self, module):
+        for hook in self.hooks:
+            module = hook.init_hook(module)
+        return module
+
+    def pre_forward(self, module, *args, **kwargs):
+        for hook in self.hooks:
+            args, kwargs = hook.pre_forward(module, *args, **kwargs)
+        return args, kwargs
+
+    def post_forward(self, module, output):
+        for hook in self.hooks:
+            output = hook.post_forward(module, output)
+        return output
+
+    def detach_hook(self, module):
+        for hook in self.hooks:
+            module = hook.detach_hook(module)
+        return module
+
+
+def add_hook_to_module(module: nn.Module, hook: ModelHook):
+    """
+    Adds a hook to a given module. This will rewrite the `forward` method of the module to include the hook, to remove
+    this behavior and restore the original `forward` method, use `remove_hook_from_module`.
+
+    <Tip warning={true}>
+
+    If the module already contains a hook, this will replace it with the new hook passed. To chain two hooks together,
+    use the `SequentialHook` class.
+
+    </Tip>
+
+    Args:
+        module (`torch.nn.Module`): The module to attach a hook to.
+        hook (`ModelHook`): The hook to attach.
+
+    Returns:
+        `torch.nn.Module`: The same module, with the hook attached (the module is modified in place, so the result can
+        be discarded).
+    """
+    if hasattr(module, "_hf_hook") and hasattr(module, "_old_forward"):
+        # If we already put some hook on this module, we replace it with the new one.
+        old_forward = module._old_forward
+    else:
+        old_forward = module.forward
+        module._old_forward = old_forward
+
+    module = hook.init_hook(module)
+    module._hf_hook = hook
+
+    @functools.wraps(old_forward)
+    def new_forward(*args, **kwargs):
+        args, kwargs = module._hf_hook.pre_forward(module, *args, **kwargs)
+        if module._hf_hook.no_grad:
+            with torch.no_grad():
+                output = old_forward(*args, **kwargs)
+        else:
+            output = old_forward(*args, **kwargs)
+        return module._hf_hook.post_forward(module, output)
+
+    module.forward = new_forward
+    return module
+
+
+def remove_hook_from_module(module: nn.Module):
+    """
+    Removes any hook attached to a module via `add_hook_to_module`.
+
+    Args:
+        module (`torch.nn.Module`): The module to attach a hook to.
+
+    Returns:
+        `torch.nn.Module`: The same module, with the hook detached (the module is modified in place, so the result can
+        be discarded).
+    """
+    if hasattr(module, "_hf_hook"):
+        module._hf_hook.detach_hook(module)
+        delattr(module, "_hf_hook")
+
+    if hasattr(module, "_old_forward"):
+        module.forward = module._old_forward
+        delattr(module, "_old_forward")
+
+    return module
+
+
+class AlignDevicesHook(ModelHook):
+    """
+    A generic `ModelHook` that ensures inputs and model weights are on the same device for the forward pass of the
+    associated module, potentially offloading the weights after the forward pass.
+
+    Args:
+        execution_device (`torch.device`, *optional*):
+            The device on which inputs and model weights should be placed before the forward pass.
+        offload (`bool`, *optional*, defauts to `False`):
+            Whether or not the weights should be offloaded after the forward pass.
+        io_same_device (`bool`, *optional*, defaults to `False`):
+            Whether or not the output should be placed on the same device as the input was.
+        weights_map (`Mapping[str, torch.Tensor]`, *optional*):
+            When the model weights are offloaded, a (potentially lazy) map from param names to the tensor values.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to include the associated module's buffers when offloading.
+        place_submodules (`bool`, *optional*, defaults to `False`):
+            Whether to place the submodules on `execution_device` during the `init_hook` event.
+    """
+
+    def __init__(
+        self,
+        execution_device: Optional[Union[int, str, torch.device]] = None,
+        offload: bool = False,
+        io_same_device: bool = False,
+        weights_map: Optional[Mapping] = None,
+        offload_buffers: bool = False,
+        place_submodules: bool = False,
+    ):
+        self.execution_device = execution_device
+        self.offload = offload
+        self.io_same_device = io_same_device
+        self.weights_map = weights_map
+        self.offload_buffers = offload_buffers
+        self.place_submodules = place_submodules
+
+        # Will contain the input device when `output_on_same_device=True`.
```suggestion
            raise ValueError("Need either a `state_dict` or a `save_folder` containing offloaded weights.")
```
@@ -0,0 +1,143 @@
+# 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 json
+import os
+from collections.abc import Mapping
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+import torch
+
+
+def offload_state_dict(save_dir: Union[str, os.PathLike], state_dict: Dict[str, torch.Tensor]):
+    """
+    Offload a state dict in a given folder.
+
+    Args:
+        save_dir (`str` or `os.PathLike`): The directory in which to offload the state dict.
+        state_dict (`Dict[str, torch.Tensor]`): The dictionary of tensors to offload.
+    """
+    os.makedirs(save_dir, exist_ok=True)
+    index = {}
+    for name, parameter in state_dict.items():
+        tensor_file = os.path.join(save_dir, f"{name}.dat")
+        array = parameter.numpy()
+        index[name] = {"dtype": str(array.dtype), "shape": list(array.shape)}
+        if array.ndim == 0:
+            array = array[None]
+        file_array = np.memmap(tensor_file, dtype=array.dtype, mode="w+", shape=array.shape)
+        file_array[:] = array[:]
+        file_array.flush()
+
+    # Update index
+    index_file = os.path.join(save_dir, "index.json")
+    if os.path.isfile(index_file):
+        with open(index_file, "r", encoding="utf-8") as f:
+            current_index = json.load(f)
+    else:
+        current_index = {}
+    current_index.update(index)
+
+    with open(index_file, "w", encoding="utf-8") as f:
+        json.dump(current_index, f, indent=2)
+
+
+class PrefixedDataset(Mapping):
+    """
+    Will access keys in a given dataset by adding a prefix.
+
+    Args:
+        dataset (`Mapping`): Any map with string keys.
+        prefix (`str`): A prefix to add when trying to access any element in the underlying dataset.
+    """
+
+    def __init__(self, dataset: Mapping, prefix: str):
+        self.dataset = dataset
+        self.prefix = prefix
+
+    def __getitem__(self, key):
+        return self.dataset[f"{self.prefix}{key}"]
+
+    def __iter__(self):
+        return iter([key for key in self.dataset if key.startswith(self.prefix)])
+
+    def __len__(self):
+        return len(self.dataset)
+
+
+class OffloadedWeightsLoader(Mapping):
+    """
+    A collection that loads weights stored in a given state dict or memory-mapped on disk.
+
+    Args:
+        state_dict (`Dict[str, torch.Tensor]`, *optional*):
+            A dictionary parameter name to tensor.
+        save_folder (`str` or `os.PathLike`, *optional*):
+            The directory in which the weights are stored (by `offload_state_dict` for instance).
+        index (`Dict`, *optional*):
+            A dictionary from weight name to their information (`dtype` and `shape`). Will default to the index saved
+            in `save_folder`.
+    """
+
+    def __init__(
+        self,
+        state_dict: Dict[str, torch.Tensor] = None,
+        save_folder: Optional[Union[str, os.PathLike]] = None,
+        index: Mapping = None,
+    ):
+        if state_dict is None and save_folder is None:
+            raise ValueError("Need either a `state_dict` or a `save_folder` containint offloaded weights.")
when this condition happens, reassigning `max_layer_size` impacts line 377 wherein second largest `max_layer_size` gets subtracted instead of largest, is this understanding correct? If yes, this isn't expected behaviour I guess.

Sample test case model:
```python
class MyCustomModel(nn.Module):
    def __init__():
        super(MyCustomModel, self).__init__()
        self.huge_layer = nn.Linear(100000, 100000) # 10 Billion params
        self.upscale_layer = nn.Linear(768, 100000) 
        self.output_layer = nn.Linear(100000, 5) 
        self.bert_model = AutoModel.from_pretrained("bert-base-uncased") # 110 Million Params

    def forward(input_ids, attention_mask):
        _, pooler_output = self.bert_model(input_ids, attention_mask, return_dict=True)
        output = torch.nn.ReLU(self.upscale_layer(pooler_output))
        output = torch.nn.ReLU(self.huge_layer(output))
        output_logits = self.output_layer(output)
        return output_logits
```
@@ -0,0 +1,568 @@
+# 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 gc
+import json
+import os
+import re
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+WEIGHTS_INDEX_NAME = "pytorch_model.bin.index.json"
+
+
+def convert_file_size_to_int(size: Union[int, str]):
+    """
+    Converts a size expressed as a string with digits an unit (like `"5MB"`) to an integer (in bytes).
+
+    Args:
+        size (`int` or `str`): The size to convert. Will be directly returned if an `int`.
+
+    Example:
+
+    ```py
+    >>> convert_file_size_to_int("1MiB")
+    1048576
+    ```
+    """
+    if isinstance(size, int):
+        return size
+    if size.upper().endswith("GIB"):
+        return int(size[:-3]) * (2**30)
+    if size.upper().endswith("MIB"):
+        return int(size[:-3]) * (2**20)
+    if size.upper().endswith("KIB"):
+        return int(size[:-3]) * (2**10)
+    if size.upper().endswith("GB"):
+        int_size = int(size[:-2]) * (10**9)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("MB"):
+        int_size = int(size[:-2]) * (10**6)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("KB"):
+        int_size = int(size[:-2]) * (10**3)
+        return int_size // 8 if size.endswith("b") else int_size
+    raise ValueError("`size` is not in a valid format. Use an integer followed by the unit, e.g., '5GB'.")
+
+
+def dtype_byte_size(dtype: torch.dtype):
+    """
+    Returns the size (in bytes) occupied by one parameter of type `dtype`.
+
+    Example:
+
+    ```py
+    >>> dtype_byte_size(torch.float32)
+    4
+    ```
+    """
+    if dtype == torch.bool:
+        return 1 / 8
+    bit_search = re.search("[^\d](\d+)$", str(dtype))
+    if bit_search is None:
+        raise ValueError(f"`dtype` is not a valid dtype: {dtype}.")
+    bit_size = int(bit_search.groups()[0])
+    return bit_size // 8
+
+
+def set_module_tensor_to_device(
+    module: nn.Module, tensor_name: str, device: Union[int, str, torch.device], value: Optional[torch.Tensor] = None
+):
+    """
+    A helper function to set a given tensor (parameter of buffer) of a module on a specific device (note that doing
+    `param.to(device)` creates a new tensor not linked to the parameter, which is why we need this function).
+
+    Args:
+        module (`torch.nn.Module`): The module in which the tensor we want to move lives.
+        param_name (`str`): The full name of the parameter/buffer.
+        device (`int`, `str` or `torch.device`): The device on which to set the tensor.
+        value (`torch.Tensor`, *optional*): The value of the tensor (useful when going from the meta device to any
+            other device).
+    """
+    # Recurse if needed
+    if "." in tensor_name:
+        splits = tensor_name.split(".")
+        for split in splits[:-1]:
+            new_module = getattr(module, split)
+            if new_module is None:
+                raise ValueError(f"{module} has no attribute {split}.")
+            module = new_module
+        tensor_name = splits[-1]
+
+    if tensor_name not in module._parameters and tensor_name not in module._buffers:
+        raise ValueError(f"{module} does not have a parameter or a buffer named {tensor_name}.")
+    is_buffer = tensor_name in module._buffers
+    old_value = getattr(module, tensor_name)
+
+    if old_value.device == torch.device("meta") and device != torch.device("meta") and value is None:
+        raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {device}.")
+
+    with torch.no_grad():
+        if value is None:
+            new_value = old_value.to(device)
+        elif isinstance(value, torch.Tensor):
+            new_value = value.to(device)
+        else:
+            new_value = torch.tensor(value, device=device)
+    if is_buffer:
+        module._buffers[tensor_name] = new_value
+    else:
+        new_value = nn.Parameter(new_value, requires_grad=old_value.requires_grad)
+        module._parameters[tensor_name] = new_value
+
+
+def named_module_tensors(module: nn.Module, include_buffers: bool = True, recurse: bool = False):
+    """
+    A helper function that gathers all the tensors (parameters + buffers) of a given module. If `include_buffers=True`
+    it's the same as doing `module.named_parameters(recurse=recurse) + module.named_buffers(recurse=recurse)`.
+
+    Args:
+        module (`torch.nn.Module`): The module we want the tensors or.
+        include_buffer (`bool`, *optional*, defaults to `True`): Whether or not to include the buffers in the result.
+        recurse (`bool`, *optional`, defaults to `False`):
+            Whether or not to go look in every submodule or just return the direct parameters and buffers.
+    """
+    for named_parameter in module.named_parameters(recurse=recurse):
+        yield named_parameter
+
+    if include_buffers:
+        for named_buffer in module.named_buffers(recurse=recurse):
+            yield named_buffer
+
+
+def find_tied_parameters(model: nn.Module, **kwargs):
+    """
+    Find the tied parameters in a given model.
+
+    Args:
+        model (`torch.nn.Module`): The model to inspect.
+
+    <Tip warning={true}>
+
+    The signature accepts keyword arguments, but they are for the recursive part of this function and you should ignore
+    them.
+
+    </Tip>
+
+    Example:
+
+
+    ```py
+    >>> from collections import OrderedDict
+    >>> import torch.nn as nn
+
+    >>> model = nn.Sequential(OrderedDict([("linear1", nn.Linear(4, 4)), ("linear2", nn.Linear(4, 4))]))
+    >>> model.linear2.weight = test_model.linear1.weight
+    >>> find_tied_parameters(test_model)
+    {'linear1.weight': 'linear2.weight'}
+    ```
+
+    Returns:
+        Dict[str, str]: A dictionary mapping tied parameter names to the name of the parameter they are tied to.
+    """
+    # Initialize result and named_parameters before recursing.
+    named_parameters = kwargs.get("named_parameters", None)
+    prefix = kwargs.get("prefix", "")
+    result = kwargs.get("result", {})
+
+    if named_parameters is None:
+        named_parameters = {n: p for n, p in model.named_parameters()}
+    else:
+        # A tied parameter will not be in the full `named_parameters` seen above but will be in the `named_parameters`
+        # of the submodule it belongs to. So while recursing we track the names that are not in the initial
+        # `named_parameters`.
+        for name, parameter in model.named_parameters():
+            full_name = name if prefix == "" else f"{prefix}.{name}"
+            if full_name not in named_parameters:
+                # When we find one, it has to be one of the existing parameters.
+                for new_name, new_param in named_parameters.items():
+                    if new_param is parameter:
+                        result[new_name] = full_name
+
+    # Once we have treated direct parameters, we move to the child modules.
+    for name, child in model.named_children():
+        child_name = name if prefix == "" else f"{prefix}.{name}"
+        find_tied_parameters(child, named_parameters=named_parameters, prefix=child_name, result=result)
+
+    return result
+
+
+def compute_module_sizes(model: nn.Module):
+    """
+    Compute the size of each submodule of a given model.
+    """
+    module_sizes = defaultdict(int)
+    for name, tensor in named_module_tensors(model, recurse=True):
+        size = tensor.numel() * dtype_byte_size(tensor.dtype)
+        name_parts = name.split(".")
+        for idx in range(len(name_parts) + 1):
+            module_sizes[".".join(name_parts[:idx])] += size
+
+    return module_sizes
+
+
+def get_max_layer_size(
+    modules: List[Tuple[str, torch.nn.Module]], module_sizes: Dict[str, int], no_split_module_classes: List[str]
+):
+    """
+    Utility function that will scan a list of named modules and return the maximum size used by one full layer. The
+    definition of a layer being:
+    - a module with no direct children (just parameters and buffers)
+    - a module whose class name is in the list `no_split_module_classes`
+
+    Args:
+        modules (`List[Tuple[str, torch.nn.Module]]`):
+            The list of named modules where we want to determine the maximum layer size.
+        module_sizes (`Dict[str, int]`):
+            A dictionary mapping each layer name to its size (as generated by `compute_module_sizes`).
+        no_split_module_classes (`List[str]`):
+            A list of class names for layers we don't want to be split.
+
+    Returns:
+        `Tuple[int, List[str]]`: The maximum size of a layer with the list of layer names realizing that maximum size.
+    """
+    max_size = 0
+    layer_names = []
+    modules_to_treat = modules.copy()
+    while len(modules_to_treat) > 0:
+        module_name, module = modules_to_treat.pop(0)
+        modules_children = list(module.named_children())
+        if len(modules_children) == 0 or module.__class__.__name__ in no_split_module_classes:
+            # No splitting this one so we compare to the max_size
+            size = module_sizes[module_name]
+            if size > max_size:
+                max_size = size
+                layer_names = [module_name]
+            elif size == max_size:
+                layer_names.append(module_name)
+        else:
+            modules_to_treat = [(f"{module_name}.{n}", v) for n, v in modules_children] + modules_to_treat
+    return max_size, layer_names
+
+
+def get_max_memory(max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None):
+    """
+    Get the maximum memory available if nothing is passed, converts string to int otherwise.
+    """
+    import psutil
+
+    if max_memory is None:
+        if not torch.cuda.is_available():
+            max_memory = {}
+        else:
+            # Make sure CUDA is initialized on each GPU to have the right memory info.
+            for i in range(torch.cuda.device_count()):
+                _ = torch.tensor([0], device=i)
+            max_memory = {i: torch.cuda.mem_get_info(i)[0] for i in range(torch.cuda.device_count())}
+        max_memory["cpu"] = psutil.virtual_memory().available
+        return max_memory
+
+    for key in max_memory:
+        if isinstance(max_memory[key], str):
+            max_memory[key] = convert_file_size_to_int(max_memory[key])
+    return max_memory
+
+
+def clean_device_map(device_map: Dict[str, Union[int, str, torch.device]], module_name: str = ""):
+    """
+    Cleans a device_map by grouping all submodules that go on the same device together.
+    """
+    # Get the value of the current module and if there is only one split across several keys, regroup it.
+    values = [v for k, v in device_map.items() if k.startswith(module_name)]
+    if len(set(values)) == 1 and len(values) > 1:
+        for k in [k for k in device_map if k.startswith(module_name)]:
+            del device_map[k]
+        device_map[module_name] = values[0]
+
+    # Recurse over the children
+    children_modules = [k for k in device_map.keys() if k.startswith(module_name) and len(k) > len(module_name)]
+    idx = len(module_name.split(".")) + 1 if len(module_name) > 0 else 1
+    children_modules = set(".".join(k.split(".")[:idx]) for k in children_modules)
+    for child in children_modules:
+        clean_device_map(device_map, module_name=child)
+
+    return device_map
+
+
+def infer_auto_device_map(
+    model: nn.Module,
+    max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None,
+    no_split_module_classes: Optional[List[str]] = None,
+):
+    """
+    Compute a device map for a given model giving priority to GPUs, then offload on CPU and finally offload to disk,
+    such that:
+    - we don't exceed the memory available of any of the GPU.
+    - if offload to the CPU is needed, there is always room left on GPU 0 to put back the layer offloaded on CPU that
+      has the largest size.
+    - if offload to the CPU is needed,we don't exceed the RAM available on the CPU.
+    - if offload to the disk is needed, there is always room left on the CPU to put back the layer offloaded on disk
+      that has the largest size.
+
+    <Tip>
+
+    All computation is done analyzing sizes and dtypes of the model parameters. As a result, the model can be on the
+    meta device (as it would if initialized within the `init_empty_weights` context manager).
+
+    </Tip>
+
+    Args:
+        model (`torch.nn.Module`): The model to analyze.
+        max_memory (`Dict`, *optional*):
+            A dictionary device identifier to maximum memory. Will default to the maximum memory available if unset.
+        no_split_module_classes (`List[str]`, *optional*):
+            A list of layer class names that should never be split across device (for instance any layer that has a
+            residual connection).
+    """
+    # Get default / clean up max_memory
+    max_memory = get_max_memory(max_memory)
+    if no_split_module_classes is None:
+        no_split_module_classes = []
+    elif not isinstance(no_split_module_classes, (list, tuple)):
+        no_split_module_classes = [no_split_module_classes]
+
+    devices = list(max_memory.keys()) + ["disk"]
+    # Devices that need to keep space for a potential offloaded layer.
+    main_devices = [devices[0], "cpu"]
+
+    module_sizes = compute_module_sizes(model)
+    tied_parameters = find_tied_parameters(model)
+
+    device_map = {}
+    current_device = 0
+    current_memory_used = 0
+
+    # Direct submodules and parameters
+    modules_to_treat = list(model.named_parameters(recurse=False)) + list(model.named_children())
+    # Initialize maximum largest layer, to know which space to keep in memory
+    max_layer_size, max_layer_names = get_max_layer_size(modules_to_treat, module_sizes, no_split_module_classes)
+
+    # Ready ? This is going to be a bit messy.
+    while len(modules_to_treat) > 0:
+        name, module = modules_to_treat.pop(0)
+        # Max size in the remaining layers may have changed since we took one, so we maybe update it.
+        max_layer_names = [n for n in max_layer_names if not n.startswith(name)]
+        if len(max_layer_names) == 0:
+            max_layer_size, max_layer_names = get_max_layer_size(
```suggestion
    # If one device and one offload, we've got one hook.
```
@@ -0,0 +1,411 @@
+# 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 functools
+from typing import Dict, Mapping, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from .utils import PrefixedDataset, find_device, named_module_tensors, send_to_device, set_module_tensor_to_device
+
+
+class ModelHook:
+    """
+    A hook that contains callbacks to be executed just before and after the forward method of a model. The difference
+    with PyTorch existing hooks is that they get passed along the kwargs.
+
+    Class attribute:
+    - **no_grad** (`bool`, *optional*, defaults to `False`) -- Whether or not to execute the actual forward pass under
+      the `torch.no_grad()` context manager.
+    """
+
+    no_grad = False
+
+    def init_hook(self, module):
+        """
+        To be executed when the hook is attached to the module.
+
+        Args:
+            module (`torch.nn.Module`): The module attached to this hook.
+        """
+        return module
+
+    def pre_forward(self, module, *args, **kwargs):
+        """
+        To be executed just before the forward method of the model.
+
+        Args:
+            module (`torch.nn.Module`): The module whose forward pass will be executed just after this event.
+            args (`Tuple[Any]`): The positional arguments passed to the module.
+            kwargs (`Dict[Str, Any]`): The keyword arguments passed to the module.
+
+        Returns:
+            `Tuple[Tuple[Any], Dict[Str, Any]]`: A tuple with the treated `args` and `kwargs`.
+        """
+        return args, kwargs
+
+    def post_forward(self, module, output):
+        """
+        To be executed just after the forward method of the model.
+
+        Args:
+            module (`torch.nn.Module`): The module whose forward pass been executed just before this event.
+            output (`Any`): The output of the module.
+
+        Returns:
+            `Any`: The processed `output`.
+        """
+        return output
+
+    def detach_hook(self, module):
+        """
+        To be executed when the hook is deached from a module.
+
+        Args:
+            module (`torch.nn.Module`): The module detached from this hook.
+        """
+        return module
+
+
+class SequentialHook(ModelHook):
+    """
+    A hook that can contain several hooks and iterates through them at each event.
+    """
+
+    def __init__(self, *hooks):
+        self.hooks = hooks
+
+    def init_hook(self, module):
+        for hook in self.hooks:
+            module = hook.init_hook(module)
+        return module
+
+    def pre_forward(self, module, *args, **kwargs):
+        for hook in self.hooks:
+            args, kwargs = hook.pre_forward(module, *args, **kwargs)
+        return args, kwargs
+
+    def post_forward(self, module, output):
+        for hook in self.hooks:
+            output = hook.post_forward(module, output)
+        return output
+
+    def detach_hook(self, module):
+        for hook in self.hooks:
+            module = hook.detach_hook(module)
+        return module
+
+
+def add_hook_to_module(module: nn.Module, hook: ModelHook):
+    """
+    Adds a hook to a given module. This will rewrite the `forward` method of the module to include the hook, to remove
+    this behavior and restore the original `forward` method, use `remove_hook_from_module`.
+
+    <Tip warning={true}>
+
+    If the module already contains a hook, this will replace it with the new hook passed. To chain two hooks together,
+    use the `SequentialHook` class.
+
+    </Tip>
+
+    Args:
+        module (`torch.nn.Module`): The module to attach a hook to.
+        hook (`ModelHook`): The hook to attach.
+
+    Returns:
+        `torch.nn.Module`: The same module, with the hook attached (the module is modified in place, so the result can
+        be discarded).
+    """
+    if hasattr(module, "_hf_hook") and hasattr(module, "_old_forward"):
+        # If we already put some hook on this module, we replace it with the new one.
+        old_forward = module._old_forward
+    else:
+        old_forward = module.forward
+        module._old_forward = old_forward
+
+    module = hook.init_hook(module)
+    module._hf_hook = hook
+
+    @functools.wraps(old_forward)
+    def new_forward(*args, **kwargs):
+        args, kwargs = module._hf_hook.pre_forward(module, *args, **kwargs)
+        if module._hf_hook.no_grad:
+            with torch.no_grad():
+                output = old_forward(*args, **kwargs)
+        else:
+            output = old_forward(*args, **kwargs)
+        return module._hf_hook.post_forward(module, output)
+
+    module.forward = new_forward
+    return module
+
+
+def remove_hook_from_module(module: nn.Module):
+    """
+    Removes any hook attached to a module via `add_hook_to_module`.
+
+    Args:
+        module (`torch.nn.Module`): The module to attach a hook to.
+
+    Returns:
+        `torch.nn.Module`: The same module, with the hook detached (the module is modified in place, so the result can
+        be discarded).
+    """
+    if hasattr(module, "_hf_hook"):
+        module._hf_hook.detach_hook(module)
+        delattr(module, "_hf_hook")
+
+    if hasattr(module, "_old_forward"):
+        module.forward = module._old_forward
+        delattr(module, "_old_forward")
+
+    return module
+
+
+class AlignDevicesHook(ModelHook):
+    """
+    A generic `ModelHook` that ensures inputs and model weights are on the same device for the forward pass of the
+    associated module, potentially offloading the weights after the forward pass.
+
+    Args:
+        execution_device (`torch.device`, *optional*):
+            The device on which inputs and model weights should be placed before the forward pass.
+        offload (`bool`, *optional*, defauts to `False`):
+            Whether or not the weights should be offloaded after the forward pass.
+        io_same_device (`bool`, *optional*, defaults to `False`):
+            Whether or not the output should be placed on the same device as the input was.
+        weights_map (`Mapping[str, torch.Tensor]`, *optional*):
+            When the model weights are offloaded, a (potentially lazy) map from param names to the tensor values.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to include the associated module's buffers when offloading.
+        place_submodules (`bool`, *optional*, defaults to `False`):
+            Whether to place the submodules on `execution_device` during the `init_hook` event.
+    """
+
+    def __init__(
+        self,
+        execution_device: Optional[Union[int, str, torch.device]] = None,
+        offload: bool = False,
+        io_same_device: bool = False,
+        weights_map: Optional[Mapping] = None,
+        offload_buffers: bool = False,
+        place_submodules: bool = False,
+    ):
+        self.execution_device = execution_device
+        self.offload = offload
+        self.io_same_device = io_same_device
+        self.weights_map = weights_map
+        self.offload_buffers = offload_buffers
+        self.place_submodules = place_submodules
+
+        # Will contain the input device when `output_on_same_device=True`.
+        self.input_device = None
+        self.param_original_devices = {}
+        self.buffer_original_devices = {}
+
+    def init_hook(self, module):
+        if not self.offload and self.execution_device is not None:
+            for name, _ in named_module_tensors(module, recurse=self.place_submodules):
+                set_module_tensor_to_device(module, name, self.execution_device)
+        elif self.offload:
+            self.original_devices = {name: param.device for name, param in named_module_tensors(module)}
+            if self.weights_map is None:
+                self.weights_map = {
+                    name: param.to("cpu")
+                    for name, param in named_module_tensors(module, include_buffers=self.offload_buffers)
+                }
+
+            for name, _ in named_module_tensors(module, include_buffers=self.offload_buffers):
+                set_module_tensor_to_device(module, name, "meta")
+            if not self.offload_buffers and self.execution_device is not None:
+                for name, _ in module.named_buffers(recurse=False):
+                    set_module_tensor_to_device(module, name, self.execution_device)
+        return module
+
+    def pre_forward(self, module, *args, **kwargs):
+        if self.io_same_device:
+            self.input_device = find_device([args, kwargs])
+        if self.offload:
+            for name, _ in named_module_tensors(module, include_buffers=self.offload_buffers):
+                set_module_tensor_to_device(module, name, self.execution_device, value=self.weights_map[name])
+
+        return send_to_device(args, self.execution_device), send_to_device(kwargs, self.execution_device)
+
+    def post_forward(self, module, output):
+        if self.offload:
+            for name, _ in named_module_tensors(module, include_buffers=self.offload_buffers):
+                set_module_tensor_to_device(module, name, "meta")
+
+        if self.io_same_device and self.input_device is not None:
+            output = send_to_device(output, self.input_device)
+
+        return output
+
+    def detach_hook(self, module):
+        if self.offload:
+            for name, device in self.original_devices.items():
+                if device != torch.device("meta"):
+                    set_module_tensor_to_device(module, name, device, value=self.weights_map.get(name, None))
+
+
+def attach_align_device_hook(
+    module: torch.nn.Module,
+    execution_device: Optional[torch.device] = None,
+    offload: bool = False,
+    weights_map: Optional[Mapping] = None,
+    offload_buffers: bool = False,
+    module_name: str = "",
+):
+    """
+    Recursively attaches `AlignDevicesHook` to all submodules of a given model that have direct parameters and/or
+    buffers.
+
+    Args:
+        module (`torch.nn.Module`):
+            The module where we want to attach the hooks.
+        execution_device (`torch.device`, *optional*):
+            The device on which inputs and model weights should be placed before the forward pass.
+        offload (`bool`, *optional*, defauts to `False`):
+            Whether or not the weights should be offloaded after the forward pass.
+        weights_map (`Mapping[str, torch.Tensor]`, *optional*):
+            When the model weights are offloaded, a (potentially lazy) map from param names to the tensor values.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to include the associated module's buffers when offloading.
+        module_name (`str`, *optional*, defaults to `""`):
+            The name of the module.
+    """
+    # Attach the hook on this module if it has any direct tensor.
+    directs = named_module_tensors(module)
+    if len(list(directs)) > 0:
+        if weights_map is not None:
+            prefix = f"{module_name}." if len(module_name) > 0 else ""
+            prefixed_weights_map = PrefixedDataset(weights_map, prefix)
+        else:
+            prefixed_weights_map = None
+        hook = AlignDevicesHook(
+            execution_device=execution_device,
+            offload=offload,
+            weights_map=prefixed_weights_map,
+            offload_buffers=offload_buffers,
+        )
+        add_hook_to_module(module, hook)
+
+    # Recurse on all children of the module.
+    for child_name, child in module.named_children():
+        child_name = f"{module_name}.{child_name}" if len(module_name) > 0 else child_name
+        attach_align_device_hook(
+            child,
+            execution_device=execution_device,
+            offload=offload,
+            weights_map=weights_map,
+            offload_buffers=offload_buffers,
+            module_name=child_name,
+        )
+
+
+def remove_hook_from_submodules(module: nn.Module):
+    """
+    Recursively removes all hooks attached on the submodules of a given model.
+
+    Args:
+        module (`torch.nn.Module`): The module on which to remove all hooks.
+    """
+    remove_hook_from_module(module)
+    for child in module.children():
+        remove_hook_from_submodules(child)
+
+
+def attach_align_device_hook_on_blocks(
+    module: nn.Module,
+    execution_device: Optional[Union[torch.device, Dict[str, torch.device]]] = None,
+    offload: Union[bool, Dict[str, bool]] = False,
+    weights_map: Mapping = None,
+    offload_buffers: bool = False,
+    module_name: str = "",
+):
+    """
+    Attaches `AlignDevicesHook` to all blocks of a given model as needed.
+
+    Args:
+        module (`torch.nn.Module`):
+            The module where we want to attach the hooks.
+        execution_device (`torch.device` or `Dict[str, torch.device]`, *optional*):
+            The device on which inputs and model weights should be placed before the forward pass. It can be one device
+            for the whole module, or a dictionary mapping module name to device.
+        offload (`bool`, *optional*, defauts to `False`):
+            Whether or not the weights should be offloaded after the forward pass. It can be one boolean for the whole
+            module, or a dictionary mapping module name to boolean.
+        weights_map (`Mapping[str, torch.Tensor]`, *optional*):
+            When the model weights are offloaded, a (potentially lazy) map from param names to the tensor values.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to include the associated module's buffers when offloading.
+        module_name (`str`, *optional*, defaults to `""`):
+            The name of the module.
+    """
+    # If one device and one offload, we've go one hook.
```suggestion
        execution_device = next(iter(model.parameters())).device
```
@@ -0,0 +1,260 @@
+# 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
+from contextlib import contextmanager
+from typing import Dict, List, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from .hooks import AlignDevicesHook, add_hook_to_module, attach_align_device_hook, attach_align_device_hook_on_blocks
+from .utils import (
+    OffloadedWeightsLoader,
+    check_device_map,
+    extract_submodules_state_dict,
+    infer_auto_device_map,
+    load_checkpoint_in_model,
+    offload_state_dict,
+)
+
+
+@contextmanager
+def init_empty_weights(include_buffers: bool = False):
+    """
+    A context manager under which models are initialized with all parameters on the meta device, therefore creating an
+    empty model. Useful when just initializing the model would blow the available RAM.
+
+    Args:
+        include_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to also put all buffers on the meta device while initializing.
+
+    Example:
+
+    ```pyton
+    import torch.nn as nn
+    from accelerate import init_empty_weights
+
+    # Initialize a model with 100 billions parameters in no time and without using any RAM.
+    with init_empty_weights():
+        tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])
+    ```
+    """
+    old_register_parameter = nn.Module.register_parameter
+    if include_buffers:
+        old_register_buffer = nn.Module.register_buffer
+
+    def register_empty_parameter(module, name, param):
+        old_register_parameter(module, name, param)
+        if param is not None:
+            module._parameters[name] = nn.Parameter(module._parameters[name].to(torch.device("meta")))
+
+    def register_empty_buffer(module, name, buffer):
+        old_register_buffer(module, name, buffer)
+        if buffer is not None:
+            module._buffers[name] = module._buffers[name].to(torch.device("meta"))
+
+    try:
+        nn.Module.register_parameter = register_empty_parameter
+        if include_buffers:
+            nn.Module.register_buffer = register_empty_buffer
+        yield
+    finally:
+        nn.Module.register_parameter = old_register_parameter
+        if include_buffers:
+            nn.Module.register_buffer = old_register_buffer
+
+
+def cpu_offload(
+    model: nn.Module,
+    execution_device: Optional[torch.device] = None,
+    offload_buffers: bool = False,
+    state_dict: Optional[Dict[str, torch.Tensor]] = None,
+):
+    """
+    Activates full CPU offload for a model. As a result, all parameters of the model will be offloaded and only one
+    copy of the state dict of the model will be kept. During the forward pass, parameters will be extracted from that
+    state dict and put on the execution device passed as they are needed, then offloaded again.
+
+    Args:
+        model (`torch.nn.Module`):
+            The model to offload.
+        execution_device (`torch.device`, *optional*):
+            The device on which the forward pass of the model will be executed (should be a GPU). Will default to the
+            model first parameter device.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+        state_dict (`Dict[str, torch.Tensor]`, *optional*):
+            The state dict of the model that will be kept on CPU.
+    """
+    if execution_device is None:
+        execution_device = next(iter(model.parameter())).device
```suggestion
        execution_device = next(iter(model.parameters())).device
```
@@ -0,0 +1,260 @@
+# 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
+from contextlib import contextmanager
+from typing import Dict, List, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from .hooks import AlignDevicesHook, add_hook_to_module, attach_align_device_hook, attach_align_device_hook_on_blocks
+from .utils import (
+    OffloadedWeightsLoader,
+    check_device_map,
+    extract_submodules_state_dict,
+    infer_auto_device_map,
+    load_checkpoint_in_model,
+    offload_state_dict,
+)
+
+
+@contextmanager
+def init_empty_weights(include_buffers: bool = False):
+    """
+    A context manager under which models are initialized with all parameters on the meta device, therefore creating an
+    empty model. Useful when just initializing the model would blow the available RAM.
+
+    Args:
+        include_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to also put all buffers on the meta device while initializing.
+
+    Example:
+
+    ```pyton
+    import torch.nn as nn
+    from accelerate import init_empty_weights
+
+    # Initialize a model with 100 billions parameters in no time and without using any RAM.
+    with init_empty_weights():
+        tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])
+    ```
+    """
+    old_register_parameter = nn.Module.register_parameter
+    if include_buffers:
+        old_register_buffer = nn.Module.register_buffer
+
+    def register_empty_parameter(module, name, param):
+        old_register_parameter(module, name, param)
+        if param is not None:
+            module._parameters[name] = nn.Parameter(module._parameters[name].to(torch.device("meta")))
+
+    def register_empty_buffer(module, name, buffer):
+        old_register_buffer(module, name, buffer)
+        if buffer is not None:
+            module._buffers[name] = module._buffers[name].to(torch.device("meta"))
+
+    try:
+        nn.Module.register_parameter = register_empty_parameter
+        if include_buffers:
+            nn.Module.register_buffer = register_empty_buffer
+        yield
+    finally:
+        nn.Module.register_parameter = old_register_parameter
+        if include_buffers:
+            nn.Module.register_buffer = old_register_buffer
+
+
+def cpu_offload(
+    model: nn.Module,
+    execution_device: Optional[torch.device] = None,
+    offload_buffers: bool = False,
+    state_dict: Optional[Dict[str, torch.Tensor]] = None,
+):
+    """
+    Activates full CPU offload for a model. As a result, all parameters of the model will be offloaded and only one
+    copy of the state dict of the model will be kept. During the forward pass, parameters will be extracted from that
+    state dict and put on the execution device passed as they are needed, then offloaded again.
+
+    Args:
+        model (`torch.nn.Module`):
+            The model to offload.
+        execution_device (`torch.device`, *optional*):
+            The device on which the forward pass of the model will be executed (should be a GPU). Will default to the
+            model first parameter device.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+        state_dict (`Dict[str, torch.Tensor]`, *optional*):
+            The state dict of the model that will be kept on CPU.
+    """
+    if execution_device is None:
+        execution_device = next(iter(model.parameter())).device
+    if state_dict is None:
+        state_dict = {n: p.to("cpu") for n, p in model.state_dict().items()}
+    attach_align_device_hook(
+        model, execution_device=execution_device, offload=True, offload_buffers=offload_buffers, weights_map=state_dict
+    )
+    add_hook_to_module(model, AlignDevicesHook(io_same_device=True))
+    return model
+
+
+def disk_offload(
+    model: nn.Module,
+    offload_dir: Union[str, os.PathLike],
+    execution_device: Optional[torch.device] = None,
+    offload_buffers: bool = False,
+):
+    """
+    Activates full disk offload for a model. As a result, all parameters of the model will be offloaded as
+    memory-mapped array in a given folder. During the forward pass, parameters will be accessed from that folder and
+    put on the execution device passed as they are needed, then offloaded again.
+
+    Args:
+        model (`torch.nn.Module`): The model to offload.
+        offload_dir (`str` or `os.PathLike`):
+            The folder in which to offload the model weights (or where the model weights are already offloaded).
+        execution_device (`torch.device`, *optional*):
+            The device on which the forward pass of the model will be executed (should be a GPU). Will default to the
+            model's first parameter device.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+    """
+    if not os.path.isdir(offload_dir) or not os.path.isfile(os.path.join(offload_dir, "index.json")):
+        offload_state_dict(offload_dir, model.state_dict())
+    if execution_device is None:
+        execution_device = next(iter(model.parameter())).device
This doesn't contain the `str` which should be accepted as `auto` is valid :)
```suggestion
    device_map: Optional[Union[str, Dict[str, Union[int, str, torch.device]]]] = None,
```
@@ -0,0 +1,260 @@
+# 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
+from contextlib import contextmanager
+from typing import Dict, List, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from .hooks import AlignDevicesHook, add_hook_to_module, attach_align_device_hook, attach_align_device_hook_on_blocks
+from .utils import (
+    OffloadedWeightsLoader,
+    check_device_map,
+    extract_submodules_state_dict,
+    infer_auto_device_map,
+    load_checkpoint_in_model,
+    offload_state_dict,
+)
+
+
+@contextmanager
+def init_empty_weights(include_buffers: bool = False):
+    """
+    A context manager under which models are initialized with all parameters on the meta device, therefore creating an
+    empty model. Useful when just initializing the model would blow the available RAM.
+
+    Args:
+        include_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to also put all buffers on the meta device while initializing.
+
+    Example:
+
+    ```pyton
+    import torch.nn as nn
+    from accelerate import init_empty_weights
+
+    # Initialize a model with 100 billions parameters in no time and without using any RAM.
+    with init_empty_weights():
+        tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])
+    ```
+    """
+    old_register_parameter = nn.Module.register_parameter
+    if include_buffers:
+        old_register_buffer = nn.Module.register_buffer
+
+    def register_empty_parameter(module, name, param):
+        old_register_parameter(module, name, param)
+        if param is not None:
+            module._parameters[name] = nn.Parameter(module._parameters[name].to(torch.device("meta")))
+
+    def register_empty_buffer(module, name, buffer):
+        old_register_buffer(module, name, buffer)
+        if buffer is not None:
+            module._buffers[name] = module._buffers[name].to(torch.device("meta"))
+
+    try:
+        nn.Module.register_parameter = register_empty_parameter
+        if include_buffers:
+            nn.Module.register_buffer = register_empty_buffer
+        yield
+    finally:
+        nn.Module.register_parameter = old_register_parameter
+        if include_buffers:
+            nn.Module.register_buffer = old_register_buffer
+
+
+def cpu_offload(
+    model: nn.Module,
+    execution_device: Optional[torch.device] = None,
+    offload_buffers: bool = False,
+    state_dict: Optional[Dict[str, torch.Tensor]] = None,
+):
+    """
+    Activates full CPU offload for a model. As a result, all parameters of the model will be offloaded and only one
+    copy of the state dict of the model will be kept. During the forward pass, parameters will be extracted from that
+    state dict and put on the execution device passed as they are needed, then offloaded again.
+
+    Args:
+        model (`torch.nn.Module`):
+            The model to offload.
+        execution_device (`torch.device`, *optional*):
+            The device on which the forward pass of the model will be executed (should be a GPU). Will default to the
+            model first parameter device.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+        state_dict (`Dict[str, torch.Tensor]`, *optional*):
+            The state dict of the model that will be kept on CPU.
+    """
+    if execution_device is None:
+        execution_device = next(iter(model.parameter())).device
+    if state_dict is None:
+        state_dict = {n: p.to("cpu") for n, p in model.state_dict().items()}
+    attach_align_device_hook(
+        model, execution_device=execution_device, offload=True, offload_buffers=offload_buffers, weights_map=state_dict
+    )
+    add_hook_to_module(model, AlignDevicesHook(io_same_device=True))
+    return model
+
+
+def disk_offload(
+    model: nn.Module,
+    offload_dir: Union[str, os.PathLike],
+    execution_device: Optional[torch.device] = None,
+    offload_buffers: bool = False,
+):
+    """
+    Activates full disk offload for a model. As a result, all parameters of the model will be offloaded as
+    memory-mapped array in a given folder. During the forward pass, parameters will be accessed from that folder and
+    put on the execution device passed as they are needed, then offloaded again.
+
+    Args:
+        model (`torch.nn.Module`): The model to offload.
+        offload_dir (`str` or `os.PathLike`):
+            The folder in which to offload the model weights (or where the model weights are already offloaded).
+        execution_device (`torch.device`, *optional*):
+            The device on which the forward pass of the model will be executed (should be a GPU). Will default to the
+            model's first parameter device.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+    """
+    if not os.path.isdir(offload_dir) or not os.path.isfile(os.path.join(offload_dir, "index.json")):
+        offload_state_dict(offload_dir, model.state_dict())
+    if execution_device is None:
+        execution_device = next(iter(model.parameter())).device
+    weights_map = OffloadedWeightsLoader(save_folder=offload_dir)
+    attach_align_device_hook(
+        model,
+        execution_device=execution_device,
+        offload=True,
+        offload_buffers=offload_buffers,
+        weights_map=weights_map,
+    )
+    add_hook_to_module(model, AlignDevicesHook(io_same_device=True))
+    return model
+
+
+def dispatch_model(
+    model: nn.Module,
+    device_map: Dict[str, Union[str, int, torch.device]],
+    main_device: Optional[torch.device] = None,
+    state_dict: Optional[Dict[str, torch.Tensor]] = None,
+    offload_dir: Union[str, os.PathLike] = None,
+    offload_buffers: bool = False,
+):
+    """
+    Dispatches a model according to a given device map. Layers of the model might be spread across GPUs, offloaded on
+    the CPU or even the disk.
+
+    Args:
+        model (`torch.nn.Module`):
+            The model to dispatch.
+        device_map (`Dict[str, Union[str, int, torch.device]]`):
+            A dictionary mapping module names in the models `state_dict` to the device they should go to. Note that
+            `"disk"` is accepted even if it's not a proper value for `torch.device`.
+        main_device (`str`, `int` or `torch.device`, *optional*):
+            The main execution device. Will default to the first device in the `device_map` different from `"cpu"` or
+            `"disk"`.
+        state_dict (`Dict[str, torch.Tensor]`, *optional*):
+            The state dict of the part of the model that will be kept on CPU.
+        offload_dir (`str` or `os.PathLike`):
+            The folder in which to offload the model weights (or where the model weights are already offloaded).
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+    """
+    # Error early if the device map is incomplete.
+    check_device_map(model, device_map)
+
+    if main_device is None:
+        main_device = [d for d in device_map.values() if d not in ["cpu", "disk"]][0]
+
+    cpu_modules = [name for name, device in device_map.items() if device == "cpu"]
+    if state_dict is None and len(cpu_modules) > 0:
+        state_dict = extract_submodules_state_dict(model.state_dict(), cpu_modules)
+
+    disk_modules = [name for name, device in device_map.items() if device == "disk"]
+    if offload_dir is None and len(disk_modules) > 0:
+        raise ValueError(
+            "We need an `offload_dir` to dispatch this model according to this `device_map`, the following submodules "
+            f"need to be offloaded: {', '.join(disk_modules)}."
+        )
+    if len(disk_modules) > 0 and (
+        not os.path.isdir(offload_dir) or not os.path.isfile(os.path.join(offload_dir, "index.json"))
+    ):
+        disk_state_dict = extract_submodules_state_dict(model.state_dict(), disk_modules)
+        offload_state_dict(offload_dir, disk_state_dict)
+
+    execution_device = {
+        name: main_device if device in ["cpu", "disk"] else device for name, device in device_map.items()
+    }
+    offload = {name: device in ["cpu", "disk"] for name, device in device_map.items()}
+    if state_dict is not None or offload_dir is not None:
+        weights_map = OffloadedWeightsLoader(state_dict=state_dict, save_folder=offload_dir)
+    else:
+        weights_map = None
+
+    attach_align_device_hook_on_blocks(
+        model,
+        execution_device=execution_device,
+        offload=offload,
+        offload_buffers=offload_buffers,
+        weights_map=weights_map,
+    )
+    return model
+
+
+def load_checkpoint_and_dispatch(
+    model: nn.Module,
+    checkpoint: Union[str, os.PathLike],
+    device_map: Optional[Dict[str, Union[int, str, torch.device]]] = None,
This `device_map` is super interesting for introspection inside the model. I'd be very interested in having this as an attribute of the model once it has been dispatched, for example as:
```py
model = load_checkpoint_and_dispatch(
    model,
   ...
)
model.device_map
``` 
@@ -0,0 +1,260 @@
+# 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
+from contextlib import contextmanager
+from typing import Dict, List, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from .hooks import AlignDevicesHook, add_hook_to_module, attach_align_device_hook, attach_align_device_hook_on_blocks
+from .utils import (
+    OffloadedWeightsLoader,
+    check_device_map,
+    extract_submodules_state_dict,
+    infer_auto_device_map,
+    load_checkpoint_in_model,
+    offload_state_dict,
+)
+
+
+@contextmanager
+def init_empty_weights(include_buffers: bool = False):
+    """
+    A context manager under which models are initialized with all parameters on the meta device, therefore creating an
+    empty model. Useful when just initializing the model would blow the available RAM.
+
+    Args:
+        include_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to also put all buffers on the meta device while initializing.
+
+    Example:
+
+    ```pyton
+    import torch.nn as nn
+    from accelerate import init_empty_weights
+
+    # Initialize a model with 100 billions parameters in no time and without using any RAM.
+    with init_empty_weights():
+        tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])
+    ```
+    """
+    old_register_parameter = nn.Module.register_parameter
+    if include_buffers:
+        old_register_buffer = nn.Module.register_buffer
+
+    def register_empty_parameter(module, name, param):
+        old_register_parameter(module, name, param)
+        if param is not None:
+            module._parameters[name] = nn.Parameter(module._parameters[name].to(torch.device("meta")))
+
+    def register_empty_buffer(module, name, buffer):
+        old_register_buffer(module, name, buffer)
+        if buffer is not None:
+            module._buffers[name] = module._buffers[name].to(torch.device("meta"))
+
+    try:
+        nn.Module.register_parameter = register_empty_parameter
+        if include_buffers:
+            nn.Module.register_buffer = register_empty_buffer
+        yield
+    finally:
+        nn.Module.register_parameter = old_register_parameter
+        if include_buffers:
+            nn.Module.register_buffer = old_register_buffer
+
+
+def cpu_offload(
+    model: nn.Module,
+    execution_device: Optional[torch.device] = None,
+    offload_buffers: bool = False,
+    state_dict: Optional[Dict[str, torch.Tensor]] = None,
+):
+    """
+    Activates full CPU offload for a model. As a result, all parameters of the model will be offloaded and only one
+    copy of the state dict of the model will be kept. During the forward pass, parameters will be extracted from that
+    state dict and put on the execution device passed as they are needed, then offloaded again.
+
+    Args:
+        model (`torch.nn.Module`):
+            The model to offload.
+        execution_device (`torch.device`, *optional*):
+            The device on which the forward pass of the model will be executed (should be a GPU). Will default to the
+            model first parameter device.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+        state_dict (`Dict[str, torch.Tensor]`, *optional*):
+            The state dict of the model that will be kept on CPU.
+    """
+    if execution_device is None:
+        execution_device = next(iter(model.parameter())).device
+    if state_dict is None:
+        state_dict = {n: p.to("cpu") for n, p in model.state_dict().items()}
+    attach_align_device_hook(
+        model, execution_device=execution_device, offload=True, offload_buffers=offload_buffers, weights_map=state_dict
+    )
+    add_hook_to_module(model, AlignDevicesHook(io_same_device=True))
+    return model
+
+
+def disk_offload(
+    model: nn.Module,
+    offload_dir: Union[str, os.PathLike],
+    execution_device: Optional[torch.device] = None,
+    offload_buffers: bool = False,
+):
+    """
+    Activates full disk offload for a model. As a result, all parameters of the model will be offloaded as
+    memory-mapped array in a given folder. During the forward pass, parameters will be accessed from that folder and
+    put on the execution device passed as they are needed, then offloaded again.
+
+    Args:
+        model (`torch.nn.Module`): The model to offload.
+        offload_dir (`str` or `os.PathLike`):
+            The folder in which to offload the model weights (or where the model weights are already offloaded).
+        execution_device (`torch.device`, *optional*):
+            The device on which the forward pass of the model will be executed (should be a GPU). Will default to the
+            model's first parameter device.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+    """
+    if not os.path.isdir(offload_dir) or not os.path.isfile(os.path.join(offload_dir, "index.json")):
+        offload_state_dict(offload_dir, model.state_dict())
+    if execution_device is None:
+        execution_device = next(iter(model.parameter())).device
+    weights_map = OffloadedWeightsLoader(save_folder=offload_dir)
+    attach_align_device_hook(
+        model,
+        execution_device=execution_device,
+        offload=True,
+        offload_buffers=offload_buffers,
+        weights_map=weights_map,
+    )
+    add_hook_to_module(model, AlignDevicesHook(io_same_device=True))
+    return model
+
+
+def dispatch_model(
+    model: nn.Module,
+    device_map: Dict[str, Union[str, int, torch.device]],
+    main_device: Optional[torch.device] = None,
+    state_dict: Optional[Dict[str, torch.Tensor]] = None,
+    offload_dir: Union[str, os.PathLike] = None,
+    offload_buffers: bool = False,
+):
+    """
+    Dispatches a model according to a given device map. Layers of the model might be spread across GPUs, offloaded on
+    the CPU or even the disk.
+
+    Args:
+        model (`torch.nn.Module`):
+            The model to dispatch.
+        device_map (`Dict[str, Union[str, int, torch.device]]`):
+            A dictionary mapping module names in the models `state_dict` to the device they should go to. Note that
+            `"disk"` is accepted even if it's not a proper value for `torch.device`.
+        main_device (`str`, `int` or `torch.device`, *optional*):
+            The main execution device. Will default to the first device in the `device_map` different from `"cpu"` or
+            `"disk"`.
+        state_dict (`Dict[str, torch.Tensor]`, *optional*):
+            The state dict of the part of the model that will be kept on CPU.
+        offload_dir (`str` or `os.PathLike`):
+            The folder in which to offload the model weights (or where the model weights are already offloaded).
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+    """
+    # Error early if the device map is incomplete.
+    check_device_map(model, device_map)
+
+    if main_device is None:
+        main_device = [d for d in device_map.values() if d not in ["cpu", "disk"]][0]
+
+    cpu_modules = [name for name, device in device_map.items() if device == "cpu"]
+    if state_dict is None and len(cpu_modules) > 0:
+        state_dict = extract_submodules_state_dict(model.state_dict(), cpu_modules)
+
+    disk_modules = [name for name, device in device_map.items() if device == "disk"]
+    if offload_dir is None and len(disk_modules) > 0:
+        raise ValueError(
+            "We need an `offload_dir` to dispatch this model according to this `device_map`, the following submodules "
+            f"need to be offloaded: {', '.join(disk_modules)}."
+        )
+    if len(disk_modules) > 0 and (
+        not os.path.isdir(offload_dir) or not os.path.isfile(os.path.join(offload_dir, "index.json"))
+    ):
+        disk_state_dict = extract_submodules_state_dict(model.state_dict(), disk_modules)
+        offload_state_dict(offload_dir, disk_state_dict)
+
+    execution_device = {
+        name: main_device if device in ["cpu", "disk"] else device for name, device in device_map.items()
+    }
+    offload = {name: device in ["cpu", "disk"] for name, device in device_map.items()}
+    if state_dict is not None or offload_dir is not None:
+        weights_map = OffloadedWeightsLoader(state_dict=state_dict, save_folder=offload_dir)
+    else:
+        weights_map = None
+
+    attach_align_device_hook_on_blocks(
+        model,
+        execution_device=execution_device,
+        offload=offload,
+        offload_buffers=offload_buffers,
+        weights_map=weights_map,
+    )
+    return model
+
+
+def load_checkpoint_and_dispatch(
+    model: nn.Module,
+    checkpoint: Union[str, os.PathLike],
+    device_map: Optional[Dict[str, Union[int, str, torch.device]]] = None,
+    max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None,
+    no_split_module_classes: Optional[List[str]] = None,
+    offload_folder: Optional[Union[str, os.PathLike]] = None,
+    offload_buffers: bool = False,
+):
+    """
+    Loads a (potentially sharded) checkpoint inside a model, potentially sending weights to a given device as they are
+    loaded and adds the various hooks that will make this model run properly (even if split across devices).
+
+    Args:
+        model (`torch.nn.Module`): The model in which we want to load a checkpoint.
+        checkpoint (`str` or `os.PathLike`):
+            The folder checkpoint to load. It can be:
+            - a path to a file containing a whole model state dict
+            - a path to a `.json` file containing the index to a sharded checkpoint
+            - a path to a folder containing a unique `.index.json` file and the shards of a checkpoint.
+        device_map (`Dict[str, Union[int, str, torch.device]]`, *optional*):
+            A map that specifies where each submodule should go. It doesn't need to be refined to each parameter/buffer
+            name, once a given module name is inside, every submodule of it will be sent to the same device.
+
+            To have Accelerate compute the most optimized `device_map` automatically, set `device_map="auto"`.
+        max_memory (`Dict`, *optional*):
+            A dictionary device identifier to maximum memory. Will default to the maximum memory available for each GPU
+            and the available CPU RAM if unset.
+        no_split_module_classes (`List[str]`, *optional*):
+            A list of layer class names that should never be split across device (for instance any layer that has a
+            residual connection).
+        offload_folder (`str` or `os.PathLike`, *optional*):
+            If the `device_map` contains any value `"disk"`, the folder where we will offload weights.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            In the layers that are offloaded on the CPU or the hard drive, whether or not to offload the buffers as
+            well as the parameters.
+    """
+    if device_map == "auto":
+        device_map = infer_auto_device_map(
+            model, max_memory=max_memory, no_split_module_classes=no_split_module_classes
+        )
This is likely the place where most memory errors will occur, if the checkpoints have not been sharded. Probably out of scope for this PR, but we could add a nice error message mentioning how to shard the checkpoint if it fails here.
@@ -0,0 +1,568 @@
+# 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 gc
+import json
+import os
+import re
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+WEIGHTS_INDEX_NAME = "pytorch_model.bin.index.json"
+
+
+def convert_file_size_to_int(size: Union[int, str]):
+    """
+    Converts a size expressed as a string with digits an unit (like `"5MB"`) to an integer (in bytes).
+
+    Args:
+        size (`int` or `str`): The size to convert. Will be directly returned if an `int`.
+
+    Example:
+
+    ```py
+    >>> convert_file_size_to_int("1MiB")
+    1048576
+    ```
+    """
+    if isinstance(size, int):
+        return size
+    if size.upper().endswith("GIB"):
+        return int(size[:-3]) * (2**30)
+    if size.upper().endswith("MIB"):
+        return int(size[:-3]) * (2**20)
+    if size.upper().endswith("KIB"):
+        return int(size[:-3]) * (2**10)
+    if size.upper().endswith("GB"):
+        int_size = int(size[:-2]) * (10**9)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("MB"):
+        int_size = int(size[:-2]) * (10**6)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("KB"):
+        int_size = int(size[:-2]) * (10**3)
+        return int_size // 8 if size.endswith("b") else int_size
+    raise ValueError("`size` is not in a valid format. Use an integer followed by the unit, e.g., '5GB'.")
+
+
+def dtype_byte_size(dtype: torch.dtype):
+    """
+    Returns the size (in bytes) occupied by one parameter of type `dtype`.
+
+    Example:
+
+    ```py
+    >>> dtype_byte_size(torch.float32)
+    4
+    ```
+    """
+    if dtype == torch.bool:
+        return 1 / 8
+    bit_search = re.search("[^\d](\d+)$", str(dtype))
+    if bit_search is None:
+        raise ValueError(f"`dtype` is not a valid dtype: {dtype}.")
+    bit_size = int(bit_search.groups()[0])
+    return bit_size // 8
+
+
+def set_module_tensor_to_device(
+    module: nn.Module, tensor_name: str, device: Union[int, str, torch.device], value: Optional[torch.Tensor] = None
+):
+    """
+    A helper function to set a given tensor (parameter of buffer) of a module on a specific device (note that doing
+    `param.to(device)` creates a new tensor not linked to the parameter, which is why we need this function).
+
+    Args:
+        module (`torch.nn.Module`): The module in which the tensor we want to move lives.
+        param_name (`str`): The full name of the parameter/buffer.
+        device (`int`, `str` or `torch.device`): The device on which to set the tensor.
+        value (`torch.Tensor`, *optional*): The value of the tensor (useful when going from the meta device to any
+            other device).
+    """
+    # Recurse if needed
+    if "." in tensor_name:
+        splits = tensor_name.split(".")
+        for split in splits[:-1]:
+            new_module = getattr(module, split)
+            if new_module is None:
+                raise ValueError(f"{module} has no attribute {split}.")
+            module = new_module
+        tensor_name = splits[-1]
+
+    if tensor_name not in module._parameters and tensor_name not in module._buffers:
+        raise ValueError(f"{module} does not have a parameter or a buffer named {tensor_name}.")
+    is_buffer = tensor_name in module._buffers
+    old_value = getattr(module, tensor_name)
+
+    if old_value.device == torch.device("meta") and device != torch.device("meta") and value is None:
+        raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {device}.")
+
+    with torch.no_grad():
+        if value is None:
+            new_value = old_value.to(device)
+        elif isinstance(value, torch.Tensor):
+            new_value = value.to(device)
+        else:
+            new_value = torch.tensor(value, device=device)
+    if is_buffer:
+        module._buffers[tensor_name] = new_value
+    else:
+        new_value = nn.Parameter(new_value, requires_grad=old_value.requires_grad)
+        module._parameters[tensor_name] = new_value
+
+
+def named_module_tensors(module: nn.Module, include_buffers: bool = True, recurse: bool = False):
+    """
+    A helper function that gathers all the tensors (parameters + buffers) of a given module. If `include_buffers=True`
+    it's the same as doing `module.named_parameters(recurse=recurse) + module.named_buffers(recurse=recurse)`.
+
+    Args:
+        module (`torch.nn.Module`): The module we want the tensors or.
+        include_buffer (`bool`, *optional*, defaults to `True`): Whether or not to include the buffers in the result.
+        recurse (`bool`, *optional`, defaults to `False`):
+            Whether or not to go look in every submodule or just return the direct parameters and buffers.
+    """
+    for named_parameter in module.named_parameters(recurse=recurse):
+        yield named_parameter
+
+    if include_buffers:
+        for named_buffer in module.named_buffers(recurse=recurse):
+            yield named_buffer
+
+
+def find_tied_parameters(model: nn.Module, **kwargs):
+    """
+    Find the tied parameters in a given model.
+
+    Args:
+        model (`torch.nn.Module`): The model to inspect.
+
+    <Tip warning={true}>
+
+    The signature accepts keyword arguments, but they are for the recursive part of this function and you should ignore
+    them.
+
+    </Tip>
+
+    Example:
+
+
+    ```py
+    >>> from collections import OrderedDict
+    >>> import torch.nn as nn
+
+    >>> model = nn.Sequential(OrderedDict([("linear1", nn.Linear(4, 4)), ("linear2", nn.Linear(4, 4))]))
+    >>> model.linear2.weight = test_model.linear1.weight
+    >>> find_tied_parameters(test_model)
+    {'linear1.weight': 'linear2.weight'}
+    ```
+
+    Returns:
+        Dict[str, str]: A dictionary mapping tied parameter names to the name of the parameter they are tied to.
+    """
+    # Initialize result and named_parameters before recursing.
+    named_parameters = kwargs.get("named_parameters", None)
+    prefix = kwargs.get("prefix", "")
+    result = kwargs.get("result", {})
+
+    if named_parameters is None:
+        named_parameters = {n: p for n, p in model.named_parameters()}
+    else:
+        # A tied parameter will not be in the full `named_parameters` seen above but will be in the `named_parameters`
+        # of the submodule it belongs to. So while recursing we track the names that are not in the initial
+        # `named_parameters`.
+        for name, parameter in model.named_parameters():
+            full_name = name if prefix == "" else f"{prefix}.{name}"
+            if full_name not in named_parameters:
+                # When we find one, it has to be one of the existing parameters.
+                for new_name, new_param in named_parameters.items():
+                    if new_param is parameter:
+                        result[new_name] = full_name
+
+    # Once we have treated direct parameters, we move to the child modules.
+    for name, child in model.named_children():
+        child_name = name if prefix == "" else f"{prefix}.{name}"
+        find_tied_parameters(child, named_parameters=named_parameters, prefix=child_name, result=result)
+
+    return result
+
+
+def compute_module_sizes(model: nn.Module):
+    """
+    Compute the size of each submodule of a given model.
+    """
+    module_sizes = defaultdict(int)
+    for name, tensor in named_module_tensors(model, recurse=True):
+        size = tensor.numel() * dtype_byte_size(tensor.dtype)
+        name_parts = name.split(".")
+        for idx in range(len(name_parts) + 1):
+            module_sizes[".".join(name_parts[:idx])] += size
+
+    return module_sizes
+
+
+def get_max_layer_size(
+    modules: List[Tuple[str, torch.nn.Module]], module_sizes: Dict[str, int], no_split_module_classes: List[str]
+):
+    """
+    Utility function that will scan a list of named modules and return the maximum size used by one full layer. The
+    definition of a layer being:
+    - a module with no direct children (just parameters and buffers)
+    - a module whose class name is in the list `no_split_module_classes`
+
+    Args:
+        modules (`List[Tuple[str, torch.nn.Module]]`):
+            The list of named modules where we want to determine the maximum layer size.
+        module_sizes (`Dict[str, int]`):
+            A dictionary mapping each layer name to its size (as generated by `compute_module_sizes`).
+        no_split_module_classes (`List[str]`):
+            A list of class names for layers we don't want to be split.
+
+    Returns:
+        `Tuple[int, List[str]]`: The maximum size of a layer with the list of layer names realizing that maximum size.
+    """
+    max_size = 0
+    layer_names = []
+    modules_to_treat = modules.copy()
+    while len(modules_to_treat) > 0:
+        module_name, module = modules_to_treat.pop(0)
+        modules_children = list(module.named_children())
+        if len(modules_children) == 0 or module.__class__.__name__ in no_split_module_classes:
+            # No splitting this one so we compare to the max_size
+            size = module_sizes[module_name]
+            if size > max_size:
+                max_size = size
+                layer_names = [module_name]
+            elif size == max_size:
+                layer_names.append(module_name)
+        else:
+            modules_to_treat = [(f"{module_name}.{n}", v) for n, v in modules_children] + modules_to_treat
+    return max_size, layer_names
+
+
+def get_max_memory(max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None):
+    """
+    Get the maximum memory available if nothing is passed, converts string to int otherwise.
+    """
+    import psutil
+
+    if max_memory is None:
+        if not torch.cuda.is_available():
+            max_memory = {}
+        else:
+            # Make sure CUDA is initialized on each GPU to have the right memory info.
+            for i in range(torch.cuda.device_count()):
+                _ = torch.tensor([0], device=i)
+            max_memory = {i: torch.cuda.mem_get_info(i)[0] for i in range(torch.cuda.device_count())}
+        max_memory["cpu"] = psutil.virtual_memory().available
+        return max_memory
+
+    for key in max_memory:
+        if isinstance(max_memory[key], str):
+            max_memory[key] = convert_file_size_to_int(max_memory[key])
+    return max_memory
+
+
+def clean_device_map(device_map: Dict[str, Union[int, str, torch.device]], module_name: str = ""):
+    """
+    Cleans a device_map by grouping all submodules that go on the same device together.
+    """
+    # Get the value of the current module and if there is only one split across several keys, regroup it.
+    values = [v for k, v in device_map.items() if k.startswith(module_name)]
+    if len(set(values)) == 1 and len(values) > 1:
+        for k in [k for k in device_map if k.startswith(module_name)]:
+            del device_map[k]
+        device_map[module_name] = values[0]
+
+    # Recurse over the children
+    children_modules = [k for k in device_map.keys() if k.startswith(module_name) and len(k) > len(module_name)]
+    idx = len(module_name.split(".")) + 1 if len(module_name) > 0 else 1
+    children_modules = set(".".join(k.split(".")[:idx]) for k in children_modules)
+    for child in children_modules:
+        clean_device_map(device_map, module_name=child)
+
+    return device_map
+
+
+def infer_auto_device_map(
+    model: nn.Module,
+    max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None,
+    no_split_module_classes: Optional[List[str]] = None,
+):
+    """
+    Compute a device map for a given model giving priority to GPUs, then offload on CPU and finally offload to disk,
+    such that:
+    - we don't exceed the memory available of any of the GPU.
+    - if offload to the CPU is needed, there is always room left on GPU 0 to put back the layer offloaded on CPU that
+      has the largest size.
+    - if offload to the CPU is needed,we don't exceed the RAM available on the CPU.
+    - if offload to the disk is needed, there is always room left on the CPU to put back the layer offloaded on disk
+      that has the largest size.
+
+    <Tip>
+
+    All computation is done analyzing sizes and dtypes of the model parameters. As a result, the model can be on the
+    meta device (as it would if initialized within the `init_empty_weights` context manager).
+
+    </Tip>
+
+    Args:
+        model (`torch.nn.Module`): The model to analyze.
+        max_memory (`Dict`, *optional*):
+            A dictionary device identifier to maximum memory. Will default to the maximum memory available if unset.
+        no_split_module_classes (`List[str]`, *optional*):
+            A list of layer class names that should never be split across device (for instance any layer that has a
+            residual connection).
+    """
+    # Get default / clean up max_memory
+    max_memory = get_max_memory(max_memory)
+    if no_split_module_classes is None:
+        no_split_module_classes = []
+    elif not isinstance(no_split_module_classes, (list, tuple)):
+        no_split_module_classes = [no_split_module_classes]
+
+    devices = list(max_memory.keys()) + ["disk"]
+    # Devices that need to keep space for a potential offloaded layer.
+    main_devices = [devices[0], "cpu"]
+
+    module_sizes = compute_module_sizes(model)
+    tied_parameters = find_tied_parameters(model)
+
+    device_map = {}
+    current_device = 0
+    current_memory_used = 0
+
+    # Direct submodules and parameters
+    modules_to_treat = list(model.named_parameters(recurse=False)) + list(model.named_children())
+    # Initialize maximum largest layer, to know which space to keep in memory
+    max_layer_size, max_layer_names = get_max_layer_size(modules_to_treat, module_sizes, no_split_module_classes)
+
+    # Ready ? This is going to be a bit messy.
+    while len(modules_to_treat) > 0:
+        name, module = modules_to_treat.pop(0)
+        # Max size in the remaining layers may have changed since we took one, so we maybe update it.
+        max_layer_names = [n for n in max_layer_names if not n.startswith(name)]
+        if len(max_layer_names) == 0:
+            max_layer_size, max_layer_names = get_max_layer_size(
+                [(n, m) for n, m in modules_to_treat if isinstance(m, torch.nn.Module)],
+                module_sizes,
+                no_split_module_classes,
+            )
+        # Assess size needed
+        module_size = module_sizes[name]
+        tied_params = [v for k, v in tied_parameters.items() if name in k]
+        # We ignore parameters that are tied when they're tied to > 1 one
+        tied_param = tied_params[0] if len(tied_params) == 1 else None
+
+        device = devices[current_device]
+        current_max_size = max_memory[device] if device != "disk" else None
+        # Reduce max size available by the largest layer.
+        if devices[current_device] in main_devices:
+            current_max_size = current_max_size - max_layer_size
+        # Case 1 -> We're too big!
+        if current_max_size is not None and current_memory_used + module_size > current_max_size:
+            # Split or not split?
+            modules_children = list(module.named_children())
+            if len(modules_children) == 0 or module.__class__.__name__ in no_split_module_classes:
+                # -> no split, we go to the next device
+                current_device += 1
+                modules_to_treat = [(name, module)] + modules_to_treat
+                current_memory_used = 0
+            else:
+                # -> split, we replace the module studied by its children + parameters
+                modules_children = list(module.named_parameters(recurse=False)) + modules_children
+                modules_to_treat = [(f"{name}.{n}", v) for n, v in modules_children] + modules_to_treat
+                # Update the max layer size.
+                max_layer_size, max_layer_names = get_max_layer_size(
+                    [(n, m) for n, m in modules_to_treat if isinstance(m, torch.nn.Module)],
+                    module_sizes,
+                    no_split_module_classes,
+                )
+
+        # Case 2, it fits! We're not entirely out of the wood though, because we may have some tied parameters.
+        elif tied_param is not None:
+            # Determine the sized occupied by this module + the module containing the tied parameter
+            tied_module_size = module_size
+            tied_module_index = [i for i, (n, _) in enumerate(modules_to_treat) if n in tied_param][0]
+            tied_module_name, tied_module = modules_to_treat[tied_module_index]
+            tied_module_size += module_sizes[tied_module_name] - module_sizes[tied_param]
+            if current_max_size is not None and current_memory_used + tied_module_size > current_max_size:
+                # Split or not split?
+                tied_module_children = list(tied_module.named_children())
+                if len(tied_module_children) == 0 or tied_module.__class__.__name__ in no_split_module_classes:
+                    # If the tied module is not split, we go to the next device
+                    current_device += 1
+                    modules_to_treat = [(name, module)] + modules_to_treat
+                    current_memory_used = 0
+                else:
+                    # Otherwise, we replace the tied module by its children.
+                    tied_module_children = list(tied_module.named_parameters(recurse=False)) + tied_module_children
+                    tied_module_children = [(f"{tied_module_name}.{n}", v) for n, v in tied_module_children]
+                    modules_to_treat = (
+                        [(name, module)]
+                        + modules_to_treat[:tied_module_index]
+                        + tied_module_children
+                        + modules_to_treat[tied_module_index + 1 :]
+                    )
+                    # Update the max layer size.
+                    max_layer_size, max_layer_names = get_max_layer_size(
+                        [(n, m) for n, m in modules_to_treat if isinstance(m, torch.nn.Module)],
+                        module_sizes,
+                        no_split_module_classes,
+                    )
+            else:
+                # We really really fit!
+                current_memory_used += tied_module_size
+                device_map[name] = devices[current_device]
+                modules_to_treat.pop(tied_module_index)
+                device_map[tied_module_name] = devices[current_device]
+        else:
+            current_memory_used += module_size
+            device_map[name] = devices[current_device]
+
+    return clean_device_map(device_map)
+
+
+def check_device_map(model: nn.Module, device_map: Dict[str, Union[int, str, torch.device]]):
+    """
+    Checks a device map covers everything in a given model.
+
+    Args:
+        model (`torch.nn.Module`): The model to check the device map against.
+        device_map (`Dict[str, Union[int, str, torch.device]]`): The device map to check.
+    """
+    all_model_tensors = [name for name, _ in model.state_dict().items()]
+    for module_name in device_map.keys():
+        all_model_tensors = [name for name in all_model_tensors if not name.startswith(module_name)]
+    if len(all_model_tensors) > 0:
+        non_covered_params = ", ".join(all_model_tensors)
+        raise ValueError(
+            f"The device_map provided does not give any device for the following parameters: {non_covered_params}"
+        )
+
+
+def load_checkpoint_in_model(
+    model: nn.Module,
+    checkpoint: Union[str, os.PathLike],
+    device_map: Optional[Dict[str, Union[int, str, torch.device]]] = None,
+    offload_folder: Optional[Union[str, os.PathLike]] = None,
+):
+    """
+    Loads a (potentially sharded) checkpoint inside a model, potentially sending weights to a given device as they are
+    loaded.
+
+    <Tip warning={true}>
+
+    Once loaded across devices, you still need to call [`dispatch_model`] on your model to make it able to run. To
+    group the checkpoint loading and dispatch in one single call, use [`load_checkpoint_and_dispatch`].
+
+    </Tip>
+
+    Args:
+        model (`torch.nn.Module`): The model in which we want to load a checkpoint.
+        checkpoint (`str` or `os.PathLike`):
+            The folder checkpoint to load. It can be:
+            - a path to a file containing a whole model state dict
+            - a path to a `.json` file containing the index to a sharded checkpoint
+            - a path to a folder containing a unique `.index.json` file and the shards of a checkpoint.
+        device_map (`Dict[str, Union[int, str, torch.device]]`, *optional*):
+            A map that specifies where each submodule should go. It doesn't need to be refined to each parameter/buffer
+            name, once a given module name is inside, every submodule of it will be sent to the same device.
+        offload_folder (`str` or `os.PathLike`, *optional*):
+            If the `device_map` contains any value `"disk"`, the folder where we will offload weights.
+    """
+    if offload_folder is None and device_map is not None and "disk" in device_map.values():
+        raise ValueError(
+            "At least one of the model submodule will be offloaded to disk, please pass along an `offload_folder`."
+        )
+
+    checkpoint_files = None
+    index_filename = None
+    if os.path.isfile(checkpoint):
+        if str(checkpoint).endswith(".json"):
+            index_filename = checkpoint
+        else:
+            checkpoint_files = [checkpoint]
+    elif os.path.isdir(checkpoint):
+        potential_index = [f for f in os.listdir(checkpoint) if f.endswith(".index.json")]
+        if len(potential_index) == 0:
+            raise ValueError(f"{checkpoint} is not a folder containing a `.index.json` file.")
+        elif len(potential_index) == 1:
+            index_filename = os.path.join(checkpoint, potential_index[0])
+        else:
+            raise ValueError(f"{checkpoint} containing mote than one `.index.json` file, delete the irrelevant ones.")
+    else:
+        raise ValueError(
+            "`checkpoint` should be the path to a file containing a whole state dict, or the index of a sharded "
+            f"checkpoint, or a folder containing a sharded checkpoint, but got {checkpoint}."
+        )
+
+    if index_filename is not None:
+        checkpoint_folder = os.path.split(index_filename)[0]
+        with open(index_filename, "r") as f:
+            index = json.loads(f.read())
+
+        if "weight_map" in index:
+            index = index["weight_map"]
+        checkpoint_files = sorted(list(set(index.values())))
+        checkpoint_files = [os.path.join(checkpoint_folder, f) for f in checkpoint_files]
+
+    # Logic for missing/unexepected keys goes here.
+
+    offload_index = {}
+    for checkpoint_file in checkpoint_files:
+        checkpoint = torch.load(checkpoint_file)
Should this default to `auto`? Right now it's optional, but not passing it results in an error:
```
raceback (most recent call last):
  File "/home/lysandre/Workspaces/Python/transformers/.env/lib/python3.9/site-packages/IPython/core/interactiveshell.py", line 3397, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-2-f756f1b4a3d4>", line 3, in <cell line: 3>
    runfile('/home/lysandre/.config/JetBrains/PyCharm2021.3/scratches/OPT/playground2.py', wdir='/home/lysandre/.config/JetBrains/PyCharm2021.3/scratches/OPT')
  File "/opt/pycharm-professional/plugins/python/helpers/pydev/_pydev_bundle/pydev_umd.py", line 198, in runfile
    pydev_imports.execfile(filename, global_vars, local_vars)  # execute the script
  File "/opt/pycharm-professional/plugins/python/helpers/pydev/_pydev_imps/_pydev_execfile.py", line 18, in execfile
    exec(compile(contents+"\n", file, 'exec'), glob, loc)
  File "/home/lysandre/.config/JetBrains/PyCharm2021.3/scratches/OPT/playground2.py", line 8, in <module>
    model = load_checkpoint_and_dispatch(
  File "/home/lysandre/Workspaces/Python/accelerate/src/accelerate/big_modeling.py", line 260, in load_checkpoint_and_dispatch
    return dispatch_model(model, device_map=device_map, offload_dir=offload_folder, offload_buffers=offload_buffers)
  File "/home/lysandre/Workspaces/Python/accelerate/src/accelerate/big_modeling.py", line 178, in dispatch_model
    check_device_map(model, device_map)
  File "/home/lysandre/Workspaces/Python/accelerate/src/accelerate/utils/modeling.py", line 451, in check_device_map
    for module_name in device_map.keys():
AttributeError: 'NoneType' object has no attribute 'keys'
```
@@ -0,0 +1,260 @@
+# 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
+from contextlib import contextmanager
+from typing import Dict, List, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from .hooks import AlignDevicesHook, add_hook_to_module, attach_align_device_hook, attach_align_device_hook_on_blocks
+from .utils import (
+    OffloadedWeightsLoader,
+    check_device_map,
+    extract_submodules_state_dict,
+    infer_auto_device_map,
+    load_checkpoint_in_model,
+    offload_state_dict,
+)
+
+
+@contextmanager
+def init_empty_weights(include_buffers: bool = False):
+    """
+    A context manager under which models are initialized with all parameters on the meta device, therefore creating an
+    empty model. Useful when just initializing the model would blow the available RAM.
+
+    Args:
+        include_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to also put all buffers on the meta device while initializing.
+
+    Example:
+
+    ```pyton
+    import torch.nn as nn
+    from accelerate import init_empty_weights
+
+    # Initialize a model with 100 billions parameters in no time and without using any RAM.
+    with init_empty_weights():
+        tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])
+    ```
+    """
+    old_register_parameter = nn.Module.register_parameter
+    if include_buffers:
+        old_register_buffer = nn.Module.register_buffer
+
+    def register_empty_parameter(module, name, param):
+        old_register_parameter(module, name, param)
+        if param is not None:
+            module._parameters[name] = nn.Parameter(module._parameters[name].to(torch.device("meta")))
+
+    def register_empty_buffer(module, name, buffer):
+        old_register_buffer(module, name, buffer)
+        if buffer is not None:
+            module._buffers[name] = module._buffers[name].to(torch.device("meta"))
+
+    try:
+        nn.Module.register_parameter = register_empty_parameter
+        if include_buffers:
+            nn.Module.register_buffer = register_empty_buffer
+        yield
+    finally:
+        nn.Module.register_parameter = old_register_parameter
+        if include_buffers:
+            nn.Module.register_buffer = old_register_buffer
+
+
+def cpu_offload(
+    model: nn.Module,
+    execution_device: Optional[torch.device] = None,
+    offload_buffers: bool = False,
+    state_dict: Optional[Dict[str, torch.Tensor]] = None,
+):
+    """
+    Activates full CPU offload for a model. As a result, all parameters of the model will be offloaded and only one
+    copy of the state dict of the model will be kept. During the forward pass, parameters will be extracted from that
+    state dict and put on the execution device passed as they are needed, then offloaded again.
+
+    Args:
+        model (`torch.nn.Module`):
+            The model to offload.
+        execution_device (`torch.device`, *optional*):
+            The device on which the forward pass of the model will be executed (should be a GPU). Will default to the
+            model first parameter device.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+        state_dict (`Dict[str, torch.Tensor]`, *optional*):
+            The state dict of the model that will be kept on CPU.
+    """
+    if execution_device is None:
+        execution_device = next(iter(model.parameter())).device
+    if state_dict is None:
+        state_dict = {n: p.to("cpu") for n, p in model.state_dict().items()}
+    attach_align_device_hook(
+        model, execution_device=execution_device, offload=True, offload_buffers=offload_buffers, weights_map=state_dict
+    )
+    add_hook_to_module(model, AlignDevicesHook(io_same_device=True))
+    return model
+
+
+def disk_offload(
+    model: nn.Module,
+    offload_dir: Union[str, os.PathLike],
+    execution_device: Optional[torch.device] = None,
+    offload_buffers: bool = False,
+):
+    """
+    Activates full disk offload for a model. As a result, all parameters of the model will be offloaded as
+    memory-mapped array in a given folder. During the forward pass, parameters will be accessed from that folder and
+    put on the execution device passed as they are needed, then offloaded again.
+
+    Args:
+        model (`torch.nn.Module`): The model to offload.
+        offload_dir (`str` or `os.PathLike`):
+            The folder in which to offload the model weights (or where the model weights are already offloaded).
+        execution_device (`torch.device`, *optional*):
+            The device on which the forward pass of the model will be executed (should be a GPU). Will default to the
+            model's first parameter device.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+    """
+    if not os.path.isdir(offload_dir) or not os.path.isfile(os.path.join(offload_dir, "index.json")):
+        offload_state_dict(offload_dir, model.state_dict())
+    if execution_device is None:
+        execution_device = next(iter(model.parameter())).device
+    weights_map = OffloadedWeightsLoader(save_folder=offload_dir)
+    attach_align_device_hook(
+        model,
+        execution_device=execution_device,
+        offload=True,
+        offload_buffers=offload_buffers,
+        weights_map=weights_map,
+    )
+    add_hook_to_module(model, AlignDevicesHook(io_same_device=True))
+    return model
+
+
+def dispatch_model(
+    model: nn.Module,
+    device_map: Dict[str, Union[str, int, torch.device]],
+    main_device: Optional[torch.device] = None,
+    state_dict: Optional[Dict[str, torch.Tensor]] = None,
+    offload_dir: Union[str, os.PathLike] = None,
+    offload_buffers: bool = False,
+):
+    """
+    Dispatches a model according to a given device map. Layers of the model might be spread across GPUs, offloaded on
+    the CPU or even the disk.
+
+    Args:
+        model (`torch.nn.Module`):
+            The model to dispatch.
+        device_map (`Dict[str, Union[str, int, torch.device]]`):
+            A dictionary mapping module names in the models `state_dict` to the device they should go to. Note that
+            `"disk"` is accepted even if it's not a proper value for `torch.device`.
+        main_device (`str`, `int` or `torch.device`, *optional*):
+            The main execution device. Will default to the first device in the `device_map` different from `"cpu"` or
+            `"disk"`.
+        state_dict (`Dict[str, torch.Tensor]`, *optional*):
+            The state dict of the part of the model that will be kept on CPU.
+        offload_dir (`str` or `os.PathLike`):
+            The folder in which to offload the model weights (or where the model weights are already offloaded).
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+    """
+    # Error early if the device map is incomplete.
+    check_device_map(model, device_map)
+
+    if main_device is None:
+        main_device = [d for d in device_map.values() if d not in ["cpu", "disk"]][0]
+
+    cpu_modules = [name for name, device in device_map.items() if device == "cpu"]
+    if state_dict is None and len(cpu_modules) > 0:
+        state_dict = extract_submodules_state_dict(model.state_dict(), cpu_modules)
+
+    disk_modules = [name for name, device in device_map.items() if device == "disk"]
+    if offload_dir is None and len(disk_modules) > 0:
+        raise ValueError(
+            "We need an `offload_dir` to dispatch this model according to this `device_map`, the following submodules "
+            f"need to be offloaded: {', '.join(disk_modules)}."
+        )
+    if len(disk_modules) > 0 and (
+        not os.path.isdir(offload_dir) or not os.path.isfile(os.path.join(offload_dir, "index.json"))
+    ):
+        disk_state_dict = extract_submodules_state_dict(model.state_dict(), disk_modules)
+        offload_state_dict(offload_dir, disk_state_dict)
+
+    execution_device = {
+        name: main_device if device in ["cpu", "disk"] else device for name, device in device_map.items()
+    }
+    offload = {name: device in ["cpu", "disk"] for name, device in device_map.items()}
+    if state_dict is not None or offload_dir is not None:
+        weights_map = OffloadedWeightsLoader(state_dict=state_dict, save_folder=offload_dir)
+    else:
+        weights_map = None
+
+    attach_align_device_hook_on_blocks(
+        model,
+        execution_device=execution_device,
+        offload=offload,
+        offload_buffers=offload_buffers,
+        weights_map=weights_map,
+    )
+    return model
+
+
+def load_checkpoint_and_dispatch(
+    model: nn.Module,
+    checkpoint: Union[str, os.PathLike],
+    device_map: Optional[Dict[str, Union[int, str, torch.device]]] = None,
`max_memory` is a dictionary, which is unordered in python <= 3.6.

Here, we're relying on the ordering of the dictionary in order to retrieve the `main_devices`. I'm open to hearing what you think; I would push to always consider dictionaries as unordered and to treat them as such, but as we're dropping python 3.6 I understand if you want to have the asumption that dictionaries are ordered.

These lists could likely be `set`s however: `devices` should theoretically not have the same device twice, and the same holds for `main_devices`.

If no GPU is available, then `devices` is a list containing `["cpu", "disk"]`. In that case, we end up with `main_devices` containing `['cpu', 'cpu']`.
@@ -0,0 +1,568 @@
+# 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 gc
+import json
+import os
+import re
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+WEIGHTS_INDEX_NAME = "pytorch_model.bin.index.json"
+
+
+def convert_file_size_to_int(size: Union[int, str]):
+    """
+    Converts a size expressed as a string with digits an unit (like `"5MB"`) to an integer (in bytes).
+
+    Args:
+        size (`int` or `str`): The size to convert. Will be directly returned if an `int`.
+
+    Example:
+
+    ```py
+    >>> convert_file_size_to_int("1MiB")
+    1048576
+    ```
+    """
+    if isinstance(size, int):
+        return size
+    if size.upper().endswith("GIB"):
+        return int(size[:-3]) * (2**30)
+    if size.upper().endswith("MIB"):
+        return int(size[:-3]) * (2**20)
+    if size.upper().endswith("KIB"):
+        return int(size[:-3]) * (2**10)
+    if size.upper().endswith("GB"):
+        int_size = int(size[:-2]) * (10**9)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("MB"):
+        int_size = int(size[:-2]) * (10**6)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("KB"):
+        int_size = int(size[:-2]) * (10**3)
+        return int_size // 8 if size.endswith("b") else int_size
+    raise ValueError("`size` is not in a valid format. Use an integer followed by the unit, e.g., '5GB'.")
+
+
+def dtype_byte_size(dtype: torch.dtype):
+    """
+    Returns the size (in bytes) occupied by one parameter of type `dtype`.
+
+    Example:
+
+    ```py
+    >>> dtype_byte_size(torch.float32)
+    4
+    ```
+    """
+    if dtype == torch.bool:
+        return 1 / 8
+    bit_search = re.search("[^\d](\d+)$", str(dtype))
+    if bit_search is None:
+        raise ValueError(f"`dtype` is not a valid dtype: {dtype}.")
+    bit_size = int(bit_search.groups()[0])
+    return bit_size // 8
+
+
+def set_module_tensor_to_device(
+    module: nn.Module, tensor_name: str, device: Union[int, str, torch.device], value: Optional[torch.Tensor] = None
+):
+    """
+    A helper function to set a given tensor (parameter of buffer) of a module on a specific device (note that doing
+    `param.to(device)` creates a new tensor not linked to the parameter, which is why we need this function).
+
+    Args:
+        module (`torch.nn.Module`): The module in which the tensor we want to move lives.
+        param_name (`str`): The full name of the parameter/buffer.
+        device (`int`, `str` or `torch.device`): The device on which to set the tensor.
+        value (`torch.Tensor`, *optional*): The value of the tensor (useful when going from the meta device to any
+            other device).
+    """
+    # Recurse if needed
+    if "." in tensor_name:
+        splits = tensor_name.split(".")
+        for split in splits[:-1]:
+            new_module = getattr(module, split)
+            if new_module is None:
+                raise ValueError(f"{module} has no attribute {split}.")
+            module = new_module
+        tensor_name = splits[-1]
+
+    if tensor_name not in module._parameters and tensor_name not in module._buffers:
+        raise ValueError(f"{module} does not have a parameter or a buffer named {tensor_name}.")
+    is_buffer = tensor_name in module._buffers
+    old_value = getattr(module, tensor_name)
+
+    if old_value.device == torch.device("meta") and device != torch.device("meta") and value is None:
+        raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {device}.")
+
+    with torch.no_grad():
+        if value is None:
+            new_value = old_value.to(device)
+        elif isinstance(value, torch.Tensor):
+            new_value = value.to(device)
+        else:
+            new_value = torch.tensor(value, device=device)
+    if is_buffer:
+        module._buffers[tensor_name] = new_value
+    else:
+        new_value = nn.Parameter(new_value, requires_grad=old_value.requires_grad)
+        module._parameters[tensor_name] = new_value
+
+
+def named_module_tensors(module: nn.Module, include_buffers: bool = True, recurse: bool = False):
+    """
+    A helper function that gathers all the tensors (parameters + buffers) of a given module. If `include_buffers=True`
+    it's the same as doing `module.named_parameters(recurse=recurse) + module.named_buffers(recurse=recurse)`.
+
+    Args:
+        module (`torch.nn.Module`): The module we want the tensors or.
+        include_buffer (`bool`, *optional*, defaults to `True`): Whether or not to include the buffers in the result.
+        recurse (`bool`, *optional`, defaults to `False`):
+            Whether or not to go look in every submodule or just return the direct parameters and buffers.
+    """
+    for named_parameter in module.named_parameters(recurse=recurse):
+        yield named_parameter
+
+    if include_buffers:
+        for named_buffer in module.named_buffers(recurse=recurse):
+            yield named_buffer
+
+
+def find_tied_parameters(model: nn.Module, **kwargs):
+    """
+    Find the tied parameters in a given model.
+
+    Args:
+        model (`torch.nn.Module`): The model to inspect.
+
+    <Tip warning={true}>
+
+    The signature accepts keyword arguments, but they are for the recursive part of this function and you should ignore
+    them.
+
+    </Tip>
+
+    Example:
+
+
+    ```py
+    >>> from collections import OrderedDict
+    >>> import torch.nn as nn
+
+    >>> model = nn.Sequential(OrderedDict([("linear1", nn.Linear(4, 4)), ("linear2", nn.Linear(4, 4))]))
+    >>> model.linear2.weight = test_model.linear1.weight
+    >>> find_tied_parameters(test_model)
+    {'linear1.weight': 'linear2.weight'}
+    ```
+
+    Returns:
+        Dict[str, str]: A dictionary mapping tied parameter names to the name of the parameter they are tied to.
+    """
+    # Initialize result and named_parameters before recursing.
+    named_parameters = kwargs.get("named_parameters", None)
+    prefix = kwargs.get("prefix", "")
+    result = kwargs.get("result", {})
+
+    if named_parameters is None:
+        named_parameters = {n: p for n, p in model.named_parameters()}
+    else:
+        # A tied parameter will not be in the full `named_parameters` seen above but will be in the `named_parameters`
+        # of the submodule it belongs to. So while recursing we track the names that are not in the initial
+        # `named_parameters`.
+        for name, parameter in model.named_parameters():
+            full_name = name if prefix == "" else f"{prefix}.{name}"
+            if full_name not in named_parameters:
+                # When we find one, it has to be one of the existing parameters.
+                for new_name, new_param in named_parameters.items():
+                    if new_param is parameter:
+                        result[new_name] = full_name
+
+    # Once we have treated direct parameters, we move to the child modules.
+    for name, child in model.named_children():
+        child_name = name if prefix == "" else f"{prefix}.{name}"
+        find_tied_parameters(child, named_parameters=named_parameters, prefix=child_name, result=result)
+
+    return result
+
+
+def compute_module_sizes(model: nn.Module):
+    """
+    Compute the size of each submodule of a given model.
+    """
+    module_sizes = defaultdict(int)
+    for name, tensor in named_module_tensors(model, recurse=True):
+        size = tensor.numel() * dtype_byte_size(tensor.dtype)
+        name_parts = name.split(".")
+        for idx in range(len(name_parts) + 1):
+            module_sizes[".".join(name_parts[:idx])] += size
+
+    return module_sizes
+
+
+def get_max_layer_size(
+    modules: List[Tuple[str, torch.nn.Module]], module_sizes: Dict[str, int], no_split_module_classes: List[str]
+):
+    """
+    Utility function that will scan a list of named modules and return the maximum size used by one full layer. The
+    definition of a layer being:
+    - a module with no direct children (just parameters and buffers)
+    - a module whose class name is in the list `no_split_module_classes`
+
+    Args:
+        modules (`List[Tuple[str, torch.nn.Module]]`):
+            The list of named modules where we want to determine the maximum layer size.
+        module_sizes (`Dict[str, int]`):
+            A dictionary mapping each layer name to its size (as generated by `compute_module_sizes`).
+        no_split_module_classes (`List[str]`):
+            A list of class names for layers we don't want to be split.
+
+    Returns:
+        `Tuple[int, List[str]]`: The maximum size of a layer with the list of layer names realizing that maximum size.
+    """
+    max_size = 0
+    layer_names = []
+    modules_to_treat = modules.copy()
+    while len(modules_to_treat) > 0:
+        module_name, module = modules_to_treat.pop(0)
+        modules_children = list(module.named_children())
+        if len(modules_children) == 0 or module.__class__.__name__ in no_split_module_classes:
+            # No splitting this one so we compare to the max_size
+            size = module_sizes[module_name]
+            if size > max_size:
+                max_size = size
+                layer_names = [module_name]
+            elif size == max_size:
+                layer_names.append(module_name)
+        else:
+            modules_to_treat = [(f"{module_name}.{n}", v) for n, v in modules_children] + modules_to_treat
+    return max_size, layer_names
+
+
+def get_max_memory(max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None):
+    """
+    Get the maximum memory available if nothing is passed, converts string to int otherwise.
+    """
+    import psutil
+
+    if max_memory is None:
+        if not torch.cuda.is_available():
+            max_memory = {}
+        else:
+            # Make sure CUDA is initialized on each GPU to have the right memory info.
+            for i in range(torch.cuda.device_count()):
+                _ = torch.tensor([0], device=i)
+            max_memory = {i: torch.cuda.mem_get_info(i)[0] for i in range(torch.cuda.device_count())}
+        max_memory["cpu"] = psutil.virtual_memory().available
+        return max_memory
+
+    for key in max_memory:
+        if isinstance(max_memory[key], str):
+            max_memory[key] = convert_file_size_to_int(max_memory[key])
+    return max_memory
+
+
+def clean_device_map(device_map: Dict[str, Union[int, str, torch.device]], module_name: str = ""):
+    """
+    Cleans a device_map by grouping all submodules that go on the same device together.
+    """
+    # Get the value of the current module and if there is only one split across several keys, regroup it.
+    values = [v for k, v in device_map.items() if k.startswith(module_name)]
+    if len(set(values)) == 1 and len(values) > 1:
+        for k in [k for k in device_map if k.startswith(module_name)]:
+            del device_map[k]
+        device_map[module_name] = values[0]
+
+    # Recurse over the children
+    children_modules = [k for k in device_map.keys() if k.startswith(module_name) and len(k) > len(module_name)]
+    idx = len(module_name.split(".")) + 1 if len(module_name) > 0 else 1
+    children_modules = set(".".join(k.split(".")[:idx]) for k in children_modules)
+    for child in children_modules:
+        clean_device_map(device_map, module_name=child)
+
+    return device_map
+
+
+def infer_auto_device_map(
+    model: nn.Module,
+    max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None,
+    no_split_module_classes: Optional[List[str]] = None,
+):
+    """
+    Compute a device map for a given model giving priority to GPUs, then offload on CPU and finally offload to disk,
+    such that:
+    - we don't exceed the memory available of any of the GPU.
+    - if offload to the CPU is needed, there is always room left on GPU 0 to put back the layer offloaded on CPU that
+      has the largest size.
+    - if offload to the CPU is needed,we don't exceed the RAM available on the CPU.
+    - if offload to the disk is needed, there is always room left on the CPU to put back the layer offloaded on disk
+      that has the largest size.
+
+    <Tip>
+
+    All computation is done analyzing sizes and dtypes of the model parameters. As a result, the model can be on the
+    meta device (as it would if initialized within the `init_empty_weights` context manager).
+
+    </Tip>
+
+    Args:
+        model (`torch.nn.Module`): The model to analyze.
+        max_memory (`Dict`, *optional*):
+            A dictionary device identifier to maximum memory. Will default to the maximum memory available if unset.
+        no_split_module_classes (`List[str]`, *optional*):
+            A list of layer class names that should never be split across device (for instance any layer that has a
+            residual connection).
+    """
+    # Get default / clean up max_memory
+    max_memory = get_max_memory(max_memory)
+    if no_split_module_classes is None:
+        no_split_module_classes = []
+    elif not isinstance(no_split_module_classes, (list, tuple)):
+        no_split_module_classes = [no_split_module_classes]
+
+    devices = list(max_memory.keys()) + ["disk"]
+    # Devices that need to keep space for a potential offloaded layer.
+    main_devices = [devices[0], "cpu"]
If I understand this correctly, this will return an empty device map if everything fits on a single device. If it receives the following device map, for example:
```py
clean_device_map({'transformer': 0, 'lm_head': 0})
```
it will return the following:
```
{'': 0}
```
which, unfortunately, cannot be digested by the code later on as keys cannot find a common radix
@@ -0,0 +1,568 @@
+# 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 gc
+import json
+import os
+import re
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+WEIGHTS_INDEX_NAME = "pytorch_model.bin.index.json"
+
+
+def convert_file_size_to_int(size: Union[int, str]):
+    """
+    Converts a size expressed as a string with digits an unit (like `"5MB"`) to an integer (in bytes).
+
+    Args:
+        size (`int` or `str`): The size to convert. Will be directly returned if an `int`.
+
+    Example:
+
+    ```py
+    >>> convert_file_size_to_int("1MiB")
+    1048576
+    ```
+    """
+    if isinstance(size, int):
+        return size
+    if size.upper().endswith("GIB"):
+        return int(size[:-3]) * (2**30)
+    if size.upper().endswith("MIB"):
+        return int(size[:-3]) * (2**20)
+    if size.upper().endswith("KIB"):
+        return int(size[:-3]) * (2**10)
+    if size.upper().endswith("GB"):
+        int_size = int(size[:-2]) * (10**9)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("MB"):
+        int_size = int(size[:-2]) * (10**6)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("KB"):
+        int_size = int(size[:-2]) * (10**3)
+        return int_size // 8 if size.endswith("b") else int_size
+    raise ValueError("`size` is not in a valid format. Use an integer followed by the unit, e.g., '5GB'.")
+
+
+def dtype_byte_size(dtype: torch.dtype):
+    """
+    Returns the size (in bytes) occupied by one parameter of type `dtype`.
+
+    Example:
+
+    ```py
+    >>> dtype_byte_size(torch.float32)
+    4
+    ```
+    """
+    if dtype == torch.bool:
+        return 1 / 8
+    bit_search = re.search("[^\d](\d+)$", str(dtype))
+    if bit_search is None:
+        raise ValueError(f"`dtype` is not a valid dtype: {dtype}.")
+    bit_size = int(bit_search.groups()[0])
+    return bit_size // 8
+
+
+def set_module_tensor_to_device(
+    module: nn.Module, tensor_name: str, device: Union[int, str, torch.device], value: Optional[torch.Tensor] = None
+):
+    """
+    A helper function to set a given tensor (parameter of buffer) of a module on a specific device (note that doing
+    `param.to(device)` creates a new tensor not linked to the parameter, which is why we need this function).
+
+    Args:
+        module (`torch.nn.Module`): The module in which the tensor we want to move lives.
+        param_name (`str`): The full name of the parameter/buffer.
+        device (`int`, `str` or `torch.device`): The device on which to set the tensor.
+        value (`torch.Tensor`, *optional*): The value of the tensor (useful when going from the meta device to any
+            other device).
+    """
+    # Recurse if needed
+    if "." in tensor_name:
+        splits = tensor_name.split(".")
+        for split in splits[:-1]:
+            new_module = getattr(module, split)
+            if new_module is None:
+                raise ValueError(f"{module} has no attribute {split}.")
+            module = new_module
+        tensor_name = splits[-1]
+
+    if tensor_name not in module._parameters and tensor_name not in module._buffers:
+        raise ValueError(f"{module} does not have a parameter or a buffer named {tensor_name}.")
+    is_buffer = tensor_name in module._buffers
+    old_value = getattr(module, tensor_name)
+
+    if old_value.device == torch.device("meta") and device != torch.device("meta") and value is None:
+        raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {device}.")
+
+    with torch.no_grad():
+        if value is None:
+            new_value = old_value.to(device)
+        elif isinstance(value, torch.Tensor):
+            new_value = value.to(device)
+        else:
+            new_value = torch.tensor(value, device=device)
+    if is_buffer:
+        module._buffers[tensor_name] = new_value
+    else:
+        new_value = nn.Parameter(new_value, requires_grad=old_value.requires_grad)
+        module._parameters[tensor_name] = new_value
+
+
+def named_module_tensors(module: nn.Module, include_buffers: bool = True, recurse: bool = False):
+    """
+    A helper function that gathers all the tensors (parameters + buffers) of a given module. If `include_buffers=True`
+    it's the same as doing `module.named_parameters(recurse=recurse) + module.named_buffers(recurse=recurse)`.
+
+    Args:
+        module (`torch.nn.Module`): The module we want the tensors or.
+        include_buffer (`bool`, *optional*, defaults to `True`): Whether or not to include the buffers in the result.
+        recurse (`bool`, *optional`, defaults to `False`):
+            Whether or not to go look in every submodule or just return the direct parameters and buffers.
+    """
+    for named_parameter in module.named_parameters(recurse=recurse):
+        yield named_parameter
+
+    if include_buffers:
+        for named_buffer in module.named_buffers(recurse=recurse):
+            yield named_buffer
+
+
+def find_tied_parameters(model: nn.Module, **kwargs):
+    """
+    Find the tied parameters in a given model.
+
+    Args:
+        model (`torch.nn.Module`): The model to inspect.
+
+    <Tip warning={true}>
+
+    The signature accepts keyword arguments, but they are for the recursive part of this function and you should ignore
+    them.
+
+    </Tip>
+
+    Example:
+
+
+    ```py
+    >>> from collections import OrderedDict
+    >>> import torch.nn as nn
+
+    >>> model = nn.Sequential(OrderedDict([("linear1", nn.Linear(4, 4)), ("linear2", nn.Linear(4, 4))]))
+    >>> model.linear2.weight = test_model.linear1.weight
+    >>> find_tied_parameters(test_model)
+    {'linear1.weight': 'linear2.weight'}
+    ```
+
+    Returns:
+        Dict[str, str]: A dictionary mapping tied parameter names to the name of the parameter they are tied to.
+    """
+    # Initialize result and named_parameters before recursing.
+    named_parameters = kwargs.get("named_parameters", None)
+    prefix = kwargs.get("prefix", "")
+    result = kwargs.get("result", {})
+
+    if named_parameters is None:
+        named_parameters = {n: p for n, p in model.named_parameters()}
+    else:
+        # A tied parameter will not be in the full `named_parameters` seen above but will be in the `named_parameters`
+        # of the submodule it belongs to. So while recursing we track the names that are not in the initial
+        # `named_parameters`.
+        for name, parameter in model.named_parameters():
+            full_name = name if prefix == "" else f"{prefix}.{name}"
+            if full_name not in named_parameters:
+                # When we find one, it has to be one of the existing parameters.
+                for new_name, new_param in named_parameters.items():
+                    if new_param is parameter:
+                        result[new_name] = full_name
+
+    # Once we have treated direct parameters, we move to the child modules.
+    for name, child in model.named_children():
+        child_name = name if prefix == "" else f"{prefix}.{name}"
+        find_tied_parameters(child, named_parameters=named_parameters, prefix=child_name, result=result)
+
+    return result
+
+
+def compute_module_sizes(model: nn.Module):
+    """
+    Compute the size of each submodule of a given model.
+    """
+    module_sizes = defaultdict(int)
+    for name, tensor in named_module_tensors(model, recurse=True):
+        size = tensor.numel() * dtype_byte_size(tensor.dtype)
+        name_parts = name.split(".")
+        for idx in range(len(name_parts) + 1):
+            module_sizes[".".join(name_parts[:idx])] += size
+
+    return module_sizes
+
+
+def get_max_layer_size(
+    modules: List[Tuple[str, torch.nn.Module]], module_sizes: Dict[str, int], no_split_module_classes: List[str]
+):
+    """
+    Utility function that will scan a list of named modules and return the maximum size used by one full layer. The
+    definition of a layer being:
+    - a module with no direct children (just parameters and buffers)
+    - a module whose class name is in the list `no_split_module_classes`
+
+    Args:
+        modules (`List[Tuple[str, torch.nn.Module]]`):
+            The list of named modules where we want to determine the maximum layer size.
+        module_sizes (`Dict[str, int]`):
+            A dictionary mapping each layer name to its size (as generated by `compute_module_sizes`).
+        no_split_module_classes (`List[str]`):
+            A list of class names for layers we don't want to be split.
+
+    Returns:
+        `Tuple[int, List[str]]`: The maximum size of a layer with the list of layer names realizing that maximum size.
+    """
+    max_size = 0
+    layer_names = []
+    modules_to_treat = modules.copy()
+    while len(modules_to_treat) > 0:
+        module_name, module = modules_to_treat.pop(0)
+        modules_children = list(module.named_children())
+        if len(modules_children) == 0 or module.__class__.__name__ in no_split_module_classes:
+            # No splitting this one so we compare to the max_size
+            size = module_sizes[module_name]
+            if size > max_size:
+                max_size = size
+                layer_names = [module_name]
+            elif size == max_size:
+                layer_names.append(module_name)
+        else:
+            modules_to_treat = [(f"{module_name}.{n}", v) for n, v in modules_children] + modules_to_treat
+    return max_size, layer_names
+
+
+def get_max_memory(max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None):
+    """
+    Get the maximum memory available if nothing is passed, converts string to int otherwise.
+    """
+    import psutil
+
+    if max_memory is None:
+        if not torch.cuda.is_available():
+            max_memory = {}
+        else:
+            # Make sure CUDA is initialized on each GPU to have the right memory info.
+            for i in range(torch.cuda.device_count()):
+                _ = torch.tensor([0], device=i)
+            max_memory = {i: torch.cuda.mem_get_info(i)[0] for i in range(torch.cuda.device_count())}
+        max_memory["cpu"] = psutil.virtual_memory().available
+        return max_memory
+
+    for key in max_memory:
+        if isinstance(max_memory[key], str):
+            max_memory[key] = convert_file_size_to_int(max_memory[key])
+    return max_memory
+
+
+def clean_device_map(device_map: Dict[str, Union[int, str, torch.device]], module_name: str = ""):
+    """
+    Cleans a device_map by grouping all submodules that go on the same device together.
+    """
+    # Get the value of the current module and if there is only one split across several keys, regroup it.
+    values = [v for k, v in device_map.items() if k.startswith(module_name)]
+    if len(set(values)) == 1 and len(values) > 1:
+        for k in [k for k in device_map if k.startswith(module_name)]:
+            del device_map[k]
+        device_map[module_name] = values[0]
Nitpick: this fails right now if we specify a checkpoint like we would do in the transformers library. For example, this fails:

```py
with init_empty_weights():
    model = GPT2LMHeadModel.from_pretrained('./gpt2')

model = load_checkpoint_and_dispatch(
    model,
    "./gpt2",
    device_map='auto'
)
```

whereas intuitively I would expect it to work. I understand if that's wanted behavior and that passing the `pytorch_model.bin` is what you want instead.
@@ -0,0 +1,568 @@
+# 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 gc
+import json
+import os
+import re
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+WEIGHTS_INDEX_NAME = "pytorch_model.bin.index.json"
+
+
+def convert_file_size_to_int(size: Union[int, str]):
+    """
+    Converts a size expressed as a string with digits an unit (like `"5MB"`) to an integer (in bytes).
+
+    Args:
+        size (`int` or `str`): The size to convert. Will be directly returned if an `int`.
+
+    Example:
+
+    ```py
+    >>> convert_file_size_to_int("1MiB")
+    1048576
+    ```
+    """
+    if isinstance(size, int):
+        return size
+    if size.upper().endswith("GIB"):
+        return int(size[:-3]) * (2**30)
+    if size.upper().endswith("MIB"):
+        return int(size[:-3]) * (2**20)
+    if size.upper().endswith("KIB"):
+        return int(size[:-3]) * (2**10)
+    if size.upper().endswith("GB"):
+        int_size = int(size[:-2]) * (10**9)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("MB"):
+        int_size = int(size[:-2]) * (10**6)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("KB"):
+        int_size = int(size[:-2]) * (10**3)
+        return int_size // 8 if size.endswith("b") else int_size
+    raise ValueError("`size` is not in a valid format. Use an integer followed by the unit, e.g., '5GB'.")
+
+
+def dtype_byte_size(dtype: torch.dtype):
+    """
+    Returns the size (in bytes) occupied by one parameter of type `dtype`.
+
+    Example:
+
+    ```py
+    >>> dtype_byte_size(torch.float32)
+    4
+    ```
+    """
+    if dtype == torch.bool:
+        return 1 / 8
+    bit_search = re.search("[^\d](\d+)$", str(dtype))
+    if bit_search is None:
+        raise ValueError(f"`dtype` is not a valid dtype: {dtype}.")
+    bit_size = int(bit_search.groups()[0])
+    return bit_size // 8
+
+
+def set_module_tensor_to_device(
+    module: nn.Module, tensor_name: str, device: Union[int, str, torch.device], value: Optional[torch.Tensor] = None
+):
+    """
+    A helper function to set a given tensor (parameter of buffer) of a module on a specific device (note that doing
+    `param.to(device)` creates a new tensor not linked to the parameter, which is why we need this function).
+
+    Args:
+        module (`torch.nn.Module`): The module in which the tensor we want to move lives.
+        param_name (`str`): The full name of the parameter/buffer.
+        device (`int`, `str` or `torch.device`): The device on which to set the tensor.
+        value (`torch.Tensor`, *optional*): The value of the tensor (useful when going from the meta device to any
+            other device).
+    """
+    # Recurse if needed
+    if "." in tensor_name:
+        splits = tensor_name.split(".")
+        for split in splits[:-1]:
+            new_module = getattr(module, split)
+            if new_module is None:
+                raise ValueError(f"{module} has no attribute {split}.")
+            module = new_module
+        tensor_name = splits[-1]
+
+    if tensor_name not in module._parameters and tensor_name not in module._buffers:
+        raise ValueError(f"{module} does not have a parameter or a buffer named {tensor_name}.")
+    is_buffer = tensor_name in module._buffers
+    old_value = getattr(module, tensor_name)
+
+    if old_value.device == torch.device("meta") and device != torch.device("meta") and value is None:
+        raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {device}.")
+
+    with torch.no_grad():
+        if value is None:
+            new_value = old_value.to(device)
+        elif isinstance(value, torch.Tensor):
+            new_value = value.to(device)
+        else:
+            new_value = torch.tensor(value, device=device)
+    if is_buffer:
+        module._buffers[tensor_name] = new_value
+    else:
+        new_value = nn.Parameter(new_value, requires_grad=old_value.requires_grad)
+        module._parameters[tensor_name] = new_value
+
+
+def named_module_tensors(module: nn.Module, include_buffers: bool = True, recurse: bool = False):
+    """
+    A helper function that gathers all the tensors (parameters + buffers) of a given module. If `include_buffers=True`
+    it's the same as doing `module.named_parameters(recurse=recurse) + module.named_buffers(recurse=recurse)`.
+
+    Args:
+        module (`torch.nn.Module`): The module we want the tensors or.
+        include_buffer (`bool`, *optional*, defaults to `True`): Whether or not to include the buffers in the result.
+        recurse (`bool`, *optional`, defaults to `False`):
+            Whether or not to go look in every submodule or just return the direct parameters and buffers.
+    """
+    for named_parameter in module.named_parameters(recurse=recurse):
+        yield named_parameter
+
+    if include_buffers:
+        for named_buffer in module.named_buffers(recurse=recurse):
+            yield named_buffer
+
+
+def find_tied_parameters(model: nn.Module, **kwargs):
+    """
+    Find the tied parameters in a given model.
+
+    Args:
+        model (`torch.nn.Module`): The model to inspect.
+
+    <Tip warning={true}>
+
+    The signature accepts keyword arguments, but they are for the recursive part of this function and you should ignore
+    them.
+
+    </Tip>
+
+    Example:
+
+
+    ```py
+    >>> from collections import OrderedDict
+    >>> import torch.nn as nn
+
+    >>> model = nn.Sequential(OrderedDict([("linear1", nn.Linear(4, 4)), ("linear2", nn.Linear(4, 4))]))
+    >>> model.linear2.weight = test_model.linear1.weight
+    >>> find_tied_parameters(test_model)
+    {'linear1.weight': 'linear2.weight'}
+    ```
+
+    Returns:
+        Dict[str, str]: A dictionary mapping tied parameter names to the name of the parameter they are tied to.
+    """
+    # Initialize result and named_parameters before recursing.
+    named_parameters = kwargs.get("named_parameters", None)
+    prefix = kwargs.get("prefix", "")
+    result = kwargs.get("result", {})
+
+    if named_parameters is None:
+        named_parameters = {n: p for n, p in model.named_parameters()}
+    else:
+        # A tied parameter will not be in the full `named_parameters` seen above but will be in the `named_parameters`
+        # of the submodule it belongs to. So while recursing we track the names that are not in the initial
+        # `named_parameters`.
+        for name, parameter in model.named_parameters():
+            full_name = name if prefix == "" else f"{prefix}.{name}"
+            if full_name not in named_parameters:
+                # When we find one, it has to be one of the existing parameters.
+                for new_name, new_param in named_parameters.items():
+                    if new_param is parameter:
+                        result[new_name] = full_name
+
+    # Once we have treated direct parameters, we move to the child modules.
+    for name, child in model.named_children():
+        child_name = name if prefix == "" else f"{prefix}.{name}"
+        find_tied_parameters(child, named_parameters=named_parameters, prefix=child_name, result=result)
+
+    return result
+
+
+def compute_module_sizes(model: nn.Module):
+    """
+    Compute the size of each submodule of a given model.
+    """
+    module_sizes = defaultdict(int)
+    for name, tensor in named_module_tensors(model, recurse=True):
+        size = tensor.numel() * dtype_byte_size(tensor.dtype)
+        name_parts = name.split(".")
+        for idx in range(len(name_parts) + 1):
+            module_sizes[".".join(name_parts[:idx])] += size
+
+    return module_sizes
+
+
+def get_max_layer_size(
+    modules: List[Tuple[str, torch.nn.Module]], module_sizes: Dict[str, int], no_split_module_classes: List[str]
+):
+    """
+    Utility function that will scan a list of named modules and return the maximum size used by one full layer. The
+    definition of a layer being:
+    - a module with no direct children (just parameters and buffers)
+    - a module whose class name is in the list `no_split_module_classes`
+
+    Args:
+        modules (`List[Tuple[str, torch.nn.Module]]`):
+            The list of named modules where we want to determine the maximum layer size.
+        module_sizes (`Dict[str, int]`):
+            A dictionary mapping each layer name to its size (as generated by `compute_module_sizes`).
+        no_split_module_classes (`List[str]`):
+            A list of class names for layers we don't want to be split.
+
+    Returns:
+        `Tuple[int, List[str]]`: The maximum size of a layer with the list of layer names realizing that maximum size.
+    """
+    max_size = 0
+    layer_names = []
+    modules_to_treat = modules.copy()
+    while len(modules_to_treat) > 0:
+        module_name, module = modules_to_treat.pop(0)
+        modules_children = list(module.named_children())
+        if len(modules_children) == 0 or module.__class__.__name__ in no_split_module_classes:
+            # No splitting this one so we compare to the max_size
+            size = module_sizes[module_name]
+            if size > max_size:
+                max_size = size
+                layer_names = [module_name]
+            elif size == max_size:
+                layer_names.append(module_name)
+        else:
+            modules_to_treat = [(f"{module_name}.{n}", v) for n, v in modules_children] + modules_to_treat
+    return max_size, layer_names
+
+
+def get_max_memory(max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None):
+    """
+    Get the maximum memory available if nothing is passed, converts string to int otherwise.
+    """
+    import psutil
+
+    if max_memory is None:
+        if not torch.cuda.is_available():
+            max_memory = {}
+        else:
+            # Make sure CUDA is initialized on each GPU to have the right memory info.
+            for i in range(torch.cuda.device_count()):
+                _ = torch.tensor([0], device=i)
+            max_memory = {i: torch.cuda.mem_get_info(i)[0] for i in range(torch.cuda.device_count())}
+        max_memory["cpu"] = psutil.virtual_memory().available
+        return max_memory
+
+    for key in max_memory:
+        if isinstance(max_memory[key], str):
+            max_memory[key] = convert_file_size_to_int(max_memory[key])
+    return max_memory
+
+
+def clean_device_map(device_map: Dict[str, Union[int, str, torch.device]], module_name: str = ""):
+    """
+    Cleans a device_map by grouping all submodules that go on the same device together.
+    """
+    # Get the value of the current module and if there is only one split across several keys, regroup it.
+    values = [v for k, v in device_map.items() if k.startswith(module_name)]
+    if len(set(values)) == 1 and len(values) > 1:
+        for k in [k for k in device_map if k.startswith(module_name)]:
+            del device_map[k]
+        device_map[module_name] = values[0]
+
+    # Recurse over the children
+    children_modules = [k for k in device_map.keys() if k.startswith(module_name) and len(k) > len(module_name)]
+    idx = len(module_name.split(".")) + 1 if len(module_name) > 0 else 1
+    children_modules = set(".".join(k.split(".")[:idx]) for k in children_modules)
+    for child in children_modules:
+        clean_device_map(device_map, module_name=child)
+
+    return device_map
+
+
+def infer_auto_device_map(
+    model: nn.Module,
+    max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None,
+    no_split_module_classes: Optional[List[str]] = None,
+):
+    """
+    Compute a device map for a given model giving priority to GPUs, then offload on CPU and finally offload to disk,
+    such that:
+    - we don't exceed the memory available of any of the GPU.
+    - if offload to the CPU is needed, there is always room left on GPU 0 to put back the layer offloaded on CPU that
+      has the largest size.
+    - if offload to the CPU is needed,we don't exceed the RAM available on the CPU.
+    - if offload to the disk is needed, there is always room left on the CPU to put back the layer offloaded on disk
+      that has the largest size.
+
+    <Tip>
+
+    All computation is done analyzing sizes and dtypes of the model parameters. As a result, the model can be on the
+    meta device (as it would if initialized within the `init_empty_weights` context manager).
+
+    </Tip>
+
+    Args:
+        model (`torch.nn.Module`): The model to analyze.
+        max_memory (`Dict`, *optional*):
+            A dictionary device identifier to maximum memory. Will default to the maximum memory available if unset.
+        no_split_module_classes (`List[str]`, *optional*):
+            A list of layer class names that should never be split across device (for instance any layer that has a
+            residual connection).
+    """
+    # Get default / clean up max_memory
+    max_memory = get_max_memory(max_memory)
+    if no_split_module_classes is None:
+        no_split_module_classes = []
+    elif not isinstance(no_split_module_classes, (list, tuple)):
+        no_split_module_classes = [no_split_module_classes]
+
+    devices = list(max_memory.keys()) + ["disk"]
+    # Devices that need to keep space for a potential offloaded layer.
+    main_devices = [devices[0], "cpu"]
+
+    module_sizes = compute_module_sizes(model)
+    tied_parameters = find_tied_parameters(model)
+
+    device_map = {}
+    current_device = 0
+    current_memory_used = 0
+
+    # Direct submodules and parameters
+    modules_to_treat = list(model.named_parameters(recurse=False)) + list(model.named_children())
+    # Initialize maximum largest layer, to know which space to keep in memory
+    max_layer_size, max_layer_names = get_max_layer_size(modules_to_treat, module_sizes, no_split_module_classes)
+
+    # Ready ? This is going to be a bit messy.
+    while len(modules_to_treat) > 0:
+        name, module = modules_to_treat.pop(0)
+        # Max size in the remaining layers may have changed since we took one, so we maybe update it.
+        max_layer_names = [n for n in max_layer_names if not n.startswith(name)]
+        if len(max_layer_names) == 0:
+            max_layer_size, max_layer_names = get_max_layer_size(
+                [(n, m) for n, m in modules_to_treat if isinstance(m, torch.nn.Module)],
+                module_sizes,
+                no_split_module_classes,
+            )
+        # Assess size needed
+        module_size = module_sizes[name]
+        tied_params = [v for k, v in tied_parameters.items() if name in k]
+        # We ignore parameters that are tied when they're tied to > 1 one
+        tied_param = tied_params[0] if len(tied_params) == 1 else None
+
+        device = devices[current_device]
+        current_max_size = max_memory[device] if device != "disk" else None
+        # Reduce max size available by the largest layer.
+        if devices[current_device] in main_devices:
+            current_max_size = current_max_size - max_layer_size
+        # Case 1 -> We're too big!
+        if current_max_size is not None and current_memory_used + module_size > current_max_size:
+            # Split or not split?
+            modules_children = list(module.named_children())
+            if len(modules_children) == 0 or module.__class__.__name__ in no_split_module_classes:
+                # -> no split, we go to the next device
+                current_device += 1
+                modules_to_treat = [(name, module)] + modules_to_treat
+                current_memory_used = 0
+            else:
+                # -> split, we replace the module studied by its children + parameters
+                modules_children = list(module.named_parameters(recurse=False)) + modules_children
+                modules_to_treat = [(f"{name}.{n}", v) for n, v in modules_children] + modules_to_treat
+                # Update the max layer size.
+                max_layer_size, max_layer_names = get_max_layer_size(
+                    [(n, m) for n, m in modules_to_treat if isinstance(m, torch.nn.Module)],
+                    module_sizes,
+                    no_split_module_classes,
+                )
+
+        # Case 2, it fits! We're not entirely out of the wood though, because we may have some tied parameters.
+        elif tied_param is not None:
+            # Determine the sized occupied by this module + the module containing the tied parameter
+            tied_module_size = module_size
+            tied_module_index = [i for i, (n, _) in enumerate(modules_to_treat) if n in tied_param][0]
+            tied_module_name, tied_module = modules_to_treat[tied_module_index]
+            tied_module_size += module_sizes[tied_module_name] - module_sizes[tied_param]
+            if current_max_size is not None and current_memory_used + tied_module_size > current_max_size:
+                # Split or not split?
+                tied_module_children = list(tied_module.named_children())
+                if len(tied_module_children) == 0 or tied_module.__class__.__name__ in no_split_module_classes:
+                    # If the tied module is not split, we go to the next device
+                    current_device += 1
+                    modules_to_treat = [(name, module)] + modules_to_treat
+                    current_memory_used = 0
+                else:
+                    # Otherwise, we replace the tied module by its children.
+                    tied_module_children = list(tied_module.named_parameters(recurse=False)) + tied_module_children
+                    tied_module_children = [(f"{tied_module_name}.{n}", v) for n, v in tied_module_children]
+                    modules_to_treat = (
+                        [(name, module)]
+                        + modules_to_treat[:tied_module_index]
+                        + tied_module_children
+                        + modules_to_treat[tied_module_index + 1 :]
+                    )
+                    # Update the max layer size.
+                    max_layer_size, max_layer_names = get_max_layer_size(
+                        [(n, m) for n, m in modules_to_treat if isinstance(m, torch.nn.Module)],
+                        module_sizes,
+                        no_split_module_classes,
+                    )
+            else:
+                # We really really fit!
+                current_memory_used += tied_module_size
+                device_map[name] = devices[current_device]
+                modules_to_treat.pop(tied_module_index)
+                device_map[tied_module_name] = devices[current_device]
+        else:
+            current_memory_used += module_size
+            device_map[name] = devices[current_device]
+
+    return clean_device_map(device_map)
+
+
+def check_device_map(model: nn.Module, device_map: Dict[str, Union[int, str, torch.device]]):
+    """
+    Checks a device map covers everything in a given model.
+
+    Args:
+        model (`torch.nn.Module`): The model to check the device map against.
+        device_map (`Dict[str, Union[int, str, torch.device]]`): The device map to check.
+    """
+    all_model_tensors = [name for name, _ in model.state_dict().items()]
+    for module_name in device_map.keys():
+        all_model_tensors = [name for name in all_model_tensors if not name.startswith(module_name)]
+    if len(all_model_tensors) > 0:
+        non_covered_params = ", ".join(all_model_tensors)
+        raise ValueError(
+            f"The device_map provided does not give any device for the following parameters: {non_covered_params}"
+        )
+
+
+def load_checkpoint_in_model(
+    model: nn.Module,
+    checkpoint: Union[str, os.PathLike],
+    device_map: Optional[Dict[str, Union[int, str, torch.device]]] = None,
+    offload_folder: Optional[Union[str, os.PathLike]] = None,
+):
+    """
+    Loads a (potentially sharded) checkpoint inside a model, potentially sending weights to a given device as they are
+    loaded.
+
+    <Tip warning={true}>
+
+    Once loaded across devices, you still need to call [`dispatch_model`] on your model to make it able to run. To
+    group the checkpoint loading and dispatch in one single call, use [`load_checkpoint_and_dispatch`].
+
+    </Tip>
+
+    Args:
+        model (`torch.nn.Module`): The model in which we want to load a checkpoint.
+        checkpoint (`str` or `os.PathLike`):
+            The folder checkpoint to load. It can be:
+            - a path to a file containing a whole model state dict
+            - a path to a `.json` file containing the index to a sharded checkpoint
+            - a path to a folder containing a unique `.index.json` file and the shards of a checkpoint.
+        device_map (`Dict[str, Union[int, str, torch.device]]`, *optional*):
+            A map that specifies where each submodule should go. It doesn't need to be refined to each parameter/buffer
+            name, once a given module name is inside, every submodule of it will be sent to the same device.
+        offload_folder (`str` or `os.PathLike`, *optional*):
+            If the `device_map` contains any value `"disk"`, the folder where we will offload weights.
+    """
+    if offload_folder is None and device_map is not None and "disk" in device_map.values():
+        raise ValueError(
+            "At least one of the model submodule will be offloaded to disk, please pass along an `offload_folder`."
+        )
+
+    checkpoint_files = None
+    index_filename = None
+    if os.path.isfile(checkpoint):
+        if str(checkpoint).endswith(".json"):
+            index_filename = checkpoint
+        else:
+            checkpoint_files = [checkpoint]
+    elif os.path.isdir(checkpoint):
If you want a reproducible code example, I've tried with `gpt2` which I can load entirely in my model's memory. For example:

```py
from accelerate import init_empty_weights, load_checkpoint_and_dispatch
from transformers import GPT2LMHeadModel
import torch

GPT2LMHeadModel.from_pretrained('gpt2').save_pretrained('./gpt2')

with init_empty_weights():
    model = GPT2LMHeadModel.from_pretrained('gpt2')

model = load_checkpoint_and_dispatch(
    model,
    "./gpt2/pytorch_model.bin",
    device_map='auto'
)

```
@@ -0,0 +1,568 @@
+# 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 gc
+import json
+import os
+import re
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+WEIGHTS_INDEX_NAME = "pytorch_model.bin.index.json"
+
+
+def convert_file_size_to_int(size: Union[int, str]):
+    """
+    Converts a size expressed as a string with digits an unit (like `"5MB"`) to an integer (in bytes).
+
+    Args:
+        size (`int` or `str`): The size to convert. Will be directly returned if an `int`.
+
+    Example:
+
+    ```py
+    >>> convert_file_size_to_int("1MiB")
+    1048576
+    ```
+    """
+    if isinstance(size, int):
+        return size
+    if size.upper().endswith("GIB"):
+        return int(size[:-3]) * (2**30)
+    if size.upper().endswith("MIB"):
+        return int(size[:-3]) * (2**20)
+    if size.upper().endswith("KIB"):
+        return int(size[:-3]) * (2**10)
+    if size.upper().endswith("GB"):
+        int_size = int(size[:-2]) * (10**9)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("MB"):
+        int_size = int(size[:-2]) * (10**6)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("KB"):
+        int_size = int(size[:-2]) * (10**3)
+        return int_size // 8 if size.endswith("b") else int_size
+    raise ValueError("`size` is not in a valid format. Use an integer followed by the unit, e.g., '5GB'.")
+
+
+def dtype_byte_size(dtype: torch.dtype):
+    """
+    Returns the size (in bytes) occupied by one parameter of type `dtype`.
+
+    Example:
+
+    ```py
+    >>> dtype_byte_size(torch.float32)
+    4
+    ```
+    """
+    if dtype == torch.bool:
+        return 1 / 8
+    bit_search = re.search("[^\d](\d+)$", str(dtype))
+    if bit_search is None:
+        raise ValueError(f"`dtype` is not a valid dtype: {dtype}.")
+    bit_size = int(bit_search.groups()[0])
+    return bit_size // 8
+
+
+def set_module_tensor_to_device(
+    module: nn.Module, tensor_name: str, device: Union[int, str, torch.device], value: Optional[torch.Tensor] = None
+):
+    """
+    A helper function to set a given tensor (parameter of buffer) of a module on a specific device (note that doing
+    `param.to(device)` creates a new tensor not linked to the parameter, which is why we need this function).
+
+    Args:
+        module (`torch.nn.Module`): The module in which the tensor we want to move lives.
+        param_name (`str`): The full name of the parameter/buffer.
+        device (`int`, `str` or `torch.device`): The device on which to set the tensor.
+        value (`torch.Tensor`, *optional*): The value of the tensor (useful when going from the meta device to any
+            other device).
+    """
+    # Recurse if needed
+    if "." in tensor_name:
+        splits = tensor_name.split(".")
+        for split in splits[:-1]:
+            new_module = getattr(module, split)
+            if new_module is None:
+                raise ValueError(f"{module} has no attribute {split}.")
+            module = new_module
+        tensor_name = splits[-1]
+
+    if tensor_name not in module._parameters and tensor_name not in module._buffers:
+        raise ValueError(f"{module} does not have a parameter or a buffer named {tensor_name}.")
+    is_buffer = tensor_name in module._buffers
+    old_value = getattr(module, tensor_name)
+
+    if old_value.device == torch.device("meta") and device != torch.device("meta") and value is None:
+        raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {device}.")
+
+    with torch.no_grad():
+        if value is None:
+            new_value = old_value.to(device)
+        elif isinstance(value, torch.Tensor):
+            new_value = value.to(device)
+        else:
+            new_value = torch.tensor(value, device=device)
+    if is_buffer:
+        module._buffers[tensor_name] = new_value
+    else:
+        new_value = nn.Parameter(new_value, requires_grad=old_value.requires_grad)
+        module._parameters[tensor_name] = new_value
+
+
+def named_module_tensors(module: nn.Module, include_buffers: bool = True, recurse: bool = False):
+    """
+    A helper function that gathers all the tensors (parameters + buffers) of a given module. If `include_buffers=True`
+    it's the same as doing `module.named_parameters(recurse=recurse) + module.named_buffers(recurse=recurse)`.
+
+    Args:
+        module (`torch.nn.Module`): The module we want the tensors or.
+        include_buffer (`bool`, *optional*, defaults to `True`): Whether or not to include the buffers in the result.
+        recurse (`bool`, *optional`, defaults to `False`):
+            Whether or not to go look in every submodule or just return the direct parameters and buffers.
+    """
+    for named_parameter in module.named_parameters(recurse=recurse):
+        yield named_parameter
+
+    if include_buffers:
+        for named_buffer in module.named_buffers(recurse=recurse):
+            yield named_buffer
+
+
+def find_tied_parameters(model: nn.Module, **kwargs):
+    """
+    Find the tied parameters in a given model.
+
+    Args:
+        model (`torch.nn.Module`): The model to inspect.
+
+    <Tip warning={true}>
+
+    The signature accepts keyword arguments, but they are for the recursive part of this function and you should ignore
+    them.
+
+    </Tip>
+
+    Example:
+
+
+    ```py
+    >>> from collections import OrderedDict
+    >>> import torch.nn as nn
+
+    >>> model = nn.Sequential(OrderedDict([("linear1", nn.Linear(4, 4)), ("linear2", nn.Linear(4, 4))]))
+    >>> model.linear2.weight = test_model.linear1.weight
+    >>> find_tied_parameters(test_model)
+    {'linear1.weight': 'linear2.weight'}
+    ```
+
+    Returns:
+        Dict[str, str]: A dictionary mapping tied parameter names to the name of the parameter they are tied to.
+    """
+    # Initialize result and named_parameters before recursing.
+    named_parameters = kwargs.get("named_parameters", None)
+    prefix = kwargs.get("prefix", "")
+    result = kwargs.get("result", {})
+
+    if named_parameters is None:
+        named_parameters = {n: p for n, p in model.named_parameters()}
+    else:
+        # A tied parameter will not be in the full `named_parameters` seen above but will be in the `named_parameters`
+        # of the submodule it belongs to. So while recursing we track the names that are not in the initial
+        # `named_parameters`.
+        for name, parameter in model.named_parameters():
+            full_name = name if prefix == "" else f"{prefix}.{name}"
+            if full_name not in named_parameters:
+                # When we find one, it has to be one of the existing parameters.
+                for new_name, new_param in named_parameters.items():
+                    if new_param is parameter:
+                        result[new_name] = full_name
+
+    # Once we have treated direct parameters, we move to the child modules.
+    for name, child in model.named_children():
+        child_name = name if prefix == "" else f"{prefix}.{name}"
+        find_tied_parameters(child, named_parameters=named_parameters, prefix=child_name, result=result)
+
+    return result
+
+
+def compute_module_sizes(model: nn.Module):
+    """
+    Compute the size of each submodule of a given model.
+    """
+    module_sizes = defaultdict(int)
+    for name, tensor in named_module_tensors(model, recurse=True):
+        size = tensor.numel() * dtype_byte_size(tensor.dtype)
+        name_parts = name.split(".")
+        for idx in range(len(name_parts) + 1):
+            module_sizes[".".join(name_parts[:idx])] += size
+
+    return module_sizes
+
+
+def get_max_layer_size(
+    modules: List[Tuple[str, torch.nn.Module]], module_sizes: Dict[str, int], no_split_module_classes: List[str]
+):
+    """
+    Utility function that will scan a list of named modules and return the maximum size used by one full layer. The
+    definition of a layer being:
+    - a module with no direct children (just parameters and buffers)
+    - a module whose class name is in the list `no_split_module_classes`
+
+    Args:
+        modules (`List[Tuple[str, torch.nn.Module]]`):
+            The list of named modules where we want to determine the maximum layer size.
+        module_sizes (`Dict[str, int]`):
+            A dictionary mapping each layer name to its size (as generated by `compute_module_sizes`).
+        no_split_module_classes (`List[str]`):
+            A list of class names for layers we don't want to be split.
+
+    Returns:
+        `Tuple[int, List[str]]`: The maximum size of a layer with the list of layer names realizing that maximum size.
+    """
+    max_size = 0
+    layer_names = []
+    modules_to_treat = modules.copy()
+    while len(modules_to_treat) > 0:
+        module_name, module = modules_to_treat.pop(0)
+        modules_children = list(module.named_children())
+        if len(modules_children) == 0 or module.__class__.__name__ in no_split_module_classes:
+            # No splitting this one so we compare to the max_size
+            size = module_sizes[module_name]
+            if size > max_size:
+                max_size = size
+                layer_names = [module_name]
+            elif size == max_size:
+                layer_names.append(module_name)
+        else:
+            modules_to_treat = [(f"{module_name}.{n}", v) for n, v in modules_children] + modules_to_treat
+    return max_size, layer_names
+
+
+def get_max_memory(max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None):
+    """
+    Get the maximum memory available if nothing is passed, converts string to int otherwise.
+    """
+    import psutil
+
+    if max_memory is None:
+        if not torch.cuda.is_available():
+            max_memory = {}
+        else:
+            # Make sure CUDA is initialized on each GPU to have the right memory info.
+            for i in range(torch.cuda.device_count()):
+                _ = torch.tensor([0], device=i)
+            max_memory = {i: torch.cuda.mem_get_info(i)[0] for i in range(torch.cuda.device_count())}
+        max_memory["cpu"] = psutil.virtual_memory().available
+        return max_memory
+
+    for key in max_memory:
+        if isinstance(max_memory[key], str):
+            max_memory[key] = convert_file_size_to_int(max_memory[key])
+    return max_memory
+
+
+def clean_device_map(device_map: Dict[str, Union[int, str, torch.device]], module_name: str = ""):
+    """
+    Cleans a device_map by grouping all submodules that go on the same device together.
+    """
+    # Get the value of the current module and if there is only one split across several keys, regroup it.
+    values = [v for k, v in device_map.items() if k.startswith(module_name)]
+    if len(set(values)) == 1 and len(values) > 1:
+        for k in [k for k in device_map if k.startswith(module_name)]:
+            del device_map[k]
+        device_map[module_name] = values[0]
This and the code below tells me that it may not have been made so that `cpu` is a `main_device` :smile: 
I'll stop bugging you with it.
@@ -0,0 +1,260 @@
+# 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
+from contextlib import contextmanager
+from typing import Dict, List, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from .hooks import AlignDevicesHook, add_hook_to_module, attach_align_device_hook, attach_align_device_hook_on_blocks
+from .utils import (
+    OffloadedWeightsLoader,
+    check_device_map,
+    extract_submodules_state_dict,
+    infer_auto_device_map,
+    load_checkpoint_in_model,
+    offload_state_dict,
+)
+
+
+@contextmanager
+def init_empty_weights(include_buffers: bool = False):
+    """
+    A context manager under which models are initialized with all parameters on the meta device, therefore creating an
+    empty model. Useful when just initializing the model would blow the available RAM.
+
+    Args:
+        include_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to also put all buffers on the meta device while initializing.
+
+    Example:
+
+    ```pyton
+    import torch.nn as nn
+    from accelerate import init_empty_weights
+
+    # Initialize a model with 100 billions parameters in no time and without using any RAM.
+    with init_empty_weights():
+        tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])
+    ```
+    """
+    old_register_parameter = nn.Module.register_parameter
+    if include_buffers:
+        old_register_buffer = nn.Module.register_buffer
+
+    def register_empty_parameter(module, name, param):
+        old_register_parameter(module, name, param)
+        if param is not None:
+            module._parameters[name] = nn.Parameter(module._parameters[name].to(torch.device("meta")))
+
+    def register_empty_buffer(module, name, buffer):
+        old_register_buffer(module, name, buffer)
+        if buffer is not None:
+            module._buffers[name] = module._buffers[name].to(torch.device("meta"))
+
+    try:
+        nn.Module.register_parameter = register_empty_parameter
+        if include_buffers:
+            nn.Module.register_buffer = register_empty_buffer
+        yield
+    finally:
+        nn.Module.register_parameter = old_register_parameter
+        if include_buffers:
+            nn.Module.register_buffer = old_register_buffer
+
+
+def cpu_offload(
+    model: nn.Module,
+    execution_device: Optional[torch.device] = None,
+    offload_buffers: bool = False,
+    state_dict: Optional[Dict[str, torch.Tensor]] = None,
+):
+    """
+    Activates full CPU offload for a model. As a result, all parameters of the model will be offloaded and only one
+    copy of the state dict of the model will be kept. During the forward pass, parameters will be extracted from that
+    state dict and put on the execution device passed as they are needed, then offloaded again.
+
+    Args:
+        model (`torch.nn.Module`):
+            The model to offload.
+        execution_device (`torch.device`, *optional*):
+            The device on which the forward pass of the model will be executed (should be a GPU). Will default to the
+            model first parameter device.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+        state_dict (`Dict[str, torch.Tensor]`, *optional*):
+            The state dict of the model that will be kept on CPU.
+    """
+    if execution_device is None:
+        execution_device = next(iter(model.parameter())).device
+    if state_dict is None:
+        state_dict = {n: p.to("cpu") for n, p in model.state_dict().items()}
+    attach_align_device_hook(
+        model, execution_device=execution_device, offload=True, offload_buffers=offload_buffers, weights_map=state_dict
+    )
+    add_hook_to_module(model, AlignDevicesHook(io_same_device=True))
+    return model
+
+
+def disk_offload(
+    model: nn.Module,
+    offload_dir: Union[str, os.PathLike],
+    execution_device: Optional[torch.device] = None,
+    offload_buffers: bool = False,
+):
+    """
+    Activates full disk offload for a model. As a result, all parameters of the model will be offloaded as
+    memory-mapped array in a given folder. During the forward pass, parameters will be accessed from that folder and
+    put on the execution device passed as they are needed, then offloaded again.
+
+    Args:
+        model (`torch.nn.Module`): The model to offload.
+        offload_dir (`str` or `os.PathLike`):
+            The folder in which to offload the model weights (or where the model weights are already offloaded).
+        execution_device (`torch.device`, *optional*):
+            The device on which the forward pass of the model will be executed (should be a GPU). Will default to the
+            model's first parameter device.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+    """
+    if not os.path.isdir(offload_dir) or not os.path.isfile(os.path.join(offload_dir, "index.json")):
+        offload_state_dict(offload_dir, model.state_dict())
+    if execution_device is None:
+        execution_device = next(iter(model.parameter())).device
+    weights_map = OffloadedWeightsLoader(save_folder=offload_dir)
+    attach_align_device_hook(
+        model,
+        execution_device=execution_device,
+        offload=True,
+        offload_buffers=offload_buffers,
+        weights_map=weights_map,
+    )
+    add_hook_to_module(model, AlignDevicesHook(io_same_device=True))
+    return model
+
+
+def dispatch_model(
+    model: nn.Module,
+    device_map: Dict[str, Union[str, int, torch.device]],
+    main_device: Optional[torch.device] = None,
+    state_dict: Optional[Dict[str, torch.Tensor]] = None,
+    offload_dir: Union[str, os.PathLike] = None,
+    offload_buffers: bool = False,
+):
+    """
+    Dispatches a model according to a given device map. Layers of the model might be spread across GPUs, offloaded on
+    the CPU or even the disk.
+
+    Args:
+        model (`torch.nn.Module`):
+            The model to dispatch.
+        device_map (`Dict[str, Union[str, int, torch.device]]`):
+            A dictionary mapping module names in the models `state_dict` to the device they should go to. Note that
+            `"disk"` is accepted even if it's not a proper value for `torch.device`.
+        main_device (`str`, `int` or `torch.device`, *optional*):
+            The main execution device. Will default to the first device in the `device_map` different from `"cpu"` or
+            `"disk"`.
+        state_dict (`Dict[str, torch.Tensor]`, *optional*):
+            The state dict of the part of the model that will be kept on CPU.
+        offload_dir (`str` or `os.PathLike`):
+            The folder in which to offload the model weights (or where the model weights are already offloaded).
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+    """
+    # Error early if the device map is incomplete.
+    check_device_map(model, device_map)
+
+    if main_device is None:
+        main_device = [d for d in device_map.values() if d not in ["cpu", "disk"]][0]
Since buffers are way smaller in most cases, I've ignored them in everything by default.
@@ -0,0 +1,260 @@
+# 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
+from contextlib import contextmanager
+from typing import Dict, List, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from .hooks import AlignDevicesHook, add_hook_to_module, attach_align_device_hook, attach_align_device_hook_on_blocks
+from .utils import (
+    OffloadedWeightsLoader,
+    check_device_map,
+    extract_submodules_state_dict,
+    infer_auto_device_map,
+    load_checkpoint_in_model,
+    offload_state_dict,
+)
+
+
+@contextmanager
+def init_empty_weights(include_buffers: bool = False):
+    """
+    A context manager under which models are initialized with all parameters on the meta device, therefore creating an
+    empty model. Useful when just initializing the model would blow the available RAM.
+
+    Args:
+        include_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to also put all buffers on the meta device while initializing.
Good catch!
@@ -0,0 +1,260 @@
+# 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
+from contextlib import contextmanager
+from typing import Dict, List, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from .hooks import AlignDevicesHook, add_hook_to_module, attach_align_device_hook, attach_align_device_hook_on_blocks
+from .utils import (
+    OffloadedWeightsLoader,
+    check_device_map,
+    extract_submodules_state_dict,
+    infer_auto_device_map,
+    load_checkpoint_in_model,
+    offload_state_dict,
+)
+
+
+@contextmanager
+def init_empty_weights(include_buffers: bool = False):
+    """
+    A context manager under which models are initialized with all parameters on the meta device, therefore creating an
+    empty model. Useful when just initializing the model would blow the available RAM.
+
+    Args:
+        include_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to also put all buffers on the meta device while initializing.
+
+    Example:
+
+    ```pyton
+    import torch.nn as nn
+    from accelerate import init_empty_weights
+
+    # Initialize a model with 100 billions parameters in no time and without using any RAM.
+    with init_empty_weights():
+        tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])
+    ```
+    """
+    old_register_parameter = nn.Module.register_parameter
+    if include_buffers:
+        old_register_buffer = nn.Module.register_buffer
+
+    def register_empty_parameter(module, name, param):
+        old_register_parameter(module, name, param)
+        if param is not None:
+            module._parameters[name] = nn.Parameter(module._parameters[name].to(torch.device("meta")))
+
+    def register_empty_buffer(module, name, buffer):
+        old_register_buffer(module, name, buffer)
+        if buffer is not None:
+            module._buffers[name] = module._buffers[name].to(torch.device("meta"))
+
+    try:
+        nn.Module.register_parameter = register_empty_parameter
+        if include_buffers:
+            nn.Module.register_buffer = register_empty_buffer
+        yield
+    finally:
+        nn.Module.register_parameter = old_register_parameter
+        if include_buffers:
+            nn.Module.register_buffer = old_register_buffer
+
+
+def cpu_offload(
+    model: nn.Module,
+    execution_device: Optional[torch.device] = None,
+    offload_buffers: bool = False,
+    state_dict: Optional[Dict[str, torch.Tensor]] = None,
+):
+    """
+    Activates full CPU offload for a model. As a result, all parameters of the model will be offloaded and only one
+    copy of the state dict of the model will be kept. During the forward pass, parameters will be extracted from that
+    state dict and put on the execution device passed as they are needed, then offloaded again.
+
+    Args:
+        model (`torch.nn.Module`):
+            The model to offload.
+        execution_device (`torch.device`, *optional*):
+            The device on which the forward pass of the model will be executed (should be a GPU). Will default to the
+            model first parameter device.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+        state_dict (`Dict[str, torch.Tensor]`, *optional*):
+            The state dict of the model that will be kept on CPU.
+    """
+    if execution_device is None:
+        execution_device = next(iter(model.parameter())).device
I'm not completely following. If we are treating the layer with the maximum number of parameters, it will be taken into account as `module_size` later on. We thus need to consider the largest layer remaining in the model (so second largest)
@@ -0,0 +1,568 @@
+# 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 gc
+import json
+import os
+import re
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+WEIGHTS_INDEX_NAME = "pytorch_model.bin.index.json"
+
+
+def convert_file_size_to_int(size: Union[int, str]):
+    """
+    Converts a size expressed as a string with digits an unit (like `"5MB"`) to an integer (in bytes).
+
+    Args:
+        size (`int` or `str`): The size to convert. Will be directly returned if an `int`.
+
+    Example:
+
+    ```py
+    >>> convert_file_size_to_int("1MiB")
+    1048576
+    ```
+    """
+    if isinstance(size, int):
+        return size
+    if size.upper().endswith("GIB"):
+        return int(size[:-3]) * (2**30)
+    if size.upper().endswith("MIB"):
+        return int(size[:-3]) * (2**20)
+    if size.upper().endswith("KIB"):
+        return int(size[:-3]) * (2**10)
+    if size.upper().endswith("GB"):
+        int_size = int(size[:-2]) * (10**9)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("MB"):
+        int_size = int(size[:-2]) * (10**6)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("KB"):
+        int_size = int(size[:-2]) * (10**3)
+        return int_size // 8 if size.endswith("b") else int_size
+    raise ValueError("`size` is not in a valid format. Use an integer followed by the unit, e.g., '5GB'.")
+
+
+def dtype_byte_size(dtype: torch.dtype):
+    """
+    Returns the size (in bytes) occupied by one parameter of type `dtype`.
+
+    Example:
+
+    ```py
+    >>> dtype_byte_size(torch.float32)
+    4
+    ```
+    """
+    if dtype == torch.bool:
+        return 1 / 8
+    bit_search = re.search("[^\d](\d+)$", str(dtype))
+    if bit_search is None:
+        raise ValueError(f"`dtype` is not a valid dtype: {dtype}.")
+    bit_size = int(bit_search.groups()[0])
+    return bit_size // 8
+
+
+def set_module_tensor_to_device(
+    module: nn.Module, tensor_name: str, device: Union[int, str, torch.device], value: Optional[torch.Tensor] = None
+):
+    """
+    A helper function to set a given tensor (parameter of buffer) of a module on a specific device (note that doing
+    `param.to(device)` creates a new tensor not linked to the parameter, which is why we need this function).
+
+    Args:
+        module (`torch.nn.Module`): The module in which the tensor we want to move lives.
+        param_name (`str`): The full name of the parameter/buffer.
+        device (`int`, `str` or `torch.device`): The device on which to set the tensor.
+        value (`torch.Tensor`, *optional*): The value of the tensor (useful when going from the meta device to any
+            other device).
+    """
+    # Recurse if needed
+    if "." in tensor_name:
+        splits = tensor_name.split(".")
+        for split in splits[:-1]:
+            new_module = getattr(module, split)
+            if new_module is None:
+                raise ValueError(f"{module} has no attribute {split}.")
+            module = new_module
+        tensor_name = splits[-1]
+
+    if tensor_name not in module._parameters and tensor_name not in module._buffers:
+        raise ValueError(f"{module} does not have a parameter or a buffer named {tensor_name}.")
+    is_buffer = tensor_name in module._buffers
+    old_value = getattr(module, tensor_name)
+
+    if old_value.device == torch.device("meta") and device != torch.device("meta") and value is None:
+        raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {device}.")
+
+    with torch.no_grad():
+        if value is None:
+            new_value = old_value.to(device)
+        elif isinstance(value, torch.Tensor):
+            new_value = value.to(device)
+        else:
+            new_value = torch.tensor(value, device=device)
+    if is_buffer:
+        module._buffers[tensor_name] = new_value
+    else:
+        new_value = nn.Parameter(new_value, requires_grad=old_value.requires_grad)
+        module._parameters[tensor_name] = new_value
+
+
+def named_module_tensors(module: nn.Module, include_buffers: bool = True, recurse: bool = False):
+    """
+    A helper function that gathers all the tensors (parameters + buffers) of a given module. If `include_buffers=True`
+    it's the same as doing `module.named_parameters(recurse=recurse) + module.named_buffers(recurse=recurse)`.
+
+    Args:
+        module (`torch.nn.Module`): The module we want the tensors or.
+        include_buffer (`bool`, *optional*, defaults to `True`): Whether or not to include the buffers in the result.
+        recurse (`bool`, *optional`, defaults to `False`):
+            Whether or not to go look in every submodule or just return the direct parameters and buffers.
+    """
+    for named_parameter in module.named_parameters(recurse=recurse):
+        yield named_parameter
+
+    if include_buffers:
+        for named_buffer in module.named_buffers(recurse=recurse):
+            yield named_buffer
+
+
+def find_tied_parameters(model: nn.Module, **kwargs):
+    """
+    Find the tied parameters in a given model.
+
+    Args:
+        model (`torch.nn.Module`): The model to inspect.
+
+    <Tip warning={true}>
+
+    The signature accepts keyword arguments, but they are for the recursive part of this function and you should ignore
+    them.
+
+    </Tip>
+
+    Example:
+
+
+    ```py
+    >>> from collections import OrderedDict
+    >>> import torch.nn as nn
+
+    >>> model = nn.Sequential(OrderedDict([("linear1", nn.Linear(4, 4)), ("linear2", nn.Linear(4, 4))]))
+    >>> model.linear2.weight = test_model.linear1.weight
+    >>> find_tied_parameters(test_model)
+    {'linear1.weight': 'linear2.weight'}
+    ```
+
+    Returns:
+        Dict[str, str]: A dictionary mapping tied parameter names to the name of the parameter they are tied to.
+    """
+    # Initialize result and named_parameters before recursing.
+    named_parameters = kwargs.get("named_parameters", None)
+    prefix = kwargs.get("prefix", "")
+    result = kwargs.get("result", {})
+
+    if named_parameters is None:
+        named_parameters = {n: p for n, p in model.named_parameters()}
+    else:
+        # A tied parameter will not be in the full `named_parameters` seen above but will be in the `named_parameters`
+        # of the submodule it belongs to. So while recursing we track the names that are not in the initial
+        # `named_parameters`.
+        for name, parameter in model.named_parameters():
+            full_name = name if prefix == "" else f"{prefix}.{name}"
+            if full_name not in named_parameters:
+                # When we find one, it has to be one of the existing parameters.
+                for new_name, new_param in named_parameters.items():
+                    if new_param is parameter:
+                        result[new_name] = full_name
+
+    # Once we have treated direct parameters, we move to the child modules.
+    for name, child in model.named_children():
+        child_name = name if prefix == "" else f"{prefix}.{name}"
+        find_tied_parameters(child, named_parameters=named_parameters, prefix=child_name, result=result)
+
+    return result
+
+
+def compute_module_sizes(model: nn.Module):
+    """
+    Compute the size of each submodule of a given model.
+    """
+    module_sizes = defaultdict(int)
+    for name, tensor in named_module_tensors(model, recurse=True):
+        size = tensor.numel() * dtype_byte_size(tensor.dtype)
+        name_parts = name.split(".")
+        for idx in range(len(name_parts) + 1):
+            module_sizes[".".join(name_parts[:idx])] += size
+
+    return module_sizes
+
+
+def get_max_layer_size(
+    modules: List[Tuple[str, torch.nn.Module]], module_sizes: Dict[str, int], no_split_module_classes: List[str]
+):
+    """
+    Utility function that will scan a list of named modules and return the maximum size used by one full layer. The
+    definition of a layer being:
+    - a module with no direct children (just parameters and buffers)
+    - a module whose class name is in the list `no_split_module_classes`
+
+    Args:
+        modules (`List[Tuple[str, torch.nn.Module]]`):
+            The list of named modules where we want to determine the maximum layer size.
+        module_sizes (`Dict[str, int]`):
+            A dictionary mapping each layer name to its size (as generated by `compute_module_sizes`).
+        no_split_module_classes (`List[str]`):
+            A list of class names for layers we don't want to be split.
+
+    Returns:
+        `Tuple[int, List[str]]`: The maximum size of a layer with the list of layer names realizing that maximum size.
+    """
+    max_size = 0
+    layer_names = []
+    modules_to_treat = modules.copy()
+    while len(modules_to_treat) > 0:
+        module_name, module = modules_to_treat.pop(0)
+        modules_children = list(module.named_children())
+        if len(modules_children) == 0 or module.__class__.__name__ in no_split_module_classes:
+            # No splitting this one so we compare to the max_size
+            size = module_sizes[module_name]
+            if size > max_size:
+                max_size = size
+                layer_names = [module_name]
+            elif size == max_size:
+                layer_names.append(module_name)
+        else:
+            modules_to_treat = [(f"{module_name}.{n}", v) for n, v in modules_children] + modules_to_treat
+    return max_size, layer_names
+
+
+def get_max_memory(max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None):
+    """
+    Get the maximum memory available if nothing is passed, converts string to int otherwise.
+    """
+    import psutil
+
+    if max_memory is None:
+        if not torch.cuda.is_available():
+            max_memory = {}
+        else:
+            # Make sure CUDA is initialized on each GPU to have the right memory info.
+            for i in range(torch.cuda.device_count()):
+                _ = torch.tensor([0], device=i)
+            max_memory = {i: torch.cuda.mem_get_info(i)[0] for i in range(torch.cuda.device_count())}
+        max_memory["cpu"] = psutil.virtual_memory().available
+        return max_memory
+
+    for key in max_memory:
+        if isinstance(max_memory[key], str):
+            max_memory[key] = convert_file_size_to_int(max_memory[key])
+    return max_memory
+
+
+def clean_device_map(device_map: Dict[str, Union[int, str, torch.device]], module_name: str = ""):
+    """
+    Cleans a device_map by grouping all submodules that go on the same device together.
+    """
+    # Get the value of the current module and if there is only one split across several keys, regroup it.
+    values = [v for k, v in device_map.items() if k.startswith(module_name)]
+    if len(set(values)) == 1 and len(values) > 1:
+        for k in [k for k in device_map if k.startswith(module_name)]:
+            del device_map[k]
+        device_map[module_name] = values[0]
+
+    # Recurse over the children
+    children_modules = [k for k in device_map.keys() if k.startswith(module_name) and len(k) > len(module_name)]
+    idx = len(module_name.split(".")) + 1 if len(module_name) > 0 else 1
+    children_modules = set(".".join(k.split(".")[:idx]) for k in children_modules)
+    for child in children_modules:
+        clean_device_map(device_map, module_name=child)
+
+    return device_map
+
+
+def infer_auto_device_map(
+    model: nn.Module,
+    max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None,
+    no_split_module_classes: Optional[List[str]] = None,
+):
+    """
+    Compute a device map for a given model giving priority to GPUs, then offload on CPU and finally offload to disk,
+    such that:
+    - we don't exceed the memory available of any of the GPU.
+    - if offload to the CPU is needed, there is always room left on GPU 0 to put back the layer offloaded on CPU that
+      has the largest size.
+    - if offload to the CPU is needed,we don't exceed the RAM available on the CPU.
+    - if offload to the disk is needed, there is always room left on the CPU to put back the layer offloaded on disk
+      that has the largest size.
+
+    <Tip>
+
+    All computation is done analyzing sizes and dtypes of the model parameters. As a result, the model can be on the
+    meta device (as it would if initialized within the `init_empty_weights` context manager).
+
+    </Tip>
+
+    Args:
+        model (`torch.nn.Module`): The model to analyze.
+        max_memory (`Dict`, *optional*):
+            A dictionary device identifier to maximum memory. Will default to the maximum memory available if unset.
+        no_split_module_classes (`List[str]`, *optional*):
+            A list of layer class names that should never be split across device (for instance any layer that has a
+            residual connection).
+    """
+    # Get default / clean up max_memory
+    max_memory = get_max_memory(max_memory)
+    if no_split_module_classes is None:
+        no_split_module_classes = []
+    elif not isinstance(no_split_module_classes, (list, tuple)):
+        no_split_module_classes = [no_split_module_classes]
+
+    devices = list(max_memory.keys()) + ["disk"]
+    # Devices that need to keep space for a potential offloaded layer.
+    main_devices = [devices[0], "cpu"]
+
+    module_sizes = compute_module_sizes(model)
+    tied_parameters = find_tied_parameters(model)
+
+    device_map = {}
+    current_device = 0
+    current_memory_used = 0
+
+    # Direct submodules and parameters
+    modules_to_treat = list(model.named_parameters(recurse=False)) + list(model.named_children())
+    # Initialize maximum largest layer, to know which space to keep in memory
+    max_layer_size, max_layer_names = get_max_layer_size(modules_to_treat, module_sizes, no_split_module_classes)
+
+    # Ready ? This is going to be a bit messy.
+    while len(modules_to_treat) > 0:
+        name, module = modules_to_treat.pop(0)
+        # Max size in the remaining layers may have changed since we took one, so we maybe update it.
+        max_layer_names = [n for n in max_layer_names if not n.startswith(name)]
+        if len(max_layer_names) == 0:
+            max_layer_size, max_layer_names = get_max_layer_size(
Assuming the modeling file is called `pytorch_model.bin` is very Transformers-centric. For Accelerate, I don't want to make this assumption.
@@ -0,0 +1,568 @@
+# 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 gc
+import json
+import os
+import re
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+WEIGHTS_INDEX_NAME = "pytorch_model.bin.index.json"
+
+
+def convert_file_size_to_int(size: Union[int, str]):
+    """
+    Converts a size expressed as a string with digits an unit (like `"5MB"`) to an integer (in bytes).
+
+    Args:
+        size (`int` or `str`): The size to convert. Will be directly returned if an `int`.
+
+    Example:
+
+    ```py
+    >>> convert_file_size_to_int("1MiB")
+    1048576
+    ```
+    """
+    if isinstance(size, int):
+        return size
+    if size.upper().endswith("GIB"):
+        return int(size[:-3]) * (2**30)
+    if size.upper().endswith("MIB"):
+        return int(size[:-3]) * (2**20)
+    if size.upper().endswith("KIB"):
+        return int(size[:-3]) * (2**10)
+    if size.upper().endswith("GB"):
+        int_size = int(size[:-2]) * (10**9)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("MB"):
+        int_size = int(size[:-2]) * (10**6)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("KB"):
+        int_size = int(size[:-2]) * (10**3)
+        return int_size // 8 if size.endswith("b") else int_size
+    raise ValueError("`size` is not in a valid format. Use an integer followed by the unit, e.g., '5GB'.")
+
+
+def dtype_byte_size(dtype: torch.dtype):
+    """
+    Returns the size (in bytes) occupied by one parameter of type `dtype`.
+
+    Example:
+
+    ```py
+    >>> dtype_byte_size(torch.float32)
+    4
+    ```
+    """
+    if dtype == torch.bool:
+        return 1 / 8
+    bit_search = re.search("[^\d](\d+)$", str(dtype))
+    if bit_search is None:
+        raise ValueError(f"`dtype` is not a valid dtype: {dtype}.")
+    bit_size = int(bit_search.groups()[0])
+    return bit_size // 8
+
+
+def set_module_tensor_to_device(
+    module: nn.Module, tensor_name: str, device: Union[int, str, torch.device], value: Optional[torch.Tensor] = None
+):
+    """
+    A helper function to set a given tensor (parameter of buffer) of a module on a specific device (note that doing
+    `param.to(device)` creates a new tensor not linked to the parameter, which is why we need this function).
+
+    Args:
+        module (`torch.nn.Module`): The module in which the tensor we want to move lives.
+        param_name (`str`): The full name of the parameter/buffer.
+        device (`int`, `str` or `torch.device`): The device on which to set the tensor.
+        value (`torch.Tensor`, *optional*): The value of the tensor (useful when going from the meta device to any
+            other device).
+    """
+    # Recurse if needed
+    if "." in tensor_name:
+        splits = tensor_name.split(".")
+        for split in splits[:-1]:
+            new_module = getattr(module, split)
+            if new_module is None:
+                raise ValueError(f"{module} has no attribute {split}.")
+            module = new_module
+        tensor_name = splits[-1]
+
+    if tensor_name not in module._parameters and tensor_name not in module._buffers:
+        raise ValueError(f"{module} does not have a parameter or a buffer named {tensor_name}.")
+    is_buffer = tensor_name in module._buffers
+    old_value = getattr(module, tensor_name)
+
+    if old_value.device == torch.device("meta") and device != torch.device("meta") and value is None:
+        raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {device}.")
+
+    with torch.no_grad():
+        if value is None:
+            new_value = old_value.to(device)
+        elif isinstance(value, torch.Tensor):
+            new_value = value.to(device)
+        else:
+            new_value = torch.tensor(value, device=device)
+    if is_buffer:
+        module._buffers[tensor_name] = new_value
+    else:
+        new_value = nn.Parameter(new_value, requires_grad=old_value.requires_grad)
+        module._parameters[tensor_name] = new_value
+
+
+def named_module_tensors(module: nn.Module, include_buffers: bool = True, recurse: bool = False):
+    """
+    A helper function that gathers all the tensors (parameters + buffers) of a given module. If `include_buffers=True`
+    it's the same as doing `module.named_parameters(recurse=recurse) + module.named_buffers(recurse=recurse)`.
+
+    Args:
+        module (`torch.nn.Module`): The module we want the tensors or.
+        include_buffer (`bool`, *optional*, defaults to `True`): Whether or not to include the buffers in the result.
+        recurse (`bool`, *optional`, defaults to `False`):
+            Whether or not to go look in every submodule or just return the direct parameters and buffers.
+    """
+    for named_parameter in module.named_parameters(recurse=recurse):
+        yield named_parameter
+
+    if include_buffers:
+        for named_buffer in module.named_buffers(recurse=recurse):
+            yield named_buffer
+
+
+def find_tied_parameters(model: nn.Module, **kwargs):
+    """
+    Find the tied parameters in a given model.
+
+    Args:
+        model (`torch.nn.Module`): The model to inspect.
+
+    <Tip warning={true}>
+
+    The signature accepts keyword arguments, but they are for the recursive part of this function and you should ignore
+    them.
+
+    </Tip>
+
+    Example:
+
+
+    ```py
+    >>> from collections import OrderedDict
+    >>> import torch.nn as nn
+
+    >>> model = nn.Sequential(OrderedDict([("linear1", nn.Linear(4, 4)), ("linear2", nn.Linear(4, 4))]))
+    >>> model.linear2.weight = test_model.linear1.weight
+    >>> find_tied_parameters(test_model)
+    {'linear1.weight': 'linear2.weight'}
+    ```
+
+    Returns:
+        Dict[str, str]: A dictionary mapping tied parameter names to the name of the parameter they are tied to.
+    """
+    # Initialize result and named_parameters before recursing.
+    named_parameters = kwargs.get("named_parameters", None)
+    prefix = kwargs.get("prefix", "")
+    result = kwargs.get("result", {})
+
+    if named_parameters is None:
+        named_parameters = {n: p for n, p in model.named_parameters()}
+    else:
+        # A tied parameter will not be in the full `named_parameters` seen above but will be in the `named_parameters`
+        # of the submodule it belongs to. So while recursing we track the names that are not in the initial
+        # `named_parameters`.
+        for name, parameter in model.named_parameters():
+            full_name = name if prefix == "" else f"{prefix}.{name}"
+            if full_name not in named_parameters:
+                # When we find one, it has to be one of the existing parameters.
+                for new_name, new_param in named_parameters.items():
+                    if new_param is parameter:
+                        result[new_name] = full_name
+
+    # Once we have treated direct parameters, we move to the child modules.
+    for name, child in model.named_children():
+        child_name = name if prefix == "" else f"{prefix}.{name}"
+        find_tied_parameters(child, named_parameters=named_parameters, prefix=child_name, result=result)
+
+    return result
+
+
+def compute_module_sizes(model: nn.Module):
+    """
+    Compute the size of each submodule of a given model.
+    """
+    module_sizes = defaultdict(int)
+    for name, tensor in named_module_tensors(model, recurse=True):
+        size = tensor.numel() * dtype_byte_size(tensor.dtype)
+        name_parts = name.split(".")
+        for idx in range(len(name_parts) + 1):
+            module_sizes[".".join(name_parts[:idx])] += size
+
+    return module_sizes
+
+
+def get_max_layer_size(
+    modules: List[Tuple[str, torch.nn.Module]], module_sizes: Dict[str, int], no_split_module_classes: List[str]
+):
+    """
+    Utility function that will scan a list of named modules and return the maximum size used by one full layer. The
+    definition of a layer being:
+    - a module with no direct children (just parameters and buffers)
+    - a module whose class name is in the list `no_split_module_classes`
+
+    Args:
+        modules (`List[Tuple[str, torch.nn.Module]]`):
+            The list of named modules where we want to determine the maximum layer size.
+        module_sizes (`Dict[str, int]`):
+            A dictionary mapping each layer name to its size (as generated by `compute_module_sizes`).
+        no_split_module_classes (`List[str]`):
+            A list of class names for layers we don't want to be split.
+
+    Returns:
+        `Tuple[int, List[str]]`: The maximum size of a layer with the list of layer names realizing that maximum size.
+    """
+    max_size = 0
+    layer_names = []
+    modules_to_treat = modules.copy()
+    while len(modules_to_treat) > 0:
+        module_name, module = modules_to_treat.pop(0)
+        modules_children = list(module.named_children())
+        if len(modules_children) == 0 or module.__class__.__name__ in no_split_module_classes:
+            # No splitting this one so we compare to the max_size
+            size = module_sizes[module_name]
+            if size > max_size:
+                max_size = size
+                layer_names = [module_name]
+            elif size == max_size:
+                layer_names.append(module_name)
+        else:
+            modules_to_treat = [(f"{module_name}.{n}", v) for n, v in modules_children] + modules_to_treat
+    return max_size, layer_names
+
+
+def get_max_memory(max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None):
+    """
+    Get the maximum memory available if nothing is passed, converts string to int otherwise.
+    """
+    import psutil
+
+    if max_memory is None:
+        if not torch.cuda.is_available():
+            max_memory = {}
+        else:
+            # Make sure CUDA is initialized on each GPU to have the right memory info.
+            for i in range(torch.cuda.device_count()):
+                _ = torch.tensor([0], device=i)
+            max_memory = {i: torch.cuda.mem_get_info(i)[0] for i in range(torch.cuda.device_count())}
+        max_memory["cpu"] = psutil.virtual_memory().available
+        return max_memory
+
+    for key in max_memory:
+        if isinstance(max_memory[key], str):
+            max_memory[key] = convert_file_size_to_int(max_memory[key])
+    return max_memory
+
+
+def clean_device_map(device_map: Dict[str, Union[int, str, torch.device]], module_name: str = ""):
+    """
+    Cleans a device_map by grouping all submodules that go on the same device together.
+    """
+    # Get the value of the current module and if there is only one split across several keys, regroup it.
+    values = [v for k, v in device_map.items() if k.startswith(module_name)]
+    if len(set(values)) == 1 and len(values) > 1:
+        for k in [k for k in device_map if k.startswith(module_name)]:
+            del device_map[k]
+        device_map[module_name] = values[0]
+
+    # Recurse over the children
+    children_modules = [k for k in device_map.keys() if k.startswith(module_name) and len(k) > len(module_name)]
+    idx = len(module_name.split(".")) + 1 if len(module_name) > 0 else 1
+    children_modules = set(".".join(k.split(".")[:idx]) for k in children_modules)
+    for child in children_modules:
+        clean_device_map(device_map, module_name=child)
+
+    return device_map
+
+
+def infer_auto_device_map(
+    model: nn.Module,
+    max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None,
+    no_split_module_classes: Optional[List[str]] = None,
+):
+    """
+    Compute a device map for a given model giving priority to GPUs, then offload on CPU and finally offload to disk,
+    such that:
+    - we don't exceed the memory available of any of the GPU.
+    - if offload to the CPU is needed, there is always room left on GPU 0 to put back the layer offloaded on CPU that
+      has the largest size.
+    - if offload to the CPU is needed,we don't exceed the RAM available on the CPU.
+    - if offload to the disk is needed, there is always room left on the CPU to put back the layer offloaded on disk
+      that has the largest size.
+
+    <Tip>
+
+    All computation is done analyzing sizes and dtypes of the model parameters. As a result, the model can be on the
+    meta device (as it would if initialized within the `init_empty_weights` context manager).
+
+    </Tip>
+
+    Args:
+        model (`torch.nn.Module`): The model to analyze.
+        max_memory (`Dict`, *optional*):
+            A dictionary device identifier to maximum memory. Will default to the maximum memory available if unset.
+        no_split_module_classes (`List[str]`, *optional*):
+            A list of layer class names that should never be split across device (for instance any layer that has a
+            residual connection).
+    """
+    # Get default / clean up max_memory
+    max_memory = get_max_memory(max_memory)
+    if no_split_module_classes is None:
+        no_split_module_classes = []
+    elif not isinstance(no_split_module_classes, (list, tuple)):
+        no_split_module_classes = [no_split_module_classes]
+
+    devices = list(max_memory.keys()) + ["disk"]
+    # Devices that need to keep space for a potential offloaded layer.
+    main_devices = [devices[0], "cpu"]
+
+    module_sizes = compute_module_sizes(model)
+    tied_parameters = find_tied_parameters(model)
+
+    device_map = {}
+    current_device = 0
+    current_memory_used = 0
+
+    # Direct submodules and parameters
+    modules_to_treat = list(model.named_parameters(recurse=False)) + list(model.named_children())
+    # Initialize maximum largest layer, to know which space to keep in memory
+    max_layer_size, max_layer_names = get_max_layer_size(modules_to_treat, module_sizes, no_split_module_classes)
+
+    # Ready ? This is going to be a bit messy.
+    while len(modules_to_treat) > 0:
+        name, module = modules_to_treat.pop(0)
+        # Max size in the remaining layers may have changed since we took one, so we maybe update it.
+        max_layer_names = [n for n in max_layer_names if not n.startswith(name)]
+        if len(max_layer_names) == 0:
+            max_layer_size, max_layer_names = get_max_layer_size(
+                [(n, m) for n, m in modules_to_treat if isinstance(m, torch.nn.Module)],
+                module_sizes,
+                no_split_module_classes,
+            )
+        # Assess size needed
+        module_size = module_sizes[name]
+        tied_params = [v for k, v in tied_parameters.items() if name in k]
+        # We ignore parameters that are tied when they're tied to > 1 one
+        tied_param = tied_params[0] if len(tied_params) == 1 else None
+
+        device = devices[current_device]
+        current_max_size = max_memory[device] if device != "disk" else None
+        # Reduce max size available by the largest layer.
+        if devices[current_device] in main_devices:
+            current_max_size = current_max_size - max_layer_size
+        # Case 1 -> We're too big!
+        if current_max_size is not None and current_memory_used + module_size > current_max_size:
+            # Split or not split?
+            modules_children = list(module.named_children())
+            if len(modules_children) == 0 or module.__class__.__name__ in no_split_module_classes:
+                # -> no split, we go to the next device
+                current_device += 1
+                modules_to_treat = [(name, module)] + modules_to_treat
+                current_memory_used = 0
+            else:
+                # -> split, we replace the module studied by its children + parameters
+                modules_children = list(module.named_parameters(recurse=False)) + modules_children
+                modules_to_treat = [(f"{name}.{n}", v) for n, v in modules_children] + modules_to_treat
+                # Update the max layer size.
+                max_layer_size, max_layer_names = get_max_layer_size(
+                    [(n, m) for n, m in modules_to_treat if isinstance(m, torch.nn.Module)],
+                    module_sizes,
+                    no_split_module_classes,
+                )
+
+        # Case 2, it fits! We're not entirely out of the wood though, because we may have some tied parameters.
+        elif tied_param is not None:
+            # Determine the sized occupied by this module + the module containing the tied parameter
+            tied_module_size = module_size
+            tied_module_index = [i for i, (n, _) in enumerate(modules_to_treat) if n in tied_param][0]
+            tied_module_name, tied_module = modules_to_treat[tied_module_index]
+            tied_module_size += module_sizes[tied_module_name] - module_sizes[tied_param]
+            if current_max_size is not None and current_memory_used + tied_module_size > current_max_size:
+                # Split or not split?
+                tied_module_children = list(tied_module.named_children())
+                if len(tied_module_children) == 0 or tied_module.__class__.__name__ in no_split_module_classes:
+                    # If the tied module is not split, we go to the next device
+                    current_device += 1
+                    modules_to_treat = [(name, module)] + modules_to_treat
+                    current_memory_used = 0
+                else:
+                    # Otherwise, we replace the tied module by its children.
+                    tied_module_children = list(tied_module.named_parameters(recurse=False)) + tied_module_children
+                    tied_module_children = [(f"{tied_module_name}.{n}", v) for n, v in tied_module_children]
+                    modules_to_treat = (
+                        [(name, module)]
+                        + modules_to_treat[:tied_module_index]
+                        + tied_module_children
+                        + modules_to_treat[tied_module_index + 1 :]
+                    )
+                    # Update the max layer size.
+                    max_layer_size, max_layer_names = get_max_layer_size(
+                        [(n, m) for n, m in modules_to_treat if isinstance(m, torch.nn.Module)],
+                        module_sizes,
+                        no_split_module_classes,
+                    )
+            else:
+                # We really really fit!
+                current_memory_used += tied_module_size
+                device_map[name] = devices[current_device]
+                modules_to_treat.pop(tied_module_index)
+                device_map[tied_module_name] = devices[current_device]
+        else:
+            current_memory_used += module_size
+            device_map[name] = devices[current_device]
+
+    return clean_device_map(device_map)
+
+
+def check_device_map(model: nn.Module, device_map: Dict[str, Union[int, str, torch.device]]):
+    """
+    Checks a device map covers everything in a given model.
+
+    Args:
+        model (`torch.nn.Module`): The model to check the device map against.
+        device_map (`Dict[str, Union[int, str, torch.device]]`): The device map to check.
+    """
+    all_model_tensors = [name for name, _ in model.state_dict().items()]
+    for module_name in device_map.keys():
+        all_model_tensors = [name for name in all_model_tensors if not name.startswith(module_name)]
+    if len(all_model_tensors) > 0:
+        non_covered_params = ", ".join(all_model_tensors)
+        raise ValueError(
+            f"The device_map provided does not give any device for the following parameters: {non_covered_params}"
+        )
+
+
+def load_checkpoint_in_model(
+    model: nn.Module,
+    checkpoint: Union[str, os.PathLike],
+    device_map: Optional[Dict[str, Union[int, str, torch.device]]] = None,
+    offload_folder: Optional[Union[str, os.PathLike]] = None,
+):
+    """
+    Loads a (potentially sharded) checkpoint inside a model, potentially sending weights to a given device as they are
+    loaded.
+
+    <Tip warning={true}>
+
+    Once loaded across devices, you still need to call [`dispatch_model`] on your model to make it able to run. To
+    group the checkpoint loading and dispatch in one single call, use [`load_checkpoint_and_dispatch`].
+
+    </Tip>
+
+    Args:
+        model (`torch.nn.Module`): The model in which we want to load a checkpoint.
+        checkpoint (`str` or `os.PathLike`):
+            The folder checkpoint to load. It can be:
+            - a path to a file containing a whole model state dict
+            - a path to a `.json` file containing the index to a sharded checkpoint
+            - a path to a folder containing a unique `.index.json` file and the shards of a checkpoint.
+        device_map (`Dict[str, Union[int, str, torch.device]]`, *optional*):
+            A map that specifies where each submodule should go. It doesn't need to be refined to each parameter/buffer
+            name, once a given module name is inside, every submodule of it will be sent to the same device.
+        offload_folder (`str` or `os.PathLike`, *optional*):
+            If the `device_map` contains any value `"disk"`, the folder where we will offload weights.
+    """
+    if offload_folder is None and device_map is not None and "disk" in device_map.values():
+        raise ValueError(
+            "At least one of the model submodule will be offloaded to disk, please pass along an `offload_folder`."
+        )
+
+    checkpoint_files = None
+    index_filename = None
+    if os.path.isfile(checkpoint):
+        if str(checkpoint).endswith(".json"):
+            index_filename = checkpoint
+        else:
+            checkpoint_files = [checkpoint]
+    elif os.path.isdir(checkpoint):
Sounds good, will likely become very intuitive after implementing it in `transformers` so happy to keep it this way :+1: 
@@ -0,0 +1,568 @@
+# 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 gc
+import json
+import os
+import re
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+WEIGHTS_INDEX_NAME = "pytorch_model.bin.index.json"
+
+
+def convert_file_size_to_int(size: Union[int, str]):
+    """
+    Converts a size expressed as a string with digits an unit (like `"5MB"`) to an integer (in bytes).
+
+    Args:
+        size (`int` or `str`): The size to convert. Will be directly returned if an `int`.
+
+    Example:
+
+    ```py
+    >>> convert_file_size_to_int("1MiB")
+    1048576
+    ```
+    """
+    if isinstance(size, int):
+        return size
+    if size.upper().endswith("GIB"):
+        return int(size[:-3]) * (2**30)
+    if size.upper().endswith("MIB"):
+        return int(size[:-3]) * (2**20)
+    if size.upper().endswith("KIB"):
+        return int(size[:-3]) * (2**10)
+    if size.upper().endswith("GB"):
+        int_size = int(size[:-2]) * (10**9)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("MB"):
+        int_size = int(size[:-2]) * (10**6)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("KB"):
+        int_size = int(size[:-2]) * (10**3)
+        return int_size // 8 if size.endswith("b") else int_size
+    raise ValueError("`size` is not in a valid format. Use an integer followed by the unit, e.g., '5GB'.")
+
+
+def dtype_byte_size(dtype: torch.dtype):
+    """
+    Returns the size (in bytes) occupied by one parameter of type `dtype`.
+
+    Example:
+
+    ```py
+    >>> dtype_byte_size(torch.float32)
+    4
+    ```
+    """
+    if dtype == torch.bool:
+        return 1 / 8
+    bit_search = re.search("[^\d](\d+)$", str(dtype))
+    if bit_search is None:
+        raise ValueError(f"`dtype` is not a valid dtype: {dtype}.")
+    bit_size = int(bit_search.groups()[0])
+    return bit_size // 8
+
+
+def set_module_tensor_to_device(
+    module: nn.Module, tensor_name: str, device: Union[int, str, torch.device], value: Optional[torch.Tensor] = None
+):
+    """
+    A helper function to set a given tensor (parameter of buffer) of a module on a specific device (note that doing
+    `param.to(device)` creates a new tensor not linked to the parameter, which is why we need this function).
+
+    Args:
+        module (`torch.nn.Module`): The module in which the tensor we want to move lives.
+        param_name (`str`): The full name of the parameter/buffer.
+        device (`int`, `str` or `torch.device`): The device on which to set the tensor.
+        value (`torch.Tensor`, *optional*): The value of the tensor (useful when going from the meta device to any
+            other device).
+    """
+    # Recurse if needed
+    if "." in tensor_name:
+        splits = tensor_name.split(".")
+        for split in splits[:-1]:
+            new_module = getattr(module, split)
+            if new_module is None:
+                raise ValueError(f"{module} has no attribute {split}.")
+            module = new_module
+        tensor_name = splits[-1]
+
+    if tensor_name not in module._parameters and tensor_name not in module._buffers:
+        raise ValueError(f"{module} does not have a parameter or a buffer named {tensor_name}.")
+    is_buffer = tensor_name in module._buffers
+    old_value = getattr(module, tensor_name)
+
+    if old_value.device == torch.device("meta") and device != torch.device("meta") and value is None:
+        raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {device}.")
+
+    with torch.no_grad():
+        if value is None:
+            new_value = old_value.to(device)
+        elif isinstance(value, torch.Tensor):
+            new_value = value.to(device)
+        else:
+            new_value = torch.tensor(value, device=device)
+    if is_buffer:
+        module._buffers[tensor_name] = new_value
+    else:
+        new_value = nn.Parameter(new_value, requires_grad=old_value.requires_grad)
+        module._parameters[tensor_name] = new_value
+
+
+def named_module_tensors(module: nn.Module, include_buffers: bool = True, recurse: bool = False):
+    """
+    A helper function that gathers all the tensors (parameters + buffers) of a given module. If `include_buffers=True`
+    it's the same as doing `module.named_parameters(recurse=recurse) + module.named_buffers(recurse=recurse)`.
+
+    Args:
+        module (`torch.nn.Module`): The module we want the tensors or.
+        include_buffer (`bool`, *optional*, defaults to `True`): Whether or not to include the buffers in the result.
+        recurse (`bool`, *optional`, defaults to `False`):
+            Whether or not to go look in every submodule or just return the direct parameters and buffers.
+    """
+    for named_parameter in module.named_parameters(recurse=recurse):
+        yield named_parameter
+
+    if include_buffers:
+        for named_buffer in module.named_buffers(recurse=recurse):
+            yield named_buffer
+
+
+def find_tied_parameters(model: nn.Module, **kwargs):
+    """
+    Find the tied parameters in a given model.
+
+    Args:
+        model (`torch.nn.Module`): The model to inspect.
+
+    <Tip warning={true}>
+
+    The signature accepts keyword arguments, but they are for the recursive part of this function and you should ignore
+    them.
+
+    </Tip>
+
+    Example:
+
+
+    ```py
+    >>> from collections import OrderedDict
+    >>> import torch.nn as nn
+
+    >>> model = nn.Sequential(OrderedDict([("linear1", nn.Linear(4, 4)), ("linear2", nn.Linear(4, 4))]))
+    >>> model.linear2.weight = test_model.linear1.weight
+    >>> find_tied_parameters(test_model)
+    {'linear1.weight': 'linear2.weight'}
+    ```
+
+    Returns:
+        Dict[str, str]: A dictionary mapping tied parameter names to the name of the parameter they are tied to.
+    """
+    # Initialize result and named_parameters before recursing.
+    named_parameters = kwargs.get("named_parameters", None)
+    prefix = kwargs.get("prefix", "")
+    result = kwargs.get("result", {})
+
+    if named_parameters is None:
+        named_parameters = {n: p for n, p in model.named_parameters()}
+    else:
+        # A tied parameter will not be in the full `named_parameters` seen above but will be in the `named_parameters`
+        # of the submodule it belongs to. So while recursing we track the names that are not in the initial
+        # `named_parameters`.
+        for name, parameter in model.named_parameters():
+            full_name = name if prefix == "" else f"{prefix}.{name}"
+            if full_name not in named_parameters:
+                # When we find one, it has to be one of the existing parameters.
+                for new_name, new_param in named_parameters.items():
+                    if new_param is parameter:
+                        result[new_name] = full_name
+
+    # Once we have treated direct parameters, we move to the child modules.
+    for name, child in model.named_children():
+        child_name = name if prefix == "" else f"{prefix}.{name}"
+        find_tied_parameters(child, named_parameters=named_parameters, prefix=child_name, result=result)
+
+    return result
+
+
+def compute_module_sizes(model: nn.Module):
+    """
+    Compute the size of each submodule of a given model.
+    """
+    module_sizes = defaultdict(int)
+    for name, tensor in named_module_tensors(model, recurse=True):
+        size = tensor.numel() * dtype_byte_size(tensor.dtype)
+        name_parts = name.split(".")
+        for idx in range(len(name_parts) + 1):
+            module_sizes[".".join(name_parts[:idx])] += size
+
+    return module_sizes
+
+
+def get_max_layer_size(
+    modules: List[Tuple[str, torch.nn.Module]], module_sizes: Dict[str, int], no_split_module_classes: List[str]
+):
+    """
+    Utility function that will scan a list of named modules and return the maximum size used by one full layer. The
+    definition of a layer being:
+    - a module with no direct children (just parameters and buffers)
+    - a module whose class name is in the list `no_split_module_classes`
+
+    Args:
+        modules (`List[Tuple[str, torch.nn.Module]]`):
+            The list of named modules where we want to determine the maximum layer size.
+        module_sizes (`Dict[str, int]`):
+            A dictionary mapping each layer name to its size (as generated by `compute_module_sizes`).
+        no_split_module_classes (`List[str]`):
+            A list of class names for layers we don't want to be split.
+
+    Returns:
+        `Tuple[int, List[str]]`: The maximum size of a layer with the list of layer names realizing that maximum size.
+    """
+    max_size = 0
+    layer_names = []
+    modules_to_treat = modules.copy()
+    while len(modules_to_treat) > 0:
+        module_name, module = modules_to_treat.pop(0)
+        modules_children = list(module.named_children())
+        if len(modules_children) == 0 or module.__class__.__name__ in no_split_module_classes:
+            # No splitting this one so we compare to the max_size
+            size = module_sizes[module_name]
+            if size > max_size:
+                max_size = size
+                layer_names = [module_name]
+            elif size == max_size:
+                layer_names.append(module_name)
+        else:
+            modules_to_treat = [(f"{module_name}.{n}", v) for n, v in modules_children] + modules_to_treat
+    return max_size, layer_names
+
+
+def get_max_memory(max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None):
+    """
+    Get the maximum memory available if nothing is passed, converts string to int otherwise.
+    """
+    import psutil
+
+    if max_memory is None:
+        if not torch.cuda.is_available():
+            max_memory = {}
+        else:
+            # Make sure CUDA is initialized on each GPU to have the right memory info.
+            for i in range(torch.cuda.device_count()):
+                _ = torch.tensor([0], device=i)
+            max_memory = {i: torch.cuda.mem_get_info(i)[0] for i in range(torch.cuda.device_count())}
+        max_memory["cpu"] = psutil.virtual_memory().available
+        return max_memory
+
+    for key in max_memory:
+        if isinstance(max_memory[key], str):
+            max_memory[key] = convert_file_size_to_int(max_memory[key])
+    return max_memory
+
+
+def clean_device_map(device_map: Dict[str, Union[int, str, torch.device]], module_name: str = ""):
+    """
+    Cleans a device_map by grouping all submodules that go on the same device together.
+    """
+    # Get the value of the current module and if there is only one split across several keys, regroup it.
+    values = [v for k, v in device_map.items() if k.startswith(module_name)]
+    if len(set(values)) == 1 and len(values) > 1:
+        for k in [k for k in device_map if k.startswith(module_name)]:
+            del device_map[k]
+        device_map[module_name] = values[0]
+
+    # Recurse over the children
+    children_modules = [k for k in device_map.keys() if k.startswith(module_name) and len(k) > len(module_name)]
+    idx = len(module_name.split(".")) + 1 if len(module_name) > 0 else 1
+    children_modules = set(".".join(k.split(".")[:idx]) for k in children_modules)
+    for child in children_modules:
+        clean_device_map(device_map, module_name=child)
+
+    return device_map
+
+
+def infer_auto_device_map(
+    model: nn.Module,
+    max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None,
+    no_split_module_classes: Optional[List[str]] = None,
+):
+    """
+    Compute a device map for a given model giving priority to GPUs, then offload on CPU and finally offload to disk,
+    such that:
+    - we don't exceed the memory available of any of the GPU.
+    - if offload to the CPU is needed, there is always room left on GPU 0 to put back the layer offloaded on CPU that
+      has the largest size.
+    - if offload to the CPU is needed,we don't exceed the RAM available on the CPU.
+    - if offload to the disk is needed, there is always room left on the CPU to put back the layer offloaded on disk
+      that has the largest size.
+
+    <Tip>
+
+    All computation is done analyzing sizes and dtypes of the model parameters. As a result, the model can be on the
+    meta device (as it would if initialized within the `init_empty_weights` context manager).
+
+    </Tip>
+
+    Args:
+        model (`torch.nn.Module`): The model to analyze.
+        max_memory (`Dict`, *optional*):
+            A dictionary device identifier to maximum memory. Will default to the maximum memory available if unset.
+        no_split_module_classes (`List[str]`, *optional*):
+            A list of layer class names that should never be split across device (for instance any layer that has a
+            residual connection).
+    """
+    # Get default / clean up max_memory
+    max_memory = get_max_memory(max_memory)
+    if no_split_module_classes is None:
+        no_split_module_classes = []
+    elif not isinstance(no_split_module_classes, (list, tuple)):
+        no_split_module_classes = [no_split_module_classes]
+
+    devices = list(max_memory.keys()) + ["disk"]
+    # Devices that need to keep space for a potential offloaded layer.
+    main_devices = [devices[0], "cpu"]
+
+    module_sizes = compute_module_sizes(model)
+    tied_parameters = find_tied_parameters(model)
+
+    device_map = {}
+    current_device = 0
+    current_memory_used = 0
+
+    # Direct submodules and parameters
+    modules_to_treat = list(model.named_parameters(recurse=False)) + list(model.named_children())
+    # Initialize maximum largest layer, to know which space to keep in memory
+    max_layer_size, max_layer_names = get_max_layer_size(modules_to_treat, module_sizes, no_split_module_classes)
+
+    # Ready ? This is going to be a bit messy.
+    while len(modules_to_treat) > 0:
+        name, module = modules_to_treat.pop(0)
+        # Max size in the remaining layers may have changed since we took one, so we maybe update it.
+        max_layer_names = [n for n in max_layer_names if not n.startswith(name)]
+        if len(max_layer_names) == 0:
+            max_layer_size, max_layer_names = get_max_layer_size(
+                [(n, m) for n, m in modules_to_treat if isinstance(m, torch.nn.Module)],
+                module_sizes,
+                no_split_module_classes,
+            )
+        # Assess size needed
+        module_size = module_sizes[name]
+        tied_params = [v for k, v in tied_parameters.items() if name in k]
+        # We ignore parameters that are tied when they're tied to > 1 one
+        tied_param = tied_params[0] if len(tied_params) == 1 else None
+
+        device = devices[current_device]
+        current_max_size = max_memory[device] if device != "disk" else None
+        # Reduce max size available by the largest layer.
+        if devices[current_device] in main_devices:
+            current_max_size = current_max_size - max_layer_size
+        # Case 1 -> We're too big!
+        if current_max_size is not None and current_memory_used + module_size > current_max_size:
+            # Split or not split?
+            modules_children = list(module.named_children())
+            if len(modules_children) == 0 or module.__class__.__name__ in no_split_module_classes:
+                # -> no split, we go to the next device
+                current_device += 1
+                modules_to_treat = [(name, module)] + modules_to_treat
+                current_memory_used = 0
+            else:
+                # -> split, we replace the module studied by its children + parameters
+                modules_children = list(module.named_parameters(recurse=False)) + modules_children
+                modules_to_treat = [(f"{name}.{n}", v) for n, v in modules_children] + modules_to_treat
+                # Update the max layer size.
+                max_layer_size, max_layer_names = get_max_layer_size(
+                    [(n, m) for n, m in modules_to_treat if isinstance(m, torch.nn.Module)],
+                    module_sizes,
+                    no_split_module_classes,
+                )
+
+        # Case 2, it fits! We're not entirely out of the wood though, because we may have some tied parameters.
+        elif tied_param is not None:
+            # Determine the sized occupied by this module + the module containing the tied parameter
+            tied_module_size = module_size
+            tied_module_index = [i for i, (n, _) in enumerate(modules_to_treat) if n in tied_param][0]
+            tied_module_name, tied_module = modules_to_treat[tied_module_index]
+            tied_module_size += module_sizes[tied_module_name] - module_sizes[tied_param]
+            if current_max_size is not None and current_memory_used + tied_module_size > current_max_size:
+                # Split or not split?
+                tied_module_children = list(tied_module.named_children())
+                if len(tied_module_children) == 0 or tied_module.__class__.__name__ in no_split_module_classes:
+                    # If the tied module is not split, we go to the next device
+                    current_device += 1
+                    modules_to_treat = [(name, module)] + modules_to_treat
+                    current_memory_used = 0
+                else:
+                    # Otherwise, we replace the tied module by its children.
+                    tied_module_children = list(tied_module.named_parameters(recurse=False)) + tied_module_children
+                    tied_module_children = [(f"{tied_module_name}.{n}", v) for n, v in tied_module_children]
+                    modules_to_treat = (
+                        [(name, module)]
+                        + modules_to_treat[:tied_module_index]
+                        + tied_module_children
+                        + modules_to_treat[tied_module_index + 1 :]
+                    )
+                    # Update the max layer size.
+                    max_layer_size, max_layer_names = get_max_layer_size(
+                        [(n, m) for n, m in modules_to_treat if isinstance(m, torch.nn.Module)],
+                        module_sizes,
+                        no_split_module_classes,
+                    )
+            else:
+                # We really really fit!
+                current_memory_used += tied_module_size
+                device_map[name] = devices[current_device]
+                modules_to_treat.pop(tied_module_index)
+                device_map[tied_module_name] = devices[current_device]
+        else:
+            current_memory_used += module_size
+            device_map[name] = devices[current_device]
+
+    return clean_device_map(device_map)
+
+
+def check_device_map(model: nn.Module, device_map: Dict[str, Union[int, str, torch.device]]):
+    """
+    Checks a device map covers everything in a given model.
+
+    Args:
+        model (`torch.nn.Module`): The model to check the device map against.
+        device_map (`Dict[str, Union[int, str, torch.device]]`): The device map to check.
+    """
+    all_model_tensors = [name for name, _ in model.state_dict().items()]
+    for module_name in device_map.keys():
+        all_model_tensors = [name for name in all_model_tensors if not name.startswith(module_name)]
+    if len(all_model_tensors) > 0:
+        non_covered_params = ", ".join(all_model_tensors)
+        raise ValueError(
+            f"The device_map provided does not give any device for the following parameters: {non_covered_params}"
+        )
+
+
+def load_checkpoint_in_model(
+    model: nn.Module,
+    checkpoint: Union[str, os.PathLike],
+    device_map: Optional[Dict[str, Union[int, str, torch.device]]] = None,
+    offload_folder: Optional[Union[str, os.PathLike]] = None,
+):
+    """
+    Loads a (potentially sharded) checkpoint inside a model, potentially sending weights to a given device as they are
+    loaded.
+
+    <Tip warning={true}>
+
+    Once loaded across devices, you still need to call [`dispatch_model`] on your model to make it able to run. To
+    group the checkpoint loading and dispatch in one single call, use [`load_checkpoint_and_dispatch`].
+
+    </Tip>
+
+    Args:
+        model (`torch.nn.Module`): The model in which we want to load a checkpoint.
+        checkpoint (`str` or `os.PathLike`):
+            The folder checkpoint to load. It can be:
+            - a path to a file containing a whole model state dict
+            - a path to a `.json` file containing the index to a sharded checkpoint
+            - a path to a folder containing a unique `.index.json` file and the shards of a checkpoint.
+        device_map (`Dict[str, Union[int, str, torch.device]]`, *optional*):
+            A map that specifies where each submodule should go. It doesn't need to be refined to each parameter/buffer
+            name, once a given module name is inside, every submodule of it will be sent to the same device.
+        offload_folder (`str` or `os.PathLike`, *optional*):
+            If the `device_map` contains any value `"disk"`, the folder where we will offload weights.
+    """
+    if offload_folder is None and device_map is not None and "disk" in device_map.values():
+        raise ValueError(
+            "At least one of the model submodule will be offloaded to disk, please pass along an `offload_folder`."
+        )
+
+    checkpoint_files = None
+    index_filename = None
+    if os.path.isfile(checkpoint):
+        if str(checkpoint).endswith(".json"):
+            index_filename = checkpoint
+        else:
+            checkpoint_files = [checkpoint]
+    elif os.path.isdir(checkpoint):
I will store it in `model.hf_device_map` to avoid conflicts with potential existing `device_map`.
@@ -0,0 +1,260 @@
+# 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
+from contextlib import contextmanager
+from typing import Dict, List, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from .hooks import AlignDevicesHook, add_hook_to_module, attach_align_device_hook, attach_align_device_hook_on_blocks
+from .utils import (
+    OffloadedWeightsLoader,
+    check_device_map,
+    extract_submodules_state_dict,
+    infer_auto_device_map,
+    load_checkpoint_in_model,
+    offload_state_dict,
+)
+
+
+@contextmanager
+def init_empty_weights(include_buffers: bool = False):
+    """
+    A context manager under which models are initialized with all parameters on the meta device, therefore creating an
+    empty model. Useful when just initializing the model would blow the available RAM.
+
+    Args:
+        include_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to also put all buffers on the meta device while initializing.
+
+    Example:
+
+    ```pyton
+    import torch.nn as nn
+    from accelerate import init_empty_weights
+
+    # Initialize a model with 100 billions parameters in no time and without using any RAM.
+    with init_empty_weights():
+        tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])
+    ```
+    """
+    old_register_parameter = nn.Module.register_parameter
+    if include_buffers:
+        old_register_buffer = nn.Module.register_buffer
+
+    def register_empty_parameter(module, name, param):
+        old_register_parameter(module, name, param)
+        if param is not None:
+            module._parameters[name] = nn.Parameter(module._parameters[name].to(torch.device("meta")))
+
+    def register_empty_buffer(module, name, buffer):
+        old_register_buffer(module, name, buffer)
+        if buffer is not None:
+            module._buffers[name] = module._buffers[name].to(torch.device("meta"))
+
+    try:
+        nn.Module.register_parameter = register_empty_parameter
+        if include_buffers:
+            nn.Module.register_buffer = register_empty_buffer
+        yield
+    finally:
+        nn.Module.register_parameter = old_register_parameter
+        if include_buffers:
+            nn.Module.register_buffer = old_register_buffer
+
+
+def cpu_offload(
+    model: nn.Module,
+    execution_device: Optional[torch.device] = None,
+    offload_buffers: bool = False,
+    state_dict: Optional[Dict[str, torch.Tensor]] = None,
+):
+    """
+    Activates full CPU offload for a model. As a result, all parameters of the model will be offloaded and only one
+    copy of the state dict of the model will be kept. During the forward pass, parameters will be extracted from that
+    state dict and put on the execution device passed as they are needed, then offloaded again.
+
+    Args:
+        model (`torch.nn.Module`):
+            The model to offload.
+        execution_device (`torch.device`, *optional*):
+            The device on which the forward pass of the model will be executed (should be a GPU). Will default to the
+            model first parameter device.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+        state_dict (`Dict[str, torch.Tensor]`, *optional*):
+            The state dict of the model that will be kept on CPU.
+    """
+    if execution_device is None:
+        execution_device = next(iter(model.parameter())).device
+    if state_dict is None:
+        state_dict = {n: p.to("cpu") for n, p in model.state_dict().items()}
+    attach_align_device_hook(
+        model, execution_device=execution_device, offload=True, offload_buffers=offload_buffers, weights_map=state_dict
+    )
+    add_hook_to_module(model, AlignDevicesHook(io_same_device=True))
+    return model
+
+
+def disk_offload(
+    model: nn.Module,
+    offload_dir: Union[str, os.PathLike],
+    execution_device: Optional[torch.device] = None,
+    offload_buffers: bool = False,
+):
+    """
+    Activates full disk offload for a model. As a result, all parameters of the model will be offloaded as
+    memory-mapped array in a given folder. During the forward pass, parameters will be accessed from that folder and
+    put on the execution device passed as they are needed, then offloaded again.
+
+    Args:
+        model (`torch.nn.Module`): The model to offload.
+        offload_dir (`str` or `os.PathLike`):
+            The folder in which to offload the model weights (or where the model weights are already offloaded).
+        execution_device (`torch.device`, *optional*):
+            The device on which the forward pass of the model will be executed (should be a GPU). Will default to the
+            model's first parameter device.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+    """
+    if not os.path.isdir(offload_dir) or not os.path.isfile(os.path.join(offload_dir, "index.json")):
+        offload_state_dict(offload_dir, model.state_dict())
+    if execution_device is None:
+        execution_device = next(iter(model.parameter())).device
+    weights_map = OffloadedWeightsLoader(save_folder=offload_dir)
+    attach_align_device_hook(
+        model,
+        execution_device=execution_device,
+        offload=True,
+        offload_buffers=offload_buffers,
+        weights_map=weights_map,
+    )
+    add_hook_to_module(model, AlignDevicesHook(io_same_device=True))
+    return model
+
+
+def dispatch_model(
+    model: nn.Module,
+    device_map: Dict[str, Union[str, int, torch.device]],
+    main_device: Optional[torch.device] = None,
+    state_dict: Optional[Dict[str, torch.Tensor]] = None,
+    offload_dir: Union[str, os.PathLike] = None,
+    offload_buffers: bool = False,
+):
+    """
+    Dispatches a model according to a given device map. Layers of the model might be spread across GPUs, offloaded on
+    the CPU or even the disk.
+
+    Args:
+        model (`torch.nn.Module`):
+            The model to dispatch.
+        device_map (`Dict[str, Union[str, int, torch.device]]`):
+            A dictionary mapping module names in the models `state_dict` to the device they should go to. Note that
+            `"disk"` is accepted even if it's not a proper value for `torch.device`.
+        main_device (`str`, `int` or `torch.device`, *optional*):
+            The main execution device. Will default to the first device in the `device_map` different from `"cpu"` or
+            `"disk"`.
+        state_dict (`Dict[str, torch.Tensor]`, *optional*):
+            The state dict of the part of the model that will be kept on CPU.
+        offload_dir (`str` or `os.PathLike`):
+            The folder in which to offload the model weights (or where the model weights are already offloaded).
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+    """
+    # Error early if the device map is incomplete.
+    check_device_map(model, device_map)
+
+    if main_device is None:
+        main_device = [d for d in device_map.values() if d not in ["cpu", "disk"]][0]
+
+    cpu_modules = [name for name, device in device_map.items() if device == "cpu"]
+    if state_dict is None and len(cpu_modules) > 0:
+        state_dict = extract_submodules_state_dict(model.state_dict(), cpu_modules)
+
+    disk_modules = [name for name, device in device_map.items() if device == "disk"]
+    if offload_dir is None and len(disk_modules) > 0:
+        raise ValueError(
+            "We need an `offload_dir` to dispatch this model according to this `device_map`, the following submodules "
+            f"need to be offloaded: {', '.join(disk_modules)}."
+        )
+    if len(disk_modules) > 0 and (
+        not os.path.isdir(offload_dir) or not os.path.isfile(os.path.join(offload_dir, "index.json"))
+    ):
+        disk_state_dict = extract_submodules_state_dict(model.state_dict(), disk_modules)
+        offload_state_dict(offload_dir, disk_state_dict)
+
+    execution_device = {
+        name: main_device if device in ["cpu", "disk"] else device for name, device in device_map.items()
+    }
+    offload = {name: device in ["cpu", "disk"] for name, device in device_map.items()}
+    if state_dict is not None or offload_dir is not None:
+        weights_map = OffloadedWeightsLoader(state_dict=state_dict, save_folder=offload_dir)
+    else:
+        weights_map = None
+
+    attach_align_device_hook_on_blocks(
+        model,
+        execution_device=execution_device,
+        offload=offload,
+        offload_buffers=offload_buffers,
+        weights_map=weights_map,
+    )
+    return model
+
+
+def load_checkpoint_and_dispatch(
+    model: nn.Module,
+    checkpoint: Union[str, os.PathLike],
+    device_map: Optional[Dict[str, Union[int, str, torch.device]]] = None,
+    max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None,
+    no_split_module_classes: Optional[List[str]] = None,
+    offload_folder: Optional[Union[str, os.PathLike]] = None,
+    offload_buffers: bool = False,
+):
+    """
+    Loads a (potentially sharded) checkpoint inside a model, potentially sending weights to a given device as they are
+    loaded and adds the various hooks that will make this model run properly (even if split across devices).
+
+    Args:
+        model (`torch.nn.Module`): The model in which we want to load a checkpoint.
+        checkpoint (`str` or `os.PathLike`):
+            The folder checkpoint to load. It can be:
+            - a path to a file containing a whole model state dict
+            - a path to a `.json` file containing the index to a sharded checkpoint
+            - a path to a folder containing a unique `.index.json` file and the shards of a checkpoint.
+        device_map (`Dict[str, Union[int, str, torch.device]]`, *optional*):
+            A map that specifies where each submodule should go. It doesn't need to be refined to each parameter/buffer
+            name, once a given module name is inside, every submodule of it will be sent to the same device.
+
+            To have Accelerate compute the most optimized `device_map` automatically, set `device_map="auto"`.
+        max_memory (`Dict`, *optional*):
+            A dictionary device identifier to maximum memory. Will default to the maximum memory available for each GPU
+            and the available CPU RAM if unset.
+        no_split_module_classes (`List[str]`, *optional*):
+            A list of layer class names that should never be split across device (for instance any layer that has a
+            residual connection).
+        offload_folder (`str` or `os.PathLike`, *optional*):
+            If the `device_map` contains any value `"disk"`, the folder where we will offload weights.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            In the layers that are offloaded on the CPU or the hard drive, whether or not to offload the buffers as
+            well as the parameters.
+    """
+    if device_map == "auto":
+        device_map = infer_auto_device_map(
+            model, max_memory=max_memory, no_split_module_classes=no_split_module_classes
+        )
We can add more defensive checks indeed, in new PRs.
@@ -0,0 +1,568 @@
+# 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 gc
+import json
+import os
+import re
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+WEIGHTS_INDEX_NAME = "pytorch_model.bin.index.json"
+
+
+def convert_file_size_to_int(size: Union[int, str]):
+    """
+    Converts a size expressed as a string with digits an unit (like `"5MB"`) to an integer (in bytes).
+
+    Args:
+        size (`int` or `str`): The size to convert. Will be directly returned if an `int`.
+
+    Example:
+
+    ```py
+    >>> convert_file_size_to_int("1MiB")
+    1048576
+    ```
+    """
+    if isinstance(size, int):
+        return size
+    if size.upper().endswith("GIB"):
+        return int(size[:-3]) * (2**30)
+    if size.upper().endswith("MIB"):
+        return int(size[:-3]) * (2**20)
+    if size.upper().endswith("KIB"):
+        return int(size[:-3]) * (2**10)
+    if size.upper().endswith("GB"):
+        int_size = int(size[:-2]) * (10**9)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("MB"):
+        int_size = int(size[:-2]) * (10**6)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("KB"):
+        int_size = int(size[:-2]) * (10**3)
+        return int_size // 8 if size.endswith("b") else int_size
+    raise ValueError("`size` is not in a valid format. Use an integer followed by the unit, e.g., '5GB'.")
+
+
+def dtype_byte_size(dtype: torch.dtype):
+    """
+    Returns the size (in bytes) occupied by one parameter of type `dtype`.
+
+    Example:
+
+    ```py
+    >>> dtype_byte_size(torch.float32)
+    4
+    ```
+    """
+    if dtype == torch.bool:
+        return 1 / 8
+    bit_search = re.search("[^\d](\d+)$", str(dtype))
+    if bit_search is None:
+        raise ValueError(f"`dtype` is not a valid dtype: {dtype}.")
+    bit_size = int(bit_search.groups()[0])
+    return bit_size // 8
+
+
+def set_module_tensor_to_device(
+    module: nn.Module, tensor_name: str, device: Union[int, str, torch.device], value: Optional[torch.Tensor] = None
+):
+    """
+    A helper function to set a given tensor (parameter of buffer) of a module on a specific device (note that doing
+    `param.to(device)` creates a new tensor not linked to the parameter, which is why we need this function).
+
+    Args:
+        module (`torch.nn.Module`): The module in which the tensor we want to move lives.
+        param_name (`str`): The full name of the parameter/buffer.
+        device (`int`, `str` or `torch.device`): The device on which to set the tensor.
+        value (`torch.Tensor`, *optional*): The value of the tensor (useful when going from the meta device to any
+            other device).
+    """
+    # Recurse if needed
+    if "." in tensor_name:
+        splits = tensor_name.split(".")
+        for split in splits[:-1]:
+            new_module = getattr(module, split)
+            if new_module is None:
+                raise ValueError(f"{module} has no attribute {split}.")
+            module = new_module
+        tensor_name = splits[-1]
+
+    if tensor_name not in module._parameters and tensor_name not in module._buffers:
+        raise ValueError(f"{module} does not have a parameter or a buffer named {tensor_name}.")
+    is_buffer = tensor_name in module._buffers
+    old_value = getattr(module, tensor_name)
+
+    if old_value.device == torch.device("meta") and device != torch.device("meta") and value is None:
+        raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {device}.")
+
+    with torch.no_grad():
+        if value is None:
+            new_value = old_value.to(device)
+        elif isinstance(value, torch.Tensor):
+            new_value = value.to(device)
+        else:
+            new_value = torch.tensor(value, device=device)
+    if is_buffer:
+        module._buffers[tensor_name] = new_value
+    else:
+        new_value = nn.Parameter(new_value, requires_grad=old_value.requires_grad)
+        module._parameters[tensor_name] = new_value
+
+
+def named_module_tensors(module: nn.Module, include_buffers: bool = True, recurse: bool = False):
+    """
+    A helper function that gathers all the tensors (parameters + buffers) of a given module. If `include_buffers=True`
+    it's the same as doing `module.named_parameters(recurse=recurse) + module.named_buffers(recurse=recurse)`.
+
+    Args:
+        module (`torch.nn.Module`): The module we want the tensors or.
+        include_buffer (`bool`, *optional*, defaults to `True`): Whether or not to include the buffers in the result.
+        recurse (`bool`, *optional`, defaults to `False`):
+            Whether or not to go look in every submodule or just return the direct parameters and buffers.
+    """
+    for named_parameter in module.named_parameters(recurse=recurse):
+        yield named_parameter
+
+    if include_buffers:
+        for named_buffer in module.named_buffers(recurse=recurse):
+            yield named_buffer
+
+
+def find_tied_parameters(model: nn.Module, **kwargs):
+    """
+    Find the tied parameters in a given model.
+
+    Args:
+        model (`torch.nn.Module`): The model to inspect.
+
+    <Tip warning={true}>
+
+    The signature accepts keyword arguments, but they are for the recursive part of this function and you should ignore
+    them.
+
+    </Tip>
+
+    Example:
+
+
+    ```py
+    >>> from collections import OrderedDict
+    >>> import torch.nn as nn
+
+    >>> model = nn.Sequential(OrderedDict([("linear1", nn.Linear(4, 4)), ("linear2", nn.Linear(4, 4))]))
+    >>> model.linear2.weight = test_model.linear1.weight
+    >>> find_tied_parameters(test_model)
+    {'linear1.weight': 'linear2.weight'}
+    ```
+
+    Returns:
+        Dict[str, str]: A dictionary mapping tied parameter names to the name of the parameter they are tied to.
+    """
+    # Initialize result and named_parameters before recursing.
+    named_parameters = kwargs.get("named_parameters", None)
+    prefix = kwargs.get("prefix", "")
+    result = kwargs.get("result", {})
+
+    if named_parameters is None:
+        named_parameters = {n: p for n, p in model.named_parameters()}
+    else:
+        # A tied parameter will not be in the full `named_parameters` seen above but will be in the `named_parameters`
+        # of the submodule it belongs to. So while recursing we track the names that are not in the initial
+        # `named_parameters`.
+        for name, parameter in model.named_parameters():
+            full_name = name if prefix == "" else f"{prefix}.{name}"
+            if full_name not in named_parameters:
+                # When we find one, it has to be one of the existing parameters.
+                for new_name, new_param in named_parameters.items():
+                    if new_param is parameter:
+                        result[new_name] = full_name
+
+    # Once we have treated direct parameters, we move to the child modules.
+    for name, child in model.named_children():
+        child_name = name if prefix == "" else f"{prefix}.{name}"
+        find_tied_parameters(child, named_parameters=named_parameters, prefix=child_name, result=result)
+
+    return result
+
+
+def compute_module_sizes(model: nn.Module):
+    """
+    Compute the size of each submodule of a given model.
+    """
+    module_sizes = defaultdict(int)
+    for name, tensor in named_module_tensors(model, recurse=True):
+        size = tensor.numel() * dtype_byte_size(tensor.dtype)
+        name_parts = name.split(".")
+        for idx in range(len(name_parts) + 1):
+            module_sizes[".".join(name_parts[:idx])] += size
+
+    return module_sizes
+
+
+def get_max_layer_size(
+    modules: List[Tuple[str, torch.nn.Module]], module_sizes: Dict[str, int], no_split_module_classes: List[str]
+):
+    """
+    Utility function that will scan a list of named modules and return the maximum size used by one full layer. The
+    definition of a layer being:
+    - a module with no direct children (just parameters and buffers)
+    - a module whose class name is in the list `no_split_module_classes`
+
+    Args:
+        modules (`List[Tuple[str, torch.nn.Module]]`):
+            The list of named modules where we want to determine the maximum layer size.
+        module_sizes (`Dict[str, int]`):
+            A dictionary mapping each layer name to its size (as generated by `compute_module_sizes`).
+        no_split_module_classes (`List[str]`):
+            A list of class names for layers we don't want to be split.
+
+    Returns:
+        `Tuple[int, List[str]]`: The maximum size of a layer with the list of layer names realizing that maximum size.
+    """
+    max_size = 0
+    layer_names = []
+    modules_to_treat = modules.copy()
+    while len(modules_to_treat) > 0:
+        module_name, module = modules_to_treat.pop(0)
+        modules_children = list(module.named_children())
+        if len(modules_children) == 0 or module.__class__.__name__ in no_split_module_classes:
+            # No splitting this one so we compare to the max_size
+            size = module_sizes[module_name]
+            if size > max_size:
+                max_size = size
+                layer_names = [module_name]
+            elif size == max_size:
+                layer_names.append(module_name)
+        else:
+            modules_to_treat = [(f"{module_name}.{n}", v) for n, v in modules_children] + modules_to_treat
+    return max_size, layer_names
+
+
+def get_max_memory(max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None):
+    """
+    Get the maximum memory available if nothing is passed, converts string to int otherwise.
+    """
+    import psutil
+
+    if max_memory is None:
+        if not torch.cuda.is_available():
+            max_memory = {}
+        else:
+            # Make sure CUDA is initialized on each GPU to have the right memory info.
+            for i in range(torch.cuda.device_count()):
+                _ = torch.tensor([0], device=i)
+            max_memory = {i: torch.cuda.mem_get_info(i)[0] for i in range(torch.cuda.device_count())}
+        max_memory["cpu"] = psutil.virtual_memory().available
+        return max_memory
+
+    for key in max_memory:
+        if isinstance(max_memory[key], str):
+            max_memory[key] = convert_file_size_to_int(max_memory[key])
+    return max_memory
+
+
+def clean_device_map(device_map: Dict[str, Union[int, str, torch.device]], module_name: str = ""):
+    """
+    Cleans a device_map by grouping all submodules that go on the same device together.
+    """
+    # Get the value of the current module and if there is only one split across several keys, regroup it.
+    values = [v for k, v in device_map.items() if k.startswith(module_name)]
+    if len(set(values)) == 1 and len(values) > 1:
+        for k in [k for k in device_map if k.startswith(module_name)]:
+            del device_map[k]
+        device_map[module_name] = values[0]
+
+    # Recurse over the children
+    children_modules = [k for k in device_map.keys() if k.startswith(module_name) and len(k) > len(module_name)]
+    idx = len(module_name.split(".")) + 1 if len(module_name) > 0 else 1
+    children_modules = set(".".join(k.split(".")[:idx]) for k in children_modules)
+    for child in children_modules:
+        clean_device_map(device_map, module_name=child)
+
+    return device_map
+
+
+def infer_auto_device_map(
+    model: nn.Module,
+    max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None,
+    no_split_module_classes: Optional[List[str]] = None,
+):
+    """
+    Compute a device map for a given model giving priority to GPUs, then offload on CPU and finally offload to disk,
+    such that:
+    - we don't exceed the memory available of any of the GPU.
+    - if offload to the CPU is needed, there is always room left on GPU 0 to put back the layer offloaded on CPU that
+      has the largest size.
+    - if offload to the CPU is needed,we don't exceed the RAM available on the CPU.
+    - if offload to the disk is needed, there is always room left on the CPU to put back the layer offloaded on disk
+      that has the largest size.
+
+    <Tip>
+
+    All computation is done analyzing sizes and dtypes of the model parameters. As a result, the model can be on the
+    meta device (as it would if initialized within the `init_empty_weights` context manager).
+
+    </Tip>
+
+    Args:
+        model (`torch.nn.Module`): The model to analyze.
+        max_memory (`Dict`, *optional*):
+            A dictionary device identifier to maximum memory. Will default to the maximum memory available if unset.
+        no_split_module_classes (`List[str]`, *optional*):
+            A list of layer class names that should never be split across device (for instance any layer that has a
+            residual connection).
+    """
+    # Get default / clean up max_memory
+    max_memory = get_max_memory(max_memory)
+    if no_split_module_classes is None:
+        no_split_module_classes = []
+    elif not isinstance(no_split_module_classes, (list, tuple)):
+        no_split_module_classes = [no_split_module_classes]
+
+    devices = list(max_memory.keys()) + ["disk"]
+    # Devices that need to keep space for a potential offloaded layer.
+    main_devices = [devices[0], "cpu"]
+
+    module_sizes = compute_module_sizes(model)
+    tied_parameters = find_tied_parameters(model)
+
+    device_map = {}
+    current_device = 0
+    current_memory_used = 0
+
+    # Direct submodules and parameters
+    modules_to_treat = list(model.named_parameters(recurse=False)) + list(model.named_children())
+    # Initialize maximum largest layer, to know which space to keep in memory
+    max_layer_size, max_layer_names = get_max_layer_size(modules_to_treat, module_sizes, no_split_module_classes)
+
+    # Ready ? This is going to be a bit messy.
+    while len(modules_to_treat) > 0:
+        name, module = modules_to_treat.pop(0)
+        # Max size in the remaining layers may have changed since we took one, so we maybe update it.
+        max_layer_names = [n for n in max_layer_names if not n.startswith(name)]
+        if len(max_layer_names) == 0:
+            max_layer_size, max_layer_names = get_max_layer_size(
+                [(n, m) for n, m in modules_to_treat if isinstance(m, torch.nn.Module)],
+                module_sizes,
+                no_split_module_classes,
+            )
+        # Assess size needed
+        module_size = module_sizes[name]
+        tied_params = [v for k, v in tied_parameters.items() if name in k]
+        # We ignore parameters that are tied when they're tied to > 1 one
+        tied_param = tied_params[0] if len(tied_params) == 1 else None
+
+        device = devices[current_device]
+        current_max_size = max_memory[device] if device != "disk" else None
+        # Reduce max size available by the largest layer.
+        if devices[current_device] in main_devices:
+            current_max_size = current_max_size - max_layer_size
+        # Case 1 -> We're too big!
+        if current_max_size is not None and current_memory_used + module_size > current_max_size:
+            # Split or not split?
+            modules_children = list(module.named_children())
+            if len(modules_children) == 0 or module.__class__.__name__ in no_split_module_classes:
+                # -> no split, we go to the next device
+                current_device += 1
+                modules_to_treat = [(name, module)] + modules_to_treat
+                current_memory_used = 0
+            else:
+                # -> split, we replace the module studied by its children + parameters
+                modules_children = list(module.named_parameters(recurse=False)) + modules_children
+                modules_to_treat = [(f"{name}.{n}", v) for n, v in modules_children] + modules_to_treat
+                # Update the max layer size.
+                max_layer_size, max_layer_names = get_max_layer_size(
+                    [(n, m) for n, m in modules_to_treat if isinstance(m, torch.nn.Module)],
+                    module_sizes,
+                    no_split_module_classes,
+                )
+
+        # Case 2, it fits! We're not entirely out of the wood though, because we may have some tied parameters.
+        elif tied_param is not None:
+            # Determine the sized occupied by this module + the module containing the tied parameter
+            tied_module_size = module_size
+            tied_module_index = [i for i, (n, _) in enumerate(modules_to_treat) if n in tied_param][0]
+            tied_module_name, tied_module = modules_to_treat[tied_module_index]
+            tied_module_size += module_sizes[tied_module_name] - module_sizes[tied_param]
+            if current_max_size is not None and current_memory_used + tied_module_size > current_max_size:
+                # Split or not split?
+                tied_module_children = list(tied_module.named_children())
+                if len(tied_module_children) == 0 or tied_module.__class__.__name__ in no_split_module_classes:
+                    # If the tied module is not split, we go to the next device
+                    current_device += 1
+                    modules_to_treat = [(name, module)] + modules_to_treat
+                    current_memory_used = 0
+                else:
+                    # Otherwise, we replace the tied module by its children.
+                    tied_module_children = list(tied_module.named_parameters(recurse=False)) + tied_module_children
+                    tied_module_children = [(f"{tied_module_name}.{n}", v) for n, v in tied_module_children]
+                    modules_to_treat = (
+                        [(name, module)]
+                        + modules_to_treat[:tied_module_index]
+                        + tied_module_children
+                        + modules_to_treat[tied_module_index + 1 :]
+                    )
+                    # Update the max layer size.
+                    max_layer_size, max_layer_names = get_max_layer_size(
+                        [(n, m) for n, m in modules_to_treat if isinstance(m, torch.nn.Module)],
+                        module_sizes,
+                        no_split_module_classes,
+                    )
+            else:
+                # We really really fit!
+                current_memory_used += tied_module_size
+                device_map[name] = devices[current_device]
+                modules_to_treat.pop(tied_module_index)
+                device_map[tied_module_name] = devices[current_device]
+        else:
+            current_memory_used += module_size
+            device_map[name] = devices[current_device]
+
+    return clean_device_map(device_map)
+
+
+def check_device_map(model: nn.Module, device_map: Dict[str, Union[int, str, torch.device]]):
+    """
+    Checks a device map covers everything in a given model.
+
+    Args:
+        model (`torch.nn.Module`): The model to check the device map against.
+        device_map (`Dict[str, Union[int, str, torch.device]]`): The device map to check.
+    """
+    all_model_tensors = [name for name, _ in model.state_dict().items()]
+    for module_name in device_map.keys():
+        all_model_tensors = [name for name in all_model_tensors if not name.startswith(module_name)]
+    if len(all_model_tensors) > 0:
+        non_covered_params = ", ".join(all_model_tensors)
+        raise ValueError(
+            f"The device_map provided does not give any device for the following parameters: {non_covered_params}"
+        )
+
+
+def load_checkpoint_in_model(
+    model: nn.Module,
+    checkpoint: Union[str, os.PathLike],
+    device_map: Optional[Dict[str, Union[int, str, torch.device]]] = None,
+    offload_folder: Optional[Union[str, os.PathLike]] = None,
+):
+    """
+    Loads a (potentially sharded) checkpoint inside a model, potentially sending weights to a given device as they are
+    loaded.
+
+    <Tip warning={true}>
+
+    Once loaded across devices, you still need to call [`dispatch_model`] on your model to make it able to run. To
+    group the checkpoint loading and dispatch in one single call, use [`load_checkpoint_and_dispatch`].
+
+    </Tip>
+
+    Args:
+        model (`torch.nn.Module`): The model in which we want to load a checkpoint.
+        checkpoint (`str` or `os.PathLike`):
+            The folder checkpoint to load. It can be:
+            - a path to a file containing a whole model state dict
+            - a path to a `.json` file containing the index to a sharded checkpoint
+            - a path to a folder containing a unique `.index.json` file and the shards of a checkpoint.
+        device_map (`Dict[str, Union[int, str, torch.device]]`, *optional*):
+            A map that specifies where each submodule should go. It doesn't need to be refined to each parameter/buffer
+            name, once a given module name is inside, every submodule of it will be sent to the same device.
+        offload_folder (`str` or `os.PathLike`, *optional*):
+            If the `device_map` contains any value `"disk"`, the folder where we will offload weights.
+    """
+    if offload_folder is None and device_map is not None and "disk" in device_map.values():
+        raise ValueError(
+            "At least one of the model submodule will be offloaded to disk, please pass along an `offload_folder`."
+        )
+
+    checkpoint_files = None
+    index_filename = None
+    if os.path.isfile(checkpoint):
+        if str(checkpoint).endswith(".json"):
+            index_filename = checkpoint
+        else:
+            checkpoint_files = [checkpoint]
+    elif os.path.isdir(checkpoint):
+        potential_index = [f for f in os.listdir(checkpoint) if f.endswith(".index.json")]
+        if len(potential_index) == 0:
+            raise ValueError(f"{checkpoint} is not a folder containing a `.index.json` file.")
+        elif len(potential_index) == 1:
+            index_filename = os.path.join(checkpoint, potential_index[0])
+        else:
+            raise ValueError(f"{checkpoint} containing mote than one `.index.json` file, delete the irrelevant ones.")
+    else:
+        raise ValueError(
+            "`checkpoint` should be the path to a file containing a whole state dict, or the index of a sharded "
+            f"checkpoint, or a folder containing a sharded checkpoint, but got {checkpoint}."
+        )
+
+    if index_filename is not None:
+        checkpoint_folder = os.path.split(index_filename)[0]
+        with open(index_filename, "r") as f:
+            index = json.loads(f.read())
+
+        if "weight_map" in index:
+            index = index["weight_map"]
+        checkpoint_files = sorted(list(set(index.values())))
+        checkpoint_files = [os.path.join(checkpoint_folder, f) for f in checkpoint_files]
+
+    # Logic for missing/unexepected keys goes here.
+
+    offload_index = {}
+    for checkpoint_file in checkpoint_files:
+        checkpoint = torch.load(checkpoint_file)
This shouldn't rely on the order (we don't care if the "main" GPU is 0, 1, etc.) we just need to make sure we don't end up with `"cpu"` twice :-)
@@ -0,0 +1,568 @@
+# 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 gc
+import json
+import os
+import re
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+WEIGHTS_INDEX_NAME = "pytorch_model.bin.index.json"
+
+
+def convert_file_size_to_int(size: Union[int, str]):
+    """
+    Converts a size expressed as a string with digits an unit (like `"5MB"`) to an integer (in bytes).
+
+    Args:
+        size (`int` or `str`): The size to convert. Will be directly returned if an `int`.
+
+    Example:
+
+    ```py
+    >>> convert_file_size_to_int("1MiB")
+    1048576
+    ```
+    """
+    if isinstance(size, int):
+        return size
+    if size.upper().endswith("GIB"):
+        return int(size[:-3]) * (2**30)
+    if size.upper().endswith("MIB"):
+        return int(size[:-3]) * (2**20)
+    if size.upper().endswith("KIB"):
+        return int(size[:-3]) * (2**10)
+    if size.upper().endswith("GB"):
+        int_size = int(size[:-2]) * (10**9)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("MB"):
+        int_size = int(size[:-2]) * (10**6)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("KB"):
+        int_size = int(size[:-2]) * (10**3)
+        return int_size // 8 if size.endswith("b") else int_size
+    raise ValueError("`size` is not in a valid format. Use an integer followed by the unit, e.g., '5GB'.")
+
+
+def dtype_byte_size(dtype: torch.dtype):
+    """
+    Returns the size (in bytes) occupied by one parameter of type `dtype`.
+
+    Example:
+
+    ```py
+    >>> dtype_byte_size(torch.float32)
+    4
+    ```
+    """
+    if dtype == torch.bool:
+        return 1 / 8
+    bit_search = re.search("[^\d](\d+)$", str(dtype))
+    if bit_search is None:
+        raise ValueError(f"`dtype` is not a valid dtype: {dtype}.")
+    bit_size = int(bit_search.groups()[0])
+    return bit_size // 8
+
+
+def set_module_tensor_to_device(
+    module: nn.Module, tensor_name: str, device: Union[int, str, torch.device], value: Optional[torch.Tensor] = None
+):
+    """
+    A helper function to set a given tensor (parameter of buffer) of a module on a specific device (note that doing
+    `param.to(device)` creates a new tensor not linked to the parameter, which is why we need this function).
+
+    Args:
+        module (`torch.nn.Module`): The module in which the tensor we want to move lives.
+        param_name (`str`): The full name of the parameter/buffer.
+        device (`int`, `str` or `torch.device`): The device on which to set the tensor.
+        value (`torch.Tensor`, *optional*): The value of the tensor (useful when going from the meta device to any
+            other device).
+    """
+    # Recurse if needed
+    if "." in tensor_name:
+        splits = tensor_name.split(".")
+        for split in splits[:-1]:
+            new_module = getattr(module, split)
+            if new_module is None:
+                raise ValueError(f"{module} has no attribute {split}.")
+            module = new_module
+        tensor_name = splits[-1]
+
+    if tensor_name not in module._parameters and tensor_name not in module._buffers:
+        raise ValueError(f"{module} does not have a parameter or a buffer named {tensor_name}.")
+    is_buffer = tensor_name in module._buffers
+    old_value = getattr(module, tensor_name)
+
+    if old_value.device == torch.device("meta") and device != torch.device("meta") and value is None:
+        raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {device}.")
+
+    with torch.no_grad():
+        if value is None:
+            new_value = old_value.to(device)
+        elif isinstance(value, torch.Tensor):
+            new_value = value.to(device)
+        else:
+            new_value = torch.tensor(value, device=device)
+    if is_buffer:
+        module._buffers[tensor_name] = new_value
+    else:
+        new_value = nn.Parameter(new_value, requires_grad=old_value.requires_grad)
+        module._parameters[tensor_name] = new_value
+
+
+def named_module_tensors(module: nn.Module, include_buffers: bool = True, recurse: bool = False):
+    """
+    A helper function that gathers all the tensors (parameters + buffers) of a given module. If `include_buffers=True`
+    it's the same as doing `module.named_parameters(recurse=recurse) + module.named_buffers(recurse=recurse)`.
+
+    Args:
+        module (`torch.nn.Module`): The module we want the tensors or.
+        include_buffer (`bool`, *optional*, defaults to `True`): Whether or not to include the buffers in the result.
+        recurse (`bool`, *optional`, defaults to `False`):
+            Whether or not to go look in every submodule or just return the direct parameters and buffers.
+    """
+    for named_parameter in module.named_parameters(recurse=recurse):
+        yield named_parameter
+
+    if include_buffers:
+        for named_buffer in module.named_buffers(recurse=recurse):
+            yield named_buffer
+
+
+def find_tied_parameters(model: nn.Module, **kwargs):
+    """
+    Find the tied parameters in a given model.
+
+    Args:
+        model (`torch.nn.Module`): The model to inspect.
+
+    <Tip warning={true}>
+
+    The signature accepts keyword arguments, but they are for the recursive part of this function and you should ignore
+    them.
+
+    </Tip>
+
+    Example:
+
+
+    ```py
+    >>> from collections import OrderedDict
+    >>> import torch.nn as nn
+
+    >>> model = nn.Sequential(OrderedDict([("linear1", nn.Linear(4, 4)), ("linear2", nn.Linear(4, 4))]))
+    >>> model.linear2.weight = test_model.linear1.weight
+    >>> find_tied_parameters(test_model)
+    {'linear1.weight': 'linear2.weight'}
+    ```
+
+    Returns:
+        Dict[str, str]: A dictionary mapping tied parameter names to the name of the parameter they are tied to.
+    """
+    # Initialize result and named_parameters before recursing.
+    named_parameters = kwargs.get("named_parameters", None)
+    prefix = kwargs.get("prefix", "")
+    result = kwargs.get("result", {})
+
+    if named_parameters is None:
+        named_parameters = {n: p for n, p in model.named_parameters()}
+    else:
+        # A tied parameter will not be in the full `named_parameters` seen above but will be in the `named_parameters`
+        # of the submodule it belongs to. So while recursing we track the names that are not in the initial
+        # `named_parameters`.
+        for name, parameter in model.named_parameters():
+            full_name = name if prefix == "" else f"{prefix}.{name}"
+            if full_name not in named_parameters:
+                # When we find one, it has to be one of the existing parameters.
+                for new_name, new_param in named_parameters.items():
+                    if new_param is parameter:
+                        result[new_name] = full_name
+
+    # Once we have treated direct parameters, we move to the child modules.
+    for name, child in model.named_children():
+        child_name = name if prefix == "" else f"{prefix}.{name}"
+        find_tied_parameters(child, named_parameters=named_parameters, prefix=child_name, result=result)
+
+    return result
+
+
+def compute_module_sizes(model: nn.Module):
+    """
+    Compute the size of each submodule of a given model.
+    """
+    module_sizes = defaultdict(int)
+    for name, tensor in named_module_tensors(model, recurse=True):
+        size = tensor.numel() * dtype_byte_size(tensor.dtype)
+        name_parts = name.split(".")
+        for idx in range(len(name_parts) + 1):
+            module_sizes[".".join(name_parts[:idx])] += size
+
+    return module_sizes
+
+
+def get_max_layer_size(
+    modules: List[Tuple[str, torch.nn.Module]], module_sizes: Dict[str, int], no_split_module_classes: List[str]
+):
+    """
+    Utility function that will scan a list of named modules and return the maximum size used by one full layer. The
+    definition of a layer being:
+    - a module with no direct children (just parameters and buffers)
+    - a module whose class name is in the list `no_split_module_classes`
+
+    Args:
+        modules (`List[Tuple[str, torch.nn.Module]]`):
+            The list of named modules where we want to determine the maximum layer size.
+        module_sizes (`Dict[str, int]`):
+            A dictionary mapping each layer name to its size (as generated by `compute_module_sizes`).
+        no_split_module_classes (`List[str]`):
+            A list of class names for layers we don't want to be split.
+
+    Returns:
+        `Tuple[int, List[str]]`: The maximum size of a layer with the list of layer names realizing that maximum size.
+    """
+    max_size = 0
+    layer_names = []
+    modules_to_treat = modules.copy()
+    while len(modules_to_treat) > 0:
+        module_name, module = modules_to_treat.pop(0)
+        modules_children = list(module.named_children())
+        if len(modules_children) == 0 or module.__class__.__name__ in no_split_module_classes:
+            # No splitting this one so we compare to the max_size
+            size = module_sizes[module_name]
+            if size > max_size:
+                max_size = size
+                layer_names = [module_name]
+            elif size == max_size:
+                layer_names.append(module_name)
+        else:
+            modules_to_treat = [(f"{module_name}.{n}", v) for n, v in modules_children] + modules_to_treat
+    return max_size, layer_names
+
+
+def get_max_memory(max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None):
+    """
+    Get the maximum memory available if nothing is passed, converts string to int otherwise.
+    """
+    import psutil
+
+    if max_memory is None:
+        if not torch.cuda.is_available():
+            max_memory = {}
+        else:
+            # Make sure CUDA is initialized on each GPU to have the right memory info.
+            for i in range(torch.cuda.device_count()):
+                _ = torch.tensor([0], device=i)
+            max_memory = {i: torch.cuda.mem_get_info(i)[0] for i in range(torch.cuda.device_count())}
+        max_memory["cpu"] = psutil.virtual_memory().available
+        return max_memory
+
+    for key in max_memory:
+        if isinstance(max_memory[key], str):
+            max_memory[key] = convert_file_size_to_int(max_memory[key])
+    return max_memory
+
+
+def clean_device_map(device_map: Dict[str, Union[int, str, torch.device]], module_name: str = ""):
+    """
+    Cleans a device_map by grouping all submodules that go on the same device together.
+    """
+    # Get the value of the current module and if there is only one split across several keys, regroup it.
+    values = [v for k, v in device_map.items() if k.startswith(module_name)]
+    if len(set(values)) == 1 and len(values) > 1:
+        for k in [k for k in device_map if k.startswith(module_name)]:
+            del device_map[k]
+        device_map[module_name] = values[0]
+
+    # Recurse over the children
+    children_modules = [k for k in device_map.keys() if k.startswith(module_name) and len(k) > len(module_name)]
+    idx = len(module_name.split(".")) + 1 if len(module_name) > 0 else 1
+    children_modules = set(".".join(k.split(".")[:idx]) for k in children_modules)
+    for child in children_modules:
+        clean_device_map(device_map, module_name=child)
+
+    return device_map
+
+
+def infer_auto_device_map(
+    model: nn.Module,
+    max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None,
+    no_split_module_classes: Optional[List[str]] = None,
+):
+    """
+    Compute a device map for a given model giving priority to GPUs, then offload on CPU and finally offload to disk,
+    such that:
+    - we don't exceed the memory available of any of the GPU.
+    - if offload to the CPU is needed, there is always room left on GPU 0 to put back the layer offloaded on CPU that
+      has the largest size.
+    - if offload to the CPU is needed,we don't exceed the RAM available on the CPU.
+    - if offload to the disk is needed, there is always room left on the CPU to put back the layer offloaded on disk
+      that has the largest size.
+
+    <Tip>
+
+    All computation is done analyzing sizes and dtypes of the model parameters. As a result, the model can be on the
+    meta device (as it would if initialized within the `init_empty_weights` context manager).
+
+    </Tip>
+
+    Args:
+        model (`torch.nn.Module`): The model to analyze.
+        max_memory (`Dict`, *optional*):
+            A dictionary device identifier to maximum memory. Will default to the maximum memory available if unset.
+        no_split_module_classes (`List[str]`, *optional*):
+            A list of layer class names that should never be split across device (for instance any layer that has a
+            residual connection).
+    """
+    # Get default / clean up max_memory
+    max_memory = get_max_memory(max_memory)
+    if no_split_module_classes is None:
+        no_split_module_classes = []
+    elif not isinstance(no_split_module_classes, (list, tuple)):
+        no_split_module_classes = [no_split_module_classes]
+
+    devices = list(max_memory.keys()) + ["disk"]
+    # Devices that need to keep space for a potential offloaded layer.
+    main_devices = [devices[0], "cpu"]
Fixed
@@ -0,0 +1,568 @@
+# 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 gc
+import json
+import os
+import re
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+WEIGHTS_INDEX_NAME = "pytorch_model.bin.index.json"
+
+
+def convert_file_size_to_int(size: Union[int, str]):
+    """
+    Converts a size expressed as a string with digits an unit (like `"5MB"`) to an integer (in bytes).
+
+    Args:
+        size (`int` or `str`): The size to convert. Will be directly returned if an `int`.
+
+    Example:
+
+    ```py
+    >>> convert_file_size_to_int("1MiB")
+    1048576
+    ```
+    """
+    if isinstance(size, int):
+        return size
+    if size.upper().endswith("GIB"):
+        return int(size[:-3]) * (2**30)
+    if size.upper().endswith("MIB"):
+        return int(size[:-3]) * (2**20)
+    if size.upper().endswith("KIB"):
+        return int(size[:-3]) * (2**10)
+    if size.upper().endswith("GB"):
+        int_size = int(size[:-2]) * (10**9)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("MB"):
+        int_size = int(size[:-2]) * (10**6)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("KB"):
+        int_size = int(size[:-2]) * (10**3)
+        return int_size // 8 if size.endswith("b") else int_size
+    raise ValueError("`size` is not in a valid format. Use an integer followed by the unit, e.g., '5GB'.")
+
+
+def dtype_byte_size(dtype: torch.dtype):
+    """
+    Returns the size (in bytes) occupied by one parameter of type `dtype`.
+
+    Example:
+
+    ```py
+    >>> dtype_byte_size(torch.float32)
+    4
+    ```
+    """
+    if dtype == torch.bool:
+        return 1 / 8
+    bit_search = re.search("[^\d](\d+)$", str(dtype))
+    if bit_search is None:
+        raise ValueError(f"`dtype` is not a valid dtype: {dtype}.")
+    bit_size = int(bit_search.groups()[0])
+    return bit_size // 8
+
+
+def set_module_tensor_to_device(
+    module: nn.Module, tensor_name: str, device: Union[int, str, torch.device], value: Optional[torch.Tensor] = None
+):
+    """
+    A helper function to set a given tensor (parameter of buffer) of a module on a specific device (note that doing
+    `param.to(device)` creates a new tensor not linked to the parameter, which is why we need this function).
+
+    Args:
+        module (`torch.nn.Module`): The module in which the tensor we want to move lives.
+        param_name (`str`): The full name of the parameter/buffer.
+        device (`int`, `str` or `torch.device`): The device on which to set the tensor.
+        value (`torch.Tensor`, *optional*): The value of the tensor (useful when going from the meta device to any
+            other device).
+    """
+    # Recurse if needed
+    if "." in tensor_name:
+        splits = tensor_name.split(".")
+        for split in splits[:-1]:
+            new_module = getattr(module, split)
+            if new_module is None:
+                raise ValueError(f"{module} has no attribute {split}.")
+            module = new_module
+        tensor_name = splits[-1]
+
+    if tensor_name not in module._parameters and tensor_name not in module._buffers:
+        raise ValueError(f"{module} does not have a parameter or a buffer named {tensor_name}.")
+    is_buffer = tensor_name in module._buffers
+    old_value = getattr(module, tensor_name)
+
+    if old_value.device == torch.device("meta") and device != torch.device("meta") and value is None:
+        raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {device}.")
+
+    with torch.no_grad():
+        if value is None:
+            new_value = old_value.to(device)
+        elif isinstance(value, torch.Tensor):
+            new_value = value.to(device)
+        else:
+            new_value = torch.tensor(value, device=device)
+    if is_buffer:
+        module._buffers[tensor_name] = new_value
+    else:
+        new_value = nn.Parameter(new_value, requires_grad=old_value.requires_grad)
+        module._parameters[tensor_name] = new_value
+
+
+def named_module_tensors(module: nn.Module, include_buffers: bool = True, recurse: bool = False):
+    """
+    A helper function that gathers all the tensors (parameters + buffers) of a given module. If `include_buffers=True`
+    it's the same as doing `module.named_parameters(recurse=recurse) + module.named_buffers(recurse=recurse)`.
+
+    Args:
+        module (`torch.nn.Module`): The module we want the tensors or.
+        include_buffer (`bool`, *optional*, defaults to `True`): Whether or not to include the buffers in the result.
+        recurse (`bool`, *optional`, defaults to `False`):
+            Whether or not to go look in every submodule or just return the direct parameters and buffers.
+    """
+    for named_parameter in module.named_parameters(recurse=recurse):
+        yield named_parameter
+
+    if include_buffers:
+        for named_buffer in module.named_buffers(recurse=recurse):
+            yield named_buffer
+
+
+def find_tied_parameters(model: nn.Module, **kwargs):
+    """
+    Find the tied parameters in a given model.
+
+    Args:
+        model (`torch.nn.Module`): The model to inspect.
+
+    <Tip warning={true}>
+
+    The signature accepts keyword arguments, but they are for the recursive part of this function and you should ignore
+    them.
+
+    </Tip>
+
+    Example:
+
+
+    ```py
+    >>> from collections import OrderedDict
+    >>> import torch.nn as nn
+
+    >>> model = nn.Sequential(OrderedDict([("linear1", nn.Linear(4, 4)), ("linear2", nn.Linear(4, 4))]))
+    >>> model.linear2.weight = test_model.linear1.weight
+    >>> find_tied_parameters(test_model)
+    {'linear1.weight': 'linear2.weight'}
+    ```
+
+    Returns:
+        Dict[str, str]: A dictionary mapping tied parameter names to the name of the parameter they are tied to.
+    """
+    # Initialize result and named_parameters before recursing.
+    named_parameters = kwargs.get("named_parameters", None)
+    prefix = kwargs.get("prefix", "")
+    result = kwargs.get("result", {})
+
+    if named_parameters is None:
+        named_parameters = {n: p for n, p in model.named_parameters()}
+    else:
+        # A tied parameter will not be in the full `named_parameters` seen above but will be in the `named_parameters`
+        # of the submodule it belongs to. So while recursing we track the names that are not in the initial
+        # `named_parameters`.
+        for name, parameter in model.named_parameters():
+            full_name = name if prefix == "" else f"{prefix}.{name}"
+            if full_name not in named_parameters:
+                # When we find one, it has to be one of the existing parameters.
+                for new_name, new_param in named_parameters.items():
+                    if new_param is parameter:
+                        result[new_name] = full_name
+
+    # Once we have treated direct parameters, we move to the child modules.
+    for name, child in model.named_children():
+        child_name = name if prefix == "" else f"{prefix}.{name}"
+        find_tied_parameters(child, named_parameters=named_parameters, prefix=child_name, result=result)
+
+    return result
+
+
+def compute_module_sizes(model: nn.Module):
+    """
+    Compute the size of each submodule of a given model.
+    """
+    module_sizes = defaultdict(int)
+    for name, tensor in named_module_tensors(model, recurse=True):
+        size = tensor.numel() * dtype_byte_size(tensor.dtype)
+        name_parts = name.split(".")
+        for idx in range(len(name_parts) + 1):
+            module_sizes[".".join(name_parts[:idx])] += size
+
+    return module_sizes
+
+
+def get_max_layer_size(
+    modules: List[Tuple[str, torch.nn.Module]], module_sizes: Dict[str, int], no_split_module_classes: List[str]
+):
+    """
+    Utility function that will scan a list of named modules and return the maximum size used by one full layer. The
+    definition of a layer being:
+    - a module with no direct children (just parameters and buffers)
+    - a module whose class name is in the list `no_split_module_classes`
+
+    Args:
+        modules (`List[Tuple[str, torch.nn.Module]]`):
+            The list of named modules where we want to determine the maximum layer size.
+        module_sizes (`Dict[str, int]`):
+            A dictionary mapping each layer name to its size (as generated by `compute_module_sizes`).
+        no_split_module_classes (`List[str]`):
+            A list of class names for layers we don't want to be split.
+
+    Returns:
+        `Tuple[int, List[str]]`: The maximum size of a layer with the list of layer names realizing that maximum size.
+    """
+    max_size = 0
+    layer_names = []
+    modules_to_treat = modules.copy()
+    while len(modules_to_treat) > 0:
+        module_name, module = modules_to_treat.pop(0)
+        modules_children = list(module.named_children())
+        if len(modules_children) == 0 or module.__class__.__name__ in no_split_module_classes:
+            # No splitting this one so we compare to the max_size
+            size = module_sizes[module_name]
+            if size > max_size:
+                max_size = size
+                layer_names = [module_name]
+            elif size == max_size:
+                layer_names.append(module_name)
+        else:
+            modules_to_treat = [(f"{module_name}.{n}", v) for n, v in modules_children] + modules_to_treat
+    return max_size, layer_names
+
+
+def get_max_memory(max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None):
+    """
+    Get the maximum memory available if nothing is passed, converts string to int otherwise.
+    """
+    import psutil
+
+    if max_memory is None:
+        if not torch.cuda.is_available():
+            max_memory = {}
+        else:
+            # Make sure CUDA is initialized on each GPU to have the right memory info.
+            for i in range(torch.cuda.device_count()):
+                _ = torch.tensor([0], device=i)
+            max_memory = {i: torch.cuda.mem_get_info(i)[0] for i in range(torch.cuda.device_count())}
+        max_memory["cpu"] = psutil.virtual_memory().available
+        return max_memory
+
+    for key in max_memory:
+        if isinstance(max_memory[key], str):
+            max_memory[key] = convert_file_size_to_int(max_memory[key])
+    return max_memory
+
+
+def clean_device_map(device_map: Dict[str, Union[int, str, torch.device]], module_name: str = ""):
+    """
+    Cleans a device_map by grouping all submodules that go on the same device together.
+    """
+    # Get the value of the current module and if there is only one split across several keys, regroup it.
+    values = [v for k, v in device_map.items() if k.startswith(module_name)]
+    if len(set(values)) == 1 and len(values) > 1:
+        for k in [k for k in device_map if k.startswith(module_name)]:
+            del device_map[k]
+        device_map[module_name] = values[0]
_The documentation is not available anymore as the PR was closed or merged._
Super cool PR! Two questions:

1. How fast is generate with parts of the weights offloaded compared to not offloaded? Guess it strongly depends on the hardware, but e.g. on a V100, for let's say `gpt2`?

2. More generally, do you think we could eventually do the following code:

```py
checkpoint = "bigscience/T0pp"
config = AutoConfig.from_pretrained(checkpoint)

# Initializes an empty shell with the model. This is instant and does not take any RAM.
with init_empty_weights():
    model = AutoModelForSeq2SeqLM.from_config(config)
# Initialize the model under the previous context manager breaks the tied weights.
model.tie_weights()
# Determine a device map that maximizes the available GPUs, you can also write your own.
# If using a different model, adjust `T5Block` to the proper class (for instance `"GPTJBlock"`)
device_map = infer_auto_device_map(model, no_split_module_classes=["T5Block"])
# Load the sharded checkpoint inside the model. This will load each part of the model on the device specified by `device_map`
load_sharded_checkpoint_in_model(model, path_to_local_t0pp, device_map=device_map)
# This will make that model that leaves on several different devices just work.
model = dispatch_model(model, device_map)
```

fully in Tranformers' `from_pretrained(...)` method with an optional flag? E.g. from a user's perspective wouldn't it be very cool to be able to just do something liike:

```py
from transformers import AutoModelForSeq2SeqLM

model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint, device_map="auto")
```

would this be desirable in a future PR?
I'll post more benchmarks toward the end but basically:
- having the model on 2 GPUs instead of one for generate does not change speed. (I do have an NVlink though, so maybe it could slower, though there is only one change of GPU per forward pass, so it's not much)
- offloading to CPU his 10-15x slower than GPU (for the part offloaded, which is why it's import to maximize GPU use before using CPU offload)
- offloading to disk is 20-30x slower than GPU.
> I'll post more benchmarks toward the end but basically:
> 
> * having the model on 2 GPUs instead of one for generate does not change speed. (I do have an NVlink though, so maybe it could slower, though there is only one change of GPU per forward pass, so it's not much)
> * offloading to CPU his 10-15x slower than GPU (for the part offloaded, which is why it's import to maximize GPU use before using CPU offload)
> * offloading to disk is 20-30x slower than GPU.

It would be interesting to compare to [deepseed inference](https://www.deepspeed.ai/tutorials/inference-tutorial/). 

I doubt it will be faster @philschmid but it won't be a black box that replaces your model.
Could support be added for TPU inference as well?
> Could support be added for TPU inference as well?

None of the things introduced here are efficient on TPUs: the model parallelism is vertical (a few layers are executed on GPU 0, then the next ones on GPU 1) which is not efficient on TPU which are more designed for horizontal parallelism, and the CPU offload would be excruciatingly slow on TPU, since the communication CPU/TPU is a huge bottleneck there.
https://api.github.com/repos/huggingface/accelerate/pulls/345
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index 559fda41c..8ff69e048 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -132,12 +132,20 @@ def get_cluster_input():
     else:
         main_training_function = "main"
 
-    num_processes = _ask_field(
-        "How many devices should be used for (potentially) distributed training? [1]:",
-        lambda x: int(x),
-        default=1,
-        error_message="Please enter an integer.",
-    )
+    if distributed_type in [DistributedType.MULTI_CPU, DistributedType.MULTI_GPU, DistributedType.TPU]:
+        machine_type = str(distributed_type).split(".")[0].replace("MULTI_", "")
+        if machine_type == "TPU":
+            machine_type += " cores"
+        else:
+            machine_type += "'s"
+        num_processes = _ask_field(
+            f"How many {machine_type} should be used for distributed training? [1]:",
+            lambda x: int(x),
+            default=1,
+            error_message="Please enter an integer.",
+        )
+    else:
+        num_processes = 1
 
     if distributed_type != DistributedType.TPU:
         mixed_precision = _ask_field(

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/344
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index cfab0f7f8..559fda41c 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -133,7 +133,7 @@ def get_cluster_input():
         main_training_function = "main"
 
     num_processes = _ask_field(
-        "How many processes in total will you use? [1]: ",
+        "How many devices should be used for (potentially) distributed training? [1]:",
         lambda x: int(x),
         default=1,
         error_message="Please enter an integer.",

```suggestion
        "How many devices should be used for (potentially) distributed training? [1]:",
```
@@ -133,7 +133,7 @@ def get_cluster_input():
         main_training_function = "main"
 
     num_processes = _ask_field(
-        "How many processes in total will you use? [1]: ",
+        "How many devices should be used for (potentially) distributed training? [1:]",
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_343). All of your documentation changes will be reflected on that endpoint.
https://api.github.com/repos/huggingface/accelerate/pulls/343
diff --git a/examples/by_feature/fsdp_with_peak_mem_tracking.py b/examples/by_feature/fsdp_with_peak_mem_tracking.py
index e17e9ff67..e39d3ede8 100644
--- a/examples/by_feature/fsdp_with_peak_mem_tracking.py
+++ b/examples/by_feature/fsdp_with_peak_mem_tracking.py
@@ -251,12 +251,13 @@ def collate_fn(examples):
             )
         )
         # Logging the peak memory usage of the GPU to the tracker
-        accelerator.log(
-            {
-                "train_total_peak_memory": tracemalloc.peaked + b2mb(tracemalloc.begin),
-            },
-            step=epoch,
-        )
+        if args.with_tracking:
+            accelerator.log(
+                {
+                    "train_total_peak_memory": tracemalloc.peaked + b2mb(tracemalloc.begin),
+                },
+                step=epoch,
+            )
 
         # New Code #
         # context manager to track the peak memory usage during the evaluation
@@ -302,12 +303,13 @@ def collate_fn(examples):
             "Total Peak Memory consumed during the eval (max): {}".format(tracemalloc.peaked + b2mb(tracemalloc.begin))
         )
         # Logging the peak memory usage of the GPU to the tracker
-        accelerator.log(
-            {
-                "eval_total_peak_memory": tracemalloc.peaked + b2mb(tracemalloc.begin),
-            },
-            step=epoch,
-        )
+        if args.with_tracking:
+            accelerator.log(
+                {
+                    "eval_total_peak_memory": tracemalloc.peaked + b2mb(tracemalloc.begin),
+                },
+                step=epoch,
+            )
 
     if args.with_tracking:
         accelerator.end_training()
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index fa4496f91..ba34587e2 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -465,17 +465,20 @@ def prepare_model(self, model):
         elif self.distributed_type == DistributedType.FSDP:
             from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP
 
-            fsdp_plugin = self.state.fsdp_plugin
-            model = FSDP(
-                model,
-                sharding_strategy=fsdp_plugin.sharding_strategy,
-                cpu_offload=fsdp_plugin.cpu_offload,
-                auto_wrap_policy=fsdp_plugin.auto_wrap_policy,
-                backward_prefetch=fsdp_plugin.backward_prefetch,
-                ignored_modules=fsdp_plugin.ignored_modules,
-            )
-            if not fsdp_plugin.cpu_offload.offload_params:
-                model.to(self.device)
+            # Check if the model is already a FSDP model due to `Manual Wrapping` and if so,
+            # don't wrap it again
+            if type(model) != FSDP:
+                fsdp_plugin = self.state.fsdp_plugin
+                model = FSDP(
+                    model,
+                    sharding_strategy=fsdp_plugin.sharding_strategy,
+                    cpu_offload=fsdp_plugin.cpu_offload,
+                    auto_wrap_policy=fsdp_plugin.auto_wrap_policy,
+                    backward_prefetch=fsdp_plugin.backward_prefetch,
+                    ignored_modules=fsdp_plugin.ignored_modules,
+                )
+                if not fsdp_plugin.cpu_offload.offload_params:
+                    model.to(self.device)
         elif self.distributed_type == DistributedType.MULTI_CPU:
             kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {}
             model = torch.nn.parallel.DistributedDataParallel(model, **kwargs)

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/342
diff --git a/docs/source/memory.mdx b/docs/source/memory.mdx
index 31a2bc66f..e578e43d7 100644
--- a/docs/source/memory.mdx
+++ b/docs/source/memory.mdx
@@ -48,4 +48,4 @@ def training_function(args):
 +   inner_training_loop()
 ```
 
-[[autodoc]] memory_utils.find_executable_batch_size
\ No newline at end of file
+[[autodoc]] utils.find_executable_batch_size
\ No newline at end of file
diff --git a/examples/by_feature/memory.py b/examples/by_feature/memory.py
index 67263de52..6f3f3e0cb 100644
--- a/examples/by_feature/memory.py
+++ b/examples/by_feature/memory.py
@@ -19,7 +19,7 @@
 from accelerate import Accelerator, DistributedType
 
 # New Code #
-from accelerate.memory_utils import find_executable_batch_size
+from accelerate.utils import find_executable_batch_size
 from datasets import load_dataset, load_metric
 from transformers import (
     AdamW,
diff --git a/src/accelerate/__init__.py b/src/accelerate/__init__.py
index 1589423be..b74c683b4 100644
--- a/src/accelerate/__init__.py
+++ b/src/accelerate/__init__.py
@@ -5,7 +5,13 @@
 __version__ = "0.8.0.dev0"
 
 from .accelerator import Accelerator
-from .kwargs_handlers import DistributedDataParallelKwargs, GradScalerKwargs, InitProcessGroupKwargs
 from .launchers import debug_launcher, notebook_launcher
-from .state import DistributedType
-from .utils import DeepSpeedPlugin, synchronize_rng_states
+from .utils import (
+    DeepSpeedPlugin,
+    DistributedDataParallelKwargs,
+    DistributedType,
+    GradScalerKwargs,
+    InitProcessGroupKwargs,
+    find_executable_batch_size,
+    synchronize_rng_states,
+)
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 7f1480436..fa4496f91 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -25,15 +25,19 @@
 
 from .checkpointing import load_accelerator_state, load_custom_state, save_accelerator_state, save_custom_state
 from .data_loader import prepare_data_loader
-from .kwargs_handlers import DistributedDataParallelKwargs, GradScalerKwargs, InitProcessGroupKwargs, KwargsHandler
 from .logging import get_logger
 from .optimizer import AcceleratedOptimizer
 from .scheduler import AcceleratedScheduler
-from .state import AcceleratorState, DistributedType, is_deepspeed_available
+from .state import AcceleratorState
 from .tracking import LOGGER_TYPE_TO_CLASS, GeneralTracker, filter_trackers
 from .utils import (
     DeepSpeedPlugin,
+    DistributedDataParallelKwargs,
+    DistributedType,
     FullyShardedDataParallelPlugin,
+    GradScalerKwargs,
+    InitProcessGroupKwargs,
+    KwargsHandler,
     LoggerType,
     PrecisionType,
     RNGType,
@@ -41,6 +45,7 @@
     extract_model_from_parallel,
     gather,
     get_pretty_name,
+    is_deepspeed_available,
     pad_across_processes,
     reduce,
     save,
@@ -51,7 +56,7 @@
 if is_deepspeed_available():
     import deepspeed
 
-    from .deepspeed_utils import DeepSpeedEngineWrapper, DeepSpeedOptimizerWrapper
+    from .utils import DeepSpeedEngineWrapper, DeepSpeedOptimizerWrapper
 
 logger = get_logger(__name__)
 
diff --git a/src/accelerate/checkpointing.py b/src/accelerate/checkpointing.py
index a94067144..b0e2c04a2 100644
--- a/src/accelerate/checkpointing.py
+++ b/src/accelerate/checkpointing.py
@@ -21,8 +21,16 @@
 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, SCHEDULER_NAME, get_pretty_name, save
+from .utils import (
+    MODEL_NAME,
+    OPTIMIZER_NAME,
+    RNG_STATE_NAME,
+    SCALER_NAME,
+    SCHEDULER_NAME,
+    get_pretty_name,
+    is_tpu_available,
+    save,
+)
 
 
 if is_tpu_available():
diff --git a/src/accelerate/commands/config/__init__.py b/src/accelerate/commands/config/__init__.py
index 5bafb3be2..11719835a 100644
--- a/src/accelerate/commands/config/__init__.py
+++ b/src/accelerate/commands/config/__init__.py
@@ -17,7 +17,7 @@
 import argparse
 import os
 
-from accelerate.state import ComputeEnvironment
+from accelerate.utils import ComputeEnvironment
 
 from .cluster import get_cluster_input
 from .config_args import cache_dir, default_config_file, default_yaml_config_file, load_config_from_file  # noqa: F401
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index f9d35ea48..cfab0f7f8 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -14,9 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from accelerate.state import ComputeEnvironment, DistributedType
-
-from ...utils import is_deepspeed_available
+from ...utils import ComputeEnvironment, DistributedType, is_deepspeed_available
 from .config_args import ClusterConfig
 from .config_utils import _ask_field, _convert_distributed_mode, _convert_yes_no_to_bool
 
diff --git a/src/accelerate/commands/config/config_args.py b/src/accelerate/commands/config/config_args.py
index 58db9b24e..fcb527aeb 100644
--- a/src/accelerate/commands/config/config_args.py
+++ b/src/accelerate/commands/config/config_args.py
@@ -21,7 +21,8 @@
 from typing import Optional, Union
 
 import yaml
-from accelerate.state import ComputeEnvironment, DistributedType, SageMakerDistributedType
+
+from ...utils import ComputeEnvironment, DistributedType, SageMakerDistributedType
 
 
 hf_cache_home = os.path.expanduser(
diff --git a/src/accelerate/commands/config/config_utils.py b/src/accelerate/commands/config/config_utils.py
index 2bc0e7ec9..fe72be267 100644
--- a/src/accelerate/commands/config/config_utils.py
+++ b/src/accelerate/commands/config/config_utils.py
@@ -14,7 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from accelerate.state import ComputeEnvironment, DistributedType, SageMakerDistributedType
+from ...utils.dataclasses import ComputeEnvironment, DistributedType, SageMakerDistributedType
 
 
 def _ask_field(input_text, convert_value=None, default=None, error_message=None):
diff --git a/src/accelerate/commands/config/sagemaker.py b/src/accelerate/commands/config/sagemaker.py
index 3b0fc78dd..f54659548 100644
--- a/src/accelerate/commands/config/sagemaker.py
+++ b/src/accelerate/commands/config/sagemaker.py
@@ -16,9 +16,8 @@
 import json
 import os
 
-from accelerate.state import ComputeEnvironment, SageMakerDistributedType
-from accelerate.utils import is_boto3_available
-
+from ...utils.dataclasses import ComputeEnvironment, SageMakerDistributedType
+from ...utils.imports import is_boto3_available
 from .config_args import SageMakerConfig
 from .config_utils import _ask_field, _convert_sagemaker_distributed_mode
 
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 9fa10da7b..e3ce3fd62 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -26,8 +26,13 @@
 
 from accelerate.commands.config import default_config_file, load_config_from_file
 from accelerate.commands.config.config_args import SageMakerConfig
-from accelerate.state import ComputeEnvironment, DistributedType
-from accelerate.utils import PrecisionType, PrepareForLaunch, is_sagemaker_available
+from accelerate.utils import (
+    ComputeEnvironment,
+    DistributedType,
+    PrecisionType,
+    PrepareForLaunch,
+    is_sagemaker_available,
+)
 
 
 def launch_command_parser(subparsers=None):
diff --git a/src/accelerate/kwargs_handlers.py b/src/accelerate/kwargs_handlers.py
deleted file mode 100644
index cf5985644..000000000
--- a/src/accelerate/kwargs_handlers.py
+++ /dev/null
@@ -1,90 +0,0 @@
-# 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 copy
-from dataclasses import dataclass
-from datetime import timedelta
-from typing import Optional
-
-
-class KwargsHandler:
-    """
-    Internal mixin that implements a `to_kwargs()` method for a dataclass.
-    """
-
-    def to_dict(self):
-        return copy.deepcopy(self.__dict__)
-
-    def to_kwargs(self):
-        """
-        Returns a dictionary containing the attributes with values different from the default of this class.
-        """
-        default_dict = self.__class__().to_dict()
-        this_dict = self.to_dict()
-        return {k: v for k, v in this_dict.items() if default_dict[k] != v}
-
-
-@dataclass
-class DistributedDataParallelKwargs(KwargsHandler):
-    """
-    Use this object in your [`Accelerator`] to customize how your model is wrapped in a
-    `torch.nn.parallel.DistributedDataParallel`. Please refer to the documentation of this
-    [wrapper](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html) for more
-    information on each argument.
-
-    <Tip warning={true}>
-
-    `gradient_as_bucket_view` is only available in PyTorch 1.7.0 and later versions.
-
-    </Tip>"""
-
-    dim: int = 0
-    broadcast_buffers: bool = True
-    bucket_cap_mb: int = 25
-    find_unused_parameters: bool = False
-    check_reduction: bool = False
-    gradient_as_bucket_view: bool = False
-
-
-@dataclass
-class GradScalerKwargs(KwargsHandler):
-    """
-    Use this object in your [`Accelerator`] to customize the behavior of mixed precision, specifically how the
-    `torch.cuda.amp.GradScaler` used is created. Please refer to the documentation of this
-    [scaler](https://pytorch.org/docs/stable/amp.html?highlight=gradscaler) for more information on each argument.
-
-    <Tip warning={true}>
-
-    `GradScaler` is only available in PyTorch 1.5.0 and later versions.
-
-    </Tip>"""
-
-    init_scale: float = 65536.0
-    growth_factor: float = 2.0
-    backoff_factor: float = 0.5
-    growth_interval: int = 2000
-    enabled: bool = True
-
-
-@dataclass
-class InitProcessGroupKwargs(KwargsHandler):
-    """
-    Use this object in your [`Accelerator`] to customize the initialization of the distributed processes. Please refer
-    to the documentation of this
-    [method](https://pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group) for more
-    information on each argument.
-    """
-
-    init_method: Optional[str] = None
-    timeout: timedelta = timedelta(seconds=1800)
diff --git a/src/accelerate/memory_utils.py b/src/accelerate/memory_utils.py
index 422ead770..eba10bb78 100644
--- a/src/accelerate/memory_utils.py
+++ b/src/accelerate/memory_utils.py
@@ -12,77 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""
-A collection of utilities for ensuring that training can always occur. Heavily influenced by the
-[toma](https://github.com/BlackHC/toma) library.
-"""
+# flake8: noqa
+# There's no way to ignore "F401 '...' imported but unused" warnings in this
+# module, but to preserve other warnings. So, don't check this module at all
 
-import functools
-import gc
-import inspect
 
-import torch
+import warnings
 
 
-def should_reduce_batch_size(exception: Exception) -> bool:
-    """
-    Checks if `exception` relates to CUDA out-of-memory, CUDNN not supported, or CPU out-of-memory
+warnings.warn(
+    "memory_utils has been reorganized to utils.memory. Import `find_executable_batchsize` from the main `__init__`: "
+    "`from accelerate import find_executable_batch_size` to avoid this warning.",
+    FutureWarning,
+)
 
-    Args:
-        exception (`Exception`):
-            An exception
-    """
-    _statements = [
-        "CUDA out of memory.",  # CUDA OOM
-        "cuDNN error: CUDNN_STATUS_NOT_SUPPORTED.",  # CUDNN SNAFU
-        "DefaultCPUAllocator: can't allocate memory",  # CPU OOM
-    ]
-    if isinstance(exception, RuntimeError) and len(exception.args) == 1:
-        return any(err in exception.args[0] for err in _statements)
-    return False
-
-
-def find_executable_batch_size(function: callable = None, starting_batch_size: int = 128):
-    """
-    A basic decorator that will try to execute `function`. If it fails from exceptions related to out-of-memory or
-    CUDNN, the batch size is cut in half and passed to `function`
-
-    `function` must take in a `batch_size` parameter as its first argument.
-
-    Args:
-        function (`callable`, *optional*):
-            A function to wrap
-        starting_batch_size (`int`, *optional*):
-            The batch size to try and fit into memory
-    """
-    if function is None:
-        return functools.partial(find_executable_batch_size, starting_batch_size=starting_batch_size)
-
-    batch_size = starting_batch_size
-
-    def decorator(*args, **kwargs):
-        nonlocal batch_size
-        gc.collect()
-        torch.cuda.empty_cache()
-        params = list(inspect.signature(function).parameters.keys())
-        # Guard against user error
-        if len(params) < (len(args) + 1):
-            arg_str = ", ".join([f"{arg}={value}" for arg, value in zip(params[1:], args[1:])])
-            raise TypeError(
-                f"Batch size was passed into `{function.__name__}` as the first argument when called."
-                f"Remove this as the decorator already does so: `{function.__name__}({arg_str})`"
-            )
-        while True:
-            if batch_size == 0:
-                raise RuntimeError("No executable batch size found, reached zero.")
-            try:
-                return function(batch_size, *args, **kwargs)
-            except Exception as e:
-                if should_reduce_batch_size(e):
-                    gc.collect()
-                    torch.cuda.empty_cache()
-                    batch_size //= 2
-                else:
-                    raise
-
-    return decorator
+from .utils.memory import find_executable_batch_size
diff --git a/src/accelerate/optimizer.py b/src/accelerate/optimizer.py
index 1bf37b7d5..dcf792592 100644
--- a/src/accelerate/optimizer.py
+++ b/src/accelerate/optimizer.py
@@ -19,8 +19,8 @@
 
 from packaging import version
 
-from .state import AcceleratorState, DistributedType, is_tpu_available
-from .utils import honor_type
+from .state import AcceleratorState
+from .utils import DistributedType, honor_type, is_tpu_available
 
 
 if is_tpu_available():
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index df2e36e9a..fbda09754 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -12,29 +12,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import importlib
 import os
 from distutils.util import strtobool
-from enum import Enum
 
 import torch
 
+from .utils import DistributedType, is_ccl_available, is_deepspeed_available, is_tpu_available
 
-try:
-    import torch_ccl  # noqa: F401
 
-    _ccl_available = True
-except ImportError:
-    _ccl_available = False
-
-
-try:
+if is_tpu_available():
     import torch_xla.core.xla_model as xm
 
-    _tpu_available = True
-except ImportError:
-    _tpu_available = False
-
 
 def get_int_from_env(env_keys, default):
     """Returns the first positive env value found in the `env_keys` list or the default."""
@@ -45,22 +33,6 @@ def get_int_from_env(env_keys, default):
     return default
 
 
-def is_ccl_available():
-    return _ccl_available
-
-
-def is_apex_available():
-    return importlib.util.find_spec("apex") is not None
-
-
-def is_tpu_available():
-    return _tpu_available
-
-
-def is_deepspeed_available():
-    return importlib.util.find_spec("deepspeed") is not None
-
-
 def parse_flag_from_env(key, default=False):
     value = os.environ.get(key, str(default))
     return strtobool(value) == 1  # As its name indicates `strtobool` actually returns an int...
@@ -71,60 +43,6 @@ def parse_choice_from_env(key, default="no"):
     return value
 
 
-class DistributedType(str, Enum):
-    """
-    Represents a type of distributed environment.
-
-    Values:
-
-        - **NO** -- Not a distributed environment, just a single process.
-        - **MULTI_CPU** -- Distributed on multiple CPU nodes.
-        - **MULTI_GPU** -- Distributed on multiple GPUs.
-        - **DEEPSPEED** -- Using DeepSpeed.
-        - **TPU** -- Distributed on TPUs.
-    """
-
-    # Subclassing str as well as Enum allows the `DistributedType` to be JSON-serializable out of the box.
-    NO = "NO"
-    MULTI_CPU = "MULTI_CPU"
-    MULTI_GPU = "MULTI_GPU"
-    DEEPSPEED = "DEEPSPEED"
-    FSDP = "FSDP"
-    TPU = "TPU"
-
-
-class SageMakerDistributedType(str, Enum):
-    """
-    Represents a type of distributed environment.
-
-    Values:
-
-        - **NO** -- Not a distributed environment, just a single process.
-        - **DATA_PARALLEL** -- using sagemaker distributed data parallelism.
-        - **MODEL_PARALLEL** -- using sagemaker distributed model parallelism.
-    """
-
-    # Subclassing str as well as Enum allows the `SageMakerDistributedType` to be JSON-serializable out of the box.
-    NO = "NO"
-    DATA_PARALLEL = "DATA_PARALLEL"
-    MODEL_PARALLEL = "MODEL_PARALLEL"
-
-
-class ComputeEnvironment(str, Enum):
-    """
-    Represents a type of the compute environment.
-
-    Values:
-
-        - **LOCAL_MACHINE** -- private/custom cluster hardware.
-        - **AMAZON_SAGEMAKER** -- Amazon SageMaker as compute environment.
-    """
-
-    # Subclassing str as well as Enum allows the `ComputeEnvironment` to be JSON-serializable out of the box.
-    LOCAL_MACHINE = "LOCAL_MACHINE"
-    AMAZON_SAGEMAKER = "AMAZON_SAGEMAKER"
-
-
 # Inspired by Alex Martelli's 'Borg'.
 class AcceleratorState:
     """
diff --git a/src/accelerate/test_utils/test_script.py b/src/accelerate/test_utils/test_script.py
index f9a99e22f..48e7bcc49 100644
--- a/src/accelerate/test_utils/test_script.py
+++ b/src/accelerate/test_utils/test_script.py
@@ -19,9 +19,9 @@
 
 from accelerate import Accelerator
 from accelerate.data_loader import prepare_data_loader
-from accelerate.state import AcceleratorState, DistributedType
+from accelerate.state import AcceleratorState
 from accelerate.test_utils import RegressionDataset, RegressionModel, are_the_same_tensors
-from accelerate.utils import gather, set_seed, synchronize_rng_states
+from accelerate.utils import DistributedType, gather, set_seed, synchronize_rng_states
 from packaging import version
 
 
diff --git a/src/accelerate/test_utils/testing.py b/src/accelerate/test_utils/testing.py
index 0b8e2250f..9379eb0b9 100644
--- a/src/accelerate/test_utils/testing.py
+++ b/src/accelerate/test_utils/testing.py
@@ -25,8 +25,8 @@
 
 import torch
 
-from ..state import AcceleratorState, is_tpu_available
-from ..utils import gather, is_comet_ml_available, is_tensorflow_available, is_wandb_available
+from ..state import AcceleratorState
+from ..utils import gather, is_comet_ml_available, is_tensorflow_available, is_tpu_available, is_wandb_available
 
 
 def parse_flag_from_env(key, default=False):
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
new file mode 100644
index 000000000..a43c7d261
--- /dev/null
+++ b/src/accelerate/utils/__init__.py
@@ -0,0 +1,61 @@
+# flake8: noqa
+# There's no way to ignore "F401 '...' imported but unused" warnings in this
+# module, but to preserve other warnings. So, don't check this module at all
+
+from .constants import MODEL_NAME, OPTIMIZER_NAME, RNG_STATE_NAME, SCALER_NAME, SCHEDULER_NAME
+from .dataclasses import (
+    ComputeEnvironment,
+    DeepSpeedPlugin,
+    DistributedDataParallelKwargs,
+    DistributedType,
+    FullyShardedDataParallelPlugin,
+    GradScalerKwargs,
+    InitProcessGroupKwargs,
+    KwargsHandler,
+    LoggerType,
+    PrecisionType,
+    RNGType,
+    SageMakerDistributedType,
+    TensorInformation,
+)
+from .imports import (
+    is_apex_available,
+    is_boto3_available,
+    is_ccl_available,
+    is_comet_ml_available,
+    is_deepspeed_available,
+    is_sagemaker_available,
+    is_tensorboard_available,
+    is_tensorflow_available,
+    is_tpu_available,
+    is_wandb_available,
+)
+from .operations import (
+    broadcast,
+    broadcast_object_list,
+    concatenate,
+    convert_outputs_to_fp32,
+    convert_to_fp32,
+    find_batch_size,
+    gather,
+    gather_object,
+    get_data_structure,
+    honor_type,
+    initialize_tensors,
+    is_tensor_information,
+    is_torch_tensor,
+    pad_across_processes,
+    recursively_apply,
+    reduce,
+    send_to_device,
+    slice_tensors,
+)
+
+
+if is_deepspeed_available():
+    from .deepspeed import DeepSpeedEngineWrapper, DeepSpeedOptimizerWrapper
+
+from .launch import PrepareForLaunch
+from .memory import find_executable_batch_size
+from .other import extract_model_from_parallel, get_pretty_name, patch_environment, save, wait_for_everyone
+from .random import set_seed, synchronize_rng_state, synchronize_rng_states
diff --git a/src/accelerate/utils/constants.py b/src/accelerate/utils/constants.py
new file mode 100644
index 000000000..c2c36a486
--- /dev/null
+++ b/src/accelerate/utils/constants.py
@@ -0,0 +1,19 @@
+# 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.
+
+SCALER_NAME = "scaler.pt"
+MODEL_NAME = "pytorch_model"
+RNG_STATE_NAME = "random_states"
+OPTIMIZER_NAME = "optimizer"
+SCHEDULER_NAME = "scheduler"
diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py
new file mode 100644
index 000000000..4a19f8b51
--- /dev/null
+++ b/src/accelerate/utils/dataclasses.py
@@ -0,0 +1,304 @@
+# 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.
+
+"""
+General namespace and dataclass related classes
+"""
+
+import copy
+import enum
+import functools
+import os
+import typing
+from dataclasses import dataclass, field
+from datetime import timedelta
+from typing import Callable, Iterable, Optional
+
+import torch
+
+
+class KwargsHandler:
+    """
+    Internal mixin that implements a `to_kwargs()` method for a dataclass.
+    """
+
+    def to_dict(self):
+        return copy.deepcopy(self.__dict__)
+
+    def to_kwargs(self):
+        """
+        Returns a dictionary containing the attributes with values different from the default of this class.
+        """
+        default_dict = self.__class__().to_dict()
+        this_dict = self.to_dict()
+        return {k: v for k, v in this_dict.items() if default_dict[k] != v}
+
+
+@dataclass
+class DistributedDataParallelKwargs(KwargsHandler):
+    """
+    Use this object in your [`Accelerator`] to customize how your model is wrapped in a
+    `torch.nn.parallel.DistributedDataParallel`. Please refer to the documentation of this
+    [wrapper](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html) for more
+    information on each argument.
+
+    <Tip warning={true}>
+
+    `gradient_as_bucket_view` is only available in PyTorch 1.7.0 and later versions.
+
+    </Tip>"""
+
+    dim: int = 0
+    broadcast_buffers: bool = True
+    bucket_cap_mb: int = 25
+    find_unused_parameters: bool = False
+    check_reduction: bool = False
+    gradient_as_bucket_view: bool = False
+
+
+@dataclass
+class GradScalerKwargs(KwargsHandler):
+    """
+    Use this object in your [`Accelerator`] to customize the behavior of mixed precision, specifically how the
+    `torch.cuda.amp.GradScaler` used is created. Please refer to the documentation of this
+    [scaler](https://pytorch.org/docs/stable/amp.html?highlight=gradscaler) for more information on each argument.
+
+    <Tip warning={true}>
+
+    `GradScaler` is only available in PyTorch 1.5.0 and later versions.
+
+    </Tip>"""
+
+    init_scale: float = 65536.0
+    growth_factor: float = 2.0
+    backoff_factor: float = 0.5
+    growth_interval: int = 2000
+    enabled: bool = True
+
+
+@dataclass
+class InitProcessGroupKwargs(KwargsHandler):
+    """
+    Use this object in your [`Accelerator`] to customize the initialization of the distributed processes. Please refer
+    to the documentation of this
+    [method](https://pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group) for more
+    information on each argument.
+    """
+
+    init_method: Optional[str] = None
+    timeout: timedelta = timedelta(seconds=1800)
+
+
+class DistributedType(str, enum.Enum):
+    """
+    Represents a type of distributed environment.
+
+    Values:
+
+        - **NO** -- Not a distributed environment, just a single process.
+        - **MULTI_CPU** -- Distributed on multiple CPU nodes.
+        - **MULTI_GPU** -- Distributed on multiple GPUs.
+        - **DEEPSPEED** -- Using DeepSpeed.
+        - **TPU** -- Distributed on TPUs.
+    """
+
+    # Subclassing str as well as Enum allows the `DistributedType` to be JSON-serializable out of the box.
+    NO = "NO"
+    MULTI_CPU = "MULTI_CPU"
+    MULTI_GPU = "MULTI_GPU"
+    DEEPSPEED = "DEEPSPEED"
+    FSDP = "FSDP"
+    TPU = "TPU"
+
+
+class SageMakerDistributedType(str, enum.Enum):
+    """
+    Represents a type of distributed environment.
+
+    Values:
+
+        - **NO** -- Not a distributed environment, just a single process.
+        - **DATA_PARALLEL** -- using sagemaker distributed data parallelism.
+        - **MODEL_PARALLEL** -- using sagemaker distributed model parallelism.
+    """
+
+    # Subclassing str as well as Enum allows the `SageMakerDistributedType` to be JSON-serializable out of the box.
+    NO = "NO"
+    DATA_PARALLEL = "DATA_PARALLEL"
+    MODEL_PARALLEL = "MODEL_PARALLEL"
+
+
+class ComputeEnvironment(str, enum.Enum):
+    """
+    Represents a type of the compute environment.
+
+    Values:
+
+        - **LOCAL_MACHINE** -- private/custom cluster hardware.
+        - **AMAZON_SAGEMAKER** -- Amazon SageMaker as compute environment.
+    """
+
+    # Subclassing str as well as Enum allows the `ComputeEnvironment` to be JSON-serializable out of the box.
+    LOCAL_MACHINE = "LOCAL_MACHINE"
+    AMAZON_SAGEMAKER = "AMAZON_SAGEMAKER"
+
+
+class EnumWithContains(enum.EnumMeta):
+    "A metaclass that adds the ability to check if `self` contains an item with the `in` operator"
+
+    def __contains__(cls, item):
+        try:
+            cls(item)
+        except ValueError:
+            return False
+        return True
+
+
+class BaseEnum(enum.Enum, metaclass=EnumWithContains):
+    "An enum class that can get the value of an item with `str(Enum.key)`"
+
+    def __str__(self):
+        return self.value
+
+    @classmethod
+    def list(cls):
+        "Method to list all the possible items in `cls`"
+        return list(map(lambda item: str(item), cls))
+
+
+class LoggerType(BaseEnum):
+    ALL = "all"
+    TENSORBOARD = "tensorboard"
+    WANDB = "wandb"
+    COMETML = "comet_ml"
+
+
+class PrecisionType(BaseEnum):
+    NO = "no"
+    FP16 = "fp16"
+    BF16 = "bf16"
+
+
+class RNGType(BaseEnum):
+    TORCH = "torch"
+    CUDA = "cuda"
+    XLA = "xla"
+    GENERATOR = "generator"
+
+
+# data classes
+
+
+@dataclass
+class TensorInformation:
+    shape: torch.Size
+    dtype: torch.dtype
+
+
+@dataclass
+class DeepSpeedPlugin:
+
+    gradient_accumulation_steps: int = field(
+        default=None, metadata={"help": "Number of steps to accumulate gradients before updating optimizer states"}
+    )
+    zero_stage: int = field(
+        default=None,
+        metadata={"help": "Possible options are 0,1,2,3; Default will be taken from environment variable"},
+    )
+    is_train_batch_min: str = field(
+        default=True,
+        metadata={"help": "If both train & eval dataloaders are specified, this will decide the train_batch_size"},
+    )
+
+    auto_opt_mapping: bool = field(
+        default=True,
+        metadata={"help": "whether to map torch.adam to deepspeed optimizer version of adam based on config"},
+    )
+
+    offload_optimizer_device: bool = field(default=None, metadata={"help": "Possible options are none|cpu|nvme"})
+
+    def __post_init__(self):
+
+        if self.gradient_accumulation_steps is None:
+            self.gradient_accumulation_steps = int(os.environ.get("GRADIENT_ACCUMULATION_STEPS", 1))
+
+        if self.zero_stage is None:
+            self.zero_stage = int(os.environ.get("DEEPSPEED_ZERO_STAGE", 2))
+
+        if self.offload_optimizer_device is None:
+            self.offload_optimizer_device = os.environ.get("DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE", "none")
+
+        self.deepspeed_config = {
+            "train_batch_size": None,
+            "gradient_accumulation_steps": self.gradient_accumulation_steps,
+            "zero_optimization": {
+                "stage": self.zero_stage,
+                "offload_optimizer": {
+                    "device": self.offload_optimizer_device,
+                },
+            },
+            "steps_per_print": float("inf"),  # this will stop deepspeed from logging @ stdout
+            "zero_allow_untested_optimizer": True,
+        }
+
+
+@dataclass
+class FullyShardedDataParallelPlugin:
+    """
+    This plugin is used to enable fully sharded data parallelism.
+    """
+
+    sharding_strategy: "typing.Any" = field(
+        default=None,
+        metadata={"help": "Possible options are [1] FULL_SHARD, [2] SHARD_GRAD_OP"},
+    )
+    backward_prefetch: "typing.Any" = field(
+        default=None,
+        metadata={"help": "Possible options are [1] BACKWARD_PRE, [2] BACKWARD_POST"},
+    )
+    auto_wrap_policy: "typing.Any" = field(
+        default=None,
+        metadata={"help": "A callable specifying a policy to recursively wrap layers with FSDP"},
+    )
+    cpu_offload: Optional[Callable] = field(
+        default=None,
+        metadata={"help": "Decides Whether to offload parameters and gradients to CPU."},
+    )
+    min_num_params: int = field(
+        default=None, metadata={"help": "FSDP's minimum number of parameters for Default Auto Wrapping."}
+    )
+    ignored_modules: Optional[Iterable[torch.nn.Module]] = field(
+        default=None,
+        metadata={"help": "A list of modules to ignore for FSDP."},
+    )
+
+    def __post_init__(self):
+        from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload, ShardingStrategy
+        from torch.distributed.fsdp.wrap import default_auto_wrap_policy
+
+        if self.sharding_strategy is None:
+            self.sharding_strategy = ShardingStrategy(int(os.environ.get("FSDP_SHARDING_STRATEGY", 1)))
+
+        if self.cpu_offload is None:
+            if os.environ.get("FSDP_OFFLOAD_PARAMS", "false") == "true":
+                self.cpu_offload = CPUOffload(offload_params=True)
+            else:
+                self.cpu_offload = CPUOffload(offload_params=False)
+
+        if self.min_num_params is None:
+            self.min_num_params = int(os.environ.get("FSDP_MIN_NUM_PARAMS", 0))
+
+        if self.auto_wrap_policy is None:
+            if self.min_num_params > 0:
+                self.auto_wrap_policy = functools.partial(default_auto_wrap_policy, min_num_params=self.min_num_params)
diff --git a/src/accelerate/deepspeed_utils.py b/src/accelerate/utils/deepspeed.py
similarity index 96%
rename from src/accelerate/deepspeed_utils.py
rename to src/accelerate/utils/deepspeed.py
index 5f261f2e3..891f6fdb8 100644
--- a/src/accelerate/deepspeed_utils.py
+++ b/src/accelerate/utils/deepspeed.py
@@ -12,8 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .optimizer import AcceleratedOptimizer
-from .state import is_apex_available, is_deepspeed_available
+from ..optimizer import AcceleratedOptimizer
+from .imports import is_apex_available, is_deepspeed_available
 
 
 if is_deepspeed_available():
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
new file mode 100644
index 000000000..3f10b83cf
--- /dev/null
+++ b/src/accelerate/utils/imports.py
@@ -0,0 +1,72 @@
+# 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 importlib
+
+
+try:
+    import torch_ccl  # noqa: F401
+
+    _ccl_available = True
+except ImportError:
+    _ccl_available = False
+
+
+try:
+    import torch_xla.core.xla_model as xm  # noqa: F401
+
+    _tpu_available = True
+except ImportError:
+    _tpu_available = False
+
+
+def is_ccl_available():
+    return _ccl_available
+
+
+def is_apex_available():
+    return importlib.util.find_spec("apex") is not None
+
+
+def is_tpu_available():
+    return _tpu_available
+
+
+def is_deepspeed_available():
+    return importlib.util.find_spec("deepspeed") is not None
+
+
+def is_tensorflow_available():
+    return importlib.util.find_spec("tensorflow") is not None
+
+
+def is_tensorboard_available():
+    return importlib.util.find_spec("tensorboard") is not None or importlib.util.find_spec("tensorboardX") is not None
+
+
+def is_wandb_available():
+    return importlib.util.find_spec("wandb") is not None
+
+
+def is_comet_ml_available():
+    return importlib.util.find_spec("comet_ml") is not None
+
+
+def is_boto3_available():
+    return importlib.util.find_spec("boto3") is not None
+
+
+def is_sagemaker_available():
+    return importlib.util.find_spec("sagemaker") is not None
diff --git a/src/accelerate/utils/launch.py b/src/accelerate/utils/launch.py
new file mode 100644
index 000000000..8ff624250
--- /dev/null
+++ b/src/accelerate/utils/launch.py
@@ -0,0 +1,55 @@
+# 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 torch
+
+from .dataclasses import DistributedType
+
+
+class PrepareForLaunch:
+    """
+    Prepare a function that will launched in a distributed setup.
+
+    Args:
+        launcher (`Callable`):
+            The function to launch.
+        distributed_type ([`~state.DistributedType`]):
+            The distributed type to prepare for.
+        debug (`bool`, *optional*, defaults to `False`):
+            Whether or not this is a debug launch.
+    """
+
+    def __init__(self, launcher, distributed_type="NO", debug=False):
+        self.launcher = launcher
+        self.distributed_type = DistributedType(distributed_type)
+        self.debug = debug
+
+    def __call__(self, index, *args):
+        if self.debug:
+            world_size = int(os.environ.get("WORLD_SIZE"))
+            rdv_file = os.environ.get("ACCELERATE_DEBUG_RDV_FILE")
+            torch.distributed.init_process_group(
+                "gloo",
+                rank=index,
+                store=torch.distributed.FileStore(rdv_file, world_size),
+                world_size=world_size,
+            )
+        elif self.distributed_type == DistributedType.MULTI_GPU or self.distributed_type == DistributedType.MULTI_CPU:
+            # Prepare the environment for torch.distributed
+            os.environ["LOCAL_RANK"] = str(index)
+            os.environ["RANK"] = str(index)
+
+        self.launcher(*args)
diff --git a/src/accelerate/utils/memory.py b/src/accelerate/utils/memory.py
new file mode 100644
index 000000000..422ead770
--- /dev/null
+++ b/src/accelerate/utils/memory.py
@@ -0,0 +1,88 @@
+# 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.
+
+"""
+A collection of utilities for ensuring that training can always occur. Heavily influenced by the
+[toma](https://github.com/BlackHC/toma) library.
+"""
+
+import functools
+import gc
+import inspect
+
+import torch
+
+
+def should_reduce_batch_size(exception: Exception) -> bool:
+    """
+    Checks if `exception` relates to CUDA out-of-memory, CUDNN not supported, or CPU out-of-memory
+
+    Args:
+        exception (`Exception`):
+            An exception
+    """
+    _statements = [
+        "CUDA out of memory.",  # CUDA OOM
+        "cuDNN error: CUDNN_STATUS_NOT_SUPPORTED.",  # CUDNN SNAFU
+        "DefaultCPUAllocator: can't allocate memory",  # CPU OOM
+    ]
+    if isinstance(exception, RuntimeError) and len(exception.args) == 1:
+        return any(err in exception.args[0] for err in _statements)
+    return False
+
+
+def find_executable_batch_size(function: callable = None, starting_batch_size: int = 128):
+    """
+    A basic decorator that will try to execute `function`. If it fails from exceptions related to out-of-memory or
+    CUDNN, the batch size is cut in half and passed to `function`
+
+    `function` must take in a `batch_size` parameter as its first argument.
+
+    Args:
+        function (`callable`, *optional*):
+            A function to wrap
+        starting_batch_size (`int`, *optional*):
+            The batch size to try and fit into memory
+    """
+    if function is None:
+        return functools.partial(find_executable_batch_size, starting_batch_size=starting_batch_size)
+
+    batch_size = starting_batch_size
+
+    def decorator(*args, **kwargs):
+        nonlocal batch_size
+        gc.collect()
+        torch.cuda.empty_cache()
+        params = list(inspect.signature(function).parameters.keys())
+        # Guard against user error
+        if len(params) < (len(args) + 1):
+            arg_str = ", ".join([f"{arg}={value}" for arg, value in zip(params[1:], args[1:])])
+            raise TypeError(
+                f"Batch size was passed into `{function.__name__}` as the first argument when called."
+                f"Remove this as the decorator already does so: `{function.__name__}({arg_str})`"
+            )
+        while True:
+            if batch_size == 0:
+                raise RuntimeError("No executable batch size found, reached zero.")
+            try:
+                return function(batch_size, *args, **kwargs)
+            except Exception as e:
+                if should_reduce_batch_size(e):
+                    gc.collect()
+                    torch.cuda.empty_cache()
+                    batch_size //= 2
+                else:
+                    raise
+
+    return decorator
diff --git a/src/accelerate/utils.py b/src/accelerate/utils/operations.py
similarity index 58%
rename from src/accelerate/utils.py
rename to src/accelerate/utils/operations.py
index 5abbf00b5..b61ae069b 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils/operations.py
@@ -1,4 +1,4 @@
-# Copyright 2021 The HuggingFace Team. All rights reserved.
+# 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.
@@ -12,172 +12,34 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import functools
-import importlib
-import os
-import random
-import typing
-from collections.abc import Mapping
-from contextlib import contextmanager
-from dataclasses import dataclass, field
-from enum import Enum, EnumMeta
+"""
+A set of basic tensor ops compatible with tpu, gpu, and multigpu
+"""
+
+
 from functools import update_wrapper
-from typing import Any, Callable, Iterable, List, Optional, Union
+from typing import Any, Mapping
 
-import numpy as np
 import torch
 from torch.distributed import ReduceOp
 
 from packaging import version
 
-from .state import AcceleratorState, DistributedType, is_deepspeed_available, is_tpu_available
+from ..state import AcceleratorState
+from .dataclasses import DistributedType, TensorInformation
+from .imports import is_tpu_available
 
 
 if is_tpu_available():
     import torch_xla.core.xla_model as xm
 
 
-def is_tensorflow_available():
-    return importlib.util.find_spec("tensorflow") is not None
-
-
-def is_tensorboard_available():
-    return importlib.util.find_spec("tensorboard") is not None or importlib.util.find_spec("tensorboardX") is not None
-
-
-def is_wandb_available():
-    return importlib.util.find_spec("wandb") is not None
-
-
-def is_comet_ml_available():
-    return importlib.util.find_spec("comet_ml") is not None
-
-
-def is_boto3_available():
-    return importlib.util.find_spec("boto3") is not None
-
-
-def is_sagemaker_available():
-    return importlib.util.find_spec("sagemaker") is not None
-
-
-if is_deepspeed_available():
-    from deepspeed import DeepSpeedEngine
-
-SCALER_NAME = "scaler.pt"
-MODEL_NAME = "pytorch_model"
-RNG_STATE_NAME = "random_states"
-OPTIMIZER_NAME = "optimizer"
-SCHEDULER_NAME = "scheduler"
-
-
-class EnumWithContains(EnumMeta):
-    "A metaclass that adds the ability to check if `self` contains an item with the `in` operator"
-
-    def __contains__(cls, item):
-        try:
-            cls(item)
-        except ValueError:
-            return False
-        return True
-
-
-class BaseEnum(Enum, metaclass=EnumWithContains):
-    "An enum class that can get the value of an item with `str(Enum.key)`"
-
-    def __str__(self):
-        return self.value
-
-    @classmethod
-    def list(cls):
-        "Method to list all the possible items in `cls`"
-        return list(map(lambda item: str(item), cls))
-
-
-class LoggerType(BaseEnum):
-    ALL = "all"
-    TENSORBOARD = "tensorboard"
-    WANDB = "wandb"
-    COMETML = "comet_ml"
-
-
-class PrecisionType(BaseEnum):
-    NO = "no"
-    FP16 = "fp16"
-    BF16 = "bf16"
-
-
-class RNGType(BaseEnum):
-    TORCH = "torch"
-    CUDA = "cuda"
-    XLA = "xla"
-    GENERATOR = "generator"
-
-
-@dataclass
-class TensorInformation:
-    shape: torch.Size
-    dtype: torch.dtype
-
+def is_torch_tensor(tensor):
+    return isinstance(tensor, torch.Tensor)
 
-def set_seed(seed: int, device_specific: bool = False):
-    """
-    Helper function for reproducible behavior to set the seed in `random`, `numpy`, `torch`.
 
-    Args:
-        seed (`int`): The seed to set.
-        device_specific (`bool`, *optional*, defaults to `False`):
-            Whether to differ the seed on each device slightly with `self.process_index`.
-    """
-    if device_specific:
-        seed += AcceleratorState().process_index
-    random.seed(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-    # ^^ safe to call this function even if cuda is not available
-    if is_tpu_available():
-        xm.set_rng_state(seed)
-
-
-def synchronize_rng_state(rng_type: Optional[RNGType] = None, generator: Optional[torch.Generator] = None):
-    # Get the proper rng state
-    if rng_type == RNGType.TORCH:
-        rng_state = torch.get_rng_state()
-    elif rng_type == RNGType.CUDA:
-        rng_state = torch.cuda.get_rng_state()
-    elif rng_type == RNGType.XLA:
-        assert is_tpu_available(), "Can't synchronize XLA seeds on an environment without TPUs."
-        rng_state = torch.tensor(xm.get_rng_state())
-    elif rng_type == RNGType.GENERATOR:
-        assert generator is not None, "Need a generator to synchronize its seed."
-        rng_state = generator.get_state()
-
-    # Broadcast the rng state from device 0 to other devices
-    state = AcceleratorState()
-    if state.distributed_type == DistributedType.TPU:
-        rng_state = xm.mesh_reduce("random_seed", rng_state, lambda x: x[0])
-    elif state.distributed_type in [DistributedType.DEEPSPEED, DistributedType.MULTI_GPU]:
-        rng_state = rng_state.to(state.device)
-        torch.distributed.broadcast(rng_state, 0)
-        rng_state = rng_state.cpu()
-    elif state.distributed_type == DistributedType.MULTI_CPU:
-        torch.distributed.broadcast(rng_state, 0)
-
-    # Set the broadcast rng state
-    if rng_type == RNGType.TORCH:
-        torch.set_rng_state(rng_state)
-    elif rng_type == RNGType.CUDA:
-        torch.cuda.set_rng_state(rng_state)
-    elif rng_type == RNGType.XLA:
-        xm.set_rng_state(rng_state.item())
-    elif rng_type == RNGType.GENERATOR:
-        generator.set_state(rng_state)
-
-
-def synchronize_rng_states(rng_types: List[Union[str, RNGType]], generator: Optional[torch.Generator] = None):
-    for rng_type in rng_types:
-        synchronize_rng_state(RNGType(rng_type), generator=generator)
+def is_tensor_information(tensor_info):
+    return isinstance(tensor_info, TensorInformation)
 
 
 def honor_type(obj, generator):
@@ -191,14 +53,6 @@ def honor_type(obj, generator):
         return type(obj)(*list(generator))
 
 
-def is_torch_tensor(tensor):
-    return isinstance(tensor, torch.Tensor)
-
-
-def is_tensor_information(tensor_info):
-    return isinstance(tensor_info, TensorInformation)
-
-
 def recursively_apply(func, data, *args, test_type=is_torch_tensor, error_on_other_type=False, **kwargs):
     """
     Recursively apply a function on a data structure that is a nested list/tuple/dictionary of a given base type.
@@ -305,73 +159,24 @@ def _initialize_tensor(tensor_info):
     return recursively_apply(_initialize_tensor, data_structure, test_type=is_tensor_information)
 
 
-def convert_to_fp32(tensor):
-    """
-    Recursively converts the elements nested list/tuple/dictionary of tensors in FP16/BF16 precision to FP32.
-
-    Args:
-        tensor (nested list/tuple/dictionary of `torch.Tensor`):
-            The data to convert from FP16/BF16 to FP32.
-
-    Returns:
-        The same data structure as `tensor` with all tensors that were in FP16/BF16 precision converted to FP32.
-    """
-
-    def _convert_to_fp32(tensor):
-        return tensor.float()
-
-    def _is_fp16_bf16_tensor(tensor):
-        return hasattr(tensor, "dtype") and (
-            tensor.dtype == torch.float16
-            or (version.parse(torch.__version__) >= version.parse("1.10") and tensor.dtype == torch.bfloat16)
-        )
-
-    return recursively_apply(_convert_to_fp32, tensor, test_type=_is_fp16_bf16_tensor)
-
-
-class ConvertOutputsToFp32:
-    """
-    Decorator to apply to a function outputing tensors (like a model forward pass) that ensures the outputs in FP16
-    precision will be convert back to FP32.
-
-    Use a class instead of a decorator because otherwise, the prepared model can no longer be pickled (issue #273).
-
-    Args:
-        model_forward (`Callable`):
-            The function which outputs we want to treat.
-
-    Returns:
-        The same function as `model_forward` but with converted outputs.
-    """
-
-    def __init__(self, model_forward):
-        self.model_forward = model_forward
-        update_wrapper(self, model_forward)
-
-    def __call__(self, *args, **kwargs):
-        return convert_to_fp32(self.model_forward(*args, **kwargs))
-
-
-convert_outputs_to_fp32 = ConvertOutputsToFp32
-
-
-def extract_model_from_parallel(model):
+def find_batch_size(data):
     """
-    Extract a model from its distributed containers.
+    Recursively finds the batch size in a nested list/tuple/dictionary of lists of tensors.
 
     Args:
-        model (`torch.nn.Module`): The model to extract.
+        data (nested list/tuple/dictionary of `torch.Tensor`): The data from which to find the batch size.
 
     Returns:
-        `torch.nn.Module`: The extracted model.
+        `int`: The batch size.
     """
-    options = (torch.nn.parallel.DistributedDataParallel, torch.nn.DataParallel)
-    if is_deepspeed_available():
-        options += (DeepSpeedEngine,)
-
-    while isinstance(model, options):
-        model = model.module
-    return model
+    if isinstance(data, (tuple, list)):
+        return find_batch_size(data[0])
+    elif isinstance(data, Mapping):
+        for k in data.keys():
+            return find_batch_size(data[k])
+    elif not isinstance(data, torch.Tensor):
+        raise TypeError(f"Can only find the batch size of tensors but got {type(data)}.")
+    return data.shape[0]
 
 
 def _tpu_gather(tensor, name="gather tensor"):
@@ -536,26 +341,6 @@ def _slice_tensor(tensor, tensor_slice):
     return recursively_apply(_slice_tensor, data, tensor_slice)
 
 
-def find_batch_size(data):
-    """
-    Recursively finds the batch size in a nested list/tuple/dictionary of lists of tensors.
-
-    Args:
-        data (nested list/tuple/dictionary of `torch.Tensor`): The data from which to find the batch size.
-
-    Returns:
-        `int`: The batch size.
-    """
-    if isinstance(data, (tuple, list)):
-        return find_batch_size(data[0])
-    elif isinstance(data, Mapping):
-        for k in data.keys():
-            return find_batch_size(data[k])
-    elif not isinstance(data, torch.Tensor):
-        raise TypeError(f"Can only find the batch size of tensors but got {type(data)}.")
-    return data.shape[0]
-
-
 def concatenate(data, dim=0):
     """
     Recursively concatenate the tensors in a nested list/tuple/dictionary of lists of tensors with the same shape.
@@ -657,198 +442,51 @@ def _reduce_across_processes(tensor, reduction="mean"):
     return recursively_apply(_reduce_across_processes, tensor, error_on_other_type=True, reduction=reduction)
 
 
-def wait_for_everyone():
-    """
-    Introduces a blocking point in the script, making sure all processes have reached this point before continuing.
-
-    <Tip warning={true}>
-
-    Make sure all processes will reach this instruction otherwise one of your processes will hang forever.
-
-    </Tip>
-    """
-    if (
-        AcceleratorState().distributed_type == DistributedType.MULTI_GPU
-        or AcceleratorState().distributed_type == DistributedType.MULTI_CPU
-        or AcceleratorState().distributed_type == DistributedType.DEEPSPEED
-    ):
-        torch.distributed.barrier()
-    elif AcceleratorState().distributed_type == DistributedType.TPU:
-        xm.rendezvous("accelerate.utils.wait_for_everyone")
-
-
-def save(obj, f):
+def convert_to_fp32(tensor):
     """
-    Save the data to disk. Use in place of `torch.save()`.
+    Recursively converts the elements nested list/tuple/dictionary of tensors in FP16/BF16 precision to FP32.
 
     Args:
-        obj: The data to save
-        f: The file (or file-like object) to use to save the data
-    """
-    if AcceleratorState().distributed_type == DistributedType.TPU:
-        xm.save(obj, f)
-    elif AcceleratorState().local_process_index == 0:
-        torch.save(obj, f)
-
+        tensor (nested list/tuple/dictionary of `torch.Tensor`):
+            The data to convert from FP16/BF16 to FP32.
 
-class PrepareForLaunch:
+    Returns:
+        The same data structure as `tensor` with all tensors that were in FP16/BF16 precision converted to FP32.
     """
-    Prepare a function that will launched in a distributed setup.
-
-    Args:
-        launcher (`Callable`):
-            The function to launch.
-        distributed_type ([`~state.DistributedType`]):
-            The distributed type to prepare for.
-        debug (`bool`, *optional*, defaults to `False`):
-            Whether or not this is a debug launch.
-    """
-
-    def __init__(self, launcher, distributed_type="NO", debug=False):
-        self.launcher = launcher
-        self.distributed_type = DistributedType(distributed_type)
-        self.debug = debug
-
-    def __call__(self, index, *args):
-        if self.debug:
-            world_size = int(os.environ.get("WORLD_SIZE"))
-            rdv_file = os.environ.get("ACCELERATE_DEBUG_RDV_FILE")
-            torch.distributed.init_process_group(
-                "gloo",
-                rank=index,
-                store=torch.distributed.FileStore(rdv_file, world_size),
-                world_size=world_size,
-            )
-        elif self.distributed_type == DistributedType.MULTI_GPU or self.distributed_type == DistributedType.MULTI_CPU:
-            # Prepare the environment for torch.distributed
-            os.environ["LOCAL_RANK"] = str(index)
-            os.environ["RANK"] = str(index)
-
-        self.launcher(*args)
-
-
-@dataclass
-class DeepSpeedPlugin:
 
-    gradient_accumulation_steps: int = field(
-        default=None, metadata={"help": "Number of steps to accumulate gradients before updating optimizer states"}
-    )
-    zero_stage: int = field(
-        default=None,
-        metadata={"help": "Possible options are 0,1,2,3; Default will be taken from environment variable"},
-    )
-    is_train_batch_min: str = field(
-        default=True,
-        metadata={"help": "If both train & eval dataloaders are specified, this will decide the train_batch_size"},
-    )
-
-    auto_opt_mapping: bool = field(
-        default=True,
-        metadata={"help": "whether to map torch.adam to deepspeed optimizer version of adam based on config"},
-    )
-
-    offload_optimizer_device: bool = field(default=None, metadata={"help": "Possible options are none|cpu|nvme"})
-
-    def __post_init__(self):
-
-        if self.gradient_accumulation_steps is None:
-            self.gradient_accumulation_steps = int(os.environ.get("GRADIENT_ACCUMULATION_STEPS", 1))
-
-        if self.zero_stage is None:
-            self.zero_stage = int(os.environ.get("DEEPSPEED_ZERO_STAGE", 2))
+    def _convert_to_fp32(tensor):
+        return tensor.float()
 
-        if self.offload_optimizer_device is None:
-            self.offload_optimizer_device = os.environ.get("DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE", "none")
+    def _is_fp16_bf16_tensor(tensor):
+        return hasattr(tensor, "dtype") and (
+            tensor.dtype == torch.float16
+            or (version.parse(torch.__version__) >= version.parse("1.10") and tensor.dtype == torch.bfloat16)
+        )
 
-        self.deepspeed_config = {
-            "train_batch_size": None,
-            "gradient_accumulation_steps": self.gradient_accumulation_steps,
-            "zero_optimization": {
-                "stage": self.zero_stage,
-                "offload_optimizer": {
-                    "device": self.offload_optimizer_device,
-                },
-            },
-            "steps_per_print": float("inf"),  # this will stop deepspeed from logging @ stdout
-            "zero_allow_untested_optimizer": True,
-        }
+    return recursively_apply(_convert_to_fp32, tensor, test_type=_is_fp16_bf16_tensor)
 
 
-@dataclass
-class FullyShardedDataParallelPlugin:
-    """
-    This plugin is used to enable fully sharded data parallelism.
+class ConvertOutputsToFp32:
     """
+    Decorator to apply to a function outputing tensors (like a model forward pass) that ensures the outputs in FP16
+    precision will be convert back to FP32.
 
-    sharding_strategy: "typing.Any" = field(
-        default=None,
-        metadata={"help": "Possible options are [1] FULL_SHARD, [2] SHARD_GRAD_OP"},
-    )
-    backward_prefetch: "typing.Any" = field(
-        default=None,
-        metadata={"help": "Possible options are [1] BACKWARD_PRE, [2] BACKWARD_POST"},
-    )
-    auto_wrap_policy: "typing.Any" = field(
-        default=None,
-        metadata={"help": "A callable specifying a policy to recursively wrap layers with FSDP"},
-    )
-    cpu_offload: Optional[Callable] = field(
-        default=None,
-        metadata={"help": "Decides Whether to offload parameters and gradients to CPU."},
-    )
-    min_num_params: int = field(
-        default=None, metadata={"help": "FSDP's minimum number of parameters for Default Auto Wrapping."}
-    )
-    ignored_modules: Optional[Iterable[torch.nn.Module]] = field(
-        default=None,
-        metadata={"help": "A list of modules to ignore for FSDP."},
-    )
-
-    def __post_init__(self):
-        from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload, ShardingStrategy
-        from torch.distributed.fsdp.wrap import default_auto_wrap_policy
-
-        if self.sharding_strategy is None:
-            self.sharding_strategy = ShardingStrategy(int(os.environ.get("FSDP_SHARDING_STRATEGY", 1)))
-
-        if self.cpu_offload is None:
-            if os.environ.get("FSDP_OFFLOAD_PARAMS", "false") == "true":
-                self.cpu_offload = CPUOffload(offload_params=True)
-            else:
-                self.cpu_offload = CPUOffload(offload_params=False)
-
-        if self.min_num_params is None:
-            self.min_num_params = int(os.environ.get("FSDP_MIN_NUM_PARAMS", 0))
-
-        if self.auto_wrap_policy is None:
-            if self.min_num_params > 0:
-                self.auto_wrap_policy = functools.partial(default_auto_wrap_policy, min_num_params=self.min_num_params)
-
+    Use a class instead of a decorator because otherwise, the prepared model can no longer be pickled (issue #273).
 
-@contextmanager
-def patch_environment(**kwargs):
-    """
-    A context manager that will add each keyword argument passed to `os.environ` and remove them when exiting.
+    Args:
+        model_forward (`Callable`):
+            The function which outputs we want to treat.
 
-    Will convert the values in `kwargs` to strings and upper-case all the keys.
+    Returns:
+        The same function as `model_forward` but with converted outputs.
     """
-    for key, value in kwargs.items():
-        os.environ[key.upper()] = str(value)
 
-    yield
+    def __init__(self, model_forward):
+        self.model_forward = model_forward
+        update_wrapper(self, model_forward)
 
-    for key in kwargs:
-        del os.environ[key.upper()]
+    def __call__(self, *args, **kwargs):
+        return convert_to_fp32(self.model_forward(*args, **kwargs))
 
 
-def get_pretty_name(obj):
-    """
-    Gets a pretty name from `obj`.
-    """
-    if not hasattr(obj, "__qualname__") and not hasattr(obj, "__name__"):
-        obj = getattr(obj, "__class__", obj)
-    if hasattr(obj, "__qualname__"):
-        return obj.__qualname__
-    if hasattr(obj, "__name__"):
-        return obj.__name__
-    return str(obj)
+convert_outputs_to_fp32 = ConvertOutputsToFp32
diff --git a/src/accelerate/utils/other.py b/src/accelerate/utils/other.py
new file mode 100644
index 000000000..35c96a21a
--- /dev/null
+++ b/src/accelerate/utils/other.py
@@ -0,0 +1,111 @@
+# 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
+from contextlib import contextmanager
+
+import torch
+
+from ..state import AcceleratorState
+from .dataclasses import DistributedType
+from .imports import is_deepspeed_available, is_tpu_available
+
+
+if is_deepspeed_available():
+    from deepspeed import DeepSpeedEngine
+
+if is_tpu_available():
+    import torch_xla.core.xla_model as xm
+
+
+def extract_model_from_parallel(model):
+    """
+    Extract a model from its distributed containers.
+
+    Args:
+        model (`torch.nn.Module`): The model to extract.
+
+    Returns:
+        `torch.nn.Module`: The extracted model.
+    """
+    options = (torch.nn.parallel.DistributedDataParallel, torch.nn.DataParallel)
+    if is_deepspeed_available():
+        options += (DeepSpeedEngine,)
+
+    while isinstance(model, options):
+        model = model.module
+    return model
+
+
+def wait_for_everyone():
+    """
+    Introduces a blocking point in the script, making sure all processes have reached this point before continuing.
+
+    <Tip warning={true}>
+
+    Make sure all processes will reach this instruction otherwise one of your processes will hang forever.
+
+    </Tip>
+    """
+    if (
+        AcceleratorState().distributed_type == DistributedType.MULTI_GPU
+        or AcceleratorState().distributed_type == DistributedType.MULTI_CPU
+        or AcceleratorState().distributed_type == DistributedType.DEEPSPEED
+    ):
+        torch.distributed.barrier()
+    elif AcceleratorState().distributed_type == DistributedType.TPU:
+        xm.rendezvous("accelerate.utils.wait_for_everyone")
+
+
+def save(obj, f):
+    """
+    Save the data to disk. Use in place of `torch.save()`.
+
+    Args:
+        obj: The data to save
+        f: The file (or file-like object) to use to save the data
+    """
+    if AcceleratorState().distributed_type == DistributedType.TPU:
+        xm.save(obj, f)
+    elif AcceleratorState().local_process_index == 0:
+        torch.save(obj, f)
+
+
+@contextmanager
+def patch_environment(**kwargs):
+    """
+    A context manager that will add each keyword argument passed to `os.environ` and remove them when exiting.
+
+    Will convert the values in `kwargs` to strings and upper-case all the keys.
+    """
+    for key, value in kwargs.items():
+        os.environ[key.upper()] = str(value)
+
+    yield
+
+    for key in kwargs:
+        del os.environ[key.upper()]
+
+
+def get_pretty_name(obj):
+    """
+    Gets a pretty name from `obj`.
+    """
+    if not hasattr(obj, "__qualname__") and not hasattr(obj, "__name__"):
+        obj = getattr(obj, "__class__", obj)
+    if hasattr(obj, "__qualname__"):
+        return obj.__qualname__
+    if hasattr(obj, "__name__"):
+        return obj.__name__
+    return str(obj)
diff --git a/src/accelerate/utils/random.py b/src/accelerate/utils/random.py
new file mode 100644
index 000000000..3958022a4
--- /dev/null
+++ b/src/accelerate/utils/random.py
@@ -0,0 +1,87 @@
+# 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 random
+from typing import List, Optional, Union
+
+import numpy as np
+import torch
+
+from ..state import AcceleratorState
+from .dataclasses import DistributedType, RNGType
+from .imports import is_tpu_available
+
+
+if is_tpu_available():
+    import torch_xla.core.xla_model as xm
+
+
+def set_seed(seed: int, device_specific: bool = False):
+    """
+    Helper function for reproducible behavior to set the seed in `random`, `numpy`, `torch`.
+
+    Args:
+        seed (`int`): The seed to set.
+        device_specific (`bool`, *optional*, defaults to `False`):
+            Whether to differ the seed on each device slightly with `self.process_index`.
+    """
+    if device_specific:
+        seed += AcceleratorState().process_index
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    # ^^ safe to call this function even if cuda is not available
+    if is_tpu_available():
+        xm.set_rng_state(seed)
+
+
+def synchronize_rng_state(rng_type: Optional[RNGType] = None, generator: Optional[torch.Generator] = None):
+    # Get the proper rng state
+    if rng_type == RNGType.TORCH:
+        rng_state = torch.get_rng_state()
+    elif rng_type == RNGType.CUDA:
+        rng_state = torch.cuda.get_rng_state()
+    elif rng_type == RNGType.XLA:
+        assert is_tpu_available(), "Can't synchronize XLA seeds on an environment without TPUs."
+        rng_state = torch.tensor(xm.get_rng_state())
+    elif rng_type == RNGType.GENERATOR:
+        assert generator is not None, "Need a generator to synchronize its seed."
+        rng_state = generator.get_state()
+
+    # Broadcast the rng state from device 0 to other devices
+    state = AcceleratorState()
+    if state.distributed_type == DistributedType.TPU:
+        rng_state = xm.mesh_reduce("random_seed", rng_state, lambda x: x[0])
+    elif state.distributed_type in [DistributedType.DEEPSPEED, DistributedType.MULTI_GPU]:
+        rng_state = rng_state.to(state.device)
+        torch.distributed.broadcast(rng_state, 0)
+        rng_state = rng_state.cpu()
+    elif state.distributed_type == DistributedType.MULTI_CPU:
+        torch.distributed.broadcast(rng_state, 0)
+
+    # Set the broadcast rng state
+    if rng_type == RNGType.TORCH:
+        torch.set_rng_state(rng_state)
+    elif rng_type == RNGType.CUDA:
+        torch.cuda.set_rng_state(rng_state)
+    elif rng_type == RNGType.XLA:
+        xm.set_rng_state(rng_state.item())
+    elif rng_type == RNGType.GENERATOR:
+        generator.set_state(rng_state)
+
+
+def synchronize_rng_states(rng_types: List[Union[str, RNGType]], generator: Optional[torch.Generator] = None):
+    for rng_type in rng_types:
+        synchronize_rng_state(RNGType(rng_type), generator=generator)
diff --git a/tests/test_kwargs_handlers.py b/tests/test_kwargs_handlers.py
index aaf15de0c..8b438d1c6 100644
--- a/tests/test_kwargs_handlers.py
+++ b/tests/test_kwargs_handlers.py
@@ -21,8 +21,8 @@
 import torch
 
 from accelerate import Accelerator, DistributedDataParallelKwargs, GradScalerKwargs
-from accelerate.kwargs_handlers import KwargsHandler
 from accelerate.test_utils import execute_subprocess_async, require_cuda, require_multi_gpu
+from accelerate.utils import KwargsHandler
 
 
 @dataclass
diff --git a/tests/test_memory_utils.py b/tests/test_memory_utils.py
index 42fa151d0..df125ea95 100644
--- a/tests/test_memory_utils.py
+++ b/tests/test_memory_utils.py
@@ -14,7 +14,7 @@
 
 import unittest
 
-from accelerate.memory_utils import find_executable_batch_size
+from accelerate.utils.memory import find_executable_batch_size
 
 
 def raise_fake_out_of_memory():
diff --git a/tests/test_sagemaker.py b/tests/test_sagemaker.py
index f2d587a39..2824493d6 100644
--- a/tests/test_sagemaker.py
+++ b/tests/test_sagemaker.py
@@ -4,7 +4,7 @@
 import pytest
 from accelerate.commands.config.config_args import SageMakerConfig
 from accelerate.commands.launch import _convert_nargs_to_dict
-from accelerate.state import ComputeEnvironment
+from accelerate.utils import ComputeEnvironment
 
 
 @dataclass

This one is rather important -> let's put it in the main init as well so that users import it directly from `accelerate` (and this way they'll be immune to the next cleanup)
@@ -19,7 +19,7 @@
 from accelerate import Accelerator, DistributedType
 
 # New Code #
-from accelerate.memory_utils import find_executable_batch_size
+from accelerate.utils import find_executable_batch_size
So as said above, let's add the batch size finder here.
@@ -5,7 +5,12 @@
 __version__ = "0.8.0.dev0"
 
 from .accelerator import Accelerator
-from .kwargs_handlers import DistributedDataParallelKwargs, GradScalerKwargs, InitProcessGroupKwargs
 from .launchers import debug_launcher, notebook_launcher
-from .state import DistributedType
-from .utils import DeepSpeedPlugin, synchronize_rng_states
+from .utils import (
+    DeepSpeedPlugin,
+    DistributedDataParallelKwargs,
+    DistributedType,
+    GradScalerKwargs,
+    InitProcessGroupKwargs,
+    synchronize_rng_states,
```suggestion
    "memory_utils has been reorganized to utils.memory. Import `find_executable_batchsize` from the main `__init__`: "
    "`from accelerate import find_executable_batch_size` to avoid this warning.",
```
@@ -0,0 +1,28 @@
+# 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.
+
+# flake8: noqa
+# There's no way to ignore "F401 '...' imported but unused" warnings in this
+# module, but to preserve other warnings. So, don't check this module at all
+
+
+import warnings
+
+
+warnings.warn(
+    "memory_utils has been reorganized to utils.memory. Import `find_executable_batchsize` from the main `__init__`: `from accelerate import find_executable_batch_size` to avoid this",
```suggestion
from .utils.memory import find_executable_batch_size
```
@@ -0,0 +1,28 @@
+# 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.
+
+# flake8: noqa
+# There's no way to ignore "F401 '...' imported but unused" warnings in this
+# module, but to preserve other warnings. So, don't check this module at all
+
+
+import warnings
+
+
+warnings.warn(
+    "memory_utils has been reorganized to utils.memory. Import `find_executable_batchsize` from the main `__init__`: `from accelerate import find_executable_batch_size` to avoid this",
+    FutureWarning,
+)
+
+from accelerate import memory_utils
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/340
diff --git a/docs/source/index.mdx b/docs/source/index.mdx
index c41ae4a0e..9d9000097 100644
--- a/docs/source/index.mdx
+++ b/docs/source/index.mdx
@@ -52,7 +52,7 @@ Changing it to work with accelerate is really easy and only adds a few lines of
 + device = accelerator.device
   my_model.to(device)
   # Pass every important object (model, optimizer, dataloader) to *accelerator.prepare*
-+ my_model, my_optimizer, my_training_dataloader = accelerate.prepare(
++ my_model, my_optimizer, my_training_dataloader = accelerator.prepare(
 +     my_model, my_optimizer, my_training_dataloader
 + )
 

_The documentation is not available anymore as the PR was closed or merged._
It's the least I can do, thank's for merging.
https://api.github.com/repos/huggingface/accelerate/pulls/339
diff --git a/examples/by_feature/tracking.py b/examples/by_feature/tracking.py
index 170211e5e..19f8e9dea 100644
--- a/examples/by_feature/tracking.py
+++ b/examples/by_feature/tracking.py
@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import argparse
-import logging
 import os
 
 import torch
@@ -30,9 +29,6 @@
 )
 
 
-logger = logging.getLogger(__name__)
-
-
 ########################################################################
 # This is a fully working simple example to use Accelerate,
 # specifically showcasing the experiment tracking capability,
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index a903386aa..7f1480436 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -26,6 +26,7 @@
 from .checkpointing import load_accelerator_state, load_custom_state, save_accelerator_state, save_custom_state
 from .data_loader import prepare_data_loader
 from .kwargs_handlers import DistributedDataParallelKwargs, GradScalerKwargs, InitProcessGroupKwargs, KwargsHandler
+from .logging import get_logger
 from .optimizer import AcceleratedOptimizer
 from .scheduler import AcceleratedScheduler
 from .state import AcceleratorState, DistributedType, is_deepspeed_available
@@ -52,10 +53,7 @@
 
     from .deepspeed_utils import DeepSpeedEngineWrapper, DeepSpeedOptimizerWrapper
 
-import logging
-
-
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 
 class Accelerator:
diff --git a/src/accelerate/checkpointing.py b/src/accelerate/checkpointing.py
index b35bec237..a94067144 100644
--- a/src/accelerate/checkpointing.py
+++ b/src/accelerate/checkpointing.py
@@ -28,10 +28,10 @@
 if is_tpu_available():
     import torch_xla.core.xla_model as xm
 
-import logging
+from .logging import get_logger
 
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 
 def save_accelerator_state(
diff --git a/src/accelerate/logging.py b/src/accelerate/logging.py
new file mode 100644
index 000000000..2128cdca7
--- /dev/null
+++ b/src/accelerate/logging.py
@@ -0,0 +1,63 @@
+# 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
+
+from .state import AcceleratorState
+
+
+class MultiProcessAdapter(logging.LoggerAdapter):
+    """
+    An adapter to assist with logging in multiprocess.
+
+    `log` takes in an additional `main_process_only` kwarg, which dictates whether it should be called on all processes
+    or only the main executed one. Default is `main_process_only=True`.
+    """
+
+    @staticmethod
+    def _should_log(main_process_only):
+        "Check if log should be performed"
+        return not main_process_only or (main_process_only and AcceleratorState().local_process_index == 0)
+
+    def log(self, level, msg, *args, **kwargs):
+        """
+        Delegates logger call after checking if we should log.
+
+        Accepts a new kwarg of `main_process_only`, which will dictate whether it will be logged across all processes
+        or only the main executed one. Default is `True` if not passed
+        """
+        main_process_only = kwargs.pop("main_process_only", True)
+        if self.isEnabledFor(level) and self._should_log(main_process_only):
+            msg, kwargs = self.process(msg, kwargs)
+            self.logger.log(level, msg, *args, **kwargs)
+
+
+def get_logger(name: str):
+    """
+    Returns a `logging.Logger` for `name` that can handle multiprocessing.
+
+    If a log should be called on all processes, pass `main_process_only=False`
+
+    E.g.
+    ```python
+    logger.info("My log", main_process_only=False)
+    logger.debug("My log", main_process_only=False)
+    ```
+
+    Args:
+        name (`str`):
+            The name for the logger, such as `__file__`
+    """
+    logger = logging.getLogger(name)
+    return MultiProcessAdapter(logger, {})
diff --git a/src/accelerate/tracking.py b/src/accelerate/tracking.py
index de3a19310..218ddf2df 100644
--- a/src/accelerate/tracking.py
+++ b/src/accelerate/tracking.py
@@ -15,11 +15,11 @@
 # Expectation:
 # Provide a project dir name, then each type of logger gets stored in project/{`logging_dir`}
 
-import logging
 import os
 from abc import ABCMeta, abstractmethod, abstractproperty
 from typing import List, Optional, Union
 
+from .logging import get_logger
 from .utils import LoggerType, is_comet_ml_available, is_tensorboard_available, is_wandb_available
 
 
@@ -41,7 +41,7 @@
     _available_trackers.append(LoggerType.COMETML)
 
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 
 def get_available_trackers():

Nice!
@@ -0,0 +1,60 @@
+# 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
+
+from .state import AcceleratorState
+
+
+class MultiProcessAdapter(logging.LoggerAdapter):
+    """
+    An adapter to assist with logging in multiprocess.
+
+    `log` takes in an additional `main_process_only` kwarg, which dictates whether it should be called on all processes
+    or only the main executed one. Default is `main_process_only=True`.
+    """
+
+    @staticmethod
+    def _should_log(main_process_only):
+        "Check if log should be performed"
+        return not main_process_only or (main_process_only and AcceleratorState().local_process_index == 0)
+
+    def log(self, level, msg, *args, **kwargs):
+        """
+        Delegates logger call after checking if we should log
+        """
+        main_process_only = kwargs.pop("main_process_only", True)
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/337
diff --git a/examples/by_feature/fsdp_with_peak_mem_tracking.py b/examples/by_feature/fsdp_with_peak_mem_tracking.py
new file mode 100644
index 000000000..e17e9ff67
--- /dev/null
+++ b/examples/by_feature/fsdp_with_peak_mem_tracking.py
@@ -0,0 +1,369 @@
+# coding=utf-8
+# Copyright 2021 The HuggingFace Inc. 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 gc
+import os
+
+import torch
+from torch.utils.data import DataLoader
+
+from accelerate import Accelerator, DistributedType
+from datasets import load_dataset, load_metric
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#   - FSDP
+#
+# This example also demonstrates the checkpointing and sharding capabilities
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+# New Code #
+# Converting Bytes to Megabytes
+def b2mb(x):
+    return int(x / 2**20)
+
+
+# New Code #
+# This context manager is used to track the peak memory usage of the process
+class TorchTracemalloc:
+    def __enter__(self):
+        gc.collect()
+        torch.cuda.empty_cache()
+        torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
+        self.begin = torch.cuda.memory_allocated()
+        return self
+
+    def __exit__(self, *exc):
+        gc.collect()
+        torch.cuda.empty_cache()
+        self.end = torch.cuda.memory_allocated()
+        self.peak = torch.cuda.max_memory_allocated()
+        self.used = b2mb(self.end - self.begin)
+        self.peaked = b2mb(self.peak - self.begin)
+        # print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")
+
+
+def training_function(config, args):
+    # Initialize accelerator
+    if args.with_tracking:
+        accelerator = Accelerator(
+            cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="wandb", logging_dir=args.logging_dir
+        )
+    else:
+        accelerator = Accelerator()
+    accelerator.print(accelerator.distributed_type)
+
+    if hasattr(args.checkpointing_steps, "isdigit"):
+        if args.checkpointing_steps == "epoch":
+            checkpointing_steps = args.checkpointing_steps
+        elif args.checkpointing_steps.isdigit():
+            checkpointing_steps = int(args.checkpointing_steps)
+        else:
+            raise ValueError(
+                f"Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed."
+            )
+    else:
+        checkpointing_steps = None
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+
+    # We need to initialize the trackers we use, and also store our configuration
+    if args.with_tracking:
+        run = os.path.split(__file__)[-1].split(".")[0]
+        if args.logging_dir:
+            run = os.path.join(args.logging_dir, run)
+            accelerator.print(run)
+        accelerator.init_trackers(run, config)
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
+    datasets = load_dataset("glue", "mrpc")
+    metric = load_metric("glue", "mrpc")
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["idx", "sentence1", "sentence2"],
+    )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    # If the batch size is too big we use gradient accumulation
+    gradient_accumulation_steps = 1
+    if batch_size > MAX_GPU_BATCH_SIZE:
+        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
+        batch_size = MAX_GPU_BATCH_SIZE
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    set_seed(seed)
+
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained(args.model_name_or_path, return_dict=True)
+    # New Code #
+    # For FSDP feature, it is highly recommended and efficient to prepare the model before creating optimizer
+    model = accelerator.prepare(model)
+
+    # Instantiate optimizer
+    # New Code #
+    # For FSDP feature, at present it doesn't support multiple parameter groups,
+    # so we need to create a single parameter group for the whole model
+    optimizer = torch.optim.AdamW(params=model.parameters(), lr=lr, weight_decay=2e-4)
+
+    # Instantiate scheduler
+    lr_scheduler = get_linear_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=10,
+        num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,
+    )
+
+    # New Code #
+    # For FSDP feature, prepare everything except the model as we have already prepared the model
+    # before creating the optimizer
+    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
+    # prepare method.
+    optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+        optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    overall_step = 0
+
+    # Potentially load in the weights and states from a previous save
+    if args.resume_from_checkpoint:
+        if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
+            accelerator.print(f"Resumed from checkpoint: {args.resume_from_checkpoint}")
+            accelerator.load_state(args.resume_from_checkpoint)
+            path = os.path.basename(args.resume_from_checkpoint)
+        else:
+            # Get the most recent checkpoint
+            dirs = [f.name for f in os.scandir(os.getcwd()) if f.is_dir()]
+            dirs.sort(key=os.path.getctime)
+            path = dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
+        # Extract `epoch_{i}` or `step_{i}`
+        training_difference = os.path.splitext(path)[0]
+
+        if "epoch" in training_difference:
+            num_epochs -= int(training_difference.replace("epoch_", ""))
+            resume_step = None
+        else:
+            resume_step = int(training_difference.replace("step_", ""))
+            num_epochs -= resume_step // len(train_dataloader)
+            # If resuming by step, we also need to know exactly how far into the DataLoader we went
+            resume_step = (num_epochs * len(train_dataloader)) - resume_step
+
+    # Now we train the model
+    for epoch in range(num_epochs):
+        # New Code #
+        # context manager to track the peak memory usage during the training epoch
+        with TorchTracemalloc() as tracemalloc:
+            model.train()
+            if args.with_tracking:
+                total_loss = 0
+            for step, batch in enumerate(train_dataloader):
+                # We need to skip steps until we reach the resumed step
+                if args.resume_from_checkpoint and epoch == 0:
+                    if resume_step is not None and step < resume_step:
+                        pass
+                # We could avoid this line since we set the accelerator with `device_placement=True`.
+                batch.to(accelerator.device)
+                outputs = model(**batch)
+                loss = outputs.loss
+                loss = loss / gradient_accumulation_steps
+                # We keep track of the loss at each epoch
+                if args.with_tracking:
+                    total_loss += loss.detach().float()
+                accelerator.backward(loss)
+                if step % gradient_accumulation_steps == 0:
+                    optimizer.step()
+                    lr_scheduler.step()
+                    optimizer.zero_grad()
+                    # accelerator.print(lr_scheduler.get_lr())
+
+                overall_step += 1
+
+                if isinstance(checkpointing_steps, int):
+                    output_dir = f"step_{overall_step}"
+                    if overall_step % checkpointing_steps == 0:
+                        if args.output_dir is not None:
+                            output_dir = os.path.join(args.output_dir, output_dir)
+                        accelerator.save_state(output_dir)
+        # New Code #
+        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
+        accelerator.print("Memory before entering the train : {}".format(b2mb(tracemalloc.begin)))
+        accelerator.print("Memory consumed at the end of the train (end-begin): {}".format(tracemalloc.used))
+        accelerator.print("Peak Memory consumed during the train (max-begin): {}".format(tracemalloc.peaked))
+        accelerator.print(
+            "Total Peak Memory consumed during the train (max): {}".format(
+                tracemalloc.peaked + b2mb(tracemalloc.begin)
+            )
+        )
+        # Logging the peak memory usage of the GPU to the tracker
+        accelerator.log(
+            {
+                "train_total_peak_memory": tracemalloc.peaked + b2mb(tracemalloc.begin),
+            },
+            step=epoch,
+        )
+
+        # New Code #
+        # context manager to track the peak memory usage during the evaluation
+        with TorchTracemalloc() as tracemalloc:
+            model.eval()
+            for step, batch in enumerate(eval_dataloader):
+                # We could avoid this line since we set the accelerator with `device_placement=True`.
+                batch.to(accelerator.device)
+                with torch.no_grad():
+                    outputs = model(**batch)
+                predictions = outputs.logits.argmax(dim=-1)
+                # It is slightly faster to call this once, than multiple times
+                predictions, references = accelerator.gather((predictions, batch["labels"]))
+                metric.add_batch(
+                    predictions=predictions,
+                    references=references,
+                )
+
+            eval_metric = metric.compute()
+            # Use accelerator.print to print only on the main process.
+            accelerator.print(f"epoch {epoch}:", eval_metric)
+            if args.with_tracking:
+                accelerator.log(
+                    {
+                        "accuracy": eval_metric["accuracy"],
+                        "f1": eval_metric["f1"],
+                        "train_loss": total_loss,
+                    },
+                    step=epoch,
+                )
+
+            if checkpointing_steps == "epoch":
+                output_dir = f"epoch_{epoch}"
+                if args.output_dir is not None:
+                    output_dir = os.path.join(args.output_dir, output_dir)
+                accelerator.save_state(output_dir)
+        # New Code #
+        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
+        accelerator.print("Memory before entering the eval : {}".format(b2mb(tracemalloc.begin)))
+        accelerator.print("Memory consumed at the end of the eval (end-begin): {}".format(tracemalloc.used))
+        accelerator.print("Peak Memory consumed during the eval (max-begin): {}".format(tracemalloc.peaked))
+        accelerator.print(
+            "Total Peak Memory consumed during the eval (max): {}".format(tracemalloc.peaked + b2mb(tracemalloc.begin))
+        )
+        # Logging the peak memory usage of the GPU to the tracker
+        accelerator.log(
+            {
+                "eval_total_peak_memory": tracemalloc.peaked + b2mb(tracemalloc.begin),
+            },
+            step=epoch,
+        )
+
+    if args.with_tracking:
+        accelerator.end_training()
+
+
+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.")
+    parser.add_argument(
+        "--checkpointing_steps",
+        type=str,
+        default=None,
+        help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.",
+    )
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        default=None,
+        help="If the training should continue from a checkpoint folder.",
+    )
+    parser.add_argument(
+        "--with_tracking",
+        action="store_true",
+        help="Whether to load in all available experiment trackers from the environment and use them for logging.",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=".",
+        help="Optional save directory where all checkpoint folders will be stored. Default is the current working directory.",
+    )
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help="Location on where to store experiment tracking logs`",
+    )
+    parser.add_argument(
+        "--model_name_or_path",
+        type=str,
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+        required=True,
+    )
+    args = parser.parse_args()
+    config = {"lr": 2e-5, "num_epochs": 3, "seed": 1, "batch_size": 16}
+    training_function(config, args)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/test_examples.py b/tests/test_examples.py
index 919667fc5..84b6adaef 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -40,7 +40,7 @@
 # Should mock `{script_name}.get_dataloaders` via:
 # @mock.patch("{script_name}.get_dataloaders", mocked_dataloaders)
 
-EXCLUDE_EXAMPLES = ["cross_validation.py", "multi_process_metrics.py", "memory.py"]
+EXCLUDE_EXAMPLES = ["cross_validation.py", "multi_process_metrics.py", "memory.py", "fsdp_with_peak_mem_tracking.py"]
 
 
 def mocked_dataloaders(accelerator, batch_size: int = 16):

This should be uncommented no? Also we can remove the print statement.
@@ -0,0 +1,378 @@
+# coding=utf-8
+# Copyright 2021 The HuggingFace Inc. 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 gc
+import os
+
+import torch
+from torch.utils.data import DataLoader
+
+from accelerate import Accelerator, DistributedType
+from datasets import load_dataset, load_metric
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#   - FSDP
+#
+# This example also demonstrates the checkpointing and sharding capabilities
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+# New Code #
+# Converting Bytes to Megabytes
+def b2mb(x):
+    return int(x / 2**20)
+
+
+# New Code #
+# This context manager is used to track the peak memory usage of the process
+class TorchTracemalloc:
+    def __enter__(self):
+        gc.collect()
+        torch.cuda.empty_cache()
+        torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
+        self.begin = torch.cuda.memory_allocated()
+        return self
+
+    def __exit__(self, *exc):
+        gc.collect()
+        torch.cuda.empty_cache()
+        self.end = torch.cuda.memory_allocated()
+        self.peak = torch.cuda.max_memory_allocated()
+        self.used = b2mb(self.end - self.begin)
+        self.peaked = b2mb(self.peak - self.begin)
+        # print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")
+
+
+def training_function(config, args):
+    # Initialize accelerator
+    if args.with_tracking:
+        accelerator = Accelerator(
+            cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="wandb", logging_dir=args.logging_dir
+        )
+    else:
+        accelerator = Accelerator()
+    accelerator.print(accelerator.distributed_type)
+
+    if hasattr(args.checkpointing_steps, "isdigit"):
+        if args.checkpointing_steps == "epoch":
+            checkpointing_steps = args.checkpointing_steps
+        elif args.checkpointing_steps.isdigit():
+            checkpointing_steps = int(args.checkpointing_steps)
+        else:
+            raise ValueError(
+                f"Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed."
+            )
+    else:
+        checkpointing_steps = None
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+
+    # We need to initialize the trackers we use, and also store our configuration
+    if args.with_tracking:
+        run = os.path.split(__file__)[-1].split(".")[0]
+        if args.logging_dir:
+            run = os.path.join(args.logging_dir, run)
+            accelerator.print(run)
+        accelerator.init_trackers(run, config)
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
+    datasets = load_dataset("glue", "mrpc")
+    metric = load_metric("glue", "mrpc")
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["idx", "sentence1", "sentence2"],
+    )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    # If the batch size is too big we use gradient accumulation
+    gradient_accumulation_steps = 1
+    if batch_size > MAX_GPU_BATCH_SIZE:
+        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
+        batch_size = MAX_GPU_BATCH_SIZE
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    set_seed(seed)
+
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained(args.model_name_or_path, return_dict=True)
+    # model = accelerator.prepare(model) #for FSDP model needs to be prepared before the optimizer
+    # accelerator.print(model)
Same, we should show the new code here?
@@ -0,0 +1,378 @@
+# coding=utf-8
+# Copyright 2021 The HuggingFace Inc. 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 gc
+import os
+
+import torch
+from torch.utils.data import DataLoader
+
+from accelerate import Accelerator, DistributedType
+from datasets import load_dataset, load_metric
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#   - FSDP
+#
+# This example also demonstrates the checkpointing and sharding capabilities
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+# New Code #
+# Converting Bytes to Megabytes
+def b2mb(x):
+    return int(x / 2**20)
+
+
+# New Code #
+# This context manager is used to track the peak memory usage of the process
+class TorchTracemalloc:
+    def __enter__(self):
+        gc.collect()
+        torch.cuda.empty_cache()
+        torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
+        self.begin = torch.cuda.memory_allocated()
+        return self
+
+    def __exit__(self, *exc):
+        gc.collect()
+        torch.cuda.empty_cache()
+        self.end = torch.cuda.memory_allocated()
+        self.peak = torch.cuda.max_memory_allocated()
+        self.used = b2mb(self.end - self.begin)
+        self.peaked = b2mb(self.peak - self.begin)
+        # print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")
+
+
+def training_function(config, args):
+    # Initialize accelerator
+    if args.with_tracking:
+        accelerator = Accelerator(
+            cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="wandb", logging_dir=args.logging_dir
+        )
+    else:
+        accelerator = Accelerator()
+    accelerator.print(accelerator.distributed_type)
+
+    if hasattr(args.checkpointing_steps, "isdigit"):
+        if args.checkpointing_steps == "epoch":
+            checkpointing_steps = args.checkpointing_steps
+        elif args.checkpointing_steps.isdigit():
+            checkpointing_steps = int(args.checkpointing_steps)
+        else:
+            raise ValueError(
+                f"Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed."
+            )
+    else:
+        checkpointing_steps = None
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+
+    # We need to initialize the trackers we use, and also store our configuration
+    if args.with_tracking:
+        run = os.path.split(__file__)[-1].split(".")[0]
+        if args.logging_dir:
+            run = os.path.join(args.logging_dir, run)
+            accelerator.print(run)
+        accelerator.init_trackers(run, config)
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
+    datasets = load_dataset("glue", "mrpc")
+    metric = load_metric("glue", "mrpc")
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["idx", "sentence1", "sentence2"],
+    )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    # If the batch size is too big we use gradient accumulation
+    gradient_accumulation_steps = 1
+    if batch_size > MAX_GPU_BATCH_SIZE:
+        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
+        batch_size = MAX_GPU_BATCH_SIZE
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    set_seed(seed)
+
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained(args.model_name_or_path, return_dict=True)
+    # model = accelerator.prepare(model) #for FSDP model needs to be prepared before the optimizer
+    # accelerator.print(model)
+
+    # Instantiate optimizer
+    no_decay = ["bias", "LayerNorm.weight"]
+    optimizer_grouped_parameters = [
+        {
+            "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
+            "weight_decay": 2e-4,
+        },
+        {
+            "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
+            "weight_decay": 0.0,
+        },
+    ]
+    # optimizer = torch.optim.AdamW(params=model.parameters(), lr=lr)
+    optimizer = torch.optim.AdamW(params=optimizer_grouped_parameters, lr=lr)
+
+    # Instantiate scheduler
+    lr_scheduler = get_linear_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=10,
+        num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,
+    )
+
+    # Prepare everything
+    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
+    # prepare method.
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+    # optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+    #     optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    # )
I thought it was more on showing how FSDP reduces memory usage. We can have this to show the proper usage for FSDP, that should be better. Should I rename the file also? 
@@ -0,0 +1,378 @@
+# coding=utf-8
+# Copyright 2021 The HuggingFace Inc. 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 gc
+import os
+
+import torch
+from torch.utils.data import DataLoader
+
+from accelerate import Accelerator, DistributedType
+from datasets import load_dataset, load_metric
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#   - FSDP
+#
+# This example also demonstrates the checkpointing and sharding capabilities
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+# New Code #
+# Converting Bytes to Megabytes
+def b2mb(x):
+    return int(x / 2**20)
+
+
+# New Code #
+# This context manager is used to track the peak memory usage of the process
+class TorchTracemalloc:
+    def __enter__(self):
+        gc.collect()
+        torch.cuda.empty_cache()
+        torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
+        self.begin = torch.cuda.memory_allocated()
+        return self
+
+    def __exit__(self, *exc):
+        gc.collect()
+        torch.cuda.empty_cache()
+        self.end = torch.cuda.memory_allocated()
+        self.peak = torch.cuda.max_memory_allocated()
+        self.used = b2mb(self.end - self.begin)
+        self.peaked = b2mb(self.peak - self.begin)
+        # print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")
+
+
+def training_function(config, args):
+    # Initialize accelerator
+    if args.with_tracking:
+        accelerator = Accelerator(
+            cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="wandb", logging_dir=args.logging_dir
+        )
+    else:
+        accelerator = Accelerator()
+    accelerator.print(accelerator.distributed_type)
+
+    if hasattr(args.checkpointing_steps, "isdigit"):
+        if args.checkpointing_steps == "epoch":
+            checkpointing_steps = args.checkpointing_steps
+        elif args.checkpointing_steps.isdigit():
+            checkpointing_steps = int(args.checkpointing_steps)
+        else:
+            raise ValueError(
+                f"Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed."
+            )
+    else:
+        checkpointing_steps = None
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+
+    # We need to initialize the trackers we use, and also store our configuration
+    if args.with_tracking:
+        run = os.path.split(__file__)[-1].split(".")[0]
+        if args.logging_dir:
+            run = os.path.join(args.logging_dir, run)
+            accelerator.print(run)
+        accelerator.init_trackers(run, config)
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
+    datasets = load_dataset("glue", "mrpc")
+    metric = load_metric("glue", "mrpc")
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["idx", "sentence1", "sentence2"],
+    )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    # If the batch size is too big we use gradient accumulation
+    gradient_accumulation_steps = 1
+    if batch_size > MAX_GPU_BATCH_SIZE:
+        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
+        batch_size = MAX_GPU_BATCH_SIZE
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    set_seed(seed)
+
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained(args.model_name_or_path, return_dict=True)
+    # model = accelerator.prepare(model) #for FSDP model needs to be prepared before the optimizer
+    # accelerator.print(model)
Both should be nice to have in the same example I think. And yes something that clearly says fsdp would be better!
@@ -0,0 +1,378 @@
+# coding=utf-8
+# Copyright 2021 The HuggingFace Inc. 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 gc
+import os
+
+import torch
+from torch.utils.data import DataLoader
+
+from accelerate import Accelerator, DistributedType
+from datasets import load_dataset, load_metric
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#   - FSDP
+#
+# This example also demonstrates the checkpointing and sharding capabilities
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+# New Code #
+# Converting Bytes to Megabytes
+def b2mb(x):
+    return int(x / 2**20)
+
+
+# New Code #
+# This context manager is used to track the peak memory usage of the process
+class TorchTracemalloc:
+    def __enter__(self):
+        gc.collect()
+        torch.cuda.empty_cache()
+        torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
+        self.begin = torch.cuda.memory_allocated()
+        return self
+
+    def __exit__(self, *exc):
+        gc.collect()
+        torch.cuda.empty_cache()
+        self.end = torch.cuda.memory_allocated()
+        self.peak = torch.cuda.max_memory_allocated()
+        self.used = b2mb(self.end - self.begin)
+        self.peaked = b2mb(self.peak - self.begin)
+        # print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")
+
+
+def training_function(config, args):
+    # Initialize accelerator
+    if args.with_tracking:
+        accelerator = Accelerator(
+            cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="wandb", logging_dir=args.logging_dir
+        )
+    else:
+        accelerator = Accelerator()
+    accelerator.print(accelerator.distributed_type)
+
+    if hasattr(args.checkpointing_steps, "isdigit"):
+        if args.checkpointing_steps == "epoch":
+            checkpointing_steps = args.checkpointing_steps
+        elif args.checkpointing_steps.isdigit():
+            checkpointing_steps = int(args.checkpointing_steps)
+        else:
+            raise ValueError(
+                f"Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed."
+            )
+    else:
+        checkpointing_steps = None
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+
+    # We need to initialize the trackers we use, and also store our configuration
+    if args.with_tracking:
+        run = os.path.split(__file__)[-1].split(".")[0]
+        if args.logging_dir:
+            run = os.path.join(args.logging_dir, run)
+            accelerator.print(run)
+        accelerator.init_trackers(run, config)
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
+    datasets = load_dataset("glue", "mrpc")
+    metric = load_metric("glue", "mrpc")
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["idx", "sentence1", "sentence2"],
+    )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    # If the batch size is too big we use gradient accumulation
+    gradient_accumulation_steps = 1
+    if batch_size > MAX_GPU_BATCH_SIZE:
+        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
+        batch_size = MAX_GPU_BATCH_SIZE
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    set_seed(seed)
+
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained(args.model_name_or_path, return_dict=True)
+    # model = accelerator.prepare(model) #for FSDP model needs to be prepared before the optimizer
+    # accelerator.print(model)
Done.
@@ -0,0 +1,378 @@
+# coding=utf-8
+# Copyright 2021 The HuggingFace Inc. 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 gc
+import os
+
+import torch
+from torch.utils.data import DataLoader
+
+from accelerate import Accelerator, DistributedType
+from datasets import load_dataset, load_metric
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#   - FSDP
+#
+# This example also demonstrates the checkpointing and sharding capabilities
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+# New Code #
+# Converting Bytes to Megabytes
+def b2mb(x):
+    return int(x / 2**20)
+
+
+# New Code #
+# This context manager is used to track the peak memory usage of the process
+class TorchTracemalloc:
+    def __enter__(self):
+        gc.collect()
+        torch.cuda.empty_cache()
+        torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
+        self.begin = torch.cuda.memory_allocated()
+        return self
+
+    def __exit__(self, *exc):
+        gc.collect()
+        torch.cuda.empty_cache()
+        self.end = torch.cuda.memory_allocated()
+        self.peak = torch.cuda.max_memory_allocated()
+        self.used = b2mb(self.end - self.begin)
+        self.peaked = b2mb(self.peak - self.begin)
+        # print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")
+
+
+def training_function(config, args):
+    # Initialize accelerator
+    if args.with_tracking:
+        accelerator = Accelerator(
+            cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="wandb", logging_dir=args.logging_dir
+        )
+    else:
+        accelerator = Accelerator()
+    accelerator.print(accelerator.distributed_type)
+
+    if hasattr(args.checkpointing_steps, "isdigit"):
+        if args.checkpointing_steps == "epoch":
+            checkpointing_steps = args.checkpointing_steps
+        elif args.checkpointing_steps.isdigit():
+            checkpointing_steps = int(args.checkpointing_steps)
+        else:
+            raise ValueError(
+                f"Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed."
+            )
+    else:
+        checkpointing_steps = None
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+
+    # We need to initialize the trackers we use, and also store our configuration
+    if args.with_tracking:
+        run = os.path.split(__file__)[-1].split(".")[0]
+        if args.logging_dir:
+            run = os.path.join(args.logging_dir, run)
+            accelerator.print(run)
+        accelerator.init_trackers(run, config)
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
+    datasets = load_dataset("glue", "mrpc")
+    metric = load_metric("glue", "mrpc")
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["idx", "sentence1", "sentence2"],
+    )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    # If the batch size is too big we use gradient accumulation
+    gradient_accumulation_steps = 1
+    if batch_size > MAX_GPU_BATCH_SIZE:
+        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
+        batch_size = MAX_GPU_BATCH_SIZE
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    set_seed(seed)
+
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained(args.model_name_or_path, return_dict=True)
+    # model = accelerator.prepare(model) #for FSDP model needs to be prepared before the optimizer
+    # accelerator.print(model)
_The documentation is not available anymore as the PR was closed or merged._
Since this is more than a "one-feature" script modifying the base, tests will pass if you add the script name to here: https://github.com/huggingface/accelerate/blob/main/tests/test_examples.py#L43
https://api.github.com/repos/huggingface/accelerate/pulls/336
diff --git a/src/accelerate/memory_utils.py b/src/accelerate/memory_utils.py
index 4ffc5d10a..422ead770 100644
--- a/src/accelerate/memory_utils.py
+++ b/src/accelerate/memory_utils.py
@@ -82,5 +82,7 @@ def decorator(*args, **kwargs):
                     gc.collect()
                     torch.cuda.empty_cache()
                     batch_size //= 2
+                else:
+                    raise
 
     return decorator
diff --git a/tests/test_memory_utils.py b/tests/test_memory_utils.py
index cc1b3e827..42fa151d0 100644
--- a/tests/test_memory_utils.py
+++ b/tests/test_memory_utils.py
@@ -80,3 +80,12 @@ def mock_training_loop_function(batch_size, arg1, arg2):
             mock_training_loop_function(128, "hello", "world")
             self.assertIn("Batch size was passed into `f`", cm.exception.args[0])
             self.assertIn("`f(arg1='hello', arg2='world')", cm.exception.args[0])
+
+    def test_any_other_error(self):
+        @find_executable_batch_size(starting_batch_size=16)
+        def mock_training_loop_function(batch_size):
+            raise ValueError("Oops, we had an error!")
+
+        with self.assertRaises(ValueError) as cm:
+            mock_training_loop_function()
+            self.assertIn("Oops, we had an error!", cm.exception.args[0])

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/335
diff --git a/src/accelerate/memory_utils.py b/src/accelerate/memory_utils.py
index 73a298caf..4ffc5d10a 100644
--- a/src/accelerate/memory_utils.py
+++ b/src/accelerate/memory_utils.py
@@ -73,6 +73,8 @@ def decorator(*args, **kwargs):
                 f"Remove this as the decorator already does so: `{function.__name__}({arg_str})`"
             )
         while True:
+            if batch_size == 0:
+                raise RuntimeError("No executable batch size found, reached zero.")
             try:
                 return function(batch_size, *args, **kwargs)
             except Exception as e:
@@ -80,7 +82,5 @@ def decorator(*args, **kwargs):
                     gc.collect()
                     torch.cuda.empty_cache()
                     batch_size //= 2
-                else:
-                    raise
 
     return decorator
diff --git a/tests/test_memory_utils.py b/tests/test_memory_utils.py
index 510d79c71..cc1b3e827 100644
--- a/tests/test_memory_utils.py
+++ b/tests/test_memory_utils.py
@@ -50,6 +50,26 @@ def mock_training_loop_function(batch_size, arg1):
         self.assertListEqual(batch_sizes, [128, 64, 32, 16, 8])
         self.assertListEqual([bs, arg1], [8, "hello"])
 
+    def test_start_zero(self):
+        @find_executable_batch_size(starting_batch_size=0)
+        def mock_training_loop_function(batch_size):
+            pass
+
+        with self.assertRaises(RuntimeError) as cm:
+            mock_training_loop_function()
+            self.assertIn("No executable batch size found, reached zero.", cm.exception.args[0])
+
+    def test_approach_zero(self):
+        @find_executable_batch_size(starting_batch_size=16)
+        def mock_training_loop_function(batch_size):
+            if batch_size > 0:
+                raise_fake_out_of_memory()
+            pass
+
+        with self.assertRaises(RuntimeError) as cm:
+            mock_training_loop_function()
+            self.assertIn("No executable batch size found, reached zero.", cm.exception.args[0])
+
     def test_verbose_guard(self):
         @find_executable_batch_size(starting_batch_size=128)
         def mock_training_loop_function(batch_size, arg1, arg2):

```suggestion
    def test_start_zero(self):
```
@@ -50,6 +50,26 @@ def mock_training_loop_function(batch_size, arg1):
         self.assertListEqual(batch_sizes, [128, 64, 32, 16, 8])
         self.assertListEqual([bs, arg1], [8, "hello"])
 
+    def test_zero(self):
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/334
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index 69464ea0b..f9d35ea48 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -64,7 +64,7 @@ def get_cluster_input():
     else:
         use_cpu = False
 
-    deepspeed_config = None
+    deepspeed_config = {}
     if distributed_type in [DistributedType.MULTI_GPU, DistributedType.NO]:
         use_deepspeed = _ask_field(
             "Do you want to use DeepSpeed? [yes/NO]: ",
@@ -78,7 +78,6 @@ def get_cluster_input():
                 is_deepspeed_available()
             ), "DeepSpeed is not installed => run `pip3 install deepspeed` or build it from source"
 
-        deepspeed_config = {}
         if distributed_type == DistributedType.DEEPSPEED:
             deepspeed_config["zero_stage"] = _ask_field(
                 "What should be your DeepSpeed's ZeRO optimization stage (0, 1, 2, 3)? [2]: ",
@@ -99,6 +98,7 @@ def get_cluster_input():
                 default=1,
             )
 
+    fsdp_config = {}
     if distributed_type in [DistributedType.MULTI_GPU]:
         use_fsdp = _ask_field(
             "Do you want to use FullyShardedDataParallel? [yes/NO]: ",
@@ -108,7 +108,6 @@ def get_cluster_input():
         )
         if use_fsdp:
             distributed_type = DistributedType.FSDP
-        fsdp_config = {}
         if distributed_type == DistributedType.FSDP:
             fsdp_config["sharding_strategy"] = _ask_field(
                 "What should be your sharding strategy ([1] FULL_SHARD, [2] SHARD_GRAD_OP)? [1]: ",

```suggestion
    fsdp_config = {}
```
Can you remove the same thing at line 112 as a result?
@@ -99,6 +99,7 @@ def get_cluster_input():
                 default=1,
             )
 
+    fsdp_config = None
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/331
diff --git a/src/accelerate/test_utils/testing.py b/src/accelerate/test_utils/testing.py
index ac6df7741..0b8e2250f 100644
--- a/src/accelerate/test_utils/testing.py
+++ b/src/accelerate/test_utils/testing.py
@@ -26,7 +26,7 @@
 import torch
 
 from ..state import AcceleratorState, is_tpu_available
-from ..utils import gather, is_tensorflow_available
+from ..utils import gather, is_comet_ml_available, is_tensorflow_available, is_wandb_available
 
 
 def parse_flag_from_env(key, default=False):
@@ -53,10 +53,51 @@ def slow(test_case):
     Decorator marking a test as slow. Slow tests are skipped by default. Set the RUN_SLOW environment variable to a
     truthy value to run them.
     """
-    if not _run_slow_tests:
-        return unittest.skip("test is slow")(test_case)
-    else:
-        return test_case
+    return unittest.skipUnless(_run_slow_tests, "test is slow")(test_case)
+
+
+def require_cuda(test_case):
+    """
+    Decorator marking a test that requires CUDA. These tests are skipped when there are no GPU available.
+    """
+    return unittest.skipUnless(torch.cuda.is_available(), "test requires a GPU")(test_case)
+
+
+def require_tpu(test_case):
+    """
+    Decorator marking a test that requires TPUs. These tests are skipped when there are no TPUs available.
+    """
+    return unittest.skipUnless(is_tpu_available(), "test requires TPU")(test_case)
+
+
+def require_multi_gpu(test_case):
+    """
+    Decorator marking a test that requires a multi-GPU setup. These tests are skipped on a machine without multiple
+    GPUs.
+    """
+    return unittest.skipUnless(torch.cuda.device_count() > 1, "test requires multiple GPUs")(test_case)
+
+
+def require_tensorflow(test_case):
+    """
+    Decorator marking a test that requires TensorFlow installed. These tests are skipped when TensorFlow isn't
+    installed
+    """
+    return unittest.skipUnless(is_tensorflow_available(), "test requires TensorFlow")(test_case)
+
+
+def require_wandb(test_case):
+    """
+    Decorator marking a test that requires wandb installed. These tests are skipped when wandb isn't installed
+    """
+    return unittest.skipUnless(is_wandb_available(), "test requires wandb")(test_case)
+
+
+def require_comet_ml(test_case):
+    """
+    Decorator marking a test that requires comet_ml installed. These tests are skipped when comet_ml isn't installed
+    """
+    return unittest.skipUnless(is_comet_ml_available(), "test requires comet_ml")(test_case)
 
 
 class TempDirTestCase(unittest.TestCase):
@@ -136,48 +177,6 @@ def are_the_same_tensors(tensor):
     return True
 
 
-def require_cuda(test_case):
-    """
-    Decorator marking a test that requires CUDA. These tests are skipped when there are no GPU available.
-    """
-    if not torch.cuda.is_available():
-        return unittest.skip("test requires a GPU")(test_case)
-    else:
-        return test_case
-
-
-def require_tpu(test_case):
-    """
-    Decorator marking a test that requires TPUs. These tests are skipped when there are no TPUs available.
-    """
-    if not is_tpu_available():
-        return unittest.skip("test requires TPU")(test_case)
-    else:
-        return test_case
-
-
-def require_multi_gpu(test_case):
-    """
-    Decorator marking a test that requires a multi-GPU setup. These tests are skipped on a machine without multiple
-    GPUs.
-    """
-    if torch.cuda.device_count() < 2:
-        return unittest.skip("test requires multiple GPUs")(test_case)
-    else:
-        return test_case
-
-
-def require_tensorflow(test_case):
-    """
-    Decorator marking a test that requires TensorFlow installed. These tests are skipped when TensorFlow isn't
-    installed
-    """
-    if not is_tensorflow_available():
-        return unittest.skip("test requires TensorFlow")(test_case)
-    else:
-        return test_case
-
-
 class _RunOutput:
     def __init__(self, returncode, stdout, stderr):
         self.returncode = returncode
diff --git a/tests/test_tracking.py b/tests/test_tracking.py
index c5b940a46..e9a2170ee 100644
--- a/tests/test_tracking.py
+++ b/tests/test_tracking.py
@@ -26,10 +26,19 @@
 
 # We use TF to parse the logs
 from accelerate import Accelerator
-from accelerate.test_utils.testing import MockingTestCase, TempDirTestCase, require_tensorflow
+from accelerate.test_utils.testing import (
+    MockingTestCase,
+    TempDirTestCase,
+    require_comet_ml,
+    require_tensorflow,
+    require_wandb,
+)
 from accelerate.tracking import CometMLTracker, GeneralTracker
-from accelerate.utils import is_tensorflow_available
-from comet_ml import OfflineExperiment
+from accelerate.utils import is_comet_ml_available, is_tensorflow_available
+
+
+if is_comet_ml_available():
+    from comet_ml import OfflineExperiment
 
 
 if is_tensorflow_available():
@@ -110,6 +119,7 @@ def test_logging_dir(self):
             _ = Accelerator(log_with="tensorboard", logging_dir=dirpath)
 
 
+@require_wandb
 @mock.patch.dict(os.environ, {"WANDB_MODE": "offline"})
 class WandBTrackingTest(TempDirTestCase, MockingTestCase):
     def setUp(self):
@@ -179,6 +189,7 @@ def offline_init(self, run_name: str, tmpdir: str):
     logger.info("Make sure to log any initial configurations with `self.store_init_configuration` before training!")
 
 
+@require_comet_ml
 @mock.patch.object(CometMLTracker, "__init__", offline_init)
 class CometMLTest(unittest.TestCase):
     @staticmethod

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/330
diff --git a/src/accelerate/commands/config/config_args.py b/src/accelerate/commands/config/config_args.py
index fadbe6211..58db9b24e 100644
--- a/src/accelerate/commands/config/config_args.py
+++ b/src/accelerate/commands/config/config_args.py
@@ -139,6 +139,13 @@ class ClusterConfig(BaseConfig):
     # args for fsdp
     fsdp_config: dict = None
 
+    def __post_init__(self):
+        if self.deepspeed_config is None:
+            self.deepspeed_config = {}
+        if self.fsdp_config is None:
+            self.fsdp_config = {}
+        return super().__post_init__()
+
 
 @dataclass
 class SageMakerConfig(BaseConfig):

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/329
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index eb3b24ce9..f62d9e7a4 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -40,6 +40,7 @@
     gather,
     get_pretty_name,
     pad_across_processes,
+    reduce,
     save,
     wait_for_everyone,
 )
@@ -556,7 +557,7 @@ def clip_grad_value_(self, parameters, clip_value):
 
     def gather(self, tensor):
         """
-        Gather the values in *tensor* accross all processes and concatenate them on the first dimension. Useful to
+        Gather the values in *tensor* across all processes and concatenate them on the first dimension. Useful to
         regroup the predictions from all processes when doing evaluation.
 
         Note:
@@ -572,6 +573,18 @@ def gather(self, tensor):
         """
         return gather(tensor)
 
+    def reduce(self, tensor: torch.Tensor, reduction="sum"):
+        """
+        Reduce the values in *tensor* across all processes based on *reduction*.
+
+        Args:
+            tensor (`torch.Tensor`):
+                The tensors to reduce across all processes.
+            reduction (`str`, *optional*, defaults to "sum"):
+                A reduction type, can be one of 'sum', 'mean', or 'none'. If 'none', will not perform any operation.
+        """
+        reduce(tensor, reduction)
+
     def pad_across_processes(self, tensor, dim=0, pad_index=0, pad_first=False):
         """
         Recursively pad the tensors in a nested list/tuple/dictionary of tensors from all devices to the same size so
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 9bcd87f5f..c2d05c162 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -24,6 +24,7 @@
 
 import numpy as np
 import torch
+from torch.distributed import ReduceOp
 
 from packaging import version
 
@@ -621,6 +622,39 @@ def _pad_across_processes(tensor, dim=0, pad_index=0, pad_first=False):
     )
 
 
+def reduce(tensor, reduction="mean"):
+    """
+    Recursively reduce the tensors in a nested list/tuple/dictionary of lists of tensors across all processes by the
+    mean of a given operation.
+
+    Args:
+        tensor (nested list/tuple/dictionary of `torch.Tensor`):
+            The data to reduce.
+        reduction (`str`, *optional*, defaults to `"mean"`):
+            A reduction method. Can be of "mean", "sum", or "none"
+
+    Returns:
+        The same data structure as `data` with all the tensors reduced.
+    """
+
+    def _reduce_across_processes(tensor, reduction="mean"):
+        state = AcceleratorState()
+        cloned_tensor = tensor.clone()
+        if state.distributed_type == DistributedType.TPU:
+            xm.all_reduce("sum", cloned_tensor)
+            return cloned_tensor
+        elif state.distributed_type in [DistributedType.DEEPSPEED, DistributedType.MULTI_GPU]:
+            torch.distributed.reduce(cloned_tensor, ReduceOp.SUM)
+            return cloned_tensor
+        else:
+            if reduction == "sum":
+                return cloned_tensor.sum()
+            else:
+                return cloned_tensor.mean()
+
+    return recursively_apply(_reduce_across_processes, tensor, error_on_other_type=True, reduction=reduction)
+
+
 def wait_for_everyone():
     """
     Introduces a blocking point in the script, making sure all processes have reached this point before continuing.

```suggestion
        reduction (`str`, *optional*, defaults to `"mean"`):
```
@@ -621,6 +622,35 @@ def _pad_across_processes(tensor, dim=0, pad_index=0, pad_first=False):
     )
 
 
+def reduce(tensor: torch.Tensor, reduction="mean"):
+    """
+    Reduce `tensor` across all processes and perform a `reduction`.
+
+    Args:
+        tensor (`torch.Tensor`):
+            The data to reduce
+        reduction (`str`, *optional*, defaults to "mean"):
```suggestion
    Reduce a tensor across all processes by the mean of a given operation.
```
The arguments are described just below, so I like to keep the general description pure English :-)
@@ -621,6 +622,35 @@ def _pad_across_processes(tensor, dim=0, pad_index=0, pad_first=False):
     )
 
 
+def reduce(tensor: torch.Tensor, reduction="mean"):
+    """
+    Reduce `tensor` across all processes and perform a `reduction`.
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/326
diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml
index 0a6d45892..1ee741653 100644
--- a/docs/source/_toctree.yml
+++ b/docs/source/_toctree.yml
@@ -23,4 +23,6 @@
     title: Checkpointing
   - local: tracking
     title: Experiment Tracking
+  - local: memory
+    title: Memory Utilities
   title: API Reference
diff --git a/docs/source/memory.mdx b/docs/source/memory.mdx
new file mode 100644
index 000000000..31a2bc66f
--- /dev/null
+++ b/docs/source/memory.mdx
@@ -0,0 +1,51 @@
+<!--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.
+-->
+
+# Memory Utilities
+
+One of the most frustrating errors when it comes to running training scripts is hitting "CUDA Out-of-Memory", 
+as the entire script needs to be restarted, progress is lost, and typically a developer would want to simply
+start their script and let it run.
+
+`Accelerate` provides a utility heavily based on [toma](https://github.com/BlackHC/toma) to give this capability.
+
+## find_executable_batch_size
+
+This algorithm operates with exponential decay, decreasing the batch size in half after each failed run on some 
+training script. To use it, restructure your training function to include an inner function that includes this wrapper, 
+and build your dataloaders inside it. At a minimum, this could look like 4 new lines of code. 
+> Note: The inner function *must* take in the batch size as the first parameter, but we do not pass one to it when called. The wrapper handles this for us
+
+```diff
+def training_function(args):
+    accelerator = Accelerator()
+    model = get_model()
+    model.to(accelerator.device)
+    optimizer = get_optimizer()
+
++   @find_executable_batch_size(starting_batch_size=args.batch_size)
++   def inner_training_loop(batch_size):
++       nonlocal model, optimizer # Ensure they can be used in our context
+        train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)
+        lr_scheduler = get_scheduler(
+            optimizer, 
+            num_training_steps=len(train_dataloader)*num_epochs
+        )
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+            model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+        )
+        train(model, optimizer, train_dataloader, lr_scheduler)
+        validate(model, eval_dataloader)
++   inner_training_loop()
+```
+
+[[autodoc]] memory_utils.find_executable_batch_size
\ No newline at end of file
diff --git a/examples/by_feature/memory.py b/examples/by_feature/memory.py
new file mode 100644
index 000000000..67263de52
--- /dev/null
+++ b/examples/by_feature/memory.py
@@ -0,0 +1,218 @@
+# 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 torch
+from torch.utils.data import DataLoader
+
+from accelerate import Accelerator, DistributedType
+
+# New Code #
+from accelerate.memory_utils import find_executable_batch_size
+from datasets import load_dataset, load_metric
+from transformers import (
+    AdamW,
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    get_linear_schedule_with_warmup,
+    set_seed,
+)
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate,
+# specifically showcasing how to ensure out-of-memory errors never
+# iterrupt training, and builds off the `nlp_example.py` script.
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#
+# New additions from the base script can be found quickly by
+# looking for the # New Code # tags
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+def get_dataloaders(accelerator: Accelerator, batch_size: int = 16):
+    """
+    Creates a set of `DataLoader`s for the `glue` dataset,
+    using "bert-base-cased" as the tokenizer.
+
+    Args:
+        accelerator (`Accelerator`):
+            An `Accelerator` object
+        batch_size (`int`, *optional*):
+            The batch size for the train and validation DataLoaders.
+    """
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    datasets = load_dataset("glue", "mrpc")
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["idx", "sentence1", "sentence2"],
+    )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    return train_dataloader, eval_dataloader
+
+
+def training_function(config, args):
+    # Initialize accelerator
+    accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    correct_bias = config["correct_bias"]
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+
+    metric = load_metric("glue", "mrpc")
+
+    # If the batch size is too big we use gradient accumulation
+    gradient_accumulation_steps = 1
+    if batch_size > MAX_GPU_BATCH_SIZE:
+        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
+        batch_size = MAX_GPU_BATCH_SIZE
+
+    set_seed(seed)
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
+
+    # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
+    # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
+    # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
+    model = model.to(accelerator.device)
+
+    # Instantiate optimizer
+    optimizer = AdamW(params=model.parameters(), lr=lr, correct_bias=correct_bias)
+
+    # New Code #
+    # We now can define an inner training loop function. It should take a batch size as the only parameter,
+    # and build the dataloaders in there.
+    # It also gets our decorator
+    @find_executable_batch_size(starting_batch_size=batch_size)
+    def inner_training_loop(batch_size):
+        # And now just move everything below under this function
+        # Ensure that anything declared outside this function is set as `nonlocal`
+        # so it is in scope
+        nonlocal model, optimizer
+        train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)
+
+        # Instantiate scheduler
+        lr_scheduler = get_linear_schedule_with_warmup(
+            optimizer=optimizer,
+            num_warmup_steps=100,
+            num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,
+        )
+
+        # Prepare everything
+        # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
+        # prepare method.
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+            model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+        )
+
+        # Now we train the model
+        for epoch in range(num_epochs):
+            model.train()
+            for step, batch in enumerate(train_dataloader):
+                # We could avoid this line since we set the accelerator with `device_placement=True`.
+                batch.to(accelerator.device)
+                outputs = model(**batch)
+                loss = outputs.loss
+                loss = loss / gradient_accumulation_steps
+                accelerator.backward(loss)
+                if step % gradient_accumulation_steps == 0:
+                    optimizer.step()
+                    lr_scheduler.step()
+                    optimizer.zero_grad()
+
+            model.eval()
+            for step, batch in enumerate(eval_dataloader):
+                # We could avoid this line since we set the accelerator with `device_placement=True`.
+                batch.to(accelerator.device)
+                with torch.no_grad():
+                    outputs = model(**batch)
+                predictions = outputs.logits.argmax(dim=-1)
+                predictions, references = accelerator.gather((predictions, batch["labels"]))
+                metric.add_batch(
+                    predictions=predictions,
+                    references=references,
+                )
+
+            eval_metric = metric.compute()
+            # Use accelerator.print to print only on the main process.
+            accelerator.print(f"epoch {epoch}:", eval_metric)
+
+    # New Code #
+    # And call it at the end with no arguments
+    # Note: You could also refactor this outside of your training loop function
+    inner_training_loop()
+
+
+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()
+    config = {"lr": 2e-5, "num_epochs": 3, "correct_bias": True, "seed": 42, "batch_size": 16}
+    training_function(config, args)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/accelerate/memory_utils.py b/src/accelerate/memory_utils.py
new file mode 100644
index 000000000..73a298caf
--- /dev/null
+++ b/src/accelerate/memory_utils.py
@@ -0,0 +1,86 @@
+# 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.
+
+"""
+A collection of utilities for ensuring that training can always occur. Heavily influenced by the
+[toma](https://github.com/BlackHC/toma) library.
+"""
+
+import functools
+import gc
+import inspect
+
+import torch
+
+
+def should_reduce_batch_size(exception: Exception) -> bool:
+    """
+    Checks if `exception` relates to CUDA out-of-memory, CUDNN not supported, or CPU out-of-memory
+
+    Args:
+        exception (`Exception`):
+            An exception
+    """
+    _statements = [
+        "CUDA out of memory.",  # CUDA OOM
+        "cuDNN error: CUDNN_STATUS_NOT_SUPPORTED.",  # CUDNN SNAFU
+        "DefaultCPUAllocator: can't allocate memory",  # CPU OOM
+    ]
+    if isinstance(exception, RuntimeError) and len(exception.args) == 1:
+        return any(err in exception.args[0] for err in _statements)
+    return False
+
+
+def find_executable_batch_size(function: callable = None, starting_batch_size: int = 128):
+    """
+    A basic decorator that will try to execute `function`. If it fails from exceptions related to out-of-memory or
+    CUDNN, the batch size is cut in half and passed to `function`
+
+    `function` must take in a `batch_size` parameter as its first argument.
+
+    Args:
+        function (`callable`, *optional*):
+            A function to wrap
+        starting_batch_size (`int`, *optional*):
+            The batch size to try and fit into memory
+    """
+    if function is None:
+        return functools.partial(find_executable_batch_size, starting_batch_size=starting_batch_size)
+
+    batch_size = starting_batch_size
+
+    def decorator(*args, **kwargs):
+        nonlocal batch_size
+        gc.collect()
+        torch.cuda.empty_cache()
+        params = list(inspect.signature(function).parameters.keys())
+        # Guard against user error
+        if len(params) < (len(args) + 1):
+            arg_str = ", ".join([f"{arg}={value}" for arg, value in zip(params[1:], args[1:])])
+            raise TypeError(
+                f"Batch size was passed into `{function.__name__}` as the first argument when called."
+                f"Remove this as the decorator already does so: `{function.__name__}({arg_str})`"
+            )
+        while True:
+            try:
+                return function(batch_size, *args, **kwargs)
+            except Exception as e:
+                if should_reduce_batch_size(e):
+                    gc.collect()
+                    torch.cuda.empty_cache()
+                    batch_size //= 2
+                else:
+                    raise
+
+    return decorator
diff --git a/tests/test_examples.py b/tests/test_examples.py
index aa78fed9f..919667fc5 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -40,7 +40,7 @@
 # Should mock `{script_name}.get_dataloaders` via:
 # @mock.patch("{script_name}.get_dataloaders", mocked_dataloaders)
 
-EXCLUDE_EXAMPLES = ["cross_validation.py", "multi_process_metrics.py"]
+EXCLUDE_EXAMPLES = ["cross_validation.py", "multi_process_metrics.py", "memory.py"]
 
 
 def mocked_dataloaders(accelerator, batch_size: int = 16):
diff --git a/tests/test_memory_utils.py b/tests/test_memory_utils.py
new file mode 100644
index 000000000..510d79c71
--- /dev/null
+++ b/tests/test_memory_utils.py
@@ -0,0 +1,62 @@
+# 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 unittest
+
+from accelerate.memory_utils import find_executable_batch_size
+
+
+def raise_fake_out_of_memory():
+    raise RuntimeError("CUDA out of memory.")
+
+
+class MemoryTest(unittest.TestCase):
+    def test_memory_implicit(self):
+        batch_sizes = []
+
+        @find_executable_batch_size(starting_batch_size=128)
+        def mock_training_loop_function(batch_size):
+            nonlocal batch_sizes
+            batch_sizes.append(batch_size)
+            if batch_size != 8:
+                raise_fake_out_of_memory()
+
+        mock_training_loop_function()
+        self.assertListEqual(batch_sizes, [128, 64, 32, 16, 8])
+
+    def test_memory_explicit(self):
+        batch_sizes = []
+
+        @find_executable_batch_size(starting_batch_size=128)
+        def mock_training_loop_function(batch_size, arg1):
+            nonlocal batch_sizes
+            batch_sizes.append(batch_size)
+            if batch_size != 8:
+                raise_fake_out_of_memory()
+            return batch_size, arg1
+
+        bs, arg1 = mock_training_loop_function("hello")
+        self.assertListEqual(batch_sizes, [128, 64, 32, 16, 8])
+        self.assertListEqual([bs, arg1], [8, "hello"])
+
+    def test_verbose_guard(self):
+        @find_executable_batch_size(starting_batch_size=128)
+        def mock_training_loop_function(batch_size, arg1, arg2):
+            if batch_size != 8:
+                raise raise_fake_out_of_memory()
+
+        with self.assertRaises(TypeError) as cm:
+            mock_training_loop_function(128, "hello", "world")
+            self.assertIn("Batch size was passed into `f`", cm.exception.args[0])
+            self.assertIn("`f(arg1='hello', arg2='world')", cm.exception.args[0])

We might hit some custom types of users here. The only thing we need to check is that there is a `batch_size` attribute and that it's an iterator, no?
@@ -0,0 +1,82 @@
+# 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.
+
+"""
+A collection of utilities for ensuring that training can always occur. Heavily influenced by the
+[toma](https://github.com/BlackHC/toma) library.
+"""
+
+import functools
+from typing import List, Union
+
+from torch.utils import DataLoader
+
+from accelerate.data_loader import DataLoaderDispatcher, DataLoaderShard
+
+
+def should_reduce_batch_size(exception: Exception) -> bool:
+    """
+    Checks if `exception` relates to CUDA out-of-memory, CUDNN not supported, or CPU out-of-memory
+
+    Args:
+        exception (`Exception`):
+            An exception
+    """
+    _statements = [
+        "CUDA out of memory.",  # CUDA OOM
+        "cuDNN error: CUDNN_STATUS_NOT_SUPPORTED.",  # CUDNN SNAFU
+        "DefaultCPUAllocator: can't allocate memory",  # CPU OOM
+    ]
+    if isinstance(exception, RuntimeError) and len(exception.args) == 1:
+        return any(err in exception.args[0] for err in _statements)
+    return False
+
+
+def memory_aware(
+    function: callable = None, dataloaders: List[Union[DataLoader, DataLoaderDispatcher, DataLoaderShard]] = None
+):
+    """
+    A decorator that will reduce the batch size in half of all `dataloaders` if the wrapped function fails from any
+    exceptions related to out-of-memory or CUDNN not supported until the function executes completely.
+
+    Args:
+        function (`callable`, *optional*):
+            A function to wrap that utilizes all declared `dataloaders`
+        dataloaders (list of `DataLoader`, [`~data_loader.DataLoaderDispatcher`], or [`~data_loader.DataLoaderShard`], *optional*):
+            A list of prepared or unprepared `DataLoaders`.
+    """
+    if not isinstance(dataloaders, (tuple, list)):
+        dataloaders = [dataloaders]
+    if not all(
+        (isinstance(dataloader, (DataLoaderDispatcher, DataLoaderShard, DataLoader)) for dataloader in dataloaders)
+    ):
+        raise TypeError(
+            "Unsupported operation attempted. One or more dataloaders passed were not of type(s) `DataLoaderDispatcher`, `DataLoaderShard`, or `torch.utils.DataLoader`"
+        )
Does this actually work? PyTorch fixes the `batch_sampler` at init, so I think this is ignore by the `dataloader`.
@@ -0,0 +1,82 @@
+# 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.
+
+"""
+A collection of utilities for ensuring that training can always occur. Heavily influenced by the
+[toma](https://github.com/BlackHC/toma) library.
+"""
+
+import functools
+from typing import List, Union
+
+from torch.utils import DataLoader
+
+from accelerate.data_loader import DataLoaderDispatcher, DataLoaderShard
+
+
+def should_reduce_batch_size(exception: Exception) -> bool:
+    """
+    Checks if `exception` relates to CUDA out-of-memory, CUDNN not supported, or CPU out-of-memory
+
+    Args:
+        exception (`Exception`):
+            An exception
+    """
+    _statements = [
+        "CUDA out of memory.",  # CUDA OOM
+        "cuDNN error: CUDNN_STATUS_NOT_SUPPORTED.",  # CUDNN SNAFU
+        "DefaultCPUAllocator: can't allocate memory",  # CPU OOM
+    ]
+    if isinstance(exception, RuntimeError) and len(exception.args) == 1:
+        return any(err in exception.args[0] for err in _statements)
+    return False
+
+
+def memory_aware(
+    function: callable = None, dataloaders: List[Union[DataLoader, DataLoaderDispatcher, DataLoaderShard]] = None
+):
+    """
+    A decorator that will reduce the batch size in half of all `dataloaders` if the wrapped function fails from any
+    exceptions related to out-of-memory or CUDNN not supported until the function executes completely.
+
+    Args:
+        function (`callable`, *optional*):
+            A function to wrap that utilizes all declared `dataloaders`
+        dataloaders (list of `DataLoader`, [`~data_loader.DataLoaderDispatcher`], or [`~data_loader.DataLoaderShard`], *optional*):
+            A list of prepared or unprepared `DataLoaders`.
+    """
+    if not isinstance(dataloaders, (tuple, list)):
+        dataloaders = [dataloaders]
+    if not all(
+        (isinstance(dataloader, (DataLoaderDispatcher, DataLoaderShard, DataLoader)) for dataloader in dataloaders)
+    ):
+        raise TypeError(
+            "Unsupported operation attempted. One or more dataloaders passed were not of type(s) `DataLoaderDispatcher`, `DataLoaderShard`, or `torch.utils.DataLoader`"
+        )
+
+    if function is None:
+        return functools.partial(memory_aware, dataloaders=dataloaders)
+
+    def decorator(*args, **kwargs):
+        while True:
+            try:
+                return function(*args, **kwargs)
+            except Exception as e:
+                if should_reduce_batch_size(e):
+                    for dataloader in dataloaders:
+                        dataloader.batch_size /= 2
Still not fond of that name ;-) What about `@find_max_batch_size`?
@@ -0,0 +1,218 @@
+# 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 torch
+from torch.utils.data import DataLoader
+
+from accelerate import Accelerator, DistributedType
+
+# New Code #
+from accelerate.memory_utils import memory_aware
+from datasets import load_dataset, load_metric
+from transformers import (
+    AdamW,
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    get_linear_schedule_with_warmup,
+    set_seed,
+)
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate,
+# specifically showcasing how to ensure out-of-memory errors never
+# iterrupt training, and builds off the `nlp_example.py` script.
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#
+# New additions from the base script can be found quickly by
+# looking for the # New Code # tags
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+def get_dataloaders(accelerator: Accelerator, batch_size: int = 16):
+    """
+    Creates a set of `DataLoader`s for the `glue` dataset,
+    using "bert-base-cased" as the tokenizer.
+
+    Args:
+        accelerator (`Accelerator`):
+            An `Accelerator` object
+        batch_size (`int`, *optional*):
+            The batch size for the train and validation DataLoaders.
+    """
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    datasets = load_dataset("glue", "mrpc")
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["idx", "sentence1", "sentence2"],
+    )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    return train_dataloader, eval_dataloader
+
+
+def training_function(config, args):
+    # Initialize accelerator
+    accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    correct_bias = config["correct_bias"]
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+
+    metric = load_metric("glue", "mrpc")
+
+    # If the batch size is too big we use gradient accumulation
+    gradient_accumulation_steps = 1
+    if batch_size > MAX_GPU_BATCH_SIZE:
+        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
+        batch_size = MAX_GPU_BATCH_SIZE
+
+    set_seed(seed)
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
+
+    # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
+    # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
+    # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
+    model = model.to(accelerator.device)
+
+    # Instantiate optimizer
+    optimizer = AdamW(params=model.parameters(), lr=lr, correct_bias=correct_bias)
+
+    # New Code #
+    # We now can define an inner training loop function. It should take a batch size as the only parameter,
+    # and build the dataloaders in there.
+    # It also gets our decorator
+    @memory_aware(starting_batch_size=batch_size)
It's cleaner to pass the string on several lines (with the `"""` you will get the new lines and indentation in the text displayed by the warning).
```suggestion
                f"Batch size was passed into `{function.__name__}` as the first argument when called."
                f"Remove this as the decorator already does so: `{function.__name__}({arg_str})`"
```
@@ -0,0 +1,87 @@
+# 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.
+
+"""
+A collection of utilities for ensuring that training can always occur. Heavily influenced by the
+[toma](https://github.com/BlackHC/toma) library.
+"""
+
+import functools
+import gc
+import inspect
+
+import torch
+
+
+def should_reduce_batch_size(exception: Exception) -> bool:
+    """
+    Checks if `exception` relates to CUDA out-of-memory, CUDNN not supported, or CPU out-of-memory
+
+    Args:
+        exception (`Exception`):
+            An exception
+    """
+    _statements = [
+        "CUDA out of memory.",  # CUDA OOM
+        "cuDNN error: CUDNN_STATUS_NOT_SUPPORTED.",  # CUDNN SNAFU
+        "DefaultCPUAllocator: can't allocate memory",  # CPU OOM
+    ]
+    if isinstance(exception, RuntimeError) and len(exception.args) == 1:
+        return any(err in exception.args[0] for err in _statements)
+    return False
+
+
+def memory_aware(function: callable = None, starting_batch_size: int = 128):
+    """
+    A basic decorator that will try to execute `function`. If it fails from exceptions related to out-of-memory or
+    CUDNN, the batch size is cut in half and passed to `function`
+
+    `function` must take in a `batch_size` parameter as its first argument.
+
+    Args:
+        function (`callable`, *optional*):
+            A function to wrap
+        starting_batch_size (`int`, *optional*):
+            The batch size to try and fit into memory
+    """
+    if function is None:
+        return functools.partial(memory_aware, starting_batch_size=starting_batch_size)
+
+    batch_size = starting_batch_size
+
+    def decorator(*args, **kwargs):
+        nonlocal batch_size
+        gc.collect()
+        torch.cuda.empty_cache()
+        params = list(inspect.signature(function).parameters.keys())
+        # Guard against user error
+        if len(params) < (len(args) + 1):
+            arg_str = ", ".join([f"{arg}={value}" for arg, value in zip(params[1:], args[1:])])
+            raise TypeError(
+                f"""Batch size was passed into `{function.__name__}` as the first argument when called.
+            Remove this as the decorator already does so:
+                `{function.__name__}({arg_str})`"""
Use something better than `f` has a function name please ;-)
@@ -0,0 +1,62 @@
+# 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 unittest
+
+from accelerate.memory_utils import memory_aware
+
+
+def raise_fake_out_of_memory():
+    raise RuntimeError("CUDA out of memory.")
+
+
+class MemoryTest(unittest.TestCase):
+    def test_memory_implicit(self):
+        batch_sizes = []
+
+        @memory_aware(starting_batch_size=128)
+        def f(batch_size):
+            nonlocal batch_sizes
+            batch_sizes.append(batch_size)
+            if batch_size != 8:
+                raise_fake_out_of_memory()
+
+        f()
Same here.
@@ -0,0 +1,62 @@
+# 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 unittest
+
+from accelerate.memory_utils import memory_aware
+
+
+def raise_fake_out_of_memory():
+    raise RuntimeError("CUDA out of memory.")
+
+
+class MemoryTest(unittest.TestCase):
+    def test_memory_implicit(self):
+        batch_sizes = []
+
+        @memory_aware(starting_batch_size=128)
+        def f(batch_size):
+            nonlocal batch_sizes
+            batch_sizes.append(batch_size)
+            if batch_size != 8:
+                raise_fake_out_of_memory()
+
+        f()
+        self.assertListEqual(batch_sizes, [128, 64, 32, 16, 8])
+
+    def test_memory_explicit(self):
+        batch_sizes = []
+
+        @memory_aware(starting_batch_size=128)
+        def f(batch_size, arg1):
Same there!
@@ -0,0 +1,62 @@
+# 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 unittest
+
+from accelerate.memory_utils import memory_aware
+
+
+def raise_fake_out_of_memory():
+    raise RuntimeError("CUDA out of memory.")
+
+
+class MemoryTest(unittest.TestCase):
+    def test_memory_implicit(self):
+        batch_sizes = []
+
+        @memory_aware(starting_batch_size=128)
+        def f(batch_size):
+            nonlocal batch_sizes
+            batch_sizes.append(batch_size)
+            if batch_size != 8:
+                raise_fake_out_of_memory()
+
+        f()
+        self.assertListEqual(batch_sizes, [128, 64, 32, 16, 8])
+
+    def test_memory_explicit(self):
+        batch_sizes = []
+
+        @memory_aware(starting_batch_size=128)
+        def f(batch_size, arg1):
+            nonlocal batch_sizes
+            batch_sizes.append(batch_size)
+            if batch_size != 8:
+                raise_fake_out_of_memory()
+            return batch_size, arg1
+
+        bs, arg1 = f("hello")
+        self.assertListEqual(batch_sizes, [128, 64, 32, 16, 8])
+        self.assertListEqual([8, "hello"], [bs, arg1])
+
+    def test_verbose_guard(self):
+        @memory_aware(starting_batch_size=128)
+        def f(batch_size, arg1, arg2):
`find_executable_batch_size` is a bit more on the nose here with what it does. WDYT?
@@ -0,0 +1,218 @@
+# 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 torch
+from torch.utils.data import DataLoader
+
+from accelerate import Accelerator, DistributedType
+
+# New Code #
+from accelerate.memory_utils import memory_aware
+from datasets import load_dataset, load_metric
+from transformers import (
+    AdamW,
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    get_linear_schedule_with_warmup,
+    set_seed,
+)
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate,
+# specifically showcasing how to ensure out-of-memory errors never
+# iterrupt training, and builds off the `nlp_example.py` script.
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#
+# New additions from the base script can be found quickly by
+# looking for the # New Code # tags
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+def get_dataloaders(accelerator: Accelerator, batch_size: int = 16):
+    """
+    Creates a set of `DataLoader`s for the `glue` dataset,
+    using "bert-base-cased" as the tokenizer.
+
+    Args:
+        accelerator (`Accelerator`):
+            An `Accelerator` object
+        batch_size (`int`, *optional*):
+            The batch size for the train and validation DataLoaders.
+    """
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    datasets = load_dataset("glue", "mrpc")
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["idx", "sentence1", "sentence2"],
+    )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    return train_dataloader, eval_dataloader
+
+
+def training_function(config, args):
+    # Initialize accelerator
+    accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    correct_bias = config["correct_bias"]
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+
+    metric = load_metric("glue", "mrpc")
+
+    # If the batch size is too big we use gradient accumulation
+    gradient_accumulation_steps = 1
+    if batch_size > MAX_GPU_BATCH_SIZE:
+        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
+        batch_size = MAX_GPU_BATCH_SIZE
+
+    set_seed(seed)
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
+
+    # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
+    # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
+    # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
+    model = model.to(accelerator.device)
+
+    # Instantiate optimizer
+    optimizer = AdamW(params=model.parameters(), lr=lr, correct_bias=correct_bias)
+
+    # New Code #
+    # We now can define an inner training loop function. It should take a batch size as the only parameter,
+    # and build the dataloaders in there.
+    # It also gets our decorator
+    @memory_aware(starting_batch_size=batch_size)
Mostly since it's a start high -> go low approach if it broke, rather than find the maximum
@@ -0,0 +1,218 @@
+# 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 torch
+from torch.utils.data import DataLoader
+
+from accelerate import Accelerator, DistributedType
+
+# New Code #
+from accelerate.memory_utils import memory_aware
+from datasets import load_dataset, load_metric
+from transformers import (
+    AdamW,
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    get_linear_schedule_with_warmup,
+    set_seed,
+)
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate,
+# specifically showcasing how to ensure out-of-memory errors never
+# iterrupt training, and builds off the `nlp_example.py` script.
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#
+# New additions from the base script can be found quickly by
+# looking for the # New Code # tags
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+def get_dataloaders(accelerator: Accelerator, batch_size: int = 16):
+    """
+    Creates a set of `DataLoader`s for the `glue` dataset,
+    using "bert-base-cased" as the tokenizer.
+
+    Args:
+        accelerator (`Accelerator`):
+            An `Accelerator` object
+        batch_size (`int`, *optional*):
+            The batch size for the train and validation DataLoaders.
+    """
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    datasets = load_dataset("glue", "mrpc")
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["idx", "sentence1", "sentence2"],
+    )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    return train_dataloader, eval_dataloader
+
+
+def training_function(config, args):
+    # Initialize accelerator
+    accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    correct_bias = config["correct_bias"]
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+
+    metric = load_metric("glue", "mrpc")
+
+    # If the batch size is too big we use gradient accumulation
+    gradient_accumulation_steps = 1
+    if batch_size > MAX_GPU_BATCH_SIZE:
+        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
+        batch_size = MAX_GPU_BATCH_SIZE
+
+    set_seed(seed)
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
+
+    # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
+    # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
+    # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
+    model = model.to(accelerator.device)
+
+    # Instantiate optimizer
+    optimizer = AdamW(params=model.parameters(), lr=lr, correct_bias=correct_bias)
+
+    # New Code #
+    # We now can define an inner training loop function. It should take a batch size as the only parameter,
+    # and build the dataloaders in there.
+    # It also gets our decorator
+    @memory_aware(starting_batch_size=batch_size)
You do end up with the max that fits on the GPU, but your name works for me too.
@@ -0,0 +1,218 @@
+# 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 torch
+from torch.utils.data import DataLoader
+
+from accelerate import Accelerator, DistributedType
+
+# New Code #
+from accelerate.memory_utils import memory_aware
+from datasets import load_dataset, load_metric
+from transformers import (
+    AdamW,
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    get_linear_schedule_with_warmup,
+    set_seed,
+)
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate,
+# specifically showcasing how to ensure out-of-memory errors never
+# iterrupt training, and builds off the `nlp_example.py` script.
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#
+# New additions from the base script can be found quickly by
+# looking for the # New Code # tags
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+def get_dataloaders(accelerator: Accelerator, batch_size: int = 16):
+    """
+    Creates a set of `DataLoader`s for the `glue` dataset,
+    using "bert-base-cased" as the tokenizer.
+
+    Args:
+        accelerator (`Accelerator`):
+            An `Accelerator` object
+        batch_size (`int`, *optional*):
+            The batch size for the train and validation DataLoaders.
+    """
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    datasets = load_dataset("glue", "mrpc")
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["idx", "sentence1", "sentence2"],
+    )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    return train_dataloader, eval_dataloader
+
+
+def training_function(config, args):
+    # Initialize accelerator
+    accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    correct_bias = config["correct_bias"]
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+
+    metric = load_metric("glue", "mrpc")
+
+    # If the batch size is too big we use gradient accumulation
+    gradient_accumulation_steps = 1
+    if batch_size > MAX_GPU_BATCH_SIZE:
+        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
+        batch_size = MAX_GPU_BATCH_SIZE
+
+    set_seed(seed)
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
+
+    # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
+    # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
+    # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
+    model = model.to(accelerator.device)
+
+    # Instantiate optimizer
+    optimizer = AdamW(params=model.parameters(), lr=lr, correct_bias=correct_bias)
+
+    # New Code #
+    # We now can define an inner training loop function. It should take a batch size as the only parameter,
+    # and build the dataloaders in there.
+    # It also gets our decorator
+    @memory_aware(starting_batch_size=batch_size)
Put the API doc at the end.
@@ -0,0 +1,51 @@
+<!--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.
+-->
+
+# Memory Utilities
+
+One of the most frustrating errors when it comes to running training scripts is hitting "CUDA Out-of-Memory", 
+as the entire script needs to be restarted, progress is lost, and typically a developer would want to simply
+start their script and let it run.
+
+`Accelerate` provides a utility heavily based on [toma](https://github.com/BlackHC/toma) to give this capability.
+
+## find_executable_batch_size
+
+[[autodoc]] memory_utils.find_executable_batch_size
Wrong name here :-)
@@ -0,0 +1,51 @@
+<!--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.
+-->
+
+# Memory Utilities
+
+One of the most frustrating errors when it comes to running training scripts is hitting "CUDA Out-of-Memory", 
+as the entire script needs to be restarted, progress is lost, and typically a developer would want to simply
+start their script and let it run.
+
+`Accelerate` provides a utility heavily based on [toma](https://github.com/BlackHC/toma) to give this capability.
+
+## find_executable_batch_size
+
+[[autodoc]] memory_utils.find_executable_batch_size
+
+This algorithm operates with exponential decay, decreasing the batch size in half after each failed run on some 
+training script. To use it, restructure your training function to include an inner function that includes this wrapper, 
+and build your dataloaders inside it. At a minimum, this could look like 4 new lines of code. 
+> Note: The inner function *must* take in the batch size as the first parameter, but we do not pass one to it when called. The wrapper handles this for us
+
+```diff
+def training_function(args):
+    accelerator = Accelerator()
+    model = get_model()
+    model.to(accelerator.device)
+    optimizer = get_optimizer()
+
++   @memory_aware(starting_batch_size=args.batch_size)
Still TODO:

- Add documentation
- Write tests
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/324
diff --git a/examples/by_feature/checkpointing.py b/examples/by_feature/checkpointing.py
index e595a8db7..ee0716b6d 100644
--- a/examples/by_feature/checkpointing.py
+++ b/examples/by_feature/checkpointing.py
@@ -166,6 +166,8 @@ def training_function(config, args):
     # New Code #
     # We need to keep track of how many total steps we have iterated over
     overall_step = 0
+    # We also need to keep track of the stating epoch so files are named properly
+    starting_epoch = 0
 
     # We need to load the checkpoint back in before training here with `load_state`
     # The total number of epochs is adjusted based on where the state is being loaded from,
@@ -184,23 +186,23 @@ def training_function(config, args):
         training_difference = os.path.splitext(path)[0]
 
         if "epoch" in training_difference:
-            num_epochs -= int(training_difference.replace("epoch_", ""))
+            starting_epoch = int(training_difference.replace("epoch_", "")) + 1
             resume_step = None
         else:
             resume_step = int(training_difference.replace("step_", ""))
-            num_epochs -= resume_step // len(train_dataloader)
-            # If resuming by step, we also need to know exactly how far into the DataLoader we went
-            resume_step = (num_epochs * len(train_dataloader)) - resume_step
+            starting_epoch = resume_step // len(train_dataloader)
+            resume_step -= starting_epoch * len(train_dataloader)
 
     # Now we train the model
-    for epoch in range(num_epochs):
+    for epoch in range(starting_epoch, num_epochs):
         model.train()
         for step, batch in enumerate(train_dataloader):
             # New Code #
             # We need to skip steps until we reach the resumed step during the first epoch
-            if args.resume_from_checkpoint and epoch == 0:
+            if args.resume_from_checkpoint and epoch == starting_epoch:
                 if resume_step is not None and step < resume_step:
-                    pass
+                    overall_step += 1
+                    continue
             # We could avoid this line since we set the accelerator with `device_placement=True`.
             batch.to(accelerator.device)
             outputs = model(**batch)
@@ -246,7 +248,7 @@ def training_function(config, args):
 
         # New Code #
         # We save the model, optimizer, lr_scheduler, and seed states by calling `save_state`
-        # These are saved to folders named `step_{overall_step}`
+        # These are saved to folders named `epoch_{epoch}`
         # Will contain files: "pytorch_model.bin", "optimizer.bin", "scheduler.bin", and "random_states.pkl"
         # If mixed precision was used, will also save a "scalar.bin" file
         if checkpointing_steps == "epoch":
diff --git a/examples/complete_cv_example.py b/examples/complete_cv_example.py
index 01033498c..80a083a8d 100644
--- a/examples/complete_cv_example.py
+++ b/examples/complete_cv_example.py
@@ -173,6 +173,10 @@ def training_function(config, args):
     model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
         model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
     )
+    # We need to keep track of how many total steps we have iterated over
+    overall_step = 0
+    # We also need to keep track of the stating epoch so files are named properly
+    starting_epoch = 0
 
     # Potentially load in the weights and states from a previous save
     if args.resume_from_checkpoint:
@@ -189,25 +193,24 @@ def training_function(config, args):
         training_difference = os.path.splitext(path)[0]
 
         if "epoch" in training_difference:
-            num_epochs -= int(training_difference.replace("epoch_", ""))
+            starting_epoch = int(training_difference.replace("epoch_", "")) + 1
             resume_step = None
         else:
             resume_step = int(training_difference.replace("step_", ""))
-            num_epochs -= resume_step // len(train_dataloader)
-            # If resuming by step, we also need to know exactly how far into the DataLoader we went
-            resume_step = (num_epochs * len(train_dataloader)) - resume_step
+            starting_epoch = resume_step // len(train_dataloader)
+            resume_step -= starting_epoch * len(train_dataloader)
 
-    overall_step = 0
     # Now we train the model
-    for epoch in range(num_epochs):
+    for epoch in range(starting_epoch, num_epochs):
         model.train()
         if args.with_tracking:
             total_loss = 0
         for step, batch in enumerate(train_dataloader):
             # We need to skip steps until we reach the resumed step
-            if args.resume_from_checkpoint and epoch == 0:
+            if args.resume_from_checkpoint and epoch == starting_epoch:
                 if resume_step is not None and step < resume_step:
-                    pass
+                    overall_step += 1
+                    continue
             # We could avoid this line since we set the accelerator with `device_placement=True`.
             batch = {k: v.to(accelerator.device) for k, v in batch.items()}
             inputs = (batch["image"] - mean) / std
diff --git a/examples/complete_nlp_example.py b/examples/complete_nlp_example.py
index a0e8d568c..1910a4f83 100644
--- a/examples/complete_nlp_example.py
+++ b/examples/complete_nlp_example.py
@@ -153,7 +153,10 @@ def collate_fn(examples):
         model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
     )
 
+    # We need to keep track of how many total steps we have iterated over
     overall_step = 0
+    # We also need to keep track of the stating epoch so files are named properly
+    starting_epoch = 0
 
     # Potentially load in the weights and states from a previous save
     if args.resume_from_checkpoint:
@@ -170,24 +173,24 @@ def collate_fn(examples):
         training_difference = os.path.splitext(path)[0]
 
         if "epoch" in training_difference:
-            num_epochs -= int(training_difference.replace("epoch_", ""))
+            starting_epoch = int(training_difference.replace("epoch_", "")) + 1
             resume_step = None
         else:
             resume_step = int(training_difference.replace("step_", ""))
-            num_epochs -= resume_step // len(train_dataloader)
-            # If resuming by step, we also need to know exactly how far into the DataLoader we went
-            resume_step = (num_epochs * len(train_dataloader)) - resume_step
+            starting_epoch = resume_step // len(train_dataloader)
+            resume_step -= starting_epoch * len(train_dataloader)
 
     # Now we train the model
-    for epoch in range(num_epochs):
+    for epoch in range(starting_epoch, num_epochs):
         model.train()
         if args.with_tracking:
             total_loss = 0
         for step, batch in enumerate(train_dataloader):
             # We need to skip steps until we reach the resumed step
-            if args.resume_from_checkpoint and epoch == 0:
+            if args.resume_from_checkpoint and epoch == starting_epoch:
                 if resume_step is not None and step < resume_step:
-                    pass
+                    overall_step += 1
+                    continue
             # We could avoid this line since we set the accelerator with `device_placement=True`.
             batch.to(accelerator.device)
             outputs = model(**batch)
diff --git a/src/accelerate/test_utils/testing.py b/src/accelerate/test_utils/testing.py
index e348e161f..ac6df7741 100644
--- a/src/accelerate/test_utils/testing.py
+++ b/src/accelerate/test_utils/testing.py
@@ -69,6 +69,8 @@ class TempDirTestCase(unittest.TestCase):
     The temporary directory location will be stored in `self.tmpdir`
     """
 
+    clear_on_setup = True
+
     @classmethod
     def setUpClass(cls):
         "Creates a `tempfile.TemporaryDirectory` and stores it in `cls.tmpdir`"
@@ -82,11 +84,12 @@ def tearDownClass(cls):
 
     def setUp(self):
         "Destroy all contents in `self.tmpdir`, but not `self.tmpdir`"
-        for path in Path(self.tmpdir).glob("**/*"):
-            if path.is_file():
-                path.unlink()
-            elif path.is_dir():
-                shutil.rmtree(path)
+        if self.clear_on_setup:
+            for path in Path(self.tmpdir).glob("**/*"):
+                if path.is_file():
+                    path.unlink()
+                elif path.is_dir():
+                    shutil.rmtree(path)
 
 
 class MockingTestCase(unittest.TestCase):
diff --git a/tests/test_examples.py b/tests/test_examples.py
index 3ecd3bf91..aa78fed9f 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -22,7 +22,7 @@
 
 from accelerate import DistributedType
 from accelerate.test_utils.examples import compare_against_test
-from accelerate.test_utils.testing import slow
+from accelerate.test_utils.testing import TempDirTestCase, slow
 from datasets import load_dataset
 from transformers import AutoTokenizer
 
@@ -158,30 +158,61 @@ def test_cv_examples(self):
         self.one_complete_example("complete_cv_example.py", False, cv_path, special_strings)
 
 
-class FeatureExamplesTests(unittest.TestCase):
+class FeatureExamplesTests(TempDirTestCase):
+    clear_on_setup = False
+
     @mock.patch("checkpointing.get_dataloaders", mocked_dataloaders)
     def test_checkpointing_by_epoch(self):
-        with tempfile.TemporaryDirectory() as tmpdir:
-            testargs = f"""
-            checkpointing.py
-            --checkpointing_steps epoch
-            --output_dir {tmpdir}
-            """.split()
+        testargs = f"""
+        checkpointing.py
+        --checkpointing_steps epoch
+        --output_dir {self.tmpdir}
+        """.split()
+        with mock.patch.object(sys, "argv", testargs):
+            checkpointing.main()
+            self.assertTrue(os.path.exists(os.path.join(self.tmpdir, "epoch_1")))
+
+    @mock.patch("checkpointing.get_dataloaders", mocked_dataloaders)
+    def test_checkpointing_by_steps(self):
+        testargs = f"""
+        checkpointing.py
+        --checkpointing_steps 2
+        --output_dir {self.tmpdir}
+        """.split()
+        with mock.patch.object(sys, "argv", testargs):
+            checkpointing.main()
+            self.assertTrue(os.path.exists(os.path.join(self.tmpdir, "step_4")))
+
+    @mock.patch("checkpointing.get_dataloaders", mocked_dataloaders)
+    def test_load_states_by_epoch(self):
+        testargs = f"""
+        checkpointing.py
+        --resume_from_checkpoint {os.path.join(self.tmpdir, "epoch_1")}
+        """.split()
+        dummy_results = {"accuracy": mock.ANY, "f1": mock.ANY}
+        with mock.patch("accelerate.Accelerator.print") as mocked_print:
             with mock.patch.object(sys, "argv", testargs):
                 checkpointing.main()
-                self.assertTrue(os.path.exists(os.path.join(tmpdir, "epoch_0")))
+            with self.assertRaises(AssertionError):
+                mocked_print.assert_any_call("epoch 0:", dummy_results)
+            with self.assertRaises(AssertionError):
+                mocked_print.assert_any_call("epoch 1:", dummy_results)
+            mocked_print.assert_any_call("epoch 2:", dummy_results)
 
     @mock.patch("checkpointing.get_dataloaders", mocked_dataloaders)
-    def test_checkpointing_by_steps(self):
-        with tempfile.TemporaryDirectory() as tmpdir:
-            testargs = f"""
-            checkpointing.py
-            --checkpointing_steps 2
-            --output_dir {tmpdir}
-            """.split()
+    def test_load_states_by_steps(self):
+        testargs = f"""
+        checkpointing.py
+        --resume_from_checkpoint {os.path.join(self.tmpdir, "step_4")}
+        """.split()
+        dummy_results = {"accuracy": mock.ANY, "f1": mock.ANY}
+        with mock.patch("accelerate.Accelerator.print") as mocked_print:
             with mock.patch.object(sys, "argv", testargs):
                 checkpointing.main()
-                self.assertTrue(os.path.exists(os.path.join(tmpdir, "step_2")))
+            with self.assertRaises(AssertionError):
+                mocked_print.assert_any_call("epoch 0:", dummy_results)
+            mocked_print.assert_any_call("epoch 1:", dummy_results)
+            mocked_print.assert_any_call("epoch 2:", dummy_results)
 
     @slow
     def test_cross_validation(self):

How are we sure we will get those values exactly?
@@ -158,30 +158,59 @@ def test_cv_examples(self):
         self.one_complete_example("complete_cv_example.py", False, cv_path, special_strings)
 
 
-class FeatureExamplesTests(unittest.TestCase):
+class FeatureExamplesTests(TempDirTestCase):
+    clear_on_setup = False
+
     @mock.patch("checkpointing.get_dataloaders", mocked_dataloaders)
     def test_checkpointing_by_epoch(self):
-        with tempfile.TemporaryDirectory() as tmpdir:
-            testargs = f"""
-            checkpointing.py
-            --checkpointing_steps epoch
-            --output_dir {tmpdir}
-            """.split()
+        testargs = f"""
+        checkpointing.py
+        --checkpointing_steps epoch
+        --output_dir {self.tmpdir}
+        """.split()
+        with mock.patch.object(sys, "argv", testargs):
+            checkpointing.main()
+            self.assertTrue(os.path.exists(os.path.join(self.tmpdir, "epoch_1")))
+
+    @mock.patch("checkpointing.get_dataloaders", mocked_dataloaders)
+    def test_checkpointing_by_steps(self):
+        testargs = f"""
+        checkpointing.py
+        --checkpointing_steps 2
+        --output_dir {self.tmpdir}
+        """.split()
+        with mock.patch.object(sys, "argv", testargs):
+            checkpointing.main()
+            self.assertTrue(os.path.exists(os.path.join(self.tmpdir, "step_4")))
+
+    @mock.patch("checkpointing.get_dataloaders", mocked_dataloaders)
+    def test_load_states_by_epoch(self):
+        testargs = f"""
+        checkpointing.py
+        --resume_from_checkpoint {os.path.join(self.tmpdir, "epoch_1")}
+        """.split()
+        with mock.patch("accelerate.Accelerator.print") as mocked_print:
             with mock.patch.object(sys, "argv", testargs):
                 checkpointing.main()
-                self.assertTrue(os.path.exists(os.path.join(tmpdir, "epoch_0")))
+            with self.assertRaises(AssertionError):
+                mocked_print.assert_any_call("epoch 0:", {"accuracy": 0.5, "f1": 0.0})
This stems from the scheduler we use in the examples, it makes it impossible for the model to train quickly so we always get an accuracy of .5 and an f1 of 0 for all of our epochs. Hence why none of these example tests check for if we get "good" accuracy, it's for independent behavior.

*But* I dug deeper and found `mock.ANY`. In this case we only care about matching the `epoch *` text, not the results. So instead we have something like this:

```python
dummy_results = {"accuracy":mock.ANY, "f1":mock.ANY}
with self.assertRaises(AssertionError):
    mocked_print.assert_any_call("epoch 0:", dummy_results)
```

Which helps me sleep much better at night
@@ -158,30 +158,59 @@ def test_cv_examples(self):
         self.one_complete_example("complete_cv_example.py", False, cv_path, special_strings)
 
 
-class FeatureExamplesTests(unittest.TestCase):
+class FeatureExamplesTests(TempDirTestCase):
+    clear_on_setup = False
+
     @mock.patch("checkpointing.get_dataloaders", mocked_dataloaders)
     def test_checkpointing_by_epoch(self):
-        with tempfile.TemporaryDirectory() as tmpdir:
-            testargs = f"""
-            checkpointing.py
-            --checkpointing_steps epoch
-            --output_dir {tmpdir}
-            """.split()
+        testargs = f"""
+        checkpointing.py
+        --checkpointing_steps epoch
+        --output_dir {self.tmpdir}
+        """.split()
+        with mock.patch.object(sys, "argv", testargs):
+            checkpointing.main()
+            self.assertTrue(os.path.exists(os.path.join(self.tmpdir, "epoch_1")))
+
+    @mock.patch("checkpointing.get_dataloaders", mocked_dataloaders)
+    def test_checkpointing_by_steps(self):
+        testargs = f"""
+        checkpointing.py
+        --checkpointing_steps 2
+        --output_dir {self.tmpdir}
+        """.split()
+        with mock.patch.object(sys, "argv", testargs):
+            checkpointing.main()
+            self.assertTrue(os.path.exists(os.path.join(self.tmpdir, "step_4")))
+
+    @mock.patch("checkpointing.get_dataloaders", mocked_dataloaders)
+    def test_load_states_by_epoch(self):
+        testargs = f"""
+        checkpointing.py
+        --resume_from_checkpoint {os.path.join(self.tmpdir, "epoch_1")}
+        """.split()
+        with mock.patch("accelerate.Accelerator.print") as mocked_print:
             with mock.patch.object(sys, "argv", testargs):
                 checkpointing.main()
-                self.assertTrue(os.path.exists(os.path.join(tmpdir, "epoch_0")))
+            with self.assertRaises(AssertionError):
+                mocked_print.assert_any_call("epoch 0:", {"accuracy": 0.5, "f1": 0.0})
You never know if your model might have learned a tiny something, so yes, much better :-)
@@ -158,30 +158,59 @@ def test_cv_examples(self):
         self.one_complete_example("complete_cv_example.py", False, cv_path, special_strings)
 
 
-class FeatureExamplesTests(unittest.TestCase):
+class FeatureExamplesTests(TempDirTestCase):
+    clear_on_setup = False
+
     @mock.patch("checkpointing.get_dataloaders", mocked_dataloaders)
     def test_checkpointing_by_epoch(self):
-        with tempfile.TemporaryDirectory() as tmpdir:
-            testargs = f"""
-            checkpointing.py
-            --checkpointing_steps epoch
-            --output_dir {tmpdir}
-            """.split()
+        testargs = f"""
+        checkpointing.py
+        --checkpointing_steps epoch
+        --output_dir {self.tmpdir}
+        """.split()
+        with mock.patch.object(sys, "argv", testargs):
+            checkpointing.main()
+            self.assertTrue(os.path.exists(os.path.join(self.tmpdir, "epoch_1")))
+
+    @mock.patch("checkpointing.get_dataloaders", mocked_dataloaders)
+    def test_checkpointing_by_steps(self):
+        testargs = f"""
+        checkpointing.py
+        --checkpointing_steps 2
+        --output_dir {self.tmpdir}
+        """.split()
+        with mock.patch.object(sys, "argv", testargs):
+            checkpointing.main()
+            self.assertTrue(os.path.exists(os.path.join(self.tmpdir, "step_4")))
+
+    @mock.patch("checkpointing.get_dataloaders", mocked_dataloaders)
+    def test_load_states_by_epoch(self):
+        testargs = f"""
+        checkpointing.py
+        --resume_from_checkpoint {os.path.join(self.tmpdir, "epoch_1")}
+        """.split()
+        with mock.patch("accelerate.Accelerator.print") as mocked_print:
             with mock.patch.object(sys, "argv", testargs):
                 checkpointing.main()
-                self.assertTrue(os.path.exists(os.path.join(tmpdir, "epoch_0")))
+            with self.assertRaises(AssertionError):
+                mocked_print.assert_any_call("epoch 0:", {"accuracy": 0.5, "f1": 0.0})
_The documentation is not available anymore as the PR was closed or merged._
@sgugger 50/50 on whether to consider these "slow" tests or not. They add ~1.5 min to the CI
https://api.github.com/repos/huggingface/accelerate/pulls/323
diff --git a/.gitignore b/.gitignore
index f41377320..7f145cd5c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -132,4 +132,7 @@ dmypy.json
 .vscode
 
 # IntelliJ
-.idea
\ No newline at end of file
+.idea
+
+# Mac .DS_Store
+.DS_Store
\ No newline at end of file
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 3073b3381..fcc7d9b1c 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -54,7 +54,7 @@ Did not find it? :( So we can act quickly on it, please follow these steps:
 * Include your **OS type and version**, the versions of **Python** and **PyTorch**.
 * A short, self-contained, code snippet that allows us to reproduce the bug in
   less than 30s;
-* Provide the with your Accelerate configuration (located by default in `~/.cache/huggingface/accelerate/default_congig.yml`)
+* Provide the with your Accelerate configuration (located by default in `~/.cache/huggingface/accelerate/default_config.yaml`)
 
 ### Do you want a new feature?
 
diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml
index 1ee741653..d35024d66 100644
--- a/docs/source/_toctree.yml
+++ b/docs/source/_toctree.yml
@@ -23,6 +23,8 @@
     title: Checkpointing
   - local: tracking
     title: Experiment Tracking
+  - local: fsdp
+    title: Fully Sharded Data Parallel
   - local: memory
     title: Memory Utilities
   title: API Reference
diff --git a/docs/source/fsdp.mdx b/docs/source/fsdp.mdx
new file mode 100644
index 000000000..e2a0210c8
--- /dev/null
+++ b/docs/source/fsdp.mdx
@@ -0,0 +1,120 @@
+<!--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.
+-->
+
+# Fully Sharded Data Parallel
+
+To accelerate training huge models on larger batch sizes, we can use a fully sharded data parallel model.
+This type of data parallel paradigm enables fitting more data and larger models by sharding the optimizer states, gradients and parameters.
+To read more about it and the benefits, check out the [Fully Sharded Data Parallel blog](https://pytorch.org/blog/introducing-pytorch-fully-sharded-data-parallel-api/).
+We have integrated the latest PyTorch's Fully Sharded Data Parallel (FSDP) training feature.
+All you need to do is enable it through the config.
+
+## How it works out the box
+
+On your machine(s) just run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked. This will generate a config file that will be used automatically to properly set the
+default options when doing
+
+```bash
+accelerate launch my_script.py --args_to_my_script
+```
+
+For instance, here is how you would run the NLP example (from the root of the repo) with FSDP enabled:
+
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config: {}
+distributed_type: FSDP
+fsdp_config:
+  min_num_params: 2000
+  offload_params: false
+  sharding_strategy: 1
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: 'no'
+num_machines: 1
+num_processes: 2
+use_cpu: false
+```
+
+```bash
+accelerate launch examples/nlp_example.py
+```
+
+Currently, `Accelerate` supports following config through the CLI:
+
+```bash
+`Sharding Strategy`: [1] FULL_SHARD, [2] SHARD_GRAD_OP
+`Min Num Params`: FSDP\'s minimum number of parameters for Default Auto Wrapping.
+`Offload Params`: Decides Whether to offload parameters and gradients to CPU.
+```
+
+## Few caveats to be aware of
+
+- PyTorch FSDP auto wraps sub-modules, flattens the parameters and shards the parameters in place.
+  Due to this, any optimizer created before model wrapping gets broken and occupies more memory.
+  Hence, it is highly recommended and efficient to prepare model before creating optimizer.
+  `Accelerate` will automatically wrap the model and create an optimizer for you in case of single model with a warning message.
+  > FSDP Warning: When using FSDP, it is efficient and recommended to call prepare for the model before creating the optimizer
+
+However, below is the recommended way to prepare model and optimizer while using FSDP:
+
+```diff
+model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
++ model = accelerator.prepare(model)
+
+optimizer = torch.optim.AdamW(params=model.parameters(), lr=lr)
+
+- model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(model,
+-        optimizer, train_dataloader, eval_dataloader, lr_scheduler
+-    )
+
++ optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
++         optimizer, train_dataloader, eval_dataloader, lr_scheduler
++        )
+
+```
+
+- In case of a single model, if you have created optimizer with multiple parameter groups and called prepare with them together,
+  then the parameter groups will be lost and the following warning is displayed:
+  > FSDP Warning: When using FSDP, several parameter groups will be conflated into
+  > a single one due to nested module wrapping and parameter flattening.
+  
+  This is because parameter groups created before wrapping will have no meaning post wrapping due parameter flattening of nested FSDP modules into 1D arrays (which can consume many layers).
+  For instance, below are the named parameters of FSDP model on GPU 0 (When using 2 GPUs. Around 55M (110M/2) params in 1D arrays as this will have the 1st shard of the parameters). 
+  Here, if one has applied no weight decay for [bias, LayerNorm.weight] named parameters of unwrapped BERT model, 
+  it can't be applied to the below FSDP wrapped model as there are no named parameters with either of those strings and 
+  the parameters of those layers are concatenated with parameters of various other layers.
+  ```
+  {
+    '_fsdp_wrapped_module.flat_param': torch.Size([494209]), 
+    '_fsdp_wrapped_module._fpw_module.bert.embeddings.word_embeddings._fsdp_wrapped_module.flat_param': torch.Size([11720448]), 
+    '_fsdp_wrapped_module._fpw_module.bert.encoder._fsdp_wrapped_module.flat_param': torch.Size([42527232])
+  }
+  ```
+
+
+- In case of multiple models, it is necessary to prepare the models before creating optimizers else it will throw an error.
+- Mixed precision is currently not supported with FSDP.
+
+For more control, users can leverage the `FullyShardedDataParallelPlugin` wherein they can specify `auto_wrap_policy`, `backward_prefetch` and `ignored_modules`.
+After creating an instance of this class, users can pass it to the Accelerator class instantiation.
+For more information on these options, please refer to the PyTorch [FullyShardedDataParallel](https://github.com/pytorch/pytorch/blob/0df2e863fbd5993a7b9e652910792bd21a516ff3/torch/distributed/fsdp/fully_sharded_data_parallel.py#L236) code.
+
+[[autodoc]] utils.FullyShardedDataParallelPlugin
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index eb3b24ce9..781330ace 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -32,6 +32,7 @@
 from .tracking import LOGGER_TYPE_TO_CLASS, GeneralTracker, filter_trackers
 from .utils import (
     DeepSpeedPlugin,
+    FullyShardedDataParallelPlugin,
     LoggerType,
     PrecisionType,
     RNGType,
@@ -80,6 +81,9 @@ class Accelerator:
         deepspeed_plugin (`DeepSpeedPlugin`, *optional*):
             Tweak your DeepSpeed related args using this argument. This argument is optional and can be configured
             directly using *accelerate config*
+        fsdp_plugin (`FullyShardedDataParallelPlugin`, *optional*):
+            Tweak your FSDP related args using this argument. This argument is optional and can be configured directly
+            using *accelerate config*
         rng_types (list of `str` or [`~utils.RNGType`]):
             The list of random number generators to synchronize at the beginning of each iteration in your prepared
             dataloaders. Should be one or several of:
@@ -127,6 +131,7 @@ def __init__(
         mixed_precision: Union[PrecisionType, str] = None,
         cpu: bool = False,
         deepspeed_plugin: DeepSpeedPlugin = None,
+        fsdp_plugin: FullyShardedDataParallelPlugin = None,
         rng_types: Optional[List[Union[str, RNGType]]] = None,
         log_with: Optional[List[Union[str, LoggerType, GeneralTracker]]] = None,
         logging_dir: Optional[Union[str, os.PathLike]] = None,
@@ -155,6 +160,15 @@ def __init__(
                 deepspeed_plugin, DeepSpeedPlugin
             ), "`deepspeed_plugin` must be a DeepSpeedPlugin object."
 
+        if os.environ.get("USE_FSDP", "false") == "true":
+            if version.parse(torch.__version__) < version.parse("1.12.0.dev20220418+cu113"):
+                raise ValueError("FSDP requires PyTorch >= 1.12.0.dev20220418+cu113")
+            if fsdp_plugin is None:  # init from env variables
+                fsdp_plugin = FullyShardedDataParallelPlugin()
+            else:
+                if not isinstance(fsdp_plugin, FullyShardedDataParallelPlugin):
+                    raise TypeError("`fsdp_plugin` must be a FullyShardedDataParallelPlugin object.")
+
         # Kwargs handlers
         self.ddp_handler = None
         self.scaler_handler = None
@@ -183,6 +197,7 @@ def __init__(
             mixed_precision=mixed_precision,
             cpu=cpu,
             deepspeed_plugin=deepspeed_plugin,
+            fsdp_plugin=fsdp_plugin,
             _from_accelerator=True,
             **kwargs,
         )
@@ -326,6 +341,44 @@ def _prepare_one(self, obj, first_pass=False):
         else:
             return obj
 
+    def _prepare_fsdp(self, *args):
+        result = []
+        for obj in args:
+            if isinstance(obj, torch.nn.Module):
+                model = obj
+                break
+        optimizers = []
+
+        self._schedulers = []
+        self._models = []
+        intermediate_result = []
+        for obj in args:
+            if isinstance(obj, torch.optim.Optimizer):
+                if len(obj.param_groups) > 1:
+                    logger.warn(
+                        "FSDP Warning: When using FSDP, several parameter groups will be conflated into "
+                        "a single one due to nested module wrapping and parameter flattening."
+                    )
+                optimizer = obj.optimizer.__class__(model.parameters(), **obj.optimizer.defaults)
+                obj = self.prepare_optimizer(optimizer)
+                optimizers.append(obj)
+            elif isinstance(obj, torch.nn.Module):
+                self._models.append(obj)
+            intermediate_result.append(obj)
+
+        for obj in intermediate_result:
+            if isinstance(obj, AcceleratedScheduler):
+                obj.optimizer = optimizers
+                for i, opt in enumerate(self._optimizers):
+                    if getattr(obj.scheduler, "optimizer", None) == opt.optimizer:
+                        obj.scheduler.optimizer = optimizers[i]
+                        obj.optimizers = [optimizers[i]]
+                        break
+                self._schedulers.append(obj)
+            result.append(obj)
+        self._optimizers = optimizers
+        return tuple(result)
+
     def prepare(self, *args):
         """
         Prepare all objects passed in `args` for distributed training and mixed precision, then return them in the same
@@ -337,6 +390,25 @@ def prepare(self, *args):
             - `torch.nn.Module`: PyTorch Module
             - `torch.optim.Optimizer`: PyTorch Optimizer
         """
+        if self.distributed_type == DistributedType.FSDP:
+            model_count = 0
+            optimizer_present = False
+            for obj in args:
+                if isinstance(obj, torch.nn.Module):
+                    model_count += 1
+                if isinstance(obj, torch.optim.Optimizer):
+                    optimizer_present = True
+            if model_count > 1 and optimizer_present:
+                raise ValueError(
+                    "For FSDP to work with multiple models (>1), "
+                    "prepare must be called for all the models before optimizers are created"
+                )
+            elif model_count == 1 and optimizer_present:
+                logger.warn(
+                    "FSDP Warning: When using FSDP, "
+                    "it is efficient and recommended to call prepare for the model before creating the optimizer"
+                )
+
         # On TPUs, putting the model on the XLA device will create new parameters, so the corresponding optimizer will
         # have parameters disconnected from the model (so no training :-( ).
         # If the model and optimizer have parameters on different devices we raise an error.
@@ -373,16 +445,33 @@ def prepare(self, *args):
                 if isinstance(obj, torch.optim.Optimizer):
                     obj._switch_parameters(mapping)
 
+        if self.distributed_type == DistributedType.FSDP and model_count == 1 and optimizer_present:
+            result = self._prepare_fsdp(*result)
+
         return result if len(result) > 1 else result[0]
 
     def prepare_model(self, model):
-        if self.device_placement:
+        if self.device_placement and self.distributed_type != DistributedType.FSDP:
             model = model.to(self.device)
         if self.distributed_type == DistributedType.MULTI_GPU:
             kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {}
             model = torch.nn.parallel.DistributedDataParallel(
                 model, device_ids=[self.local_process_index], output_device=self.local_process_index, **kwargs
             )
+        elif self.distributed_type == DistributedType.FSDP:
+            from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP
+
+            fsdp_plugin = self.state.fsdp_plugin
+            model = FSDP(
+                model,
+                sharding_strategy=fsdp_plugin.sharding_strategy,
+                cpu_offload=fsdp_plugin.cpu_offload,
+                auto_wrap_policy=fsdp_plugin.auto_wrap_policy,
+                backward_prefetch=fsdp_plugin.backward_prefetch,
+                ignored_modules=fsdp_plugin.ignored_modules,
+            )
+            if not fsdp_plugin.cpu_offload.offload_params:
+                model.to(self.device)
         elif self.distributed_type == DistributedType.MULTI_CPU:
             kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {}
             model = torch.nn.parallel.DistributedDataParallel(model, **kwargs)
@@ -544,6 +633,12 @@ def clip_grad_norm_(self, parameters, max_norm, norm_type=2):
         """
         Should be used in place of `torch.nn.utils.clip_grad_norm_`.
         """
+        if self.distributed_type == DistributedType.FSDP:
+            parameters = [p for p in parameters]
+            for model in self._models:
+                if parameters == [p for p in model.parameters()]:
+                    model.clip_grad_norm_(max_norm, norm_type)
+                    return
         self.unscale_gradients()
         torch.nn.utils.clip_grad_norm_(parameters, max_norm, norm_type=norm_type)
 
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index d72680bea..69464ea0b 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -99,6 +99,34 @@ def get_cluster_input():
                 default=1,
             )
 
+    if distributed_type in [DistributedType.MULTI_GPU]:
+        use_fsdp = _ask_field(
+            "Do you want to use FullyShardedDataParallel? [yes/NO]: ",
+            _convert_yes_no_to_bool,
+            default=False,
+            error_message="Please enter yes or no.",
+        )
+        if use_fsdp:
+            distributed_type = DistributedType.FSDP
+        fsdp_config = {}
+        if distributed_type == DistributedType.FSDP:
+            fsdp_config["sharding_strategy"] = _ask_field(
+                "What should be your sharding strategy ([1] FULL_SHARD, [2] SHARD_GRAD_OP)? [1]: ",
+                lambda x: int(x),
+                default=1,
+            )
+            fsdp_config["offload_params"] = _ask_field(
+                "Do you want to offload parameters and gradients to CPU? [yes/NO]: ",
+                _convert_yes_no_to_bool,
+                default=False,
+                error_message="Please enter yes or no.",
+            )
+            fsdp_config["min_num_params"] = _ask_field(
+                "What should be your FSDP's minimum number of parameters for Default Auto Wrapping Policy? [1e8]: ",
+                lambda x: int(x),
+                default=1e8,
+            )
+
     if distributed_type == DistributedType.TPU:
         main_training_function = _ask_field(
             "What is the name of the function in your script that should be launched in all parallel scripts? [main]: ",
@@ -134,5 +162,6 @@ def get_cluster_input():
         main_process_port=main_process_port,
         main_training_function=main_training_function,
         deepspeed_config=deepspeed_config,
+        fsdp_config=fsdp_config,
         use_cpu=use_cpu,
     )
diff --git a/src/accelerate/commands/config/config_args.py b/src/accelerate/commands/config/config_args.py
index 5af912fc3..fadbe6211 100644
--- a/src/accelerate/commands/config/config_args.py
+++ b/src/accelerate/commands/config/config_args.py
@@ -136,6 +136,8 @@ class ClusterConfig(BaseConfig):
 
     # args for deepspeed_plugin
     deepspeed_config: dict = None
+    # args for fsdp
+    fsdp_config: dict = None
 
 
 @dataclass
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index f2cab1ccb..9fa10da7b 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -51,6 +51,30 @@ def launch_command_parser(subparsers=None):
         action="store_true",
         help="Whether to use deepspeed.",
     )
+    parser.add_argument(
+        "--use_fsdp",
+        default=False,
+        action="store_true",
+        help="Whether to use fsdp.",
+    )
+    parser.add_argument(
+        "--offload_params",
+        default="false",
+        type=str,
+        help="Decides Whether (true|false) to offload parameters and gradients to CPU. (useful only when `use_fsdp` flag is passed).",
+    )
+    parser.add_argument(
+        "--min_num_params",
+        type=int,
+        default=1e8,
+        help="FSDP's minimum number of parameters for Default Auto Wrapping. (useful only when `use_fsdp` flag is passed).",
+    )
+    parser.add_argument(
+        "--sharding_strategy",
+        type=int,
+        default=1,
+        help="FSDP's Sharding Strategy. (useful only when `use_fsdp` flag is passed).",
+    )
     parser.add_argument(
         "--tpu", default=False, action="store_true", help="Whether or not this should launch a TPU training."
     )
@@ -225,7 +249,11 @@ def multi_gpu_launcher(args):
         mixed_precision = "fp16"
 
     current_env["MIXED_PRECISION"] = str(mixed_precision)
-
+    if args.use_fsdp:
+        current_env["USE_FSDP"] = "true"
+        current_env["FSDP_OFFLOAD_PARAMS"] = str(args.offload_params).lower()
+        current_env["FSDP_MIN_NUM_PARAMS"] = str(args.min_num_params)
+        current_env["FSDP_SHARDING_STRATEGY"] = str(args.sharding_strategy)
     process = subprocess.Popen(cmd, env=current_env)
     process.wait()
     if process.returncode != 0:
@@ -433,8 +461,8 @@ def sagemaker_launcher(sagemaker_config: SageMakerConfig, args):
 
 def launch_command(args):
     # Sanity checks
-    if sum([args.multi_gpu, args.tpu, args.use_deepspeed]) > 1:
-        raise ValueError("You can only pick one between `--multi_gpu`, `--use_deepspeed`, `--tpu`.")
+    if sum([args.multi_gpu, args.tpu, args.use_deepspeed, args.use_fsdp]) > 1:
+        raise ValueError("You can only pick one between `--multi_gpu`, `--use_deepspeed`, `--tpu`, `--use_fsdp`.")
 
     defaults = None
     # Get the default from the config file.
@@ -444,6 +472,7 @@ def launch_command(args):
             args.use_deepspeed = defaults.distributed_type == DistributedType.DEEPSPEED
             args.multi_gpu = defaults.distributed_type == DistributedType.MULTI_GPU
             args.tpu = defaults.distributed_type == DistributedType.TPU
+            args.use_fsdp = defaults.distributed_type == DistributedType.FSDP
         if defaults.compute_environment == ComputeEnvironment.LOCAL_MACHINE:
             # Update args with the defaults
             for name, attr in defaults.__dict__.items():
@@ -451,6 +480,8 @@ def launch_command(args):
                     for k in defaults.deepspeed_config:
                         if getattr(args, k) is None:
                             setattr(args, k, defaults.deepspeed_config[k])
+                    for k in defaults.fsdp_config:
+                        setattr(args, k, defaults.fsdp_config[k])
                     continue
 
                 # Those args are handled separately
@@ -472,6 +503,8 @@ def launch_command(args):
     # Use the proper launcher
     if args.use_deepspeed and not args.cpu:
         deepspeed_launcher(args)
+    elif args.use_fsdp and not args.cpu:
+        multi_gpu_launcher(args)
     elif args.multi_gpu and not args.cpu:
         multi_gpu_launcher(args)
     elif args.tpu and not args.cpu:
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index b0521c7eb..df2e36e9a 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -89,6 +89,7 @@ class DistributedType(str, Enum):
     MULTI_CPU = "MULTI_CPU"
     MULTI_GPU = "MULTI_GPU"
     DEEPSPEED = "DEEPSPEED"
+    FSDP = "FSDP"
     TPU = "TPU"
 
 
@@ -149,6 +150,7 @@ def __init__(
         mixed_precision: str = None,
         cpu: bool = False,
         deepspeed_plugin=None,
+        fsdp_plugin=None,
         _from_accelerator: bool = False,
         **kwargs,
     ):
@@ -206,7 +208,13 @@ def __init__(
                 self.mixed_precision = (
                     parse_choice_from_env("MIXED_PRECISION", "no") if mixed_precision is None else mixed_precision
                 )
-
+                if os.environ.get("USE_FSDP", "false") == "true":
+                    self.distributed_type = DistributedType.FSDP
+                    if self.mixed_precision != "no":
+                        raise ValueError(
+                            "Mixed precision is currently not supported for FSDP. Please set `mixed_precision` to `no`."
+                        )
+                    self.fsdp_plugin = fsdp_plugin
             elif get_int_from_env(["PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE", "WORLD_SIZE"], 1) > 1:
                 self.distributed_type = DistributedType.MULTI_CPU
                 if is_ccl_available() and get_int_from_env(["CCL_WORKER_COUNT"], 0) > 0:
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 9bcd87f5f..8101fb7ec 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -12,15 +12,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import functools
 import importlib
 import os
 import random
+import typing
 from collections.abc import Mapping
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 from enum import Enum, EnumMeta
 from functools import update_wrapper
-from typing import Any, List, Optional, Union
+from typing import Any, Callable, Iterable, List, Optional, Union
 
 import numpy as np
 import torch
@@ -738,6 +740,57 @@ def __post_init__(self):
         }
 
 
+@dataclass
+class FullyShardedDataParallelPlugin:
+    """
+    This plugin is used to enable fully sharded data parallelism.
+    """
+
+    sharding_strategy: "typing.Any" = field(
+        default=None,
+        metadata={"help": "Possible options are [1] FULL_SHARD, [2] SHARD_GRAD_OP"},
+    )
+    backward_prefetch: "typing.Any" = field(
+        default=None,
+        metadata={"help": "Possible options are [1] BACKWARD_PRE, [2] BACKWARD_POST"},
+    )
+    auto_wrap_policy: "typing.Any" = field(
+        default=None,
+        metadata={"help": "A callable specifying a policy to recursively wrap layers with FSDP"},
+    )
+    cpu_offload: Optional[Callable] = field(
+        default=None,
+        metadata={"help": "Decides Whether to offload parameters and gradients to CPU."},
+    )
+    min_num_params: int = field(
+        default=None, metadata={"help": "FSDP's minimum number of parameters for Default Auto Wrapping."}
+    )
+    ignored_modules: Optional[Iterable[torch.nn.Module]] = field(
+        default=None,
+        metadata={"help": "A list of modules to ignore for FSDP."},
+    )
+
+    def __post_init__(self):
+        from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload, ShardingStrategy
+        from torch.distributed.fsdp.wrap import default_auto_wrap_policy
+
+        if self.sharding_strategy is None:
+            self.sharding_strategy = ShardingStrategy(int(os.environ.get("FSDP_SHARDING_STRATEGY", 1)))
+
+        if self.cpu_offload is None:
+            if os.environ.get("FSDP_OFFLOAD_PARAMS", "false") == "true":
+                self.cpu_offload = CPUOffload(offload_params=True)
+            else:
+                self.cpu_offload = CPUOffload(offload_params=False)
+
+        if self.min_num_params is None:
+            self.min_num_params = int(os.environ.get("FSDP_MIN_NUM_PARAMS", 0))
+
+        if self.auto_wrap_policy is None:
+            if self.min_num_params > 0:
+                self.auto_wrap_policy = functools.partial(default_auto_wrap_policy, min_num_params=self.min_num_params)
+
+
 @contextmanager
 def patch_environment(**kwargs):
     """

This is a first step but it assumes there is only one model and one optimizer with all the model parameters. We can call this method before the model is prepared and prepare it there, which would allow us to identify which parameters are use in which optimizers and then properly re-instantiate them with the right parameters/parameters groups.

Even better would be to replace the parameters in the existing optimizers without re-instantiating as we are losing parameter groups etc.


@@ -348,6 +349,32 @@ def _prepare_one(self, obj, first_pass=False):
         else:
             return obj
 
+    def _prepare_fsdp(self, *args):
+        result = []
+        for obj in args:
+            if isinstance(obj, torch.nn.Module):
+                model = obj
+                break
+        optimizers = []
+
+        self._schedulers = []
+        for obj in args:
+            if isinstance(obj, torch.optim.Optimizer):
+                optimizer = obj.optimizer.__class__(model.parameters(), **obj.optimizer.defaults)
+                obj = self.prepare_optimizer(optimizer)
+                optimizers.append(obj)
+            elif isinstance(obj, AcceleratedScheduler):
+                obj.optimizer = optimizers
+                for i, opt in enumerate(self._optimizers):
+                    if getattr(obj.scheduler, "optimizer", None) == opt.optimizer:
+                        obj.scheduler.optimizer = optimizers[i]
+                        obj.optimizers = [optimizers[i]]
+                        break
+                self._schedulers.append(obj)
+            result.append(obj)
+        self._optimizers = optimizers
+        return tuple(result)
Is this still necessary?
@@ -359,6 +386,19 @@ def prepare(self, *args):
             - `torch.nn.Module`: PyTorch Module
             - `torch.optim.Optimizer`: PyTorch Optimizer
         """
+        if self.distributed_type == DistributedType.FSDP:
+            model_count = 0
+            optimizer_present = False
+            for obj in args:
+                if isinstance(obj, torch.nn.Module):
+                    model_count += 1
+                if isinstance(obj, torch.optim.Optimizer):
+                    optimizer_present = True
+            if model_count > 1 and optimizer_present:
+                raise ValueError(
+                    "For FSDP to work, prepare must be called for all the models before optimizers are created"
+                )
Those two should also be added to the state, and we should get them from it here.
@@ -405,6 +448,21 @@ def prepare_model(self, model):
             model = torch.nn.parallel.DistributedDataParallel(
                 model, device_ids=[self.local_process_index], output_device=self.local_process_index, **kwargs
             )
+        elif self.distributed_type == DistributedType.FSDP:
+            from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload
+            from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP
+            from torch.distributed.fsdp.wrap import default_auto_wrap_policy
+
+            auto_wrap_policy = None
+            cpu_offload = None
+            if os.environ.get("FSDP_OFFLOAD_PARAMS", "false") == "true":
+                cpu_offload = CPUOffload(offload_params=True)
+
+            if int(os.environ.get("FSDP_MIN_NUM_PARAMS", "0")) > 0:
+                auto_wrap_policy = functools.partial(
+                    default_auto_wrap_policy, min_num_params=int(os.environ["FSDP_MIN_NUM_PARAMS"])
+                )
Was it not out in 1.11 ?. Let's use the `base_version` attribute to compare to 1.12.0 without anything added ;-)
@@ -472,6 +499,10 @@ def launch_command(args):
     # Use the proper launcher
     if args.use_deepspeed and not args.cpu:
         deepspeed_launcher(args)
+    elif args.use_fsdp and not args.cpu:
+        if version.parse(torch.__version__) < version.parse("1.12.0.dev20220418+cu113"):
+            raise ValueError("FSDP requires PyTorch >= 1.12.0.dev20220418+cu113")
We should throw an error if we have mixed_precision different from "no" instead of silently ignoring what the user said in their config.
@@ -203,10 +204,13 @@ def __init__(
                 self.local_process_index = int(os.environ.get("LOCAL_RANK", -1))
                 self.device = torch.device("cuda", self.local_process_index)
                 torch.cuda.set_device(self.device)
-                self.mixed_precision = (
-                    parse_choice_from_env("MIXED_PRECISION", "no") if mixed_precision is None else mixed_precision
-                )
-
+                if os.environ.get("USE_FSDP", "false") == "true":
+                    self.distributed_type = DistributedType.FSDP
+                    self.mixed_precision = "no"
Hello, I've updated the code as per yesterday's offline discussion. 
@@ -348,6 +349,32 @@ def _prepare_one(self, obj, first_pass=False):
         else:
             return obj
 
+    def _prepare_fsdp(self, *args):
+        result = []
+        for obj in args:
+            if isinstance(obj, torch.nn.Module):
+                model = obj
+                break
+        optimizers = []
+
+        self._schedulers = []
+        for obj in args:
+            if isinstance(obj, torch.optim.Optimizer):
+                optimizer = obj.optimizer.__class__(model.parameters(), **obj.optimizer.defaults)
+                obj = self.prepare_optimizer(optimizer)
+                optimizers.append(obj)
+            elif isinstance(obj, AcceleratedScheduler):
+                obj.optimizer = optimizers
+                for i, opt in enumerate(self._optimizers):
+                    if getattr(obj.scheduler, "optimizer", None) == opt.optimizer:
+                        obj.scheduler.optimizer = optimizers[i]
+                        obj.optimizers = [optimizers[i]]
+                        break
+                self._schedulers.append(obj)
+            result.append(obj)
+        self._optimizers = optimizers
+        return tuple(result)
Yes, currently we will only be able to handle the case of single model and hence throwing error in case of multiple models. 
@@ -359,6 +386,19 @@ def prepare(self, *args):
             - `torch.nn.Module`: PyTorch Module
             - `torch.optim.Optimizer`: PyTorch Optimizer
         """
+        if self.distributed_type == DistributedType.FSDP:
+            model_count = 0
+            optimizer_present = False
+            for obj in args:
+                if isinstance(obj, torch.nn.Module):
+                    model_count += 1
+                if isinstance(obj, torch.optim.Optimizer):
+                    optimizer_present = True
+            if model_count > 1 and optimizer_present:
+                raise ValueError(
+                    "For FSDP to work, prepare must be called for all the models before optimizers are created"
+                )
Done. Created FullyShardedDataParallelPlugin for handling the FSDP config.
@@ -405,6 +448,21 @@ def prepare_model(self, model):
             model = torch.nn.parallel.DistributedDataParallel(
                 model, device_ids=[self.local_process_index], output_device=self.local_process_index, **kwargs
             )
+        elif self.distributed_type == DistributedType.FSDP:
+            from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload
+            from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP
+            from torch.distributed.fsdp.wrap import default_auto_wrap_policy
+
+            auto_wrap_policy = None
+            cpu_offload = None
+            if os.environ.get("FSDP_OFFLOAD_PARAMS", "false") == "true":
+                cpu_offload = CPUOffload(offload_params=True)
+
+            if int(os.environ.get("FSDP_MIN_NUM_PARAMS", "0")) > 0:
+                auto_wrap_policy = functools.partial(
+                    default_auto_wrap_policy, min_num_params=int(os.environ["FSDP_MIN_NUM_PARAMS"])
+                )
As mentioned offline, at present the model saving/checkpointing is supported through nightly version only.
@@ -472,6 +499,10 @@ def launch_command(args):
     # Use the proper launcher
     if args.use_deepspeed and not args.cpu:
         deepspeed_launcher(args)
+    elif args.use_fsdp and not args.cpu:
+        if version.parse(torch.__version__) < version.parse("1.12.0.dev20220418+cu113"):
+            raise ValueError("FSDP requires PyTorch >= 1.12.0.dev20220418+cu113")
Done.
@@ -203,10 +204,13 @@ def __init__(
                 self.local_process_index = int(os.environ.get("LOCAL_RANK", -1))
                 self.device = torch.device("cuda", self.local_process_index)
                 torch.cuda.set_device(self.device)
-                self.mixed_precision = (
-                    parse_choice_from_env("MIXED_PRECISION", "no") if mixed_precision is None else mixed_precision
-                )
-
+                if os.environ.get("USE_FSDP", "false") == "true":
+                    self.distributed_type = DistributedType.FSDP
+                    self.mixed_precision = "no"
For new code, we try to avoid assert, so please use a test and rasie a `TypeError` here.
@@ -155,6 +160,16 @@ def __init__(
                 deepspeed_plugin, DeepSpeedPlugin
             ), "`deepspeed_plugin` must be a DeepSpeedPlugin object."
 
+        if os.environ.get("USE_FSDP", "false") == "true":
+            if version.parse(torch.__version__) < version.parse("1.12.0.dev20220418+cu113"):
+                raise ValueError("FSDP requires PyTorch >= 1.12.0.dev20220418+cu113")
+            if fsdp_plugin is None:  # init from env variables
+                fsdp_plugin = FullyShardedDataParallelPlugin()
+            else:
+                assert isinstance(
+                    fsdp_plugin, FullyShardedDataParallelPlugin
+                ), "`fsdp_plugin` must be a FullyShardedDataParallelPlugin object."
Split to respect the 119 char limit please :-)
@@ -326,6 +342,43 @@ def _prepare_one(self, obj, first_pass=False):
         else:
             return obj
 
+    def _prepare_fsdp(self, *args):
+        result = []
+        for obj in args:
+            if isinstance(obj, torch.nn.Module):
+                model = obj
+                break
+        optimizers = []
+
+        self._schedulers = []
+        self._models = []
+        intermediate_result = []
+        for obj in args:
+            if isinstance(obj, torch.optim.Optimizer):
+                if len(obj.param_groups) > 1:
+                    logger.warn(
+                        "FSDP Warning: When using FSDP, several parameter groups will be conflated into a single one due to nested module wrapping and parameter flattening."
The error message could be made clearer then.
@@ -359,6 +386,19 @@ def prepare(self, *args):
             - `torch.nn.Module`: PyTorch Module
             - `torch.optim.Optimizer`: PyTorch Optimizer
         """
+        if self.distributed_type == DistributedType.FSDP:
+            model_count = 0
+            optimizer_present = False
+            for obj in args:
+                if isinstance(obj, torch.nn.Module):
+                    model_count += 1
+                if isinstance(obj, torch.optim.Optimizer):
+                    optimizer_present = True
+            if model_count > 1 and optimizer_present:
+                raise ValueError(
+                    "For FSDP to work, prepare must be called for all the models before optimizers are created"
+                )
There are two consecutive calls to `to` here in this case, would be better to only do one.
@@ -373,16 +443,33 @@ def prepare(self, *args):
                 if isinstance(obj, torch.optim.Optimizer):
                     obj._switch_parameters(mapping)
 
+        if self.distributed_type == DistributedType.FSDP:
+            result = self._prepare_fsdp(*result)
+
         return result if len(result) > 1 else result[0]
 
     def prepare_model(self, model):
-        if self.device_placement:
+        if self.device_placement and self.distributed_type != DistributedType.FSDP:
             model = model.to(self.device)
         if self.distributed_type == DistributedType.MULTI_GPU:
             kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {}
             model = torch.nn.parallel.DistributedDataParallel(
                 model, device_ids=[self.local_process_index], output_device=self.local_process_index, **kwargs
             )
+        elif self.distributed_type == DistributedType.FSDP:
+            from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP
+
+            fsdp_plugin = self.state.fsdp_plugin
+            model = FSDP(
+                model,
+                sharding_strategy=fsdp_plugin.sharding_strategy,
+                cpu_offload=fsdp_plugin.cpu_offload,
+                auto_wrap_policy=fsdp_plugin.auto_wrap_policy,
+                backward_prefetch=fsdp_plugin.backward_prefetch,
+                ignored_modules=fsdp_plugin.ignored_modules,
+            ).to(self.device)
+            if fsdp_plugin.cpu_offload.offload_params:
+                model.to("cpu")
```suggestion
                            "Mixed precision is currently not supported for FSDP. Please set `mixed_precision` to `no`."
```
@@ -206,7 +208,13 @@ def __init__(
                 self.mixed_precision = (
                     parse_choice_from_env("MIXED_PRECISION", "no") if mixed_precision is None else mixed_precision
                 )
-
+                if os.environ.get("USE_FSDP", "false") == "true":
+                    self.distributed_type = DistributedType.FSDP
+                    if self.mixed_precision != "no":
+                        raise ValueError(
+                            "Mixed precision is currently not supported for FSDPS. Please set `mixed_precision` to `no`."
Why add this here, is this used anywhere? The dataclass in itself should be enough to contain the full config?
@@ -738,6 +739,68 @@ def __post_init__(self):
         }
 
 
+@dataclass
+class FullyShardedDataParallelPlugin:
+    """
+    This plugin is used to enable fully sharded data parallelism.
+    """
+
+    from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch, CPUOffload, ShardingStrategy
+
+    sharding_strategy: Optional[ShardingStrategy] = field(
+        default=None,
+        metadata={"help": "Possible options are {}".format(ShardingStrategy.__members__)},
+    )
+    backward_prefetch: Optional[BackwardPrefetch] = field(
+        default=None,
+        metadata={"help": "Possible options are {}".format(BackwardPrefetch.__members__)},
+    )
+    auto_wrap_policy: Optional[Callable] = field(
+        default=None,
+        metadata={"help": "A callable specifying a policy to recursively wrap layers with FSDP"},
+    )
+    cpu_offload: Optional[CPUOffload] = field(
+        default=None,
+        metadata={"help": "Decides Whether to offload parameters and gradients to CPU."},
+    )
+    min_num_params: int = field(
+        default=None, metadata={"help": "FSDP's minimum number of parameters for Default Auto Wrapping."}
+    )
+    ignored_modules: Optional[Iterable[torch.nn.Module]] = field(
+        default=None,
+        metadata={"help": "A list of modules to ignore for FSDP."},
+    )
+
+    def __post_init__(self):
+        from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload, ShardingStrategy
+        from torch.distributed.fsdp.wrap import default_auto_wrap_policy
+
+        if self.sharding_strategy is None:
+            self.sharding_strategy = ShardingStrategy(int(os.environ.get("FSDP_SHARDING_STRATEGY", 1)))
+
+        if self.cpu_offload is None:
+            if os.environ.get("FSDP_OFFLOAD_PARAMS", "false") == "true":
+                self.cpu_offload = CPUOffload(offload_params=True)
+            else:
+                self.cpu_offload = CPUOffload(offload_params=False)
+
+        if self.min_num_params is None:
+            self.min_num_params = int(os.environ.get("FSDP_MIN_NUM_PARAMS", 0))
+
+        if self.auto_wrap_policy is None:
+            if self.min_num_params > 0:
+                self.auto_wrap_policy = functools.partial(default_auto_wrap_policy, min_num_params=self.min_num_params)
+
+        self.fsdp_config = {
+            "sharding_strategy": str(self.sharding_strategy),
+            "cpu_offload": self.cpu_offload.offload_params,
+            "min_num_params": self.min_num_params,
+            "is_default_auto_wrap_policy": self.auto_wrap_policy.__name__ == "default_auto_wrap_policy",
+            "ignored_modules_present": self.ignored_modules is not None,
+            "backward_prefetch_present": self.backward_prefetch is not None,
+        }
As said offline this is the line making the CI red.
@@ -738,6 +739,68 @@ def __post_init__(self):
         }
 
 
+@dataclass
+class FullyShardedDataParallelPlugin:
+    """
+    This plugin is used to enable fully sharded data parallelism.
+    """
+
+    from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch, CPUOffload, ShardingStrategy
Got it. Thanks for pointing this out. Fixed this.
@@ -738,6 +739,68 @@ def __post_init__(self):
         }
 
 
+@dataclass
+class FullyShardedDataParallelPlugin:
+    """
+    This plugin is used to enable fully sharded data parallelism.
+    """
+
+    from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch, CPUOffload, ShardingStrategy
It was used sort of "__repr__" to be printed when printing the accelerator state. removed this as the dataclass should be the full config.
@@ -738,6 +739,68 @@ def __post_init__(self):
         }
 
 
+@dataclass
+class FullyShardedDataParallelPlugin:
+    """
+    This plugin is used to enable fully sharded data parallelism.
+    """
+
+    from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch, CPUOffload, ShardingStrategy
+
+    sharding_strategy: Optional[ShardingStrategy] = field(
+        default=None,
+        metadata={"help": "Possible options are {}".format(ShardingStrategy.__members__)},
+    )
+    backward_prefetch: Optional[BackwardPrefetch] = field(
+        default=None,
+        metadata={"help": "Possible options are {}".format(BackwardPrefetch.__members__)},
+    )
+    auto_wrap_policy: Optional[Callable] = field(
+        default=None,
+        metadata={"help": "A callable specifying a policy to recursively wrap layers with FSDP"},
+    )
+    cpu_offload: Optional[CPUOffload] = field(
+        default=None,
+        metadata={"help": "Decides Whether to offload parameters and gradients to CPU."},
+    )
+    min_num_params: int = field(
+        default=None, metadata={"help": "FSDP's minimum number of parameters for Default Auto Wrapping."}
+    )
+    ignored_modules: Optional[Iterable[torch.nn.Module]] = field(
+        default=None,
+        metadata={"help": "A list of modules to ignore for FSDP."},
+    )
+
+    def __post_init__(self):
+        from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload, ShardingStrategy
+        from torch.distributed.fsdp.wrap import default_auto_wrap_policy
+
+        if self.sharding_strategy is None:
+            self.sharding_strategy = ShardingStrategy(int(os.environ.get("FSDP_SHARDING_STRATEGY", 1)))
+
+        if self.cpu_offload is None:
+            if os.environ.get("FSDP_OFFLOAD_PARAMS", "false") == "true":
+                self.cpu_offload = CPUOffload(offload_params=True)
+            else:
+                self.cpu_offload = CPUOffload(offload_params=False)
+
+        if self.min_num_params is None:
+            self.min_num_params = int(os.environ.get("FSDP_MIN_NUM_PARAMS", 0))
+
+        if self.auto_wrap_policy is None:
+            if self.min_num_params > 0:
+                self.auto_wrap_policy = functools.partial(default_auto_wrap_policy, min_num_params=self.min_num_params)
+
+        self.fsdp_config = {
+            "sharding_strategy": str(self.sharding_strategy),
+            "cpu_offload": self.cpu_offload.offload_params,
+            "min_num_params": self.min_num_params,
+            "is_default_auto_wrap_policy": self.auto_wrap_policy.__name__ == "default_auto_wrap_policy",
+            "ignored_modules_present": self.ignored_modules is not None,
+            "backward_prefetch_present": self.backward_prefetch is not None,
+        }
Done.
@@ -373,16 +443,33 @@ def prepare(self, *args):
                 if isinstance(obj, torch.optim.Optimizer):
                     obj._switch_parameters(mapping)
 
+        if self.distributed_type == DistributedType.FSDP:
+            result = self._prepare_fsdp(*result)
+
         return result if len(result) > 1 else result[0]
 
     def prepare_model(self, model):
-        if self.device_placement:
+        if self.device_placement and self.distributed_type != DistributedType.FSDP:
             model = model.to(self.device)
         if self.distributed_type == DistributedType.MULTI_GPU:
             kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {}
             model = torch.nn.parallel.DistributedDataParallel(
                 model, device_ids=[self.local_process_index], output_device=self.local_process_index, **kwargs
             )
+        elif self.distributed_type == DistributedType.FSDP:
+            from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP
+
+            fsdp_plugin = self.state.fsdp_plugin
+            model = FSDP(
+                model,
+                sharding_strategy=fsdp_plugin.sharding_strategy,
+                cpu_offload=fsdp_plugin.cpu_offload,
+                auto_wrap_policy=fsdp_plugin.auto_wrap_policy,
+                backward_prefetch=fsdp_plugin.backward_prefetch,
+                ignored_modules=fsdp_plugin.ignored_modules,
+            ).to(self.device)
+            if fsdp_plugin.cpu_offload.offload_params:
+                model.to("cpu")
Done.
@@ -326,6 +342,43 @@ def _prepare_one(self, obj, first_pass=False):
         else:
             return obj
 
+    def _prepare_fsdp(self, *args):
+        result = []
+        for obj in args:
+            if isinstance(obj, torch.nn.Module):
+                model = obj
+                break
+        optimizers = []
+
+        self._schedulers = []
+        self._models = []
+        intermediate_result = []
+        for obj in args:
+            if isinstance(obj, torch.optim.Optimizer):
+                if len(obj.param_groups) > 1:
+                    logger.warn(
+                        "FSDP Warning: When using FSDP, several parameter groups will be conflated into a single one due to nested module wrapping and parameter flattening."
Done.
@@ -155,6 +160,16 @@ def __init__(
                 deepspeed_plugin, DeepSpeedPlugin
             ), "`deepspeed_plugin` must be a DeepSpeedPlugin object."
 
+        if os.environ.get("USE_FSDP", "false") == "true":
+            if version.parse(torch.__version__) < version.parse("1.12.0.dev20220418+cu113"):
+                raise ValueError("FSDP requires PyTorch >= 1.12.0.dev20220418+cu113")
+            if fsdp_plugin is None:  # init from env variables
+                fsdp_plugin = FullyShardedDataParallelPlugin()
+            else:
+                assert isinstance(
+                    fsdp_plugin, FullyShardedDataParallelPlugin
+                ), "`fsdp_plugin` must be a FullyShardedDataParallelPlugin object."
Done.
@@ -359,6 +386,19 @@ def prepare(self, *args):
             - `torch.nn.Module`: PyTorch Module
             - `torch.optim.Optimizer`: PyTorch Optimizer
         """
+        if self.distributed_type == DistributedType.FSDP:
+            model_count = 0
+            optimizer_present = False
+            for obj in args:
+                if isinstance(obj, torch.nn.Module):
+                    model_count += 1
+                if isinstance(obj, torch.optim.Optimizer):
+                    optimizer_present = True
+            if model_count > 1 and optimizer_present:
+                raise ValueError(
+                    "For FSDP to work, prepare must be called for all the models before optimizers are created"
+                )
Can we add the new guide without changing the indentation of the whole file? ;-)
@@ -1,26 +1,28 @@
-- sections: 
-  - local: index
-    title: 🤗 Accelerate
-  - local: quicktour
-    title: Quick tour
-  - local: installation
-    title: Installation
+- sections:
+    - local: index
+      title: 🤗 Accelerate
+    - local: quicktour
+      title: Quick tour
+    - local: installation
+      title: Installation
   title: Get started
 - sections:
-  - local: sagemaker
-    title: Amazon SageMaker
+    - local: sagemaker
+      title: Amazon SageMaker
   title: Guides
 - sections:
-  - local: accelerator
-    title: Accelerator
-  - local: launcher
-    title: Notebook Launcher
-  - local: kwargs
-    title: Kwargs Handlers
-  - local: internal
-    title: Internals
-  - local: checkpoint
-    title: Checkpointing
-  - local: tracking
-    title: Experiment Tracking
+    - local: accelerator
+      title: Accelerator
+    - local: launcher
+      title: Notebook Launcher
+    - local: kwargs
+      title: Kwargs Handlers
+    - local: internal
+      title: Internals
+    - local: checkpoint
+      title: Checkpointing
+    - local: tracking
+      title: Experiment Tracking
+    - local: fsdp
+      title: Fully Sharded Data Parallel
I would put this before the API doc (so before the line [[autodoc]] above).
@@ -0,0 +1,107 @@
+<!--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.
+-->
+
+# Fully Sharded Data Parallel
+
+To accelerate training huge models on larger batch sizes, we can use a fully sharded data parallel model.
+This type of data parallel paradigm enables fitting more data and larger models by sharding the optimizer states, gradients and parameters.
+To read more about it and the benefits, check out the [Fully Sharded Data Parallel blog](https://pytorch.org/blog/introducing-pytorch-fully-sharded-data-parallel-api/).
+We have integrated the latest PyTorch's Fully Sharded Data Parallel (FSDP) training feature.
+All you need to do is enable it through the config.
+
+## How it works out the box
+
+On your machine(s) just run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked. This will generate a config file that will be used automatically to properly set the
+default options when doing
+
+```bash
+accelerate launch my_script.py --args_to_my_script
+```
+
+For instance, here is how you would run the NLP example (from the root of the repo) with FSDP enabled:
+
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config: {}
+distributed_type: FSDP
+fsdp_config:
+  min_num_params: 2000
+  offload_params: false
+  sharding_strategy: 1
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: 'no'
+num_machines: 1
+num_processes: 2
+use_cpu: false
+```
+
+```bash
+accelerate launch examples/nlp_example.py
+```
+
+Currently, `Accelerate` supports following config through the CLI:
+
+```bash
+`Sharding Strategy`: [1] FULL_SHARD, [2] SHARD_GRAD_OP
+`Min Num Params`: FSDP\'s minimum number of parameters for Default Auto Wrapping.
+`Offload Params`: Decides Whether to offload parameters and gradients to CPU.
+```
+
+For more control, users can leverage the `FullyShardedDataParallelPlugin` wherein they can specify `auto_wrap_policy`, `backward_prefetch` and `ignored_modules`.
+After creating an instance of this class, users can pass it to the Accelerator class instantiation.
+For more information on these options, please refer to the PyTorch [FullyShardedDataParallel](https://github.com/pytorch/pytorch/blob/0df2e863fbd5993a7b9e652910792bd21a516ff3/torch/distributed/fsdp/fully_sharded_data_parallel.py#L236) code.
+
+[[autodoc]] utils.FullyShardedDataParallelPlugin
+
+## Few caveats to be aware of:
You should highlight that FSDP throws away any parameter groups anyway.
@@ -0,0 +1,107 @@
+<!--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.
+-->
+
+# Fully Sharded Data Parallel
+
+To accelerate training huge models on larger batch sizes, we can use a fully sharded data parallel model.
+This type of data parallel paradigm enables fitting more data and larger models by sharding the optimizer states, gradients and parameters.
+To read more about it and the benefits, check out the [Fully Sharded Data Parallel blog](https://pytorch.org/blog/introducing-pytorch-fully-sharded-data-parallel-api/).
+We have integrated the latest PyTorch's Fully Sharded Data Parallel (FSDP) training feature.
+All you need to do is enable it through the config.
+
+## How it works out the box
+
+On your machine(s) just run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked. This will generate a config file that will be used automatically to properly set the
+default options when doing
+
+```bash
+accelerate launch my_script.py --args_to_my_script
+```
+
+For instance, here is how you would run the NLP example (from the root of the repo) with FSDP enabled:
+
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config: {}
+distributed_type: FSDP
+fsdp_config:
+  min_num_params: 2000
+  offload_params: false
+  sharding_strategy: 1
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: 'no'
+num_machines: 1
+num_processes: 2
+use_cpu: false
+```
+
+```bash
+accelerate launch examples/nlp_example.py
+```
+
+Currently, `Accelerate` supports following config through the CLI:
+
+```bash
+`Sharding Strategy`: [1] FULL_SHARD, [2] SHARD_GRAD_OP
+`Min Num Params`: FSDP\'s minimum number of parameters for Default Auto Wrapping.
+`Offload Params`: Decides Whether to offload parameters and gradients to CPU.
+```
+
+For more control, users can leverage the `FullyShardedDataParallelPlugin` wherein they can specify `auto_wrap_policy`, `backward_prefetch` and `ignored_modules`.
+After creating an instance of this class, users can pass it to the Accelerator class instantiation.
+For more information on these options, please refer to the PyTorch [FullyShardedDataParallel](https://github.com/pytorch/pytorch/blob/0df2e863fbd5993a7b9e652910792bd21a516ff3/torch/distributed/fsdp/fully_sharded_data_parallel.py#L236) code.
+
+[[autodoc]] utils.FullyShardedDataParallelPlugin
+
+## Few caveats to be aware of:
+
+- PyTorch FSDP auto wraps sub-modules, flattens the parameters and shards the parameters in place.
+  Due to this, any optimizer created before model wrapping gets broken and occupies more memory.
+  Hence, it is highly recommended and efficient to prepare model before creating optimizer.
+  `Accelerate` will automatically wrap the model and create an optimizer for you in case of single model with a warning message.
+  > FSDP Warning: When using FSDP, it is efficient and recommended to call prepare for the model before creating the optimizer
+
+However, below is the recommended way to prepare model and optimizer while using FSDP:
+
+```diff
+model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
++ model = accelerator.prepare(model)
+
+optimizer = torch.optim.AdamW(params=model.parameters(), lr=lr)
+
+- model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(model,
+-        optimizer, train_dataloader, eval_dataloader, lr_scheduler
+-    )
+
++ optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
++         optimizer, train_dataloader, eval_dataloader, lr_scheduler
++        )
+
+```
+
+- In case of a single model, if you have created optimizer with multiple optimizer groups and called prepare with them together,
+  then the optimizer groups will be lost and the following warning is displayed:
+
+  > FSDP Warning: When using FSDP, several parameter groups will be conflated into
+  > a single one due to nested module wrapping and parameter flattening.
Done.
@@ -0,0 +1,107 @@
+<!--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.
+-->
+
+# Fully Sharded Data Parallel
+
+To accelerate training huge models on larger batch sizes, we can use a fully sharded data parallel model.
+This type of data parallel paradigm enables fitting more data and larger models by sharding the optimizer states, gradients and parameters.
+To read more about it and the benefits, check out the [Fully Sharded Data Parallel blog](https://pytorch.org/blog/introducing-pytorch-fully-sharded-data-parallel-api/).
+We have integrated the latest PyTorch's Fully Sharded Data Parallel (FSDP) training feature.
+All you need to do is enable it through the config.
+
+## How it works out the box
+
+On your machine(s) just run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked. This will generate a config file that will be used automatically to properly set the
+default options when doing
+
+```bash
+accelerate launch my_script.py --args_to_my_script
+```
+
+For instance, here is how you would run the NLP example (from the root of the repo) with FSDP enabled:
+
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config: {}
+distributed_type: FSDP
+fsdp_config:
+  min_num_params: 2000
+  offload_params: false
+  sharding_strategy: 1
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: 'no'
+num_machines: 1
+num_processes: 2
+use_cpu: false
+```
+
+```bash
+accelerate launch examples/nlp_example.py
+```
+
+Currently, `Accelerate` supports following config through the CLI:
+
+```bash
+`Sharding Strategy`: [1] FULL_SHARD, [2] SHARD_GRAD_OP
+`Min Num Params`: FSDP\'s minimum number of parameters for Default Auto Wrapping.
+`Offload Params`: Decides Whether to offload parameters and gradients to CPU.
+```
+
+For more control, users can leverage the `FullyShardedDataParallelPlugin` wherein they can specify `auto_wrap_policy`, `backward_prefetch` and `ignored_modules`.
+After creating an instance of this class, users can pass it to the Accelerator class instantiation.
+For more information on these options, please refer to the PyTorch [FullyShardedDataParallel](https://github.com/pytorch/pytorch/blob/0df2e863fbd5993a7b9e652910792bd21a516ff3/torch/distributed/fsdp/fully_sharded_data_parallel.py#L236) code.
+
+[[autodoc]] utils.FullyShardedDataParallelPlugin
+
+## Few caveats to be aware of:
Done.
@@ -1,26 +1,28 @@
-- sections: 
-  - local: index
-    title: 🤗 Accelerate
-  - local: quicktour
-    title: Quick tour
-  - local: installation
-    title: Installation
+- sections:
+    - local: index
+      title: 🤗 Accelerate
+    - local: quicktour
+      title: Quick tour
+    - local: installation
+      title: Installation
   title: Get started
 - sections:
-  - local: sagemaker
-    title: Amazon SageMaker
+    - local: sagemaker
+      title: Amazon SageMaker
   title: Guides
 - sections:
-  - local: accelerator
-    title: Accelerator
-  - local: launcher
-    title: Notebook Launcher
-  - local: kwargs
-    title: Kwargs Handlers
-  - local: internal
-    title: Internals
-  - local: checkpoint
-    title: Checkpointing
-  - local: tracking
-    title: Experiment Tracking
+    - local: accelerator
+      title: Accelerator
+    - local: launcher
+      title: Notebook Launcher
+    - local: kwargs
+      title: Kwargs Handlers
+    - local: internal
+      title: Internals
+    - local: checkpoint
+      title: Checkpointing
+    - local: tracking
+      title: Experiment Tracking
+    - local: fsdp
+      title: Fully Sharded Data Parallel
Done.
@@ -0,0 +1,107 @@
+<!--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.
+-->
+
+# Fully Sharded Data Parallel
+
+To accelerate training huge models on larger batch sizes, we can use a fully sharded data parallel model.
+This type of data parallel paradigm enables fitting more data and larger models by sharding the optimizer states, gradients and parameters.
+To read more about it and the benefits, check out the [Fully Sharded Data Parallel blog](https://pytorch.org/blog/introducing-pytorch-fully-sharded-data-parallel-api/).
+We have integrated the latest PyTorch's Fully Sharded Data Parallel (FSDP) training feature.
+All you need to do is enable it through the config.
+
+## How it works out the box
+
+On your machine(s) just run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked. This will generate a config file that will be used automatically to properly set the
+default options when doing
+
+```bash
+accelerate launch my_script.py --args_to_my_script
+```
+
+For instance, here is how you would run the NLP example (from the root of the repo) with FSDP enabled:
+
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config: {}
+distributed_type: FSDP
+fsdp_config:
+  min_num_params: 2000
+  offload_params: false
+  sharding_strategy: 1
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: 'no'
+num_machines: 1
+num_processes: 2
+use_cpu: false
+```
+
+```bash
+accelerate launch examples/nlp_example.py
+```
+
+Currently, `Accelerate` supports following config through the CLI:
+
+```bash
+`Sharding Strategy`: [1] FULL_SHARD, [2] SHARD_GRAD_OP
+`Min Num Params`: FSDP\'s minimum number of parameters for Default Auto Wrapping.
+`Offload Params`: Decides Whether to offload parameters and gradients to CPU.
+```
+
+For more control, users can leverage the `FullyShardedDataParallelPlugin` wherein they can specify `auto_wrap_policy`, `backward_prefetch` and `ignored_modules`.
+After creating an instance of this class, users can pass it to the Accelerator class instantiation.
+For more information on these options, please refer to the PyTorch [FullyShardedDataParallel](https://github.com/pytorch/pytorch/blob/0df2e863fbd5993a7b9e652910792bd21a516ff3/torch/distributed/fsdp/fully_sharded_data_parallel.py#L236) code.
+
+[[autodoc]] utils.FullyShardedDataParallelPlugin
+
+## Few caveats to be aware of:
+
+- PyTorch FSDP auto wraps sub-modules, flattens the parameters and shards the parameters in place.
+  Due to this, any optimizer created before model wrapping gets broken and occupies more memory.
+  Hence, it is highly recommended and efficient to prepare model before creating optimizer.
+  `Accelerate` will automatically wrap the model and create an optimizer for you in case of single model with a warning message.
+  > FSDP Warning: When using FSDP, it is efficient and recommended to call prepare for the model before creating the optimizer
+
+However, below is the recommended way to prepare model and optimizer while using FSDP:
+
+```diff
+model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
++ model = accelerator.prepare(model)
+
+optimizer = torch.optim.AdamW(params=model.parameters(), lr=lr)
+
+- model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(model,
+-        optimizer, train_dataloader, eval_dataloader, lr_scheduler
+-    )
+
++ optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
++         optimizer, train_dataloader, eval_dataloader, lr_scheduler
++        )
+
+```
+
+- In case of a single model, if you have created optimizer with multiple optimizer groups and called prepare with them together,
+  then the optimizer groups will be lost and the following warning is displayed:
+
+  > FSDP Warning: When using FSDP, several parameter groups will be conflated into
+  > a single one due to nested module wrapping and parameter flattening.
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/321
diff --git a/docs/source/tracking.mdx b/docs/source/tracking.mdx
index 959317887..6d4ca437c 100644
--- a/docs/source/tracking.mdx
+++ b/docs/source/tracking.mdx
@@ -28,12 +28,12 @@ Currently `Accelerate` supports three trackers out-of-the-box:
 
 To use any of them, pass in the selected type(s) to the `log_with` parameter in [`Accelerate`]:
 ```python
-from accelerate import Accelerate
+from accelerate import Accelerator
 from accelerate.utils import LoggerType
 
-accelerator = Accelerate(log_with="all")  # For all available trackers in the environment
-accelerator = Accelerate(log_with="wandb")
-accelerator = Accelerate(log_with=["wandb", LoggerType.TENSORBOARD])
+accelerator = Accelerator(log_with="all")  # For all available trackers in the environment
+accelerator = Accelerator(log_with="wandb")
+accelerator = Accelerator(log_with=["wandb", LoggerType.TENSORBOARD])
 ```
 
 At the start of your experiment [`~Accelerator.init_trackers`] should be used to setup your project, and potentially add any experiment hyperparameters to be logged:
@@ -123,14 +123,14 @@ be used with the API:
 
 ```python
 tracker = MyCustomTracker("some_run_name")
-accelerator = Accelerate(log_with=tracker)
+accelerator = Accelerator(log_with=tracker)
 ```
 
 These also can be mixed with existing trackers, including with `"all"`:
 
 ```python
 tracker = MyCustomTracker("some_run_name")
-accelerator = Accelerate(log_with=[tracker, "all"])
+accelerator = Accelerator(log_with=[tracker, "all"])
 ```
 
 ## When a wrapper cannot work

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/320
diff --git a/examples/by_feature/multi_process_metrics.py b/examples/by_feature/multi_process_metrics.py
new file mode 100644
index 000000000..abba1a746
--- /dev/null
+++ b/examples/by_feature/multi_process_metrics.py
@@ -0,0 +1,215 @@
+# coding=utf-8
+# Copyright 2022 The HuggingFace Inc. 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 torch
+from torch.utils.data import DataLoader
+
+from accelerate import Accelerator, DistributedType
+from datasets import load_dataset, load_metric
+from transformers import (
+    AdamW,
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    get_linear_schedule_with_warmup,
+    set_seed,
+)
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate,
+# specifically showcasing how to properly calculate the metrics on the
+# validation dataset when in a distributed system, and builds off the
+# `nlp_example.py` script.
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#
+# To help focus on the differences in the code, building `DataLoaders`
+# was refactored into its own function.
+# New additions from the base script can be found quickly by
+# looking for the # New Code # tags
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+def get_dataloaders(accelerator: Accelerator, batch_size: int = 16):
+    """
+    Creates a set of `DataLoader`s for the `glue` dataset,
+    using "bert-base-cased" as the tokenizer.
+
+    Args:
+        accelerator (`Accelerator`):
+            An `Accelerator` object
+        batch_size (`int`, *optional*):
+            The batch size for the train and validation DataLoaders.
+    """
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    datasets = load_dataset("glue", "mrpc")
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["idx", "sentence1", "sentence2"],
+    )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    return train_dataloader, eval_dataloader
+
+
+def training_function(config, args):
+    # Initialize accelerator
+    accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    correct_bias = config["correct_bias"]
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+
+    metric = load_metric("glue", "mrpc")
+
+    # If the batch size is too big we use gradient accumulation
+    gradient_accumulation_steps = 1
+    if batch_size > MAX_GPU_BATCH_SIZE:
+        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
+        batch_size = MAX_GPU_BATCH_SIZE
+
+    set_seed(seed)
+    train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
+
+    # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
+    # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
+    # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
+    model = model.to(accelerator.device)
+
+    # Instantiate optimizer
+    optimizer = AdamW(params=model.parameters(), lr=lr, correct_bias=correct_bias)
+
+    # Instantiate scheduler
+    lr_scheduler = get_linear_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=100,
+        num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,
+    )
+
+    # Prepare everything
+    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
+    # prepare method.
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    # Now we train the model
+    for epoch in range(num_epochs):
+        model.train()
+        for step, batch in enumerate(train_dataloader):
+            # We could avoid this line since we set the accelerator with `device_placement=True`.
+            batch.to(accelerator.device)
+            outputs = model(**batch)
+            loss = outputs.loss
+            loss = loss / gradient_accumulation_steps
+            accelerator.backward(loss)
+            if step % gradient_accumulation_steps == 0:
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+
+        model.eval()
+        samples_seen = 0
+        for step, batch in enumerate(eval_dataloader):
+            # We could avoid this line since we set the accelerator with `device_placement=True`.
+            batch.to(accelerator.device)
+            with torch.no_grad():
+                outputs = model(**batch)
+            predictions = outputs.logits.argmax(dim=-1)
+            predictions, references = accelerator.gather((predictions, batch["labels"]))
+            # New Code #
+            # First we check if it's a distributed system
+            if accelerator.num_processes > 1:
+                # Then see if we're on the last batch of our eval dataloader
+                if step == len(eval_dataloader):
+                    # Last batch needs to be truncated on distributed systems as it contains additional samples
+                    predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
+                    references = references[: len(eval_dataloader.dataset) - samples_seen]
+                else:
+                    # Otherwise we add the number of samples seen
+                    samples_seen += references.shape[0]
+            metric.add_batch(
+                predictions=predictions,
+                references=references,
+            )
+
+        eval_metric = metric.compute()
+        # Use accelerator.print to print only on the main process.
+        accelerator.print(f"epoch {epoch}:", eval_metric)
+
+
+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()
+    config = {"lr": 2e-5, "num_epochs": 3, "correct_bias": True, "seed": 42, "batch_size": 16}
+    training_function(config, args)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/test_examples.py b/tests/test_examples.py
index 354412aa5..3ecd3bf91 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -33,13 +33,14 @@
 if SRC_DIRS is not None:
     import checkpointing
     import cross_validation
+    import multi_process_metrics
     import tracking
 
 # DataLoaders built from `test_samples/MRPC` for quick testing
 # Should mock `{script_name}.get_dataloaders` via:
 # @mock.patch("{script_name}.get_dataloaders", mocked_dataloaders)
 
-EXCLUDE_EXAMPLES = ["cross_validation.py"]
+EXCLUDE_EXAMPLES = ["cross_validation.py", "multi_process_metrics.py"]
 
 
 def mocked_dataloaders(accelerator, batch_size: int = 16):
@@ -182,18 +183,6 @@ def test_checkpointing_by_steps(self):
                 checkpointing.main()
                 self.assertTrue(os.path.exists(os.path.join(tmpdir, "step_2")))
 
-    @mock.patch("tracking.get_dataloaders", mocked_dataloaders)
-    def test_tracking(self):
-        with tempfile.TemporaryDirectory() as tmpdir:
-            testargs = f"""
-            tracking.py
-            --with_tracking
-            --logging_dir {tmpdir}
-            """.split()
-            with mock.patch.object(sys, "argv", testargs):
-                tracking.main()
-                self.assertTrue(os.path.exists(os.path.join(tmpdir, "tracking")))
-
     @slow
     def test_cross_validation(self):
         testargs = """
@@ -205,3 +194,21 @@ def test_cross_validation(self):
                 cross_validation.main()
                 call = mocked_print.mock_calls[-1]
                 self.assertGreaterEqual(call.args[1]["accuracy"], 0.75)
+
+    @mock.patch("multi_process_metrics.get_dataloaders", mocked_dataloaders)
+    def test_multi_process_metrics(self):
+        testargs = ["multi_process_metrics.py"]
+        with mock.patch.object(sys, "argv", testargs):
+            multi_process_metrics.main()
+
+    @mock.patch("tracking.get_dataloaders", mocked_dataloaders)
+    def test_tracking(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            testargs = f"""
+            tracking.py
+            --with_tracking
+            --logging_dir {tmpdir}
+            """.split()
+            with mock.patch.object(sys, "argv", testargs):
+                tracking.main()
+                self.assertTrue(os.path.exists(os.path.join(tmpdir, "tracking")))

Should be named `multi_process_metrics` as the problem arises as soon as you have several processes, not several nodes.
@@ -0,0 +1,215 @@
+# coding=utf-8
```suggestion
# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
```
@@ -0,0 +1,215 @@
+# coding=utf-8
+# Copyright 2021 The HuggingFace Inc. team. All rights reserved.
_The documentation is not available anymore as the PR was closed or merged._
Note: This is just an initial to make sure the format and whatnot looks right and then all the other examples will follow suite :) 
https://api.github.com/repos/huggingface/accelerate/pulls/319
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index a79012d5c..eb3b24ce9 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -725,6 +725,13 @@ def free_memory(self):
         gc.collect()
         torch.cuda.empty_cache()
 
+    def clear(self):
+        """
+        Alias for [`Accelerate.free_memory`], releases all references to the internal objects stored and call the
+        garbage collector. You should call this method between two trainings with different models/optimizers.
+        """
+        self.free_memory()
+
     def _get_named_parameters(self, *args):
         named_parameters = {}
         for obj in args:

```suggestion
        Alias for [`Accelerate.free_memory`], releases all references to the internal objects stored and call the garbage
```
@@ -725,6 +725,13 @@ def free_memory(self):
         gc.collect()
         torch.cuda.empty_cache()
 
+    def clear(self):
+        """
+        Alias for `Accelerate.free_memory`, releases all references to the internal objects stored and call the garbage
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/318
diff --git a/examples/by_feature/README.md b/examples/by_feature/README.md
index 40ffa3c1e..f0db09937 100644
--- a/examples/by_feature/README.md
+++ b/examples/by_feature/README.md
@@ -54,3 +54,15 @@ These arguments should be added at the end of any method for starting the python
 ```bash
 accelerate launch ./tracking.py --with_tracking
 ```
+
+### Cross Validation (`cross_validation.py`)
+
+- Shows how to use `Accelerator.free_memory` and run cross validation efficiently with `datasets`.
+- Arguments available:
+  - `num_folds`, the number of folds the training dataset should be split into.
+
+These arguments should be added at the end of any method for starting the python script (such as `python`, `accelerate launch`, `python -m torch.distributed.launch`), such as:
+
+```bash
+accelerate launch ./cross_validation.py --num_folds 2
+```
diff --git a/examples/by_feature/cross_validation.py b/examples/by_feature/cross_validation.py
new file mode 100644
index 000000000..9fdc3fe72
--- /dev/null
+++ b/examples/by_feature/cross_validation.py
@@ -0,0 +1,275 @@
+# coding=utf-8
+# Copyright 2022 The HuggingFace Inc. 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
+from typing import List
+
+import numpy as np
+import torch
+from torch.utils.data import DataLoader
+
+from accelerate import Accelerator, DistributedType
+from datasets import DatasetDict, load_dataset, load_metric
+
+# New Code #
+# We'll be using StratifiedKFold for this example
+from sklearn.model_selection import StratifiedKFold
+from transformers import (
+    AdamW,
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    get_linear_schedule_with_warmup,
+    set_seed,
+)
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate,
+# specifically showcasing how to perform Cross Validation,
+# and builds off the `nlp_example.py` script.
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#
+# To help focus on the differences in the code, building `DataLoaders`
+# was refactored into its own function.
+# New additions from the base script can be found quickly by
+# looking for the # New Code # tags
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+# New Code #
+# We need a different `get_dataloaders` function that will build dataloaders by indexs
+
+
+def get_fold_dataloaders(
+    accelerator: Accelerator, dataset: DatasetDict, train_idxs: List[int], valid_idxs: List[int], batch_size: int = 16
+):
+    """
+    Gets a set of train, valid, and test dataloaders for a particular fold
+
+    Args:
+        accelerator (`Accelerator`):
+            The main `Accelerator` object
+        train_idxs (list of `int`):
+            The split indicies for the training dataset
+        valid_idxs (list of `int`):
+            The split indicies for the validation dataset
+        batch_size (`int`):
+            The size of the minibatch. Default is 16
+    """
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    datasets = DatasetDict(
+        {
+            "train": dataset["train"].select(train_idxs),
+            "validation": dataset["train"].select(valid_idxs),
+            "test": dataset["validation"],
+        }
+    )
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["idx", "sentence1", "sentence2"],
+    )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    test_dataloader = DataLoader(
+        tokenized_datasets["test"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    return train_dataloader, eval_dataloader, test_dataloader
+
+
+def training_function(config, args):
+    # New Code #
+    test_labels = None
+    test_predictions = []
+    # Download the dataset
+    datasets = load_dataset("glue", "mrpc")
+    # Create our splits
+    kfold = StratifiedKFold(n_splits=int(args.num_folds))
+    # Initialize accelerator
+    accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    correct_bias = config["correct_bias"]
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+
+    metric = load_metric("glue", "mrpc")
+
+    # If the batch size is too big we use gradient accumulation
+    gradient_accumulation_steps = 1
+    if batch_size > MAX_GPU_BATCH_SIZE:
+        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
+        batch_size = MAX_GPU_BATCH_SIZE
+
+    set_seed(seed)
+
+    # New Code #
+    # Create our folds:
+    folds = kfold.split(np.zeros(datasets["train"].num_rows), datasets["train"]["label"])
+
+    # Iterate over them
+    for train_idxs, valid_idxs in folds:
+        train_dataloader, eval_dataloader, test_dataloader = get_fold_dataloaders(
+            accelerator,
+            datasets,
+            train_idxs,
+            valid_idxs,
+        )
+        if test_labels is None:
+            test_labels = datasets["validation"]["label"]
+        # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+        model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
+
+        # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
+        # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
+        # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
+        model = model.to(accelerator.device)
+
+        # Instantiate optimizer
+        optimizer = AdamW(params=model.parameters(), lr=lr, correct_bias=correct_bias)
+
+        # Instantiate scheduler
+        lr_scheduler = get_linear_schedule_with_warmup(
+            optimizer=optimizer,
+            num_warmup_steps=100,
+            num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,
+        )
+
+        # Prepare everything
+        # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
+        # prepare method.
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+            model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+        )
+
+        # Now we train the model
+        for epoch in range(num_epochs):
+            model.train()
+            for step, batch in enumerate(train_dataloader):
+                # We could avoid this line since we set the accelerator with `device_placement=True`.
+                batch.to(accelerator.device)
+                outputs = model(**batch)
+                loss = outputs.loss
+                loss = loss / gradient_accumulation_steps
+                accelerator.backward(loss)
+                if step % gradient_accumulation_steps == 0:
+                    optimizer.step()
+                    lr_scheduler.step()
+                    optimizer.zero_grad()
+
+            model.eval()
+            for step, batch in enumerate(eval_dataloader):
+                # We could avoid this line since we set the accelerator with `device_placement=True`.
+                batch.to(accelerator.device)
+                with torch.no_grad():
+                    outputs = model(**batch)
+                predictions = outputs.logits.argmax(dim=-1)
+                predictions, references = accelerator.gather((predictions, batch["labels"]))
+                metric.add_batch(
+                    predictions=predictions,
+                    references=references,
+                )
+
+            eval_metric = metric.compute()
+            # Use accelerator.print to print only on the main process.
+            accelerator.print(f"epoch {epoch}:", eval_metric)
+
+        # New Code #
+        # We also run predictions on the test set at the very end
+        fold_predictions = []
+        for step, batch in enumerate(test_dataloader):
+            # We could avoid this line since we set the accelerator with `device_placement=True`.
+            batch.to(accelerator.device)
+            with torch.no_grad():
+                outputs = model(**batch)
+            predictions = outputs.logits
+            predictions, references = accelerator.gather((predictions, batch["labels"]))
+            fold_predictions.append(predictions.cpu())
+            metric.add_batch(
+                predictions=predictions.argmax(dim=-1),
+                references=references,
+            )
+        test_metric = metric.compute()
+        # Use accelerator.print to print only on the main process.
+        test_predictions.append(torch.cat(fold_predictions, dim=0))
+        # We now need to release all our memory and get rid of the current model, optimizer, etc
+        accelerator.free_memory()
+    # New Code #
+    # Finally we check the accuracy of our folded results:
+    preds = torch.stack(test_predictions, dim=0).sum(dim=0).div(int(config["n_splits"])).argmax(dim=-1)
+    test_metric = metric.compute(predictions=preds, references=test_labels)
+    accelerator.print("Average test metrics from all folds:", test_metric)
+
+
+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.")
+    # New Code #
+    parser.add_argument("--num_folds", type=int, default=3, help="The number of splits to perform across the dataset")
+    args = parser.parse_args()
+    config = {"lr": 2e-5, "num_epochs": 3, "correct_bias": True, "seed": 42, "batch_size": 16}
+    training_function(config, args)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/accelerate/test_utils/testing.py b/src/accelerate/test_utils/testing.py
index 6a9e66d6a..e348e161f 100644
--- a/src/accelerate/test_utils/testing.py
+++ b/src/accelerate/test_utils/testing.py
@@ -18,6 +18,7 @@
 import sys
 import tempfile
 import unittest
+from distutils.util import strtobool
 from pathlib import Path
 from typing import List, Union
 from unittest import mock
@@ -28,6 +29,36 @@
 from ..utils import gather, is_tensorflow_available
 
 
+def parse_flag_from_env(key, default=False):
+    try:
+        value = os.environ[key]
+    except KeyError:
+        # KEY isn't set, default to `default`.
+        _value = default
+    else:
+        # KEY is set, convert it to True or False.
+        try:
+            _value = strtobool(value)
+        except ValueError:
+            # More values are supported, but let's keep the message simple.
+            raise ValueError(f"If set, {key} must be yes or no.")
+    return _value
+
+
+_run_slow_tests = parse_flag_from_env("RUN_SLOW", default=False)
+
+
+def slow(test_case):
+    """
+    Decorator marking a test as slow. Slow tests are skipped by default. Set the RUN_SLOW environment variable to a
+    truthy value to run them.
+    """
+    if not _run_slow_tests:
+        return unittest.skip("test is slow")(test_case)
+    else:
+        return test_case
+
+
 class TempDirTestCase(unittest.TestCase):
     """
     A TestCase class that keeps a single `tempfile.TemporaryDirectory` open for the duration of the class, wipes its
diff --git a/tests/test_examples.py b/tests/test_examples.py
index d60d7c08f..354412aa5 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -22,6 +22,7 @@
 
 from accelerate import DistributedType
 from accelerate.test_utils.examples import compare_against_test
+from accelerate.test_utils.testing import slow
 from datasets import load_dataset
 from transformers import AutoTokenizer
 
@@ -31,12 +32,15 @@
 
 if SRC_DIRS is not None:
     import checkpointing
+    import cross_validation
     import tracking
 
 # DataLoaders built from `test_samples/MRPC` for quick testing
 # Should mock `{script_name}.get_dataloaders` via:
 # @mock.patch("{script_name}.get_dataloaders", mocked_dataloaders)
 
+EXCLUDE_EXAMPLES = ["cross_validation.py"]
+
 
 def mocked_dataloaders(accelerator, batch_size: int = 16):
     tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
@@ -117,21 +121,22 @@ def one_complete_example(
         by_feature_path = os.path.abspath(os.path.join("examples", "by_feature"))
         examples_path = os.path.abspath("examples")
         for item in os.listdir(by_feature_path):
-            item_path = os.path.join(by_feature_path, item)
-            if os.path.isfile(item_path) and ".py" in item_path:
-                with self.subTest(
-                    tested_script=complete_file_name,
-                    feature_script=item,
-                    tested_section="main()" if parser_only else "training_function()",
-                ):
-                    diff = compare_against_test(
-                        os.path.join(examples_path, complete_file_name), item_path, parser_only, secondary_filename
-                    )
-                    diff = "\n".join(diff)
-                    if special_strings is not None:
-                        for string in special_strings:
-                            diff = diff.replace(string, "")
-                    self.assertEqual(diff, "")
+            if item not in EXCLUDE_EXAMPLES:
+                item_path = os.path.join(by_feature_path, item)
+                if os.path.isfile(item_path) and ".py" in item_path:
+                    with self.subTest(
+                        tested_script=complete_file_name,
+                        feature_script=item,
+                        tested_section="main()" if parser_only else "training_function()",
+                    ):
+                        diff = compare_against_test(
+                            os.path.join(examples_path, complete_file_name), item_path, parser_only, secondary_filename
+                        )
+                        diff = "\n".join(diff)
+                        if special_strings is not None:
+                            for string in special_strings:
+                                diff = diff.replace(string, "")
+                        self.assertEqual(diff, "")
 
     def test_nlp_examples(self):
         self.one_complete_example("complete_nlp_example.py", True)
@@ -188,3 +193,15 @@ def test_tracking(self):
             with mock.patch.object(sys, "argv", testargs):
                 tracking.main()
                 self.assertTrue(os.path.exists(os.path.join(tmpdir, "tracking")))
+
+    @slow
+    def test_cross_validation(self):
+        testargs = """
+        cross_validation.py
+        --num_folds 2
+        """.split()
+        with mock.patch.object(sys, "argv", testargs):
+            with mock.patch("accelerate.Accelerator.print") as mocked_print:
+                cross_validation.main()
+                call = mocked_print.mock_calls[-1]
+                self.assertGreaterEqual(call.args[1]["accuracy"], 0.75)

<s>Still todo is update the script docs, but code wise it's good</s> 

Done!
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/317
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 514b3a4f4..a79012d5c 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -29,7 +29,7 @@
 from .optimizer import AcceleratedOptimizer
 from .scheduler import AcceleratedScheduler
 from .state import AcceleratorState, DistributedType, is_deepspeed_available
-from .tracking import CometMLTracker, GeneralTracker, TensorBoardTracker, WandBTracker, get_available_trackers
+from .tracking import LOGGER_TYPE_TO_CLASS, GeneralTracker, filter_trackers
 from .utils import (
     DeepSpeedPlugin,
     LoggerType,
@@ -39,9 +39,6 @@
     extract_model_from_parallel,
     gather,
     get_pretty_name,
-    is_comet_ml_available,
-    is_tensorboard_available,
-    is_wandb_available,
     pad_across_processes,
     save,
     wait_for_everyone,
@@ -132,32 +129,13 @@ def __init__(
         deepspeed_plugin: DeepSpeedPlugin = None,
         rng_types: Optional[List[Union[str, RNGType]]] = None,
         log_with: Optional[List[Union[str, LoggerType, GeneralTracker]]] = None,
-        logging_dir: Optional[Union[str, os.PathLike]] = "",
+        logging_dir: Optional[Union[str, os.PathLike]] = None,
         dispatch_batches: Optional[bool] = None,
         step_scheduler_with_optimizer: bool = True,
         kwargs_handlers: Optional[List[KwargsHandler]] = None,
     ):
-        loggers = []
-        if log_with is not None:
-            if not isinstance(log_with, (list, tuple)):
-                log_with = [log_with]
-                logger.debug(f"{log_with}")
-            if "all" in log_with or LoggerType.ALL in log_with:
-                loggers = [o for o in log_with if issubclass(type(o), GeneralTracker)] + get_available_trackers()
-            else:
-                for log_type in log_with:
-                    if log_type not in LoggerType and not issubclass(type(log_type), GeneralTracker):
-                        raise ValueError(
-                            f"Unsupported logging capability: {log_type}. Choose between {LoggerType.list()}"
-                        )
-                    if issubclass(type(log_type), GeneralTracker):
-                        loggers.append(log_type)
-                    else:
-                        log_type = LoggerType(log_type)
-                        if log_type not in loggers:
-                            loggers.append(log_type)
-        self.log_with = loggers
         self.logging_dir = logging_dir
+        self.log_with = filter_trackers(log_with, self.logging_dir)
 
         if mixed_precision is not None:
             mixed_precision = str(mixed_precision)
@@ -644,12 +622,13 @@ def init_trackers(self, project_name: str, config: Optional[dict] = None):
             if issubclass(type(tracker), GeneralTracker):
                 # Custom trackers are already initialized
                 self.trackers.append(tracker)
-            elif str(tracker).lower() == "tensorboard" and is_tensorboard_available():
-                self.trackers.append(TensorBoardTracker(project_name, self.logging_dir))
-            elif str(tracker).lower() == "wandb" and is_wandb_available():
-                self.trackers.append(WandBTracker(project_name))
-            elif str(tracker).lower() == "comet_ml" and is_comet_ml_available():
-                self.trackers.append(CometMLTracker(project_name))
+            else:
+                tracker_init = LOGGER_TYPE_TO_CLASS[str(tracker)]
+                if getattr(tracker_init, "requires_logging_directory"):
+                    # We can skip this check since it was done in `__init__`
+                    self.trackers.append(tracker_init(project_name, self.logging_dir))
+                else:
+                    self.trackers.append(tracker_init(project_name))
         if config is not None:
             for tracker in self.trackers:
                 tracker.store_init_configuration(config)
diff --git a/src/accelerate/tracking.py b/src/accelerate/tracking.py
index 8d496a3d1..de3a19310 100644
--- a/src/accelerate/tracking.py
+++ b/src/accelerate/tracking.py
@@ -17,8 +17,8 @@
 
 import logging
 import os
-from abc import ABCMeta, abstractmethod
-from typing import Optional, Union
+from abc import ABCMeta, abstractmethod, abstractproperty
+from typing import List, Optional, Union
 
 from .utils import LoggerType, is_comet_ml_available, is_tensorboard_available, is_wandb_available
 
@@ -54,6 +54,13 @@ class GeneralTracker(object, metaclass=ABCMeta):
     A base Tracker class to be used for all logging integration implementations.
     """
 
+    @abstractproperty
+    def requires_logging_directory(self):
+        """
+        Whether the logger requires a directory to store their logs. Should either return `True` or `False`.
+        """
+        pass
+
     @abstractmethod
     def store_init_configuration(self, values: dict):
         """
@@ -100,7 +107,9 @@ class TensorBoardTracker(GeneralTracker):
             Location for TensorBoard logs to be stored.
     """
 
-    def __init__(self, run_name: str, logging_dir: Optional[Union[str, os.PathLike]] = ""):
+    requires_logging_directory = True
+
+    def __init__(self, run_name: str, logging_dir: Optional[Union[str, os.PathLike]]):
         self.run_name = run_name
         self.logging_dir = os.path.join(logging_dir, run_name)
         self.writer = tensorboard.SummaryWriter(self.logging_dir)
@@ -157,6 +166,8 @@ class WandBTracker(GeneralTracker):
             The name of the experiment run.
     """
 
+    requires_logging_directory = False
+
     def __init__(self, run_name: str):
         self.run_name = run_name
         self.run = wandb.init(self.run_name)
@@ -209,6 +220,8 @@ class CometMLTracker(GeneralTracker):
             The name of the experiment run.
     """
 
+    requires_logging_directory = False
+
     def __init__(self, run_name: str):
         self.run_name = run_name
         self.writer = Experiment(project_name=run_name)
@@ -250,3 +263,59 @@ def finish(self):
         """
         self.writer.end()
         logger.info("CometML run closed")
+
+
+LOGGER_TYPE_TO_CLASS = {"tensorboard": TensorBoardTracker, "wandb": WandBTracker, "comet_ml": CometMLTracker}
+
+
+def filter_trackers(
+    log_with: List[Union[str, LoggerType, GeneralTracker]], logging_dir: Union[str, os.PathLike] = None
+):
+    """
+    Takes in a list of potential tracker types and checks that:
+        - The tracker wanted is available in that environment
+        - Filters out repeats of tracker types
+        - If `all` is in `log_with`, will return all trackers in the environment
+        - If a tracker requires a `logging_dir`, ensures that `logging_dir` is not `None`
+
+    Args:
+        log_with (list of `str`, [`~utils.LoggerType`] or [`~tracking.GeneralTracker`], *optional*):
+            A list of loggers to be setup for experiment tracking. Should be one or several of:
+
+            - `"all"`
+            - `"tensorboard"`
+            - `"wandb"`
+            - `"comet_ml"`
+            If `"all`" is selected, will pick up all available trackers in the environment and intialize them. Can also
+            accept implementations of `GeneralTracker` for custom trackers, and can be combined with `"all"`.
+        logging_dir (`str`, `os.PathLike`, *optional*):
+            A path to a directory for storing logs of locally-compatible loggers.
+    """
+    loggers = []
+    if log_with is not None:
+        if not isinstance(log_with, (list, tuple)):
+            log_with = [log_with]
+            logger.debug(f"{log_with}")
+        if "all" in log_with or LoggerType.ALL in log_with:
+            loggers = [o for o in log_with if issubclass(type(o), GeneralTracker)] + get_available_trackers()
+        else:
+            for log_type in log_with:
+                if log_type not in LoggerType and not issubclass(type(log_type), GeneralTracker):
+                    raise ValueError(f"Unsupported logging capability: {log_type}. Choose between {LoggerType.list()}")
+                if issubclass(type(log_type), GeneralTracker):
+                    loggers.append(log_type)
+                else:
+                    log_type = LoggerType(log_type)
+                    if log_type not in loggers:
+                        if log_type in get_available_trackers():
+                            tracker_init = LOGGER_TYPE_TO_CLASS[str(log_type)]
+                            if getattr(tracker_init, "requires_logging_directory"):
+                                if logging_dir is None:
+                                    raise ValueError(
+                                        f"Logging with `{str(log_type)}` requires a `logging_dir` to be passed in."
+                                    )
+                            loggers.append(log_type)
+                        else:
+                            logger.info(f"Tried adding logger {log_type}, but package is unavailable in the system.")
+
+    return loggers
diff --git a/tests/test_tracking.py b/tests/test_tracking.py
index b3974afe2..c5b940a46 100644
--- a/tests/test_tracking.py
+++ b/tests/test_tracking.py
@@ -103,6 +103,12 @@ def test_log(self):
         self.assertEqual(iteration, values["iteration"])
         self.assertEqual(my_text, values["my_text"])
 
+    def test_logging_dir(self):
+        with self.assertRaisesRegex(ValueError, "Logging with `tensorboard` requires a `logging_dir`"):
+            _ = Accelerator(log_with="tensorboard")
+        with tempfile.TemporaryDirectory() as dirpath:
+            _ = Accelerator(log_with="tensorboard", logging_dir=dirpath)
+
 
 @mock.patch.dict(os.environ, {"WANDB_MODE": "offline"})
 class WandBTrackingTest(TempDirTestCase, MockingTestCase):
@@ -237,6 +243,8 @@ class MyCustomTracker(GeneralTracker):
         "some_string",
     ]
 
+    requires_logging_directory = False
+
     def __init__(self, dir: str):
         self.f = open(f"{dir}/log.csv", "w+")
         self.writer = csv.DictWriter(self.f, fieldnames=self._col_names)

Why not use a class attribute directly? We could set a default then (I don't think this actually defaults to False).
@@ -54,6 +54,14 @@ class GeneralTracker(object, metaclass=ABCMeta):
     A base Tracker class to be used for all logging integration implementations.
     """
 
+    @abstractproperty
+    def requires_logging_directory(self):
+        """
+        Whether the logger requires a directory to store their logs. Should either return `True` or `False`. The
+        default behavior is `False`
+        """
+        pass
Very nice!
@@ -644,12 +653,13 @@ def init_trackers(self, project_name: str, config: Optional[dict] = None):
             if issubclass(type(tracker), GeneralTracker):
                 # Custom trackers are already initialized
                 self.trackers.append(tracker)
-            elif str(tracker).lower() == "tensorboard" and is_tensorboard_available():
-                self.trackers.append(TensorBoardTracker(project_name, self.logging_dir))
-            elif str(tracker).lower() == "wandb" and is_wandb_available():
-                self.trackers.append(WandBTracker(project_name))
-            elif str(tracker).lower() == "comet_ml" and is_comet_ml_available():
-                self.trackers.append(CometMLTracker(project_name))
+            else:
This won't right. (This is also evidence of me going back and forth with it). Having it like this IMO forces implementations to think about how they want the behavior to go, and have it explicit in the source code for users interested in the code to see. 

For example, forcing it to be an abstract property/function, W&B could eventually make it a conditional where if it's in `offline_mode` (by checking the env var), then it would be True, otherwise False.

But, I'm not against just making it a normal class attr directly
@@ -54,6 +54,14 @@ class GeneralTracker(object, metaclass=ABCMeta):
     A base Tracker class to be used for all logging integration implementations.
     """
 
+    @abstractproperty
+    def requires_logging_directory(self):
+        """
+        Whether the logger requires a directory to store their logs. Should either return `True` or `False`. The
+        default behavior is `False`
+        """
+        pass
In that case, you should remove the default behavior mention in the docstring, as there is no default :-)
@@ -54,6 +54,14 @@ class GeneralTracker(object, metaclass=ABCMeta):
     A base Tracker class to be used for all logging integration implementations.
     """
 
+    @abstractproperty
+    def requires_logging_directory(self):
+        """
+        Whether the logger requires a directory to store their logs. Should either return `True` or `False`. The
+        default behavior is `False`
+        """
+        pass
Much cleaner!
@@ -132,32 +129,13 @@ def __init__(
         deepspeed_plugin: DeepSpeedPlugin = None,
         rng_types: Optional[List[Union[str, RNGType]]] = None,
         log_with: Optional[List[Union[str, LoggerType, GeneralTracker]]] = None,
-        logging_dir: Optional[Union[str, os.PathLike]] = "",
+        logging_dir: Optional[Union[str, os.PathLike]] = None,
         dispatch_batches: Optional[bool] = None,
         step_scheduler_with_optimizer: bool = True,
         kwargs_handlers: Optional[List[KwargsHandler]] = None,
     ):
-        loggers = []
-        if log_with is not None:
-            if not isinstance(log_with, (list, tuple)):
-                log_with = [log_with]
-                logger.debug(f"{log_with}")
-            if "all" in log_with or LoggerType.ALL in log_with:
-                loggers = [o for o in log_with if issubclass(type(o), GeneralTracker)] + get_available_trackers()
-            else:
-                for log_type in log_with:
-                    if log_type not in LoggerType and not issubclass(type(log_type), GeneralTracker):
-                        raise ValueError(
-                            f"Unsupported logging capability: {log_type}. Choose between {LoggerType.list()}"
-                        )
-                    if issubclass(type(log_type), GeneralTracker):
-                        loggers.append(log_type)
-                    else:
-                        log_type = LoggerType(log_type)
-                        if log_type not in loggers:
-                            loggers.append(log_type)
-        self.log_with = loggers
         self.logging_dir = logging_dir
+        self.log_with = filter_trackers(log_with, self.logging_dir)
_The documentation is not available anymore as the PR was closed or merged._
@sgugger pinging for a rereview just to make sure the refactor seems sound to you 😄 
https://api.github.com/repos/huggingface/accelerate/pulls/316
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index e18686e12..514b3a4f4 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -440,7 +440,10 @@ def _prepare_deepspeed(self, *args):
             batch_size_per_device * deepspeed_plugin.gradient_accumulation_steps * self.num_processes
         )
 
-        result = [self._prepare_one(obj) if isinstance(obj, torch.utils.data.DataLoader) else obj for obj in args]
+        result = [
+            self._prepare_one(obj, first_pass=True) if isinstance(obj, torch.utils.data.DataLoader) else obj
+            for obj in args
+        ]
 
         model = None
         optimizer = None

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/315
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index c213a9167..5020c6cc4 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -14,4 +14,17 @@ jobs:
     - name: Install Python dependencies
       run: pip install setuptools==59.5.0; pip install -e .[test,test_trackers]
     - name: Run Tests
-      run: make test
\ No newline at end of file
+      run: make test
+      
+  test_examples:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v2
+    - name: Set up Python 3.6
+      uses: actions/setup-python@v2
+      with:
+        python-version: 3.6
+    - name: Install Python dependencies
+      run: pip install setuptools==59.5.0; pip install -e .[test] tensorboard
+    - name: Run Tests
+      run: make test_examples
\ No newline at end of file
diff --git a/Makefile b/Makefile
index e5b557f10..db7893565 100644
--- a/Makefile
+++ b/Makefile
@@ -25,4 +25,7 @@ style:
 	
 # Run tests for the library
 test:
-	python -m pytest -n auto --dist=loadfile -s -v ./tests/
\ No newline at end of file
+	python -m pytest -n auto --dist=loadfile -s -v ./tests/ --ignore=./tests/test_examples.py
+
+test_examples:
+	python -m pytest -n auto --dist=loadfile -s -v ./tests/test_examples.py
diff --git a/examples/by_feature/checkpointing.py b/examples/by_feature/checkpointing.py
index 35205107d..e595a8db7 100644
--- a/examples/by_feature/checkpointing.py
+++ b/examples/by_feature/checkpointing.py
@@ -115,14 +115,17 @@ def training_function(config, args):
 
     # New Code #
     # Parse out whether we are saving every epoch or after a certain number of batches
-    if args.checkpointing_steps == "epoch":
-        checkpointing_steps = args.checkpointing_steps
-    elif args.checkpointing_steps.isdigit():
-        checkpointing_steps = int(args.checkpointing_steps)
+    if hasattr(args.checkpointing_steps, "isdigit"):
+        if args.checkpointing_steps == "epoch":
+            checkpointing_steps = args.checkpointing_steps
+        elif args.checkpointing_steps.isdigit():
+            checkpointing_steps = int(args.checkpointing_steps)
+        else:
+            raise ValueError(
+                f"Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed."
+            )
     else:
-        raise ValueError(
-            f"Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed."
-        )
+        checkpointing_steps = None
 
     set_seed(seed)
 
@@ -162,21 +165,29 @@ def training_function(config, args):
 
     # New Code #
     # We need to keep track of how many total steps we have iterated over
-    if isinstance(checkpointing_steps, int):
-        overall_step = 0
+    overall_step = 0
 
     # We need to load the checkpoint back in before training here with `load_state`
     # The total number of epochs is adjusted based on where the state is being loaded from,
     # as we assume continuation of the same training script
     if args.resume_from_checkpoint:
-        accelerator.print(f"Resuming from checkpoint: {args.resume_from_checkpoint}")
-        accelerator.load_state(args.resume_from_checkpoint)
-
-        if "epoch" in args.resume_from_checkpoint:
-            num_epochs -= int(args.resume_from_checkpoint.replace("epoch_", ""))
+        if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
+            accelerator.print(f"Resumed from checkpoint: {args.resume_from_checkpoint}")
+            accelerator.load_state(args.resume_from_checkpoint)
+            path = os.path.basename(args.resume_from_checkpoint)
+        else:
+            # Get the most recent checkpoint
+            dirs = [f.name for f in os.scandir(os.getcwd()) if f.is_dir()]
+            dirs.sort(key=os.path.getctime)
+            path = dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
+        # Extract `epoch_{i}` or `step_{i}`
+        training_difference = os.path.splitext(path)[0]
+
+        if "epoch" in training_difference:
+            num_epochs -= int(training_difference.replace("epoch_", ""))
             resume_step = None
         else:
-            resume_step = int(args.resume_from_checkpoint.replace("step_", ""))
+            resume_step = int(training_difference.replace("step_", ""))
             num_epochs -= resume_step // len(train_dataloader)
             # If resuming by step, we also need to know exactly how far into the DataLoader we went
             resume_step = (num_epochs * len(train_dataloader)) - resume_step
@@ -200,6 +211,8 @@ def training_function(config, args):
                 optimizer.step()
                 lr_scheduler.step()
                 optimizer.zero_grad()
+            # New Code #
+            overall_step += 1
 
             # New Code #
             # We save the model, optimizer, lr_scheduler, and seed states by calling `save_state`
@@ -237,7 +250,7 @@ def training_function(config, args):
         # Will contain files: "pytorch_model.bin", "optimizer.bin", "scheduler.bin", and "random_states.pkl"
         # If mixed precision was used, will also save a "scalar.bin" file
         if checkpointing_steps == "epoch":
-            output_dir = f"epoch_{num_epochs}"
+            output_dir = f"epoch_{epoch}"
             if args.output_dir is not None:
                 output_dir = os.path.join(args.output_dir, output_dir)
             accelerator.save_state(output_dir)
@@ -258,7 +271,7 @@ def main():
     parser.add_argument(
         "--checkpointing_steps",
         type=str,
-        default="epoch",
+        default=None,
         help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.",
     )
     parser.add_argument(
diff --git a/examples/by_feature/tracking.py b/examples/by_feature/tracking.py
index ca11fa04b..170211e5e 100644
--- a/examples/by_feature/tracking.py
+++ b/examples/by_feature/tracking.py
@@ -13,6 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import argparse
+import logging
+import os
 
 import torch
 from torch.utils.data import DataLoader
@@ -28,6 +30,9 @@
 )
 
 
+logger = logging.getLogger(__name__)
+
+
 ########################################################################
 # This is a fully working simple example to use Accelerate,
 # specifically showcasing the experiment tracking capability,
@@ -110,7 +115,9 @@ def training_function(config, args):
     # Note: If using a custom `Tracker` class, should be passed in here such as:
     # >>> log_with = ["all", MyCustomTrackerClassInstance()]
     if args.with_tracking:
-        accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="all")
+        accelerator = Accelerator(
+            cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="all", logging_dir=args.logging_dir
+        )
     else:
         accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
@@ -158,7 +165,10 @@ def training_function(config, args):
     # New Code #
     # We need to initalize the trackers we use. Overall configurations can also be stored
     if args.with_tracking:
-        accelerator.init_trackers("accelerate_glue_with_tracking", config)
+        run = os.path.split(__file__)[-1].split(".")[0]
+        if args.logging_dir:
+            run = os.path.join(args.logging_dir, run)
+        accelerator.init_trackers(run, config)
 
     # Now we train the model
     for epoch in range(num_epochs):
@@ -173,7 +183,8 @@ def training_function(config, args):
             outputs = model(**batch)
             loss = outputs.loss
             # New Code #
-            total_loss += loss.detach().float()
+            if args.with_tracking:
+                total_loss += loss.detach().float()
             loss = loss / gradient_accumulation_steps
             accelerator.backward(loss)
             if step % gradient_accumulation_steps == 0:
@@ -202,14 +213,21 @@ def training_function(config, args):
         # New Code #
         # To actually log, we call `Accelerator.log`
         # The values passed can be of `str`, `int`, or `float`
-        accelerator.log(
-            {"accuracy": eval_metric["accuracy"], "f1": eval_metric["f1"], "train_loss": total_loss, "epoch": epoch}
-        )
+        if args.with_tracking:
+            accelerator.log(
+                {
+                    "accuracy": eval_metric["accuracy"],
+                    "f1": eval_metric["f1"],
+                    "train_loss": total_loss,
+                    "epoch": epoch,
+                }
+            )
 
     # New Code #
     # When a run is finished, you should call `accelerator.end_training()`
     # to close all of the open trackers
-    accelerator.end_training()
+    if args.with_tracking:
+        accelerator.end_training()
 
 
 def main():
@@ -229,6 +247,12 @@ def main():
         action="store_true",
         help="Whether to load in all available experiment trackers from the environment and use them for logging.",
     )
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help="Location on where to store experiment tracking logs`",
+    )
     args = parser.parse_args()
     config = {"lr": 2e-5, "num_epochs": 3, "correct_bias": True, "seed": 42, "batch_size": 16}
     training_function(config, args)
diff --git a/examples/complete_cv_example.py b/examples/complete_cv_example.py
index e62568a39..01033498c 100644
--- a/examples/complete_cv_example.py
+++ b/examples/complete_cv_example.py
@@ -74,16 +74,12 @@ def __getitem__(self, idx):
 def training_function(config, args):
     # Initialize accelerator
     if args.with_tracking:
-        accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="all")
+        accelerator = Accelerator(
+            cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="all", logging_dir=args.logging_dir
+        )
     else:
-        accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision)
+        accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
 
-    if hasattr(args.checkpointing_steps, "isdigit"):
-        checkpointing_steps = args.checkpointing_steps
-        if args.checkpointing_steps.isdigit():
-            checkpointing_steps = int(args.checkpointing_steps)
-    else:
-        checkpointing_steps = None
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
     lr = config["lr"]
     num_epochs = int(config["num_epochs"])
@@ -93,9 +89,25 @@ def training_function(config, args):
     if not isinstance(image_size, (list, tuple)):
         image_size = (image_size, image_size)
 
+    # Parse out whether we are saving every epoch or after a certain number of batches
+    if hasattr(args.checkpointing_steps, "isdigit"):
+        if args.checkpointing_steps == "epoch":
+            checkpointing_steps = args.checkpointing_steps
+        elif args.checkpointing_steps.isdigit():
+            checkpointing_steps = int(args.checkpointing_steps)
+        else:
+            raise ValueError(
+                f"Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed."
+            )
+    else:
+        checkpointing_steps = None
+
     # We need to initialize the trackers we use, and also store our configuration
     if args.with_tracking:
-        accelerator.init_trackers("cv_example", config)
+        run = os.path.split(__file__)[-1].split(".")[0]
+        if args.logging_dir:
+            run = os.path.join(args.logging_dir, run)
+        accelerator.init_trackers(run, config)
 
     # Grab all the image filenames
     file_names = [os.path.join(args.data_dir, fname) for fname in os.listdir(args.data_dir) if fname.endswith(".jpg")]
@@ -163,24 +175,27 @@ def training_function(config, args):
     )
 
     # Potentially load in the weights and states from a previous save
-    state_restored = True
     if args.resume_from_checkpoint:
         if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
             accelerator.print(f"Resumed from checkpoint: {args.resume_from_checkpoint}")
             accelerator.load_state(args.resume_from_checkpoint)
-            resume_step = None
+            path = os.path.basename(args.resume_from_checkpoint)
         else:
             # Get the most recent checkpoint
             dirs = [f.name for f in os.scandir(os.getcwd()) if f.is_dir()]
             dirs.sort(key=os.path.getctime)
             path = dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
-            if "epoch" in path.name:
-                num_epochs -= int(path.name.replace("epoch_", ""))
-            else:
-                resume_step = int(path.name.replace("step_", ""))
-                num_epochs -= resume_step // len(train_dataloader)
-                resume_step = (num_epochs * len(train_dataloader)) - resume_step
-                state_restored = False
+        # Extract `epoch_{i}` or `step_{i}`
+        training_difference = os.path.splitext(path)[0]
+
+        if "epoch" in training_difference:
+            num_epochs -= int(training_difference.replace("epoch_", ""))
+            resume_step = None
+        else:
+            resume_step = int(training_difference.replace("step_", ""))
+            num_epochs -= resume_step // len(train_dataloader)
+            # If resuming by step, we also need to know exactly how far into the DataLoader we went
+            resume_step = (num_epochs * len(train_dataloader)) - resume_step
 
     overall_step = 0
     # Now we train the model
@@ -190,8 +205,9 @@ def training_function(config, args):
             total_loss = 0
         for step, batch in enumerate(train_dataloader):
             # We need to skip steps until we reach the resumed step
-            if args.resume_from_checkpoint and epoch == 0 and step < resume_step:
-                continue
+            if args.resume_from_checkpoint and epoch == 0:
+                if resume_step is not None and step < resume_step:
+                    pass
             # We could avoid this line since we set the accelerator with `device_placement=True`.
             batch = {k: v.to(accelerator.device) for k, v in batch.items()}
             inputs = (batch["image"] - mean) / std
@@ -206,32 +222,40 @@ def training_function(config, args):
             optimizer.zero_grad()
             overall_step += 1
             if isinstance(checkpointing_steps, int):
+                output_dir = f"step_{overall_step}"
                 if overall_step % checkpointing_steps == 0:
-                    accelerator.save_state(f"step_{overall_step}")
-        if state_restored:
-            model.eval()
-            accurate = 0
-            num_elems = 0
-            for step, batch in enumerate(eval_dataloader):
-                # We could avoid this line since we set the accelerator with `device_placement=True`.
-                batch = {k: v.to(accelerator.device) for k, v in batch.items()}
-                inputs = (batch["image"] - mean) / std
-                with torch.no_grad():
-                    outputs = model(inputs)
-                predictions = outputs.argmax(dim=-1)
-                accurate_preds = accelerator.gather(predictions) == accelerator.gather(batch["label"])
-                num_elems += accurate_preds.shape[0]
-                accurate += accurate_preds.long().sum()
-
-            eval_metric = accurate.item() / num_elems
-            # Use accelerator.print to print only on the main process.
-            accelerator.print(f"epoch {epoch}: {100 * eval_metric:.2f}")
-            if args.with_tracking:
-                accelerator.log(
-                    {"accuracy": 100 * eval_metric, "total_loss": total_loss, "epoch": epoch}, step=overall_step
-                )
-            if args.checkpointing_steps == "epoch":
-                accelerator.save_state(f"epoch_{epoch}")
+                    if args.output_dir is not None:
+                        output_dir = os.path.join(args.output_dir, output_dir)
+                    accelerator.save_state(output_dir)
+        model.eval()
+        accurate = 0
+        num_elems = 0
+        for step, batch in enumerate(eval_dataloader):
+            # We could avoid this line since we set the accelerator with `device_placement=True`.
+            batch = {k: v.to(accelerator.device) for k, v in batch.items()}
+            inputs = (batch["image"] - mean) / std
+            with torch.no_grad():
+                outputs = model(inputs)
+            predictions = outputs.argmax(dim=-1)
+            accurate_preds = accelerator.gather(predictions) == accelerator.gather(batch["label"])
+            num_elems += accurate_preds.shape[0]
+            accurate += accurate_preds.long().sum()
+
+        eval_metric = accurate.item() / num_elems
+        # Use accelerator.print to print only on the main process.
+        accelerator.print(f"epoch {epoch}: {100 * eval_metric:.2f}")
+        if args.with_tracking:
+            accelerator.log(
+                {"accuracy": 100 * eval_metric, "total_loss": total_loss, "epoch": epoch}, step=overall_step
+            )
+        if checkpointing_steps == "epoch":
+            output_dir = f"epoch_{epoch}"
+            if args.output_dir is not None:
+                output_dir = os.path.join(args.output_dir, output_dir)
+            accelerator.save_state(output_dir)
+
+    if args.with_tracking:
+        accelerator.end_training()
 
 
 def main():
@@ -254,6 +278,12 @@ def main():
         default=None,
         help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.",
     )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=".",
+        help="Optional save directory where all checkpoint folders will be stored. Default is the current working directory.",
+    )
     parser.add_argument(
         "--resume_from_checkpoint",
         type=str,
@@ -262,9 +292,15 @@ def main():
     )
     parser.add_argument(
         "--with_tracking",
-        required=False,
+        action="store_true",
         help="Whether to load in all available experiment trackers from the environment and use them for logging.",
     )
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help="Location on where to store experiment tracking logs`",
+    )
     args = parser.parse_args()
     config = {"lr": 3e-2, "num_epochs": 3, "seed": 42, "batch_size": 64, "image_size": 224}
     training_function(config, args)
diff --git a/examples/complete_nlp_example.py b/examples/complete_nlp_example.py
index 57125aa29..a0e8d568c 100644
--- a/examples/complete_nlp_example.py
+++ b/examples/complete_nlp_example.py
@@ -55,14 +55,21 @@
 def training_function(config, args):
     # Initialize accelerator
     if args.with_tracking:
-        accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="all")
+        accelerator = Accelerator(
+            cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="all", logging_dir=args.logging_dir
+        )
     else:
-        accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision)
+        accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
 
     if hasattr(args.checkpointing_steps, "isdigit"):
-        checkpointing_steps = args.checkpointing_steps
-        if args.checkpointing_steps.isdigit():
+        if args.checkpointing_steps == "epoch":
+            checkpointing_steps = args.checkpointing_steps
+        elif args.checkpointing_steps.isdigit():
             checkpointing_steps = int(args.checkpointing_steps)
+        else:
+            raise ValueError(
+                f"Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed."
+            )
     else:
         checkpointing_steps = None
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
@@ -74,7 +81,10 @@ def training_function(config, args):
 
     # We need to initialize the trackers we use, and also store our configuration
     if args.with_tracking:
-        accelerator.init_trackers("nlp_example", config)
+        run = os.path.split(__file__)[-1].split(".")[0]
+        if args.logging_dir:
+            run = os.path.join(args.logging_dir, run)
+        accelerator.init_trackers(run, config)
 
     tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
     datasets = load_dataset("glue", "mrpc")
@@ -143,27 +153,31 @@ def collate_fn(examples):
         model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
     )
 
+    overall_step = 0
+
     # Potentially load in the weights and states from a previous save
-    state_restored = True
     if args.resume_from_checkpoint:
         if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
             accelerator.print(f"Resumed from checkpoint: {args.resume_from_checkpoint}")
             accelerator.load_state(args.resume_from_checkpoint)
-            resume_step = None
+            path = os.path.basename(args.resume_from_checkpoint)
         else:
             # Get the most recent checkpoint
             dirs = [f.name for f in os.scandir(os.getcwd()) if f.is_dir()]
             dirs.sort(key=os.path.getctime)
             path = dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
-            if "epoch" in path.name:
-                num_epochs -= int(path.name.replace("epoch_", ""))
-            else:
-                resume_step = int(path.name.replace("step_", ""))
-                num_epochs -= resume_step // len(train_dataloader)
-                resume_step = (num_epochs * len(train_dataloader)) - resume_step
-                state_restored = False
+        # Extract `epoch_{i}` or `step_{i}`
+        training_difference = os.path.splitext(path)[0]
+
+        if "epoch" in training_difference:
+            num_epochs -= int(training_difference.replace("epoch_", ""))
+            resume_step = None
+        else:
+            resume_step = int(training_difference.replace("step_", ""))
+            num_epochs -= resume_step // len(train_dataloader)
+            # If resuming by step, we also need to know exactly how far into the DataLoader we went
+            resume_step = (num_epochs * len(train_dataloader)) - resume_step
 
-    overall_step = 0
     # Now we train the model
     for epoch in range(num_epochs):
         model.train()
@@ -171,8 +185,9 @@ def collate_fn(examples):
             total_loss = 0
         for step, batch in enumerate(train_dataloader):
             # We need to skip steps until we reach the resumed step
-            if args.resume_from_checkpoint and epoch == 0 and step < resume_step:
-                continue
+            if args.resume_from_checkpoint and epoch == 0:
+                if resume_step is not None and step < resume_step:
+                    pass
             # We could avoid this line since we set the accelerator with `device_placement=True`.
             batch.to(accelerator.device)
             outputs = model(**batch)
@@ -190,42 +205,51 @@ def collate_fn(examples):
             overall_step += 1
 
             if isinstance(checkpointing_steps, int):
+                output_dir = f"step_{overall_step}"
                 if overall_step % checkpointing_steps == 0:
-                    accelerator.save_state(f"step_{overall_step}")
-        if state_restored:
-            model.eval()
-            for step, batch in enumerate(eval_dataloader):
-                # We could avoid this line since we set the accelerator with `device_placement=True`.
-                batch.to(accelerator.device)
-                with torch.no_grad():
-                    outputs = model(**batch)
-                predictions = outputs.logits.argmax(dim=-1)
-                metric.add_batch(
-                    predictions=accelerator.gather(predictions),
-                    references=accelerator.gather(batch["labels"]),
-                )
-
-            eval_metric = metric.compute()
-            # Use accelerator.print to print only on the main process.
-            accelerator.print(f"epoch {epoch}:", eval_metric)
-            if args.with_tracking:
-                accelerator.log(
-                    {
-                        "accuracy": eval_metric["accuracy"],
-                        "f1": eval_metric["f1"],
-                        "total_loss": total_loss,
-                        "epoch": epoch,
-                    },
-                    step=overall_step,
-                )
+                    if args.output_dir is not None:
+                        output_dir = os.path.join(args.output_dir, output_dir)
+                    accelerator.save_state(output_dir)
+
+        model.eval()
+        for step, batch in enumerate(eval_dataloader):
+            # We could avoid this line since we set the accelerator with `device_placement=True`.
+            batch.to(accelerator.device)
+            with torch.no_grad():
+                outputs = model(**batch)
+            predictions = outputs.logits.argmax(dim=-1)
+            # It is slightly faster to call this once, than multiple times
+            predictions, references = accelerator.gather((predictions, batch["labels"]))
+            metric.add_batch(
+                predictions=predictions,
+                references=references,
+            )
+
+        eval_metric = metric.compute()
+        # Use accelerator.print to print only on the main process.
+        accelerator.print(f"epoch {epoch}:", eval_metric)
+        if args.with_tracking:
+            accelerator.log(
+                {
+                    "accuracy": eval_metric["accuracy"],
+                    "f1": eval_metric["f1"],
+                    "train_loss": total_loss,
+                    "epoch": epoch,
+                }
+            )
+
+        if checkpointing_steps == "epoch":
+            output_dir = f"epoch_{epoch}"
+            if args.output_dir is not None:
+                output_dir = os.path.join(args.output_dir, output_dir)
+            accelerator.save_state(output_dir)
 
-            if args.checkpointing_steps == "epoch":
-                accelerator.save_state(f"epoch_{epoch}")
+    if args.with_tracking:
+        accelerator.end_training()
 
 
 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,
@@ -250,9 +274,21 @@ def main():
     )
     parser.add_argument(
         "--with_tracking",
-        required=False,
+        action="store_true",
         help="Whether to load in all available experiment trackers from the environment and use them for logging.",
     )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=".",
+        help="Optional save directory where all checkpoint folders will be stored. Default is the current working directory.",
+    )
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help="Location on where to store experiment tracking logs`",
+    )
     args = parser.parse_args()
     config = {"lr": 2e-5, "num_epochs": 3, "correct_bias": True, "seed": 42, "batch_size": 16}
     training_function(config, args)
diff --git a/examples/cv_example.py b/examples/cv_example.py
index b14b067be..a80b42d13 100644
--- a/examples/cv_example.py
+++ b/examples/cv_example.py
@@ -73,7 +73,7 @@ def __getitem__(self, idx):
 
 def training_function(config, args):
     # Initialize accelerator
-    accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mix_precision)
+    accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mix_precision)
 
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
     lr = config["lr"]
@@ -185,7 +185,6 @@ def training_function(config, args):
 def main():
     parser = argparse.ArgumentParser(description="Simple example of training script.")
     parser.add_argument("--data_dir", required=True, help="The data folder on disk.")
-    parser.add_argument("--fp16", action="store_true", help="If passed, will use FP16 training.")
     parser.add_argument(
         "--mixed_precision",
         type=str,
@@ -195,6 +194,12 @@ def main():
         "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
         "and an Nvidia Ampere GPU.",
     )
+    parser.add_argument(
+        "--checkpointing_steps",
+        type=str,
+        default=None,
+        help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.",
+    )
     parser.add_argument("--cpu", action="store_true", help="If passed, will train on the CPU.")
     args = parser.parse_args()
     config = {"lr": 3e-2, "num_epochs": 3, "seed": 42, "batch_size": 64, "image_size": 224}
diff --git a/examples/nlp_example.py b/examples/nlp_example.py
index 915d35807..87cfd7698 100644
--- a/examples/nlp_example.py
+++ b/examples/nlp_example.py
@@ -49,21 +49,19 @@
 EVAL_BATCH_SIZE = 32
 
 
-def training_function(config, args):
-    # Initialize accelerator
-    accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
-    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
-    lr = config["lr"]
-    num_epochs = int(config["num_epochs"])
-    correct_bias = config["correct_bias"]
-    seed = int(config["seed"])
-    batch_size = int(config["batch_size"])
-
+def get_dataloaders(accelerator: Accelerator, batch_size: int = 16):
+    """
+    Creates a set of `DataLoader`s for the `glue` dataset,
+    using "bert-base-cased" as the tokenizer.
+
+    Args:
+        accelerator (`Accelerator`):
+            An `Accelerator` object
+        batch_size (`int`, *optional*):
+            The batch size for the train and validation DataLoaders.
+    """
     tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
     datasets = load_dataset("glue", "mrpc")
-    metric = load_metric("glue", "mrpc")
-
-    set_seed(seed)
 
     def tokenize_function(examples):
         # max_length=None => use the model max length (it's actually the default)
@@ -81,12 +79,6 @@ def tokenize_function(examples):
     # transformers library
     tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
 
-    # If the batch size is too big we use gradient accumulation
-    gradient_accumulation_steps = 1
-    if batch_size > MAX_GPU_BATCH_SIZE:
-        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
-        batch_size = MAX_GPU_BATCH_SIZE
-
     def collate_fn(examples):
         # On TPU it's best to pad everything to the same length or training will be very slow.
         if accelerator.distributed_type == DistributedType.TPU:
@@ -101,6 +93,29 @@ def collate_fn(examples):
         tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
     )
 
+    return train_dataloader, eval_dataloader
+
+
+def training_function(config, args):
+    # Initialize accelerator
+    accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    correct_bias = config["correct_bias"]
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+
+    metric = load_metric("glue", "mrpc")
+
+    # If the batch size is too big we use gradient accumulation
+    gradient_accumulation_steps = 1
+    if batch_size > MAX_GPU_BATCH_SIZE:
+        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
+        batch_size = MAX_GPU_BATCH_SIZE
+
+    set_seed(seed)
+    train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)
     # Instantiate the model (we build the model here so that the seed also control new weights initialization)
     model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
 
@@ -148,9 +163,10 @@ def collate_fn(examples):
             with torch.no_grad():
                 outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1)
+            predictions, references = accelerator.gather((predictions, batch["labels"]))
             metric.add_batch(
-                predictions=accelerator.gather(predictions),
-                references=accelerator.gather(batch["labels"]),
+                predictions=predictions,
+                references=references,
             )
 
         eval_metric = metric.compute()
diff --git a/setup.py b/setup.py
index 0ce1e5af4..7e26759d1 100644
--- a/setup.py
+++ b/setup.py
@@ -20,7 +20,12 @@
 extras["docs"] = []
 extras["test"] = [
     "pytest",
-    "pytest-xdist"
+    "pytest-xdist",
+    "pytest-subtests",
+    "datasets",
+    "transformers",
+    "scipy",
+    "sklearn"
 ]
 extras["test_trackers"] = ["wandb", "comet-ml", "tensorflow"]
 extras["dev"] = extras["quality"] + extras["test"]
diff --git a/src/accelerate/test_utils/examples.py b/src/accelerate/test_utils/examples.py
new file mode 100644
index 000000000..4e4092c0e
--- /dev/null
+++ b/src/accelerate/test_utils/examples.py
@@ -0,0 +1,139 @@
+#!/usr/bin/env python
+
+# 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.
+"""
+A collection of utilities for comparing `examples/complete_*_example.py` scripts with the capabilities inside of each
+`examples/by_feature` example. `compare_against_test` is the main function that should be used when testing, while the
+others are used to either get the code that matters, or to preprocess them (such as stripping comments)
+"""
+
+import os
+from typing import List
+
+
+def get_function_contents_by_name(lines: List[str], name: str):
+    """
+    Extracts a function from `lines` of segmented source code with the name `name`.
+
+    Args:
+        lines (`List[str]`):
+            Source code of a script seperated by line.
+        name (`str`):
+            The name of the function to extract. Should be either `training_function` or `main`
+    """
+    if name != "training_function" and name != "main":
+        raise ValueError(f"Incorrect function name passed: {name}, choose either 'main' or 'training_function'")
+    good_lines, found_start = [], False
+    for line in lines:
+        if not found_start and f"def {name}" in line:
+            found_start = True
+            good_lines.append(line)
+            continue
+        if found_start:
+            if name == "training_function" and "def main" in line:
+                return good_lines
+            if name == "main" and "if __name__" in line:
+                return good_lines
+            good_lines.append(line)
+
+
+def clean_lines(lines: List[str]):
+    """
+    Filters `lines` and removes any entries that start with a comment ('#') or is just a newline ('\n')
+
+    Args:
+        lines (`List[str]`):
+            Source code of a script seperated by line.
+    """
+    return [line for line in lines if not line.lstrip().startswith("#") and line != "\n"]
+
+
+def compare_against_test(base_filename: str, feature_filename: str, parser_only: bool, secondary_filename: str = None):
+    """
+    Tests whether the additional code inside of `feature_filename` was implemented in `base_filename`. This should be
+    used when testing to see if `complete_*_.py` examples have all of the implementations from each of the
+    `examples/by_feature/*` scripts.
+
+    It utilizes `nlp_example.py` to extract out all of the repeated training code, so that only the new additional code
+    is examined and checked. If something *other* than `nlp_example.py` should be used, such as `cv_example.py` for the
+    `complete_cv_example.py` script, it should be passed in for the `secondary_filename` parameter.
+
+    Args:
+        base_filename (`str` or `os.PathLike`):
+            The filepath of a single "complete" example script to test, such as `examples/complete_cv_example.py`
+        feature_filename (`str` or `os.PathLike`):
+            The filepath of a single feature example script. The contents of this script are checked to see if they
+            exist in `base_filename`
+        parser_only (`bool`):
+            Whether to compare only the `main()` sections in both files, or to compare the contents of
+            `training_loop()`
+        secondary_filename (`str`, *optional*):
+            A potential secondary filepath that should be included in the check. This function extracts the base
+            functionalities off of "examples/nlp_example.py", so if `base_filename` is a script other than
+            `complete_nlp_example.py`, the template script should be included here. Such as `examples/cv_example.py`
+    """
+    with open(base_filename, "r") as f:
+        base_file_contents = f.readlines()
+    with open(os.path.abspath(os.path.join("examples", "nlp_example.py")), "r") as f:
+        full_file_contents = f.readlines()
+    with open(feature_filename, "r") as f:
+        feature_file_contents = f.readlines()
+    if secondary_filename is not None:
+        with open(secondary_filename, "r") as f:
+            secondary_file_contents = f.readlines()
+
+    # This is our base, we remove all the code from here in our `full_filename` and `feature_filename` to find the new content
+    if parser_only:
+        base_file_func = clean_lines(get_function_contents_by_name(base_file_contents, "main"))
+        full_file_func = clean_lines(get_function_contents_by_name(full_file_contents, "main"))
+        feature_file_func = clean_lines(get_function_contents_by_name(feature_file_contents, "main"))
+        if secondary_filename is not None:
+            secondary_file_func = clean_lines(get_function_contents_by_name(secondary_file_contents, "main"))
+    else:
+        base_file_func = clean_lines(get_function_contents_by_name(base_file_contents, "training_function"))
+        full_file_func = clean_lines(get_function_contents_by_name(full_file_contents, "training_function"))
+        feature_file_func = clean_lines(get_function_contents_by_name(feature_file_contents, "training_function"))
+        if secondary_filename is not None:
+            secondary_file_func = clean_lines(
+                get_function_contents_by_name(secondary_file_contents, "training_function")
+            )
+
+    _dl_line = "train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)\n"
+
+    # Specific code in our script that differs from the full version, aka what is new
+    new_feature_code = []
+    passed_idxs = []  # We keep track of the idxs just in case it's a repeated statement
+    for i, line in enumerate(feature_file_func):
+        if i not in passed_idxs:
+            if (line not in full_file_func) and (line.lstrip() != _dl_line):
+                new_feature_code.append(line)
+                passed_idxs.append(i)
+
+    # Extract out just the new parts from the full_file_training_func
+    new_full_example_parts = []
+    passed_idxs = []  # We keep track of the idxs just in case it's a repeated statement
+    for i, line in enumerate(base_file_func):
+        if i not in passed_idxs:
+            if (line not in full_file_func) and (line.lstrip() != _dl_line):
+                new_full_example_parts.append(line)
+                passed_idxs.append(i)
+
+    # Finally, get the overall diff
+    diff_from_example = [line for line in new_feature_code if line not in new_full_example_parts]
+    if secondary_filename is not None:
+        diff_from_two = [line for line in full_file_contents if line not in secondary_file_func]
+        diff_from_example = [line for line in diff_from_example if line not in diff_from_two]
+
+    return diff_from_example
diff --git a/tests/test_examples.py b/tests/test_examples.py
new file mode 100644
index 000000000..d60d7c08f
--- /dev/null
+++ b/tests/test_examples.py
@@ -0,0 +1,190 @@
+# 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 sys
+import tempfile
+import unittest
+from unittest import mock
+
+from torch.utils.data import DataLoader
+
+from accelerate import DistributedType
+from accelerate.test_utils.examples import compare_against_test
+from datasets import load_dataset
+from transformers import AutoTokenizer
+
+
+SRC_DIRS = [os.path.abspath(os.path.join("examples", "by_feature"))]
+sys.path.extend(SRC_DIRS)
+
+if SRC_DIRS is not None:
+    import checkpointing
+    import tracking
+
+# DataLoaders built from `test_samples/MRPC` for quick testing
+# Should mock `{script_name}.get_dataloaders` via:
+# @mock.patch("{script_name}.get_dataloaders", mocked_dataloaders)
+
+
+def mocked_dataloaders(accelerator, batch_size: int = 16):
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    data_files = {"train": "tests/test_samples/MRPC/train.csv", "validation": "tests/test_samples/MRPC/dev.csv"}
+    datasets = load_dataset("csv", data_files=data_files)
+    label_list = datasets["train"].unique("label")
+
+    label_to_id = {v: i for i, v in enumerate(label_list)}
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(
+            examples["sentence1"], examples["sentence2"], truncation=True, max_length=None, padding="max_length"
+        )
+        if "label" in examples:
+            outputs["labels"] = [label_to_id[l] for l in examples["label"]]
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["sentence1", "sentence2", "label"],
+    )
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=2)
+    eval_dataloader = DataLoader(tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=1)
+
+    return train_dataloader, eval_dataloader
+
+
+class ExampleDifferenceTests(unittest.TestCase):
+    """
+    This TestCase checks that all of the `complete_*` scripts contain all of the
+    information found in the `by_feature` scripts, line for line. If one fails,
+    then a complete example does not contain all of the features in the features
+    scripts, and should be updated.
+
+    Each example script should be a single test (such as `test_nlp_example`),
+    and should run `one_complete_example` twice: once with `parser_only=True`,
+    and the other with `parser_only=False`. This is so that when the test
+    failures are returned to the user, they understand if the discrepancy lies in
+    the `main` function, or the `training_loop` function. Otherwise it will be
+    unclear.
+
+    Also, if there are any expected differences between the base script used and
+    `complete_nlp_example.py` (the canonical base script), these should be included in
+    `special_strings`. These would be differences in how something is logged, print statements,
+    etc (such as calls to `Accelerate.log()`)
+    """
+
+    def one_complete_example(
+        self, complete_file_name: str, parser_only: bool, secondary_filename: str = None, special_strings: list = None
+    ):
+        """
+        Tests a single `complete` example against all of the implemented `by_feature` scripts
+
+        Args:
+            complete_file_name (`str`):
+                The filename of a complete example
+            parser_only (`bool`):
+                Whether to look at the main training function, or the argument parser
+            secondary_filename (`str`, *optional*):
+                A potential secondary base file to strip all script information not relevant for checking,
+                such as "cv_example.py" when testing "complete_cv_example.py"
+            special_strings (`list`, *optional*):
+                A list of strings to potentially remove before checking no differences are left. These should be
+                diffs that are file specific, such as different logging variations between files.
+        """
+        self.maxDiff = None
+        by_feature_path = os.path.abspath(os.path.join("examples", "by_feature"))
+        examples_path = os.path.abspath("examples")
+        for item in os.listdir(by_feature_path):
+            item_path = os.path.join(by_feature_path, item)
+            if os.path.isfile(item_path) and ".py" in item_path:
+                with self.subTest(
+                    tested_script=complete_file_name,
+                    feature_script=item,
+                    tested_section="main()" if parser_only else "training_function()",
+                ):
+                    diff = compare_against_test(
+                        os.path.join(examples_path, complete_file_name), item_path, parser_only, secondary_filename
+                    )
+                    diff = "\n".join(diff)
+                    if special_strings is not None:
+                        for string in special_strings:
+                            diff = diff.replace(string, "")
+                    self.assertEqual(diff, "")
+
+    def test_nlp_examples(self):
+        self.one_complete_example("complete_nlp_example.py", True)
+        self.one_complete_example("complete_nlp_example.py", False)
+
+    def test_cv_examples(self):
+        cv_path = os.path.abspath(os.path.join("examples", "cv_example.py"))
+        special_strings = [
+            " " * 16 + "{\n\n",
+            " " * 18 + '"accuracy": eval_metric["accuracy"],\n\n',
+            " " * 18 + '"f1": eval_metric["f1"],\n\n',
+            " " * 18 + '"train_loss": total_loss,\n\n',
+            " " * 18 + '"epoch": epoch,\n\n',
+            " " * 16 + "}\n",
+            " " * 8,
+        ]
+        self.one_complete_example("complete_cv_example.py", True, cv_path, special_strings)
+        self.one_complete_example("complete_cv_example.py", False, cv_path, special_strings)
+
+
+class FeatureExamplesTests(unittest.TestCase):
+    @mock.patch("checkpointing.get_dataloaders", mocked_dataloaders)
+    def test_checkpointing_by_epoch(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            testargs = f"""
+            checkpointing.py
+            --checkpointing_steps epoch
+            --output_dir {tmpdir}
+            """.split()
+            with mock.patch.object(sys, "argv", testargs):
+                checkpointing.main()
+                self.assertTrue(os.path.exists(os.path.join(tmpdir, "epoch_0")))
+
+    @mock.patch("checkpointing.get_dataloaders", mocked_dataloaders)
+    def test_checkpointing_by_steps(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            testargs = f"""
+            checkpointing.py
+            --checkpointing_steps 2
+            --output_dir {tmpdir}
+            """.split()
+            with mock.patch.object(sys, "argv", testargs):
+                checkpointing.main()
+                self.assertTrue(os.path.exists(os.path.join(tmpdir, "step_2")))
+
+    @mock.patch("tracking.get_dataloaders", mocked_dataloaders)
+    def test_tracking(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            testargs = f"""
+            tracking.py
+            --with_tracking
+            --logging_dir {tmpdir}
+            """.split()
+            with mock.patch.object(sys, "argv", testargs):
+                tracking.main()
+                self.assertTrue(os.path.exists(os.path.join(tmpdir, "tracking")))
diff --git a/tests/test_samples/MRPC/dev.csv b/tests/test_samples/MRPC/dev.csv
new file mode 100644
index 000000000..96beccda9
--- /dev/null
+++ b/tests/test_samples/MRPC/dev.csv
@@ -0,0 +1,7 @@
+label,sentence1,sentence2
+equivalent,He said the foodservice pie business doesn 't fit the company 's long-term growth strategy .,""" The foodservice pie business does not fit our long-term growth strategy ."
+not_equivalent,Magnarelli said Racicot hated the Iraqi regime and looked forward to using his long years of training in the war .,"His wife said he was "" 100 percent behind George Bush "" and looked forward to using his years of training in the war ."
+not_equivalent,"The dollar was at 116.92 yen against the yen , flat on the session , and at 1.2891 against the Swiss franc , also flat .","The dollar was at 116.78 yen JPY = , virtually flat on the session , and at 1.2871 against the Swiss franc CHF = , down 0.1 percent ."
+equivalent,The AFL-CIO is waiting until October to decide if it will endorse a candidate .,The AFL-CIO announced Wednesday that it will decide in October whether to endorse a candidate before the primaries .
+not_equivalent,No dates have been set for the civil or the criminal trial .,"No dates have been set for the criminal or civil cases , but Shanley has pleaded not guilty ."
+equivalent,Wal-Mart said it would check all of its million-plus domestic workers to ensure they were legally employed .,It has also said it would review all of its domestic employees more than 1 million to ensure they have legal status .
diff --git a/tests/test_samples/MRPC/train.csv b/tests/test_samples/MRPC/train.csv
new file mode 100644
index 000000000..96beccda9
--- /dev/null
+++ b/tests/test_samples/MRPC/train.csv
@@ -0,0 +1,7 @@
+label,sentence1,sentence2
+equivalent,He said the foodservice pie business doesn 't fit the company 's long-term growth strategy .,""" The foodservice pie business does not fit our long-term growth strategy ."
+not_equivalent,Magnarelli said Racicot hated the Iraqi regime and looked forward to using his long years of training in the war .,"His wife said he was "" 100 percent behind George Bush "" and looked forward to using his years of training in the war ."
+not_equivalent,"The dollar was at 116.92 yen against the yen , flat on the session , and at 1.2891 against the Swiss franc , also flat .","The dollar was at 116.78 yen JPY = , virtually flat on the session , and at 1.2871 against the Swiss franc CHF = , down 0.1 percent ."
+equivalent,The AFL-CIO is waiting until October to decide if it will endorse a candidate .,The AFL-CIO announced Wednesday that it will decide in October whether to endorse a candidate before the primaries .
+not_equivalent,No dates have been set for the civil or the criminal trial .,"No dates have been set for the criminal or civil cases , but Shanley has pleaded not guilty ."
+equivalent,Wal-Mart said it would check all of its million-plus domestic workers to ensure they were legally employed .,It has also said it would review all of its domestic employees more than 1 million to ensure they have legal status .

Why not always do this?
@@ -200,6 +200,8 @@ def training_function(config, args):
                 optimizer.step()
                 lr_scheduler.step()
                 optimizer.zero_grad()
+            if isinstance(checkpointing_steps, int):
+                overall_step += 1
It's only needed if we save via checkpointing steps rather than epoch, so didn't want people to assume we *always* need to do that. 

(Which means a comment is needed!)
@@ -200,6 +200,8 @@ def training_function(config, args):
                 optimizer.step()
                 lr_scheduler.step()
                 optimizer.zero_grad()
+            if isinstance(checkpointing_steps, int):
+                overall_step += 1
Yeah I get the variable won't be used if we don't checkpoint with steps, but it doesn't hurt to always have it (and would save one line of code).
@@ -200,6 +200,8 @@ def training_function(config, args):
                 optimizer.step()
                 lr_scheduler.step()
                 optimizer.zero_grad()
+            if isinstance(checkpointing_steps, int):
+                overall_step += 1
```suggestion
# Copyright 2022 The HuggingFace Team. All rights reserved.
```
@@ -0,0 +1,132 @@
+#!/usr/bin/env python
+
+# Copyright 2021 The HuggingFace Team. All rights reserved.
A quick intro here would be useful.
@@ -0,0 +1,132 @@
+#!/usr/bin/env python
+
+# 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.
+
```suggestion
def get_train_func(lines: list):
```
If we use type annotations, better use `List[str]` here.
@@ -0,0 +1,132 @@
+#!/usr/bin/env python
+
+# 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 logging
+import os
+
+
+logger = logging.getLogger(__name__)
+
+
+def get_train_func(lines: list):
Sounds like it could be refactored with the previous one and an extra arg?
@@ -0,0 +1,132 @@
+#!/usr/bin/env python
+
+# 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 logging
+import os
+
+
+logger = logging.getLogger(__name__)
+
+
+def get_train_func(lines: list):
+    """
+    Finds the main training function from inside segmented source code.
+
+    Args:
+        lines (`list`):
+            Source code of a script
+    """
+    good_lines, found_start = [], False
+    for line in lines:
+        if not found_start and "def training_function" in line:
+            found_start = True
+            good_lines.append(line)
+            continue
+        if found_start:
+            if "def main" in line:
+                return good_lines
+            good_lines.append(line)
+
+
+def get_main_func(lines: list):
Those do not match the signature. Also this could be expanded a little bit, I'm not sure I fully understand.
@@ -0,0 +1,132 @@
+#!/usr/bin/env python
+
+# 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 logging
+import os
+
+
+logger = logging.getLogger(__name__)
+
+
+def get_train_func(lines: list):
+    """
+    Finds the main training function from inside segmented source code.
+
+    Args:
+        lines (`list`):
+            Source code of a script
+    """
+    good_lines, found_start = [], False
+    for line in lines:
+        if not found_start and "def training_function" in line:
+            found_start = True
+            good_lines.append(line)
+            continue
+        if found_start:
+            if "def main" in line:
+                return good_lines
+            good_lines.append(line)
+
+
+def get_main_func(lines: list):
+    """
+    Finds the main function from inside segmented source code
+
+    Args:
+        lines (`list`):
+            Source code of a script
+    """
+    good_lines, found_start = [], False
+    for line in lines:
+        if not found_start and "def main" in line:
+            found_start = True
+            good_lines.append(line)
+            continue
+        if found_start:
+            if "if __name__" in line:
+                return good_lines
+            good_lines.append(line)
+
+
+def clean_lines(lines: list):
+    """
+    Filters `lines` and removes any entries that start with a comment ('#') or is just a newline ('\n')
+
+    Args:
+        lines (`list`):
+            Source code of a script
+    """
+    return [line for line in lines if not line.lstrip().startswith("#") and line != "\n"]
+
+
+def compare_against_test(base_filename: str, feature_filename: str, parser_only: bool, secondary_filename: str = None):
+    """
+    Checks whether the content aligned in `test_filename` is included inside of `full_filename`.
```suggestion
# Copyright 2022 The HuggingFace Team. All rights reserved.
```
@@ -0,0 +1,176 @@
+# Copyright 2021 The HuggingFace Team. All rights reserved.
It seems we always use the method with both flags, should we just remove that arg and put the two tests inside?
@@ -0,0 +1,176 @@
+# 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 sys
+import tempfile
+import unittest
+from unittest import mock
+
+from torch.utils.data import DataLoader
+
+from accelerate import DistributedType
+from accelerate.test_utils.examples import compare_against_test
+from datasets import load_dataset
+from transformers import AutoTokenizer
+
+
+SRC_DIRS = [os.path.abspath(os.path.join("examples", "by_feature"))]
+sys.path.extend(SRC_DIRS)
+
+if SRC_DIRS is not None:
+    import checkpointing
+    import tracking
+
+# DataLoaders built from `test_samples/MRPC` for quick testing
+# Should mock `{script_name}.get_dataloaders` via:
+# @mock.patch("{script_name}.get_dataloaders", mocked_dataloaders)
+
+
+def mocked_dataloaders(accelerator, batch_size: int = 16):
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    data_files = {"train": "tests/test_samples/MRPC/train.csv", "validation": "tests/test_samples/MRPC/dev.csv"}
+    datasets = load_dataset("csv", data_files=data_files)
+    label_list = datasets["train"].unique("label")
+
+    label_to_id = {v: i for i, v in enumerate(label_list)}
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(
+            examples["sentence1"], examples["sentence2"], truncation=True, max_length=None, padding="max_length"
+        )
+        if "label" in examples:
+            outputs["labels"] = [label_to_id[l] for l in examples["label"]]
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["sentence1", "sentence2", "label"],
+    )
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=2)
+    eval_dataloader = DataLoader(tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=1)
+
+    return train_dataloader, eval_dataloader
+
+
+class ExampleDifferenceTests(unittest.TestCase):
+    """
+    This TestCase checks that all of the `complete_*` scripts contain all of the
+    information found in the `by_feature` scripts, line for line.
+    """
+
+    def one_complete_example(
+        self, complete_file_name: str, parser_only: bool, secondary_filename: str = None, special_strings: list = None
+    ):
+        """
+        Tests a single `complete` example against all of the implemented `by_feature` scripts
+
+        Args:
+            complete_file_name (`str`):
+                The filename of a complete example
+            parser_only (`bool`):
+                Whether to look at the main training function, or the argument parser
+            secondary_filename (`str`, *optional*):
+                A potential secondary base file to strip all script information not relevant for checking,
+                such as "cv_example.py" when testing "complete_cv_example.py"
+            special_strings (`list`, *optional*):
+                A list of strings to potentially remove before checking no differences are left. These should be
+                diffs that are file specific, such as different logging variations between files.
+        """
+        self.maxDiff = None
+        by_feature_path = os.path.abspath(os.path.join("examples", "by_feature"))
+        examples_path = os.path.abspath("examples")
+        for item in os.listdir(by_feature_path):
+            item_path = os.path.join(by_feature_path, item)
+            if os.path.isfile(item_path) and ".py" in item_path:
+                with self.subTest(
+                    tested_script=complete_file_name,
+                    feature_script=item,
+                    tested_section="main()" if parser_only else "training_function()",
+                ):
+                    diff = compare_against_test(
+                        os.path.join(examples_path, complete_file_name), item_path, parser_only, secondary_filename
+                    )
+                    diff = "\n".join(diff)
+                    if special_strings is not None:
+                        for string in special_strings:
+                            diff = diff.replace(string, "")
+                    self.assertEqual(diff, "")
+
+    def test_nlp_examples(self):
+        self.one_complete_example("complete_nlp_example.py", True)
+        self.one_complete_example("complete_nlp_example.py", False)
+
I'm making a note for this inside of the documentation, but the reasoning for separation is it lets the test failure be more readable as to what section failed, rather than one complete error.
Notice the `tested_section` part
E.g.:

```bash
_________ ExampleDifferenceTests.test_cv_example (feature_script='tracking.py', tested_script='complete_cv_example.py', tested_section='training_function()') _________
```
vs:
```bash
_________ ExampleDifferenceTests.test_cv_example (feature_script='tracking.py', tested_script='complete_cv_example.py', tested_section='main()') _________
```

(This is a `pytest-subtest` hack)
@@ -0,0 +1,176 @@
+# 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 sys
+import tempfile
+import unittest
+from unittest import mock
+
+from torch.utils.data import DataLoader
+
+from accelerate import DistributedType
+from accelerate.test_utils.examples import compare_against_test
+from datasets import load_dataset
+from transformers import AutoTokenizer
+
+
+SRC_DIRS = [os.path.abspath(os.path.join("examples", "by_feature"))]
+sys.path.extend(SRC_DIRS)
+
+if SRC_DIRS is not None:
+    import checkpointing
+    import tracking
+
+# DataLoaders built from `test_samples/MRPC` for quick testing
+# Should mock `{script_name}.get_dataloaders` via:
+# @mock.patch("{script_name}.get_dataloaders", mocked_dataloaders)
+
+
+def mocked_dataloaders(accelerator, batch_size: int = 16):
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    data_files = {"train": "tests/test_samples/MRPC/train.csv", "validation": "tests/test_samples/MRPC/dev.csv"}
+    datasets = load_dataset("csv", data_files=data_files)
+    label_list = datasets["train"].unique("label")
+
+    label_to_id = {v: i for i, v in enumerate(label_list)}
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(
+            examples["sentence1"], examples["sentence2"], truncation=True, max_length=None, padding="max_length"
+        )
+        if "label" in examples:
+            outputs["labels"] = [label_to_id[l] for l in examples["label"]]
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["sentence1", "sentence2", "label"],
+    )
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=2)
+    eval_dataloader = DataLoader(tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=1)
+
+    return train_dataloader, eval_dataloader
+
+
+class ExampleDifferenceTests(unittest.TestCase):
+    """
+    This TestCase checks that all of the `complete_*` scripts contain all of the
+    information found in the `by_feature` scripts, line for line.
+    """
+
+    def one_complete_example(
+        self, complete_file_name: str, parser_only: bool, secondary_filename: str = None, special_strings: list = None
+    ):
+        """
+        Tests a single `complete` example against all of the implemented `by_feature` scripts
+
+        Args:
+            complete_file_name (`str`):
+                The filename of a complete example
+            parser_only (`bool`):
+                Whether to look at the main training function, or the argument parser
+            secondary_filename (`str`, *optional*):
+                A potential secondary base file to strip all script information not relevant for checking,
+                such as "cv_example.py" when testing "complete_cv_example.py"
+            special_strings (`list`, *optional*):
+                A list of strings to potentially remove before checking no differences are left. These should be
+                diffs that are file specific, such as different logging variations between files.
+        """
+        self.maxDiff = None
+        by_feature_path = os.path.abspath(os.path.join("examples", "by_feature"))
+        examples_path = os.path.abspath("examples")
+        for item in os.listdir(by_feature_path):
+            item_path = os.path.join(by_feature_path, item)
+            if os.path.isfile(item_path) and ".py" in item_path:
+                with self.subTest(
+                    tested_script=complete_file_name,
+                    feature_script=item,
+                    tested_section="main()" if parser_only else "training_function()",
+                ):
+                    diff = compare_against_test(
+                        os.path.join(examples_path, complete_file_name), item_path, parser_only, secondary_filename
+                    )
+                    diff = "\n".join(diff)
+                    if special_strings is not None:
+                        for string in special_strings:
+                            diff = diff.replace(string, "")
+                    self.assertEqual(diff, "")
+
+    def test_nlp_examples(self):
+        self.one_complete_example("complete_nlp_example.py", True)
+        self.one_complete_example("complete_nlp_example.py", False)
+
Should go inside triple quotes.
@@ -0,0 +1,143 @@
+#!/usr/bin/env python
+
+# 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.
+
+# A collection of utilities for comparing `examples/complete_*_example.py` scripts with the capabilities inside of each
+# `examples/by_feature` example. `compare_against_test` is the main function that should be used when testing,
+# while the others are used to either get the code that matters, or to preprocess them (such as stripping comments)
_The documentation is not available anymore as the PR was closed or merged._
I've come up with a solution to make sure that also the main `complete_` examples will raise an error if they differ from the bits inside of `by_feature`. Now eventually it would be good to have an exclusion list of files for when we *wouldn't* expect to see some feature added (such as say cross validation). But we'll cross that bridge when we get there. In the meantime here is the report generated back to the user from pytest:

(Note: `pytest-subtests` was added as a dep to make sure we could use `TestCase.subTest` and make this code clean and efficient)
```bash
======================================================================= short test summary info =======================================================================
SUBFAIL tests/test_examples.py::ExampleDifferenceTests::test_complete_cv_example_body - AssertionError: 14 != 0
SUBFAIL tests/test_examples.py::ExampleDifferenceTests::test_complete_cv_example_body - AssertionError: 27 != 0
SUBFAIL tests/test_examples.py::ExampleDifferenceTests::test_complete_cv_example_parser - AssertionError: 1 != 0
SUBFAIL tests/test_examples.py::ExampleDifferenceTests::test_complete_cv_example_parser - AssertionError: 4 != 0
SUBFAIL tests/test_examples.py::ExampleDifferenceTests::test_complete_nlp_example_parser - AssertionError: 1 != 0
SUBFAIL tests/test_examples.py::ExampleDifferenceTests::test_complete_nlp_example_parser - AssertionError: 1 != 0
=============================================================== 6 failed, 4 passed, 1 warning in 1.64s ================================================================
```

{{ Removed full trace, see the following message for an example }} 
Slightly tweaked how they look, I've now included it posting the source code diffs. I believe this is a must as it will tell the user exactly what parts were missing from the full example:

```bash
============================================================================== FAILURES ===============================================================================
_________ ExampleDifferenceTests.test_cv_example (feature_script='tracking.py', tested_script='complete_cv_example.py', tested_section='training_function()') _________

self = <test_examples.ExampleDifferenceTests testMethod=test_cv_example>, complete_file_name = 'complete_cv_example.py', parser_only = False

    def one_complete_example(self, complete_file_name: str, parser_only: bool):
        """
        Tests a single `complete` example against all of the implemented `by_feature` scripts
    
        Args:
            complete_file_name (`str`):
                The filename of a complete example
            parser_only (`bool`):
                Whether to look at the main training function, or the argument parser
        """
        self.maxDiff = None
        by_feature_path = os.path.abspath(os.path.join("examples", "by_feature"))
        examples_path = os.path.abspath("examples")
        for item in os.listdir(by_feature_path):
            item_path = os.path.join(by_feature_path, item)
            if os.path.isfile(item_path) and ".py" in item_path:
                with self.subTest(tested_script=complete_file_name, feature_script=item, tested_section="main()" if parser_only else "training_function()"):
                    diff = compare_against_test(
                        os.path.join(examples_path, "nlp_example.py"),
                        os.path.join(examples_path, complete_file_name),
                        item_path,
                        parser_only
                    )
>                   self.assertEqual('\n'.join(diff), '')
E                   AssertionError: '        accelerator = Accelerator(cpu=arg[693 chars]()\n' != ''
E                   -         accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="all")
E                   - 
E                   -         accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
E                   - 
E                   -         accelerator.init_trackers("nlp_example", config)
E                   - 
E                   -             predictions, references = accelerator.gather((predictions, batch["labels"]))
E                   - 
E                   -                 predictions=predictions,
E                   - 
E                   -                 references=references,
E                   - 
E                   -             accelerator.log(
E                   - 
E                   -                 {
E                   - 
E                   -                     "accuracy": eval_metric["accuracy"],
E                   - 
E                   -                     "f1": eval_metric["f1"],
E                   - 
E                   -                     "train_loss": total_loss,
E                   - 
E                   -                     "epoch": epoch,
E                   - 
E                   -                 }
E                   - 
E                   -         accelerator.end_training()

tests/test_examples.py:107: AssertionError
```
https://api.github.com/repos/huggingface/accelerate/pulls/313
diff --git a/examples/README.md b/examples/README.md
index 115384a42..bb248e9d1 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -186,10 +186,20 @@ To run it in each of these various modes, use the following commands:
 
 ## Finer Examples
 
-While the first two scripts are extremely barebones when it comes to what you can do with accelerate, there are two complementary scripts for each: `complete_cv_example.py` and `complete_nlp_example.py`.
+While the first two scripts are extremely barebones when it comes to what you can do with accelerate, more advanced features are documented in two other locations.
 
-These two serve to document what else the `Accelerate` library is able to do, outside of just the basic training loop. They each have the following additional arguments:
+### `by_feature` examples
+
+These scripts are *individual* examples highlighting one particular feature or use-case within Accelerate. They all stem from the [nlp_example.py](./nlp_example.py) script, and any changes or modifications is denoted with a `# New Code #` comment.
+
+Read the README.md file located in the `by_feature` folder for more information.
+
+### `complete_*` examples
+
+These two scripts contain *every* single feature currently available in Accelerate in one place, as one giant script.
+
+New arguments that can be passed include:
 
 - `checkpointing_steps`, whether the various states should be saved at the end of every `n` steps, or `"epoch"` for each epoch. States are then saved to folders named `step_{n}` or `epoch_{n}`
 - `resume_from_checkpoint`, should be used if you want to resume training off of a previous call to the script and passed a `checkpointing_steps` to it.
-- `with_tracking`, should be used if you want to log the training run using all available experiment trackers in your environment. Currently supported trackers include TensorBoard, Weights and Biases, and CometML.
\ No newline at end of file
+- `with_tracking`, should be used if you want to log the training run using all available experiment trackers in your environment. Currently supported trackers include TensorBoard, Weights and Biases, and CometML.
diff --git a/examples/by_feature/README.md b/examples/by_feature/README.md
new file mode 100644
index 000000000..40ffa3c1e
--- /dev/null
+++ b/examples/by_feature/README.md
@@ -0,0 +1,56 @@
+# What are these scripts?
+
+All scripts in this folder originate from the `nlp_example.py` file, as it is a very simplistic NLP training example using Accelerate with zero extra features.
+
+From there, each further script adds in just **one** feature of Accelerate, showing how you can quickly modify your own scripts to implement these capabilities.
+
+A full example with all of these parts integrated together can be found in the `complete_nlp_example.py` script and `complete_cv_example.py` script.
+
+Adjustments to each script from the base `nlp_example.py` file can be found quickly by searching for "# New Code #"
+
+## Example Scripts by Feature and their Arguments
+
+### Base Example (`../nlp_example.py`)
+
+- Shows how to use `Accelerator` in an extremely simplistic PyTorch training loop
+- Arguments available:
+  - `mixed_precision`, whether to use mixed precision. ("no", "fp16", or "bf16")
+  - `cpu`, whether to train using only the CPU. (yes/no/1/0)
+
+All following scripts also accept these arguments in addition to their added ones.
+
+These arguments should be added at the end of any method for starting the python script (such as `python`, `accelerate launch`, `python -m torch.distributed.launch`), such as:
+
+```bash
+accelerate launch ../nlp_example.py --mixed_precision fp16 --cpu 0
+```
+
+### Checkpointing and Resuming Training (`checkpointing.py`)
+
+- Shows how to use `Accelerator.save_state` and `Accelerator.load_state` to save or continue training
+- **It is assumed you are continuing off the same training script**
+- Arguments available:
+  - `checkpointing_steps`, after how many steps the various states should be saved. ("epoch", 1, 2, ...)
+  - `output_dir`, where saved state folders should be saved to, default is current working directory
+  - `resume_from_checkpoint`, what checkpoint folder to resume from. ("epoch_0", "step_22", ...)
+
+These arguments should be added at the end of any method for starting the python script (such as `python`, `accelerate launch`, `python -m torch.distributed.launch`), such as:
+
+(Note, `resume_from_checkpoint` assumes that we've ran the script for one epoch with the `--checkpointing_steps epoch` flag)
+
+```bash
+accelerate launch ./checkpointing.py --checkpointing_steps epoch output_dir "checkpointing_tutorial" --resume_from_checkpoint "checkpointing_tutorial/epoch_0"
+```
+
+### Experiment Tracking (`tracking.py`)
+
+- Shows how to use `Accelerate.init_trackers` and `Accelerator.log`
+- Can be used with Weights and Biases, TensorBoard, or CometML.
+- Arguments available:
+  - `with_tracking`, whether to load in all available experiment trackers from the environment.
+
+These arguments should be added at the end of any method for starting the python script (such as `python`, `accelerate launch`, `python -m torch.distributed.launch`), such as:
+
+```bash
+accelerate launch ./tracking.py --with_tracking
+```
diff --git a/examples/by_feature/checkpointing.py b/examples/by_feature/checkpointing.py
new file mode 100644
index 000000000..35205107d
--- /dev/null
+++ b/examples/by_feature/checkpointing.py
@@ -0,0 +1,282 @@
+# coding=utf-8
+# Copyright 2021 The HuggingFace Inc. 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 torch
+from torch.utils.data import DataLoader
+
+from accelerate import Accelerator, DistributedType
+from datasets import load_dataset, load_metric
+from transformers import (
+    AdamW,
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    get_linear_schedule_with_warmup,
+    set_seed,
+)
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate,
+# specifically showcasing the checkpointing capability,
+# and builds off the `nlp_example.py` script.
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#
+# To help focus on the differences in the code, building `DataLoaders`
+# was refactored into its own function.
+# New additions from the base script can be found quickly by
+# looking for the # New Code # tags
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+def get_dataloaders(accelerator: Accelerator, batch_size: int = 16):
+    """
+    Creates a set of `DataLoader`s for the `glue` dataset,
+    using "bert-base-cased" as the tokenizer.
+
+    Args:
+        accelerator (`Accelerator`):
+            An `Accelerator` object
+        batch_size (`int`, *optional*):
+            The batch size for the train and validation DataLoaders.
+    """
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    datasets = load_dataset("glue", "mrpc")
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["idx", "sentence1", "sentence2"],
+    )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    return train_dataloader, eval_dataloader
+
+
+def training_function(config, args):
+    # Initialize accelerator
+    accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    correct_bias = config["correct_bias"]
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+
+    # New Code #
+    # Parse out whether we are saving every epoch or after a certain number of batches
+    if args.checkpointing_steps == "epoch":
+        checkpointing_steps = args.checkpointing_steps
+    elif args.checkpointing_steps.isdigit():
+        checkpointing_steps = int(args.checkpointing_steps)
+    else:
+        raise ValueError(
+            f"Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed."
+        )
+
+    set_seed(seed)
+
+    train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)
+    metric = load_metric("glue", "mrpc")
+
+    # If the batch size is too big we use gradient accumulation
+    gradient_accumulation_steps = 1
+    if batch_size > MAX_GPU_BATCH_SIZE:
+        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
+        batch_size = MAX_GPU_BATCH_SIZE
+
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
+
+    # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
+    # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
+    # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
+    model = model.to(accelerator.device)
+
+    # Instantiate optimizer
+    optimizer = AdamW(params=model.parameters(), lr=lr, correct_bias=correct_bias)
+
+    # Instantiate scheduler
+    lr_scheduler = get_linear_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=100,
+        num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,
+    )
+
+    # Prepare everything
+    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
+    # prepare method.
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    # New Code #
+    # We need to keep track of how many total steps we have iterated over
+    if isinstance(checkpointing_steps, int):
+        overall_step = 0
+
+    # We need to load the checkpoint back in before training here with `load_state`
+    # The total number of epochs is adjusted based on where the state is being loaded from,
+    # as we assume continuation of the same training script
+    if args.resume_from_checkpoint:
+        accelerator.print(f"Resuming from checkpoint: {args.resume_from_checkpoint}")
+        accelerator.load_state(args.resume_from_checkpoint)
+
+        if "epoch" in args.resume_from_checkpoint:
+            num_epochs -= int(args.resume_from_checkpoint.replace("epoch_", ""))
+            resume_step = None
+        else:
+            resume_step = int(args.resume_from_checkpoint.replace("step_", ""))
+            num_epochs -= resume_step // len(train_dataloader)
+            # If resuming by step, we also need to know exactly how far into the DataLoader we went
+            resume_step = (num_epochs * len(train_dataloader)) - resume_step
+
+    # Now we train the model
+    for epoch in range(num_epochs):
+        model.train()
+        for step, batch in enumerate(train_dataloader):
+            # New Code #
+            # We need to skip steps until we reach the resumed step during the first epoch
+            if args.resume_from_checkpoint and epoch == 0:
+                if resume_step is not None and step < resume_step:
+                    pass
+            # We could avoid this line since we set the accelerator with `device_placement=True`.
+            batch.to(accelerator.device)
+            outputs = model(**batch)
+            loss = outputs.loss
+            loss = loss / gradient_accumulation_steps
+            accelerator.backward(loss)
+            if step % gradient_accumulation_steps == 0:
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+
+            # New Code #
+            # We save the model, optimizer, lr_scheduler, and seed states by calling `save_state`
+            # These are saved to folders named `step_{overall_step}`
+            # Will contain files: "pytorch_model.bin", "optimizer.bin", "scheduler.bin", and "random_states.pkl"
+            # If mixed precision was used, will also save a "scalar.bin" file
+            if isinstance(checkpointing_steps, int):
+                output_dir = f"step_{overall_step}"
+                if overall_step % checkpointing_steps == 0:
+                    if args.output_dir is not None:
+                        output_dir = os.path.join(args.output_dir, output_dir)
+                    accelerator.save_state(output_dir)
+
+        model.eval()
+        for step, batch in enumerate(eval_dataloader):
+            # We could avoid this line since we set the accelerator with `device_placement=True` (the default).
+            batch.to(accelerator.device)
+            with torch.no_grad():
+                outputs = model(**batch)
+            predictions = outputs.logits.argmax(dim=-1)
+            # It is slightly faster to call this once, than multiple times
+            predictions, references = accelerator.gather((predictions, batch["labels"]))
+            metric.add_batch(
+                predictions=predictions,
+                references=references,
+            )
+
+        eval_metric = metric.compute()
+        # Use accelerator.print to print only on the main process.
+        accelerator.print(f"epoch {epoch}:", eval_metric)
+
+        # New Code #
+        # We save the model, optimizer, lr_scheduler, and seed states by calling `save_state`
+        # These are saved to folders named `step_{overall_step}`
+        # Will contain files: "pytorch_model.bin", "optimizer.bin", "scheduler.bin", and "random_states.pkl"
+        # If mixed precision was used, will also save a "scalar.bin" file
+        if checkpointing_steps == "epoch":
+            output_dir = f"epoch_{num_epochs}"
+            if args.output_dir is not None:
+                output_dir = os.path.join(args.output_dir, output_dir)
+            accelerator.save_state(output_dir)
+
+
+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.")
+    parser.add_argument(
+        "--checkpointing_steps",
+        type=str,
+        default="epoch",
+        help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=".",
+        help="Optional save directory where all checkpoint folders will be stored. Default is the current working directory.",
+    )
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        default=None,
+        help="If the training should continue from a checkpoint folder.",
+    )
+    args = parser.parse_args()
+    config = {"lr": 2e-5, "num_epochs": 3, "correct_bias": True, "seed": 42, "batch_size": 16}
+    training_function(config, args)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/by_feature/tracking.py b/examples/by_feature/tracking.py
new file mode 100644
index 000000000..ca11fa04b
--- /dev/null
+++ b/examples/by_feature/tracking.py
@@ -0,0 +1,238 @@
+# coding=utf-8
+# Copyright 2021 The HuggingFace Inc. 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 torch
+from torch.utils.data import DataLoader
+
+from accelerate import Accelerator, DistributedType
+from datasets import load_dataset, load_metric
+from transformers import (
+    AdamW,
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    get_linear_schedule_with_warmup,
+    set_seed,
+)
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate,
+# specifically showcasing the experiment tracking capability,
+# and builds off the `nlp_example.py` script.
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#
+# To help focus on the differences in the code, building `DataLoaders`
+# was refactored into its own function.
+# New additions from the base script can be found quickly by
+# looking for the # New Code # tags
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+def get_dataloaders(accelerator: Accelerator, batch_size: int = 16):
+    """
+    Creates a set of `DataLoader`s for the `glue` dataset,
+    using "bert-base-cased" as the tokenizer.
+
+    Args:
+        accelerator (`Accelerator`):
+            An `Accelerator` object
+        batch_size (`int`, *optional*):
+            The batch size for the train and validation DataLoaders.
+    """
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    datasets = load_dataset("glue", "mrpc")
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["idx", "sentence1", "sentence2"],
+    )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    return train_dataloader, eval_dataloader
+
+
+def training_function(config, args):
+    # Initialize Accelerator
+
+    # New Code #
+    # We pass in "all" to `log_with` to grab all available trackers in the environment
+    # Note: If using a custom `Tracker` class, should be passed in here such as:
+    # >>> log_with = ["all", MyCustomTrackerClassInstance()]
+    if args.with_tracking:
+        accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="all")
+    else:
+        accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    correct_bias = config["correct_bias"]
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+    set_seed(seed)
+
+    train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)
+    metric = load_metric("glue", "mrpc")
+
+    # If the batch size is too big we use gradient accumulation
+    gradient_accumulation_steps = 1
+    if batch_size > MAX_GPU_BATCH_SIZE:
+        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
+        batch_size = MAX_GPU_BATCH_SIZE
+
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
+
+    # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
+    # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
+    # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
+    model = model.to(accelerator.device)
+
+    # Instantiate optimizer
+    optimizer = AdamW(params=model.parameters(), lr=lr, correct_bias=correct_bias)
+
+    # Instantiate scheduler
+    lr_scheduler = get_linear_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=100,
+        num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,
+    )
+
+    # Prepare everything
+    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
+    # prepare method.
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    # New Code #
+    # We need to initalize the trackers we use. Overall configurations can also be stored
+    if args.with_tracking:
+        accelerator.init_trackers("accelerate_glue_with_tracking", config)
+
+    # Now we train the model
+    for epoch in range(num_epochs):
+        model.train()
+        # New Code #
+        # For our tracking example, we will log the total loss of each epoch
+        if args.with_tracking:
+            total_loss = 0
+        for step, batch in enumerate(train_dataloader):
+            # We could avoid this line since we set the accelerator with `device_placement=True`.
+            batch.to(accelerator.device)
+            outputs = model(**batch)
+            loss = outputs.loss
+            # New Code #
+            total_loss += loss.detach().float()
+            loss = loss / gradient_accumulation_steps
+            accelerator.backward(loss)
+            if step % gradient_accumulation_steps == 0:
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+
+        model.eval()
+        for step, batch in enumerate(eval_dataloader):
+            # We could avoid this line since we set the accelerator with `device_placement=True` (the default).
+            batch.to(accelerator.device)
+            with torch.no_grad():
+                outputs = model(**batch)
+            predictions = outputs.logits.argmax(dim=-1)
+            # It is slightly faster to call this once, than multiple times
+            predictions, references = accelerator.gather((predictions, batch["labels"]))
+            metric.add_batch(
+                predictions=predictions,
+                references=references,
+            )
+
+        eval_metric = metric.compute()
+        # Use accelerator.print to print only on the main process.
+        accelerator.print(f"epoch {epoch}:", eval_metric)
+
+        # New Code #
+        # To actually log, we call `Accelerator.log`
+        # The values passed can be of `str`, `int`, or `float`
+        accelerator.log(
+            {"accuracy": eval_metric["accuracy"], "f1": eval_metric["f1"], "train_loss": total_loss, "epoch": epoch}
+        )
+
+    # New Code #
+    # When a run is finished, you should call `accelerator.end_training()`
+    # to close all of the open trackers
+    accelerator.end_training()
+
+
+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.")
+    parser.add_argument(
+        "--with_tracking",
+        action="store_true",
+        help="Whether to load in all available experiment trackers from the environment and use them for logging.",
+    )
+    args = parser.parse_args()
+    config = {"lr": 2e-5, "num_epochs": 3, "correct_bias": True, "seed": 42, "batch_size": 16}
+    training_function(config, args)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/cv_example.py b/examples/cv_example.py
index a40fe6bef..b14b067be 100644
--- a/examples/cv_example.py
+++ b/examples/cv_example.py
@@ -166,7 +166,7 @@ def training_function(config, args):
         model.eval()
         accurate = 0
         num_elems = 0
-        for step, batch in enumerate(eval_dataloader):
+        for _, batch in enumerate(eval_dataloader):
             # We could avoid this line since we set the accelerator with `device_placement=True`.
             batch = {k: v.to(accelerator.device) for k, v in batch.items()}
             inputs = (batch["image"] - mean) / std
diff --git a/examples/nlp_example.py b/examples/nlp_example.py
index 8f0d25541..915d35807 100644
--- a/examples/nlp_example.py
+++ b/examples/nlp_example.py
@@ -51,8 +51,7 @@
 
 def training_function(config, args):
     # Initialize accelerator
-    accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision)
-
+    accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
     lr = config["lr"]
     num_epochs = int(config["num_epochs"])
@@ -64,6 +63,8 @@ def training_function(config, args):
     datasets = load_dataset("glue", "mrpc")
     metric = load_metric("glue", "mrpc")
 
+    set_seed(seed)
+
     def tokenize_function(examples):
         # max_length=None => use the model max length (it's actually the default)
         outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
@@ -100,8 +101,6 @@ def collate_fn(examples):
         tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
     )
 
-    set_seed(seed)
-
     # Instantiate the model (we build the model here so that the seed also control new weights initialization)
     model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
 
@@ -161,7 +160,6 @@ def collate_fn(examples):
 
 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,

Just as a note, I noticed we set_seed *before* building the dataloaders, hence why this change. 
@@ -64,6 +63,8 @@ def training_function(config, args):
     datasets = load_dataset("glue", "mrpc")
     metric = load_metric("glue", "mrpc")
 
+    set_seed(seed)
+
_The documentation is not available anymore as the PR was closed or merged._
A note on `transformers` examples:

> This should be decided on a case-by-case basis, as to whether the new feature is important enough to put in those scripts.
https://api.github.com/repos/huggingface/accelerate/pulls/312
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index 15eaaecb6..d72680bea 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -54,9 +54,10 @@ def get_cluster_input():
 
     if distributed_type == DistributedType.NO:
         use_cpu = _ask_field(
-            "Do you want to run your training on CPU only (even if a GPU is available)? [no]:",
-            lambda x: bool(x),
+            "Do you want to run your training on CPU only (even if a GPU is available)? [yes/NO]:",
+            _convert_yes_no_to_bool,
             default=False,
+            error_message="Please enter yes or no.",
         )
     elif distributed_type == DistributedType.MULTI_CPU:
         use_cpu = True
diff --git a/tests/test_tracking.py b/tests/test_tracking.py
index b958dbbfc..b3974afe2 100644
--- a/tests/test_tracking.py
+++ b/tests/test_tracking.py
@@ -159,10 +159,10 @@ def test_log(self):
                 break
         # Check HPS through careful parsing and cleaning
         cleaned_log = re.sub(r"[\x00-\x1f]+", " ", content.decode("utf8", "ignore"))
-        self.assertEqual(self.get_value_from_log("total_loss", cleaned_log), "0.1")
-        self.assertEqual(self.get_value_from_log("iteration", cleaned_log), "1")
-        self.assertEqual(self.get_value_from_log("my_text", cleaned_log), "some_value")
-        self.assertEqual(self.get_value_from_log("_step", cleaned_log), "0")
+        self.assertTrue("0.1" in self.get_value_from_log("total_loss", cleaned_log))
+        self.assertTrue("1" in self.get_value_from_log("iteration", cleaned_log))
+        self.assertTrue("some_value" in self.get_value_from_log("my_text", cleaned_log))
+        self.assertTrue("0" in self.get_value_from_log("_step", cleaned_log))
 
 
 # Comet has a special `OfflineExperiment` we need to use for testing

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/311
diff --git a/README.md b/README.md
index 3e21a3ae1..a3fc74cef 100644
--- a/README.md
+++ b/README.md
@@ -168,7 +168,7 @@ mpirun -np 2 python examples/nlp_example.py
 
 ## Launching training using DeepSpeed
 
-🤗 Accelerate supports training on single/multiple GPUs using DeepSpeed. to use it, you don't need to change anything in your training code; you can set everything using just `accelerate config`. However, if you desire to tweak your DeepSpeed related args from your python script, we provide you the `DeepSpeedPlugin`.
+🤗 Accelerate supports training on single/multiple GPUs using DeepSpeed. To use it, you don't need to change anything in your training code; you can set everything using just `accelerate config`. However, if you desire to tweak your DeepSpeed related args from your python script, we provide you the `DeepSpeedPlugin`.
 
 ```python
 from accelerator import Accelerator, DeepSpeedPlugin
diff --git a/docs/source/quicktour.mdx b/docs/source/quicktour.mdx
index 2b405e8ec..1c62f05f5 100644
--- a/docs/source/quicktour.mdx
+++ b/docs/source/quicktour.mdx
@@ -398,10 +398,6 @@ DeepSpeed support is experimental, so the underlying API will evolve in the near
 breaking changes. In particular, 🤗 Accelerate does not support DeepSpeed config you have written yourself yet, this
 will be added in a next version.
 
-One main caveat for the DeepSpeed integration is that the DeepSpeed launcher always passes a `local_rank` variable to
-the training script, so your training script should accept it (whether you launch training with the DeepSpeed launcher
-or `accelerate launch`).
-
 <Tip warning={true}>
 
 The [`notebook_launcher`] does not support the DeepSpeed integration yet.
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 0cbd66884..f2cab1ccb 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -233,7 +233,7 @@ def multi_gpu_launcher(args):
 
 
 def deepspeed_launcher(args):
-    cmd = ["deepspeed"]
+    cmd = ["deepspeed", "--no_local_rank"]
     if args.num_machines > 1:
         cmd.extend(
             [

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/309
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 6e0742e02..e18686e12 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -284,9 +284,10 @@ def use_fp16(self):
     @property
     def mixed_precision(self):
         if self.distributed_type == DistributedType.DEEPSPEED:
-            if self.state.deepspeed_plugin.deepspeed_config["fp16"]["enabled"]:
+            config = self.state.deepspeed_plugin.deepspeed_config
+            if config.get("fp16", {}).get("enabled", False):
                 mixed_precision = "fp16"
-            elif self.state.deepspeed_plugin.deepspeed_config["bf16"]["enabled"]:
+            elif config.get("bf16", {}).get("enabled", False):
                 mixed_precision = "bf16"
             else:
                 mixed_precision = "no"

```suggestion
            if config.get("fp16", {}).get("enabled", False):
```
@@ -284,9 +284,10 @@ def use_fp16(self):
     @property
     def mixed_precision(self):
         if self.distributed_type == DistributedType.DEEPSPEED:
-            if self.state.deepspeed_plugin.deepspeed_config["fp16"]["enabled"]:
+            config = self.state.deepspeed_plugin.deepspeed_config
+            if "fp16" in config and config["fp16"].get("enabled", False):
```suggestion
            elif config.get("bf16", {}).get("enabled", False):
```
@@ -284,9 +284,10 @@ def use_fp16(self):
     @property
     def mixed_precision(self):
         if self.distributed_type == DistributedType.DEEPSPEED:
-            if self.state.deepspeed_plugin.deepspeed_config["fp16"]["enabled"]:
+            config = self.state.deepspeed_plugin.deepspeed_config
+            if "fp16" in config and config["fp16"].get("enabled", False):
                 mixed_precision = "fp16"
-            elif self.state.deepspeed_plugin.deepspeed_config["bf16"]["enabled"]:
+            elif "bf16" in config and config["bf16"].get("enabled", False):
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/308
diff --git a/examples/README.md b/examples/README.md
index 19995fc0a..115384a42 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -183,3 +183,13 @@ To run it in each of these various modes, use the following commands:
         ```
     * In PyTorch:
         Add an `xmp.spawn` line in your script as you usually do.
+
+## Finer Examples
+
+While the first two scripts are extremely barebones when it comes to what you can do with accelerate, there are two complementary scripts for each: `complete_cv_example.py` and `complete_nlp_example.py`.
+
+These two serve to document what else the `Accelerate` library is able to do, outside of just the basic training loop. They each have the following additional arguments:
+
+- `checkpointing_steps`, whether the various states should be saved at the end of every `n` steps, or `"epoch"` for each epoch. States are then saved to folders named `step_{n}` or `epoch_{n}`
+- `resume_from_checkpoint`, should be used if you want to resume training off of a previous call to the script and passed a `checkpointing_steps` to it.
+- `with_tracking`, should be used if you want to log the training run using all available experiment trackers in your environment. Currently supported trackers include TensorBoard, Weights and Biases, and CometML.
\ No newline at end of file
diff --git a/examples/complete_cv_example.py b/examples/complete_cv_example.py
new file mode 100644
index 000000000..e62568a39
--- /dev/null
+++ b/examples/complete_cv_example.py
@@ -0,0 +1,274 @@
+# coding=utf-8
+# Copyright 2021 The HuggingFace Inc. 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 re
+
+import numpy as np
+import torch
+from torch.optim.lr_scheduler import OneCycleLR
+from torch.utils.data import DataLoader, Dataset
+
+import PIL
+from accelerate import Accelerator
+from timm import create_model
+from torchvision.transforms import Compose, RandomResizedCrop, Resize, ToTensor
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate
+#
+# This example trains a ResNet50 on the Oxford-IIT Pet Dataset
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+
+# Function to get the label from the filename
+def extract_label(fname):
+    stem = fname.split(os.path.sep)[-1]
+    return re.search(r"^(.*)_\d+\.jpg$", stem).groups()[0]
+
+
+class PetsDataset(Dataset):
+    def __init__(self, file_names, image_transform=None, label_to_id=None):
+        self.file_names = file_names
+        self.image_transform = image_transform
+        self.label_to_id = label_to_id
+
+    def __len__(self):
+        return len(self.file_names)
+
+    def __getitem__(self, idx):
+        fname = self.file_names[idx]
+        raw_image = PIL.Image.open(fname)
+        image = raw_image.convert("RGB")
+        if self.image_transform is not None:
+            image = self.image_transform(image)
+        label = extract_label(fname)
+        if self.label_to_id is not None:
+            label = self.label_to_id[label]
+        return {"image": image, "label": label}
+
+
+def training_function(config, args):
+    # Initialize accelerator
+    if args.with_tracking:
+        accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="all")
+    else:
+        accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision)
+
+    if hasattr(args.checkpointing_steps, "isdigit"):
+        checkpointing_steps = args.checkpointing_steps
+        if args.checkpointing_steps.isdigit():
+            checkpointing_steps = int(args.checkpointing_steps)
+    else:
+        checkpointing_steps = None
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+    image_size = config["image_size"]
+    if not isinstance(image_size, (list, tuple)):
+        image_size = (image_size, image_size)
+
+    # We need to initialize the trackers we use, and also store our configuration
+    if args.with_tracking:
+        accelerator.init_trackers("cv_example", config)
+
+    # Grab all the image filenames
+    file_names = [os.path.join(args.data_dir, fname) for fname in os.listdir(args.data_dir) if fname.endswith(".jpg")]
+
+    # Build the label correspondences
+    all_labels = [extract_label(fname) for fname in file_names]
+    id_to_label = list(set(all_labels))
+    id_to_label.sort()
+    label_to_id = {lbl: i for i, lbl in enumerate(id_to_label)}
+
+    # Set the seed before splitting the data.
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+    # Split our filenames between train and validation
+    random_perm = np.random.permutation(len(file_names))
+    cut = int(0.8 * len(file_names))
+    train_split = random_perm[:cut]
+    eval_split = random_perm[cut:]
+
+    # For training we use a simple RandomResizedCrop
+    train_tfm = Compose([RandomResizedCrop(image_size, scale=(0.5, 1.0)), ToTensor()])
+    train_dataset = PetsDataset(
+        [file_names[i] for i in train_split], image_transform=train_tfm, label_to_id=label_to_id
+    )
+
+    # For evaluation, we use a deterministic Resize
+    eval_tfm = Compose([Resize(image_size), ToTensor()])
+    eval_dataset = PetsDataset([file_names[i] for i in eval_split], image_transform=eval_tfm, label_to_id=label_to_id)
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size, num_workers=4)
+    eval_dataloader = DataLoader(eval_dataset, shuffle=False, batch_size=batch_size, num_workers=4)
+
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = create_model("resnet50d", pretrained=True, num_classes=len(label_to_id))
+
+    # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
+    # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
+    # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
+    model = model.to(accelerator.device)
+
+    # Freezing the base model
+    for param in model.parameters():
+        param.requires_grad = False
+    for param in model.get_classifier().parameters():
+        param.requires_grad = True
+
+    # We normalize the batches of images to be a bit faster.
+    mean = torch.tensor(model.default_cfg["mean"])[None, :, None, None].to(accelerator.device)
+    std = torch.tensor(model.default_cfg["std"])[None, :, None, None].to(accelerator.device)
+
+    # Instantiate optimizer
+    optimizer = torch.optim.Adam(params=model.parameters(), lr=lr / 25)
+
+    # Instantiate learning rate scheduler
+    lr_scheduler = OneCycleLR(optimizer=optimizer, max_lr=lr, epochs=num_epochs, steps_per_epoch=len(train_dataloader))
+
+    # Prepare everything
+    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
+    # prepare method.
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    # Potentially load in the weights and states from a previous save
+    state_restored = True
+    if args.resume_from_checkpoint:
+        if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
+            accelerator.print(f"Resumed from checkpoint: {args.resume_from_checkpoint}")
+            accelerator.load_state(args.resume_from_checkpoint)
+            resume_step = None
+        else:
+            # Get the most recent checkpoint
+            dirs = [f.name for f in os.scandir(os.getcwd()) if f.is_dir()]
+            dirs.sort(key=os.path.getctime)
+            path = dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
+            if "epoch" in path.name:
+                num_epochs -= int(path.name.replace("epoch_", ""))
+            else:
+                resume_step = int(path.name.replace("step_", ""))
+                num_epochs -= resume_step // len(train_dataloader)
+                resume_step = (num_epochs * len(train_dataloader)) - resume_step
+                state_restored = False
+
+    overall_step = 0
+    # Now we train the model
+    for epoch in range(num_epochs):
+        model.train()
+        if args.with_tracking:
+            total_loss = 0
+        for step, batch in enumerate(train_dataloader):
+            # We need to skip steps until we reach the resumed step
+            if args.resume_from_checkpoint and epoch == 0 and step < resume_step:
+                continue
+            # We could avoid this line since we set the accelerator with `device_placement=True`.
+            batch = {k: v.to(accelerator.device) for k, v in batch.items()}
+            inputs = (batch["image"] - mean) / std
+            outputs = model(inputs)
+            loss = torch.nn.functional.cross_entropy(outputs, batch["label"])
+            # We keep track of the loss at each epoch
+            if args.with_tracking:
+                total_loss += loss.detach().float()
+            accelerator.backward(loss)
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
+            overall_step += 1
+            if isinstance(checkpointing_steps, int):
+                if overall_step % checkpointing_steps == 0:
+                    accelerator.save_state(f"step_{overall_step}")
+        if state_restored:
+            model.eval()
+            accurate = 0
+            num_elems = 0
+            for step, batch in enumerate(eval_dataloader):
+                # We could avoid this line since we set the accelerator with `device_placement=True`.
+                batch = {k: v.to(accelerator.device) for k, v in batch.items()}
+                inputs = (batch["image"] - mean) / std
+                with torch.no_grad():
+                    outputs = model(inputs)
+                predictions = outputs.argmax(dim=-1)
+                accurate_preds = accelerator.gather(predictions) == accelerator.gather(batch["label"])
+                num_elems += accurate_preds.shape[0]
+                accurate += accurate_preds.long().sum()
+
+            eval_metric = accurate.item() / num_elems
+            # Use accelerator.print to print only on the main process.
+            accelerator.print(f"epoch {epoch}: {100 * eval_metric:.2f}")
+            if args.with_tracking:
+                accelerator.log(
+                    {"accuracy": 100 * eval_metric, "total_loss": total_loss, "epoch": epoch}, step=overall_step
+                )
+            if args.checkpointing_steps == "epoch":
+                accelerator.save_state(f"epoch_{epoch}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Simple example of training script.")
+    parser.add_argument("--data_dir", required=True, help="The data folder on disk.")
+    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.")
+    parser.add_argument(
+        "--checkpointing_steps",
+        type=str,
+        default=None,
+        help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.",
+    )
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        default=None,
+        help="If the training should continue from a checkpoint folder.",
+    )
+    parser.add_argument(
+        "--with_tracking",
+        required=False,
+        help="Whether to load in all available experiment trackers from the environment and use them for logging.",
+    )
+    args = parser.parse_args()
+    config = {"lr": 3e-2, "num_epochs": 3, "seed": 42, "batch_size": 64, "image_size": 224}
+    training_function(config, args)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/complete_nlp_example.py b/examples/complete_nlp_example.py
new file mode 100644
index 000000000..57125aa29
--- /dev/null
+++ b/examples/complete_nlp_example.py
@@ -0,0 +1,262 @@
+# coding=utf-8
+# Copyright 2021 The HuggingFace Inc. 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 torch
+from torch.utils.data import DataLoader
+
+from accelerate import Accelerator, DistributedType
+from datasets import load_dataset, load_metric
+from transformers import (
+    AdamW,
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    get_linear_schedule_with_warmup,
+    set_seed,
+)
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#
+# This example also demonstrates the checkpointing and sharding capabilities
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+def training_function(config, args):
+    # Initialize accelerator
+    if args.with_tracking:
+        accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="all")
+    else:
+        accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision)
+
+    if hasattr(args.checkpointing_steps, "isdigit"):
+        checkpointing_steps = args.checkpointing_steps
+        if args.checkpointing_steps.isdigit():
+            checkpointing_steps = int(args.checkpointing_steps)
+    else:
+        checkpointing_steps = None
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    correct_bias = config["correct_bias"]
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+
+    # We need to initialize the trackers we use, and also store our configuration
+    if args.with_tracking:
+        accelerator.init_trackers("nlp_example", config)
+
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    datasets = load_dataset("glue", "mrpc")
+    metric = load_metric("glue", "mrpc")
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["idx", "sentence1", "sentence2"],
+    )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    # If the batch size is too big we use gradient accumulation
+    gradient_accumulation_steps = 1
+    if batch_size > MAX_GPU_BATCH_SIZE:
+        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
+        batch_size = MAX_GPU_BATCH_SIZE
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    set_seed(seed)
+
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
+
+    # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
+    # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
+    # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
+    model = model.to(accelerator.device)
+
+    # Instantiate optimizer
+    optimizer = AdamW(params=model.parameters(), lr=lr, correct_bias=correct_bias)
+
+    # Instantiate scheduler
+    lr_scheduler = get_linear_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=100,
+        num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,
+    )
+
+    # Prepare everything
+    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
+    # prepare method.
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    # Potentially load in the weights and states from a previous save
+    state_restored = True
+    if args.resume_from_checkpoint:
+        if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
+            accelerator.print(f"Resumed from checkpoint: {args.resume_from_checkpoint}")
+            accelerator.load_state(args.resume_from_checkpoint)
+            resume_step = None
+        else:
+            # Get the most recent checkpoint
+            dirs = [f.name for f in os.scandir(os.getcwd()) if f.is_dir()]
+            dirs.sort(key=os.path.getctime)
+            path = dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
+            if "epoch" in path.name:
+                num_epochs -= int(path.name.replace("epoch_", ""))
+            else:
+                resume_step = int(path.name.replace("step_", ""))
+                num_epochs -= resume_step // len(train_dataloader)
+                resume_step = (num_epochs * len(train_dataloader)) - resume_step
+                state_restored = False
+
+    overall_step = 0
+    # Now we train the model
+    for epoch in range(num_epochs):
+        model.train()
+        if args.with_tracking:
+            total_loss = 0
+        for step, batch in enumerate(train_dataloader):
+            # We need to skip steps until we reach the resumed step
+            if args.resume_from_checkpoint and epoch == 0 and step < resume_step:
+                continue
+            # We could avoid this line since we set the accelerator with `device_placement=True`.
+            batch.to(accelerator.device)
+            outputs = model(**batch)
+            loss = outputs.loss
+            loss = loss / gradient_accumulation_steps
+            # We keep track of the loss at each epoch
+            if args.with_tracking:
+                total_loss += loss.detach().float()
+            accelerator.backward(loss)
+            if step % gradient_accumulation_steps == 0:
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+
+            overall_step += 1
+
+            if isinstance(checkpointing_steps, int):
+                if overall_step % checkpointing_steps == 0:
+                    accelerator.save_state(f"step_{overall_step}")
+        if state_restored:
+            model.eval()
+            for step, batch in enumerate(eval_dataloader):
+                # We could avoid this line since we set the accelerator with `device_placement=True`.
+                batch.to(accelerator.device)
+                with torch.no_grad():
+                    outputs = model(**batch)
+                predictions = outputs.logits.argmax(dim=-1)
+                metric.add_batch(
+                    predictions=accelerator.gather(predictions),
+                    references=accelerator.gather(batch["labels"]),
+                )
+
+            eval_metric = metric.compute()
+            # Use accelerator.print to print only on the main process.
+            accelerator.print(f"epoch {epoch}:", eval_metric)
+            if args.with_tracking:
+                accelerator.log(
+                    {
+                        "accuracy": eval_metric["accuracy"],
+                        "f1": eval_metric["f1"],
+                        "total_loss": total_loss,
+                        "epoch": epoch,
+                    },
+                    step=overall_step,
+                )
+
+            if args.checkpointing_steps == "epoch":
+                accelerator.save_state(f"epoch_{epoch}")
+
+
+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.")
+    parser.add_argument(
+        "--checkpointing_steps",
+        type=str,
+        default=None,
+        help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.",
+    )
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        default=None,
+        help="If the training should continue from a checkpoint folder.",
+    )
+    parser.add_argument(
+        "--with_tracking",
+        required=False,
+        help="Whether to load in all available experiment trackers from the environment and use them for logging.",
+    )
+    args = parser.parse_args()
+    config = {"lr": 2e-5, "num_epochs": 3, "correct_bias": True, "seed": 42, "batch_size": 16}
+    training_function(config, args)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/cv_example.py b/examples/cv_example.py
index 5f2adb7e4..a40fe6bef 100644
--- a/examples/cv_example.py
+++ b/examples/cv_example.py
@@ -149,10 +149,6 @@ def training_function(config, args):
         model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
     )
 
-    # Instantiate learning rate scheduler after preparing the training dataloader as the prepare method
-    # may change its length.
-    lr_scheduler = OneCycleLR(optimizer=optimizer, max_lr=lr, epochs=num_epochs, steps_per_epoch=len(train_dataloader))
-
     # Now we train the model
     for epoch in range(num_epochs):
         model.train()

This is the NLP example ;-)
@@ -51,15 +51,28 @@
 
 def training_function(config, args):
     # Initialize accelerator
-    accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision)
-
+    if args.with_tracking:
+        accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="all")
+    else:
+        accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision)
+
+    if hasattr(args.checkpointing_steps, "isdigit"):
+        checkpointing_steps = args.checkpointing_steps
+        if args.checkpointing_steps.isdigit():
+            checkpointing_steps = int(args.checkpointing_steps)
+    else:
+        checkpointing_steps = None
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
     lr = config["lr"]
     num_epochs = int(config["num_epochs"])
     correct_bias = config["correct_bias"]
     seed = int(config["seed"])
     batch_size = int(config["batch_size"])
 
+    # We need to initialize the trackers we use, and also store our configuration
+    if args.with_tracking:
+        accelerator.init_trackers("cv_example", config)
```suggestion
        help="Whether to load in all available experiment trackers from the environment and use them for logging.",
```
@@ -172,6 +214,19 @@ def main():
         "and an Nvidia Ampere GPU.",
     )
     parser.add_argument("--cpu", action="store_true", help="If passed, will train on the CPU.")
+    parser.add_argument(
+        "--checkpointing_steps",
+        required=False,
+        help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.",
+    )
+    parser.add_argument(
+        "--resume_from_checkpoint", required=False, help="If the training should continue from a checkpoint folder."
+    )
+    parser.add_argument(
+        "--with_tracking",
+        required=False,
+        help="Whether to load in all available experiment trackers from the enviornment and use them for logging.",
I'd put `type=str` and `default=None`here instead of `required=False`
@@ -172,6 +214,19 @@ def main():
         "and an Nvidia Ampere GPU.",
     )
     parser.add_argument("--cpu", action="store_true", help="If passed, will train on the CPU.")
+    parser.add_argument(
+        "--checkpointing_steps",
+        required=False,
I'd put `type=str` and `default=None`here
@@ -172,6 +214,19 @@ def main():
         "and an Nvidia Ampere GPU.",
     )
     parser.add_argument("--cpu", action="store_true", help="If passed, will train on the CPU.")
+    parser.add_argument(
+        "--checkpointing_steps",
+        required=False,
+        help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.",
+    )
+    parser.add_argument(
+        "--resume_from_checkpoint", required=False, help="If the training should continue from a checkpoint folder."
There is a difference between this argument being None or "". In the second case we want to resume from the last checkpoint ideally.
@@ -127,21 +140,38 @@ def collate_fn(examples):
         model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
     )
 
+    # Potentially load in the weights and states from a previous save
+    if args.resume_from_checkpoint:
When resuming a training, we need to skip some steps.
@@ -127,21 +140,38 @@ def collate_fn(examples):
         model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
     )
 
+    # Potentially load in the weights and states from a previous save
+    if args.resume_from_checkpoint:
+        accelerator.print(f"Resumed from checkpoint: {args.resume_from_checkpoint}")
+        accelerator.load_state(args.resume_from_checkpoint)
+
+    overall_step = 0
Did you mean `-=` here?
@@ -0,0 +1,262 @@
+# coding=utf-8
+# Copyright 2021 The HuggingFace Inc. 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 torch
+from torch.utils.data import DataLoader
+
+from accelerate import Accelerator, DistributedType
+from datasets import load_dataset, load_metric
+from transformers import (
+    AdamW,
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    get_linear_schedule_with_warmup,
+    set_seed,
+)
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#
+# This example also demonstrates the checkpointing and sharding capabilities
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+def training_function(config, args):
+    # Initialize accelerator
+    if args.with_tracking:
+        accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="all")
+    else:
+        accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision)
+
+    if hasattr(args.checkpointing_steps, "isdigit"):
+        checkpointing_steps = args.checkpointing_steps
+        if args.checkpointing_steps.isdigit():
+            checkpointing_steps = int(args.checkpointing_steps)
+    else:
+        checkpointing_steps = None
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    correct_bias = config["correct_bias"]
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+
+    # We need to initialize the trackers we use, and also store our configuration
+    if args.with_tracking:
+        accelerator.init_trackers("nlp_example", config)
+
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    datasets = load_dataset("glue", "mrpc")
+    metric = load_metric("glue", "mrpc")
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["idx", "sentence1", "sentence2"],
+    )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    # If the batch size is too big we use gradient accumulation
+    gradient_accumulation_steps = 1
+    if batch_size > MAX_GPU_BATCH_SIZE:
+        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
+        batch_size = MAX_GPU_BATCH_SIZE
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    set_seed(seed)
+
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
+
+    # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
+    # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
+    # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
+    model = model.to(accelerator.device)
+
+    # Instantiate optimizer
+    optimizer = AdamW(params=model.parameters(), lr=lr, correct_bias=correct_bias)
+
+    # Instantiate scheduler
+    lr_scheduler = get_linear_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=100,
+        num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,
+    )
+
+    # Prepare everything
+    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
+    # prepare method.
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    # Potentially load in the weights and states from a previous save
+    state_restored = True
+    if args.resume_from_checkpoint:
+        if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
+            accelerator.print(f"Resumed from checkpoint: {args.resume_from_checkpoint}")
+            accelerator.load_state(args.resume_from_checkpoint)
+            resume_step = None
+        else:
+            # Get the most recent checkpoint
+            dirs = [f.name for f in os.scandir(os.getcwd()) if f.is_dir()]
+            dirs.sort(key=os.path.getctime)
+            path = dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
+            if "epoch" in path.name:
+                num_epochs = -int(path.name.replace("epoch_", ""))
Same here, this is weird.
@@ -0,0 +1,262 @@
+# coding=utf-8
+# Copyright 2021 The HuggingFace Inc. 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 torch
+from torch.utils.data import DataLoader
+
+from accelerate import Accelerator, DistributedType
+from datasets import load_dataset, load_metric
+from transformers import (
+    AdamW,
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    get_linear_schedule_with_warmup,
+    set_seed,
+)
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#
+# This example also demonstrates the checkpointing and sharding capabilities
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+def training_function(config, args):
+    # Initialize accelerator
+    if args.with_tracking:
+        accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="all")
+    else:
+        accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision)
+
+    if hasattr(args.checkpointing_steps, "isdigit"):
+        checkpointing_steps = args.checkpointing_steps
+        if args.checkpointing_steps.isdigit():
+            checkpointing_steps = int(args.checkpointing_steps)
+    else:
+        checkpointing_steps = None
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    correct_bias = config["correct_bias"]
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+
+    # We need to initialize the trackers we use, and also store our configuration
+    if args.with_tracking:
+        accelerator.init_trackers("nlp_example", config)
+
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    datasets = load_dataset("glue", "mrpc")
+    metric = load_metric("glue", "mrpc")
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["idx", "sentence1", "sentence2"],
+    )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    # If the batch size is too big we use gradient accumulation
+    gradient_accumulation_steps = 1
+    if batch_size > MAX_GPU_BATCH_SIZE:
+        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
+        batch_size = MAX_GPU_BATCH_SIZE
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    set_seed(seed)
+
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
+
+    # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
+    # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
+    # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
+    model = model.to(accelerator.device)
+
+    # Instantiate optimizer
+    optimizer = AdamW(params=model.parameters(), lr=lr, correct_bias=correct_bias)
+
+    # Instantiate scheduler
+    lr_scheduler = get_linear_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=100,
+        num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,
+    )
+
+    # Prepare everything
+    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
+    # prepare method.
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    # Potentially load in the weights and states from a previous save
+    state_restored = True
+    if args.resume_from_checkpoint:
+        if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
+            accelerator.print(f"Resumed from checkpoint: {args.resume_from_checkpoint}")
+            accelerator.load_state(args.resume_from_checkpoint)
+            resume_step = None
+        else:
+            # Get the most recent checkpoint
+            dirs = [f.name for f in os.scandir(os.getcwd()) if f.is_dir()]
+            dirs.sort(key=os.path.getctime)
+            path = dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
+            if "epoch" in path.name:
+                num_epochs = -int(path.name.replace("epoch_", ""))
+            else:
+                resume_step = int(path.name.replace("step_", ""))
+                num_epochs = -resume_step // len(train_dataloader)
```suggestion
            if args.resume_from_checkpoint and epoch == 0 and step < resume_step:
                continue
```
and this needs a comment on what we are doing here.
@@ -0,0 +1,262 @@
+# coding=utf-8
+# Copyright 2021 The HuggingFace Inc. 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 torch
+from torch.utils.data import DataLoader
+
+from accelerate import Accelerator, DistributedType
+from datasets import load_dataset, load_metric
+from transformers import (
+    AdamW,
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    get_linear_schedule_with_warmup,
+    set_seed,
+)
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#
+# This example also demonstrates the checkpointing and sharding capabilities
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+def training_function(config, args):
+    # Initialize accelerator
+    if args.with_tracking:
+        accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="all")
+    else:
+        accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision)
+
+    if hasattr(args.checkpointing_steps, "isdigit"):
+        checkpointing_steps = args.checkpointing_steps
+        if args.checkpointing_steps.isdigit():
+            checkpointing_steps = int(args.checkpointing_steps)
+    else:
+        checkpointing_steps = None
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    correct_bias = config["correct_bias"]
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+
+    # We need to initialize the trackers we use, and also store our configuration
+    if args.with_tracking:
+        accelerator.init_trackers("nlp_example", config)
+
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    datasets = load_dataset("glue", "mrpc")
+    metric = load_metric("glue", "mrpc")
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["idx", "sentence1", "sentence2"],
+    )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    # If the batch size is too big we use gradient accumulation
+    gradient_accumulation_steps = 1
+    if batch_size > MAX_GPU_BATCH_SIZE:
+        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
+        batch_size = MAX_GPU_BATCH_SIZE
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    set_seed(seed)
+
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
+
+    # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
+    # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
+    # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
+    model = model.to(accelerator.device)
+
+    # Instantiate optimizer
+    optimizer = AdamW(params=model.parameters(), lr=lr, correct_bias=correct_bias)
+
+    # Instantiate scheduler
+    lr_scheduler = get_linear_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=100,
+        num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,
+    )
+
+    # Prepare everything
+    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
+    # prepare method.
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    # Potentially load in the weights and states from a previous save
+    state_restored = True
+    if args.resume_from_checkpoint:
+        if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
+            accelerator.print(f"Resumed from checkpoint: {args.resume_from_checkpoint}")
+            accelerator.load_state(args.resume_from_checkpoint)
+            resume_step = None
+        else:
+            # Get the most recent checkpoint
+            dirs = [f.name for f in os.scandir(os.getcwd()) if f.is_dir()]
+            dirs.sort(key=os.path.getctime)
+            path = dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
+            if "epoch" in path.name:
+                num_epochs = -int(path.name.replace("epoch_", ""))
+            else:
+                resume_step = int(path.name.replace("step_", ""))
+                num_epochs = -resume_step // len(train_dataloader)
+                resume_step = (num_epochs * len(train_dataloader)) - resume_step
+                state_restored = False
+
+    overall_step = 0
+    # Now we train the model
+    for epoch in range(num_epochs):
+        model.train()
+        if args.with_tracking:
+            total_loss = 0
+        for step, batch in enumerate(train_dataloader):
+            if args.resume_from_checkpoint and epoch == 0:
+                for _ in range(resume_step - 1):
+                    pass
Then no need for an else branch
@@ -0,0 +1,262 @@
+# coding=utf-8
+# Copyright 2021 The HuggingFace Inc. 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 torch
+from torch.utils.data import DataLoader
+
+from accelerate import Accelerator, DistributedType
+from datasets import load_dataset, load_metric
+from transformers import (
+    AdamW,
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    get_linear_schedule_with_warmup,
+    set_seed,
+)
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#
+# This example also demonstrates the checkpointing and sharding capabilities
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+def training_function(config, args):
+    # Initialize accelerator
+    if args.with_tracking:
+        accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="all")
+    else:
+        accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision)
+
+    if hasattr(args.checkpointing_steps, "isdigit"):
+        checkpointing_steps = args.checkpointing_steps
+        if args.checkpointing_steps.isdigit():
+            checkpointing_steps = int(args.checkpointing_steps)
+    else:
+        checkpointing_steps = None
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    correct_bias = config["correct_bias"]
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+
+    # We need to initialize the trackers we use, and also store our configuration
+    if args.with_tracking:
+        accelerator.init_trackers("nlp_example", config)
+
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    datasets = load_dataset("glue", "mrpc")
+    metric = load_metric("glue", "mrpc")
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["idx", "sentence1", "sentence2"],
+    )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    # If the batch size is too big we use gradient accumulation
+    gradient_accumulation_steps = 1
+    if batch_size > MAX_GPU_BATCH_SIZE:
+        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
+        batch_size = MAX_GPU_BATCH_SIZE
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    set_seed(seed)
+
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
+
+    # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
+    # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
+    # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
+    model = model.to(accelerator.device)
+
+    # Instantiate optimizer
+    optimizer = AdamW(params=model.parameters(), lr=lr, correct_bias=correct_bias)
+
+    # Instantiate scheduler
+    lr_scheduler = get_linear_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=100,
+        num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,
+    )
+
+    # Prepare everything
+    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
+    # prepare method.
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    # Potentially load in the weights and states from a previous save
+    state_restored = True
+    if args.resume_from_checkpoint:
+        if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
+            accelerator.print(f"Resumed from checkpoint: {args.resume_from_checkpoint}")
+            accelerator.load_state(args.resume_from_checkpoint)
+            resume_step = None
+        else:
+            # Get the most recent checkpoint
+            dirs = [f.name for f in os.scandir(os.getcwd()) if f.is_dir()]
+            dirs.sort(key=os.path.getctime)
+            path = dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
+            if "epoch" in path.name:
+                num_epochs = -int(path.name.replace("epoch_", ""))
+            else:
+                resume_step = int(path.name.replace("step_", ""))
+                num_epochs = -resume_step // len(train_dataloader)
+                resume_step = (num_epochs * len(train_dataloader)) - resume_step
+                state_restored = False
+
+    overall_step = 0
+    # Now we train the model
+    for epoch in range(num_epochs):
+        model.train()
+        if args.with_tracking:
+            total_loss = 0
+        for step, batch in enumerate(train_dataloader):
+            if args.resume_from_checkpoint and epoch == 0:
+                for _ in range(resume_step - 1):
+                    pass
+            else:
Good catch!
@@ -149,10 +149,6 @@ def training_function(config, args):
         model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
     )
 
-    # Instantiate learning rate scheduler after preparing the training dataloader as the prepare method
-    # may change its length.
-    lr_scheduler = OneCycleLR(optimizer=optimizer, max_lr=lr, epochs=num_epochs, steps_per_epoch=len(train_dataloader))
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/307
diff --git a/src/accelerate/tracking.py b/src/accelerate/tracking.py
index 65a5caec2..8d496a3d1 100644
--- a/src/accelerate/tracking.py
+++ b/src/accelerate/tracking.py
@@ -243,3 +243,10 @@ def log(self, values: dict, step: Optional[int] = None):
             self.writer.set_step(step)
         self.writer.log_others(values)
         logger.info("Successfully logged to CometML")
+
+    def finish(self):
+        """
+        Closes `comet-ml` writer
+        """
+        self.writer.end()
+        logger.info("CometML run closed")
diff --git a/tests/test_tracking.py b/tests/test_tracking.py
index 62c423ef7..b958dbbfc 100644
--- a/tests/test_tracking.py
+++ b/tests/test_tracking.py
@@ -13,11 +13,13 @@
 # limitations under the License.
 
 import csv
+import json
 import logging
 import os
 import re
 import tempfile
 import unittest
+import zipfile
 from pathlib import Path
 from typing import Optional
 from unittest import mock
@@ -25,8 +27,9 @@
 # We use TF to parse the logs
 from accelerate import Accelerator
 from accelerate.test_utils.testing import MockingTestCase, TempDirTestCase, require_tensorflow
-from accelerate.tracking import GeneralTracker
+from accelerate.tracking import CometMLTracker, GeneralTracker
 from accelerate.utils import is_tensorflow_available
+from comet_ml import OfflineExperiment
 
 
 if is_tensorflow_available():
@@ -162,6 +165,66 @@ def test_log(self):
         self.assertEqual(self.get_value_from_log("_step", cleaned_log), "0")
 
 
+# Comet has a special `OfflineExperiment` we need to use for testing
+def offline_init(self, run_name: str, tmpdir: str):
+    self.run_name = run_name
+    self.writer = OfflineExperiment(project_name=run_name, offline_directory=tmpdir)
+    logger.info(f"Initialized offline CometML project {self.run_name}")
+    logger.info("Make sure to log any initial configurations with `self.store_init_configuration` before training!")
+
+
+@mock.patch.object(CometMLTracker, "__init__", offline_init)
+class CometMLTest(unittest.TestCase):
+    @staticmethod
+    def get_value_from_key(log_list, key: str, is_param: bool = False):
+        "Extracts `key` from Comet `log`"
+        for log in log_list:
+            j = json.loads(log)["payload"]
+            if is_param and "param" in j.keys():
+                if j["param"]["paramName"] == key:
+                    return j["param"]["paramValue"]
+            if "log_other" in j.keys():
+                if j["log_other"]["key"] == key:
+                    return j["log_other"]["val"]
+
+    def test_init_trackers(self):
+        with tempfile.TemporaryDirectory() as d:
+            tracker = CometMLTracker("test_project_with_config", d)
+            accelerator = Accelerator(log_with=tracker)
+            config = {"num_iterations": 12, "learning_rate": 1e-2, "some_boolean": False, "some_string": "some_value"}
+            accelerator.init_trackers(None, config)
+            accelerator.end_training()
+            log = os.listdir(d)[0]  # Comet is nice, it's just a zip file here
+            # We parse the raw logs
+            p = os.path.join(d, log)
+            archive = zipfile.ZipFile(p, "r")
+            log = archive.open("messages.json").read().decode("utf-8")
+        list_of_json = log.split("\n")[:-1]
+        self.assertEqual(self.get_value_from_key(list_of_json, "num_iterations", True), 12)
+        self.assertEqual(self.get_value_from_key(list_of_json, "learning_rate", True), 0.01)
+        self.assertEqual(self.get_value_from_key(list_of_json, "some_boolean", True), False)
+        self.assertEqual(self.get_value_from_key(list_of_json, "some_string", True), "some_value")
+
+    def test_log(self):
+        with tempfile.TemporaryDirectory() as d:
+            tracker = CometMLTracker("test_project_with_config", d)
+            accelerator = Accelerator(log_with=tracker)
+            accelerator.init_trackers(None)
+            values = {"total_loss": 0.1, "iteration": 1, "my_text": "some_value"}
+            accelerator.log(values, step=0)
+            accelerator.end_training()
+            log = os.listdir(d)[0]  # Comet is nice, it's just a zip file here
+            # We parse the raw logs
+            p = os.path.join(d, log)
+            archive = zipfile.ZipFile(p, "r")
+            log = archive.open("messages.json").read().decode("utf-8")
+        list_of_json = log.split("\n")[:-1]
+        self.assertEqual(self.get_value_from_key(list_of_json, "curr_step", True), 0)
+        self.assertEqual(self.get_value_from_key(list_of_json, "total_loss"), 0.1)
+        self.assertEqual(self.get_value_from_key(list_of_json, "iteration"), 1)
+        self.assertEqual(self.get_value_from_key(list_of_json, "my_text"), "some_value")
+
+
 class MyCustomTracker(GeneralTracker):
     "Basic tracker that writes to a csv for testing"
     _col_names = [

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/306
diff --git a/docs/source/tracking.mdx b/docs/source/tracking.mdx
index 256e2848b..959317887 100644
--- a/docs/source/tracking.mdx
+++ b/docs/source/tracking.mdx
@@ -106,7 +106,7 @@ from typing import Optional
 import wandb
 
 
-class WandBTracker(GeneralTracker):
+class MyCustomTracker(GeneralTracker):
     def __init__(self, run_name: str):
         self.run_name = run_name
         wandb.init(self.run_name)
@@ -118,6 +118,21 @@ class WandBTracker(GeneralTracker):
         wandb.log(values, step=step)
 ```
 
+When you are ready to build your `Accelerator` object, pass in an **instance** of your tracker to [`~Accelerator.log_with`] to have it automatically
+be used with the API:
+
+```python
+tracker = MyCustomTracker("some_run_name")
+accelerator = Accelerate(log_with=tracker)
+```
+
+These also can be mixed with existing trackers, including with `"all"`:
+
+```python
+tracker = MyCustomTracker("some_run_name")
+accelerator = Accelerate(log_with=[tracker, "all"])
+```
+
 ## When a wrapper cannot work
 
 If a library has an API that does not follow a strict `.log` with an overall dictionary such as Neptune.AI, logging can be done manually under an `if accelerator.is_main_process` statement:
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 67c981de3..6e0742e02 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -29,7 +29,7 @@
 from .optimizer import AcceleratedOptimizer
 from .scheduler import AcceleratedScheduler
 from .state import AcceleratorState, DistributedType, is_deepspeed_available
-from .tracking import CometMLTracker, TensorBoardTracker, WandBTracker, get_available_trackers
+from .tracking import CometMLTracker, GeneralTracker, TensorBoardTracker, WandBTracker, get_available_trackers
 from .utils import (
     DeepSpeedPlugin,
     LoggerType,
@@ -94,15 +94,15 @@ class Accelerator:
               dataloader) or of the iterable dataset (if it exists) if the underlying dataset is of that type.
 
             Will default to `["torch"]` for PyTorch versions <=1.5.1 and `["generator"]` for PyTorch versions >= 1.6.
-        log_with (list of `str` or [`~utils.LoggerType`], *optional*):
+        log_with (list of `str`, [`~utils.LoggerType`] or [`~tracking.GeneralTracker`], *optional*):
             A list of loggers to be setup for experiment tracking. Should be one or several of:
 
             - `"all"`
             - `"tensorboard"`
             - `"wandb"`
             - `"comet_ml"`
-
-            If `"all`" is selected, will pick up all available trackers in the environment and intialize them.
+            If `"all`" is selected, will pick up all available trackers in the environment and intialize them. Can also
+            accept implementations of `GeneralTracker` for custom trackers, and can be combined with `"all"`.
         logging_dir (`str`, `os.PathLike`, *optional*):
             A path to a directory for storing logs of locally-compatible loggers.
         dispatch_batches (`bool`, *optional*):
@@ -131,26 +131,32 @@ def __init__(
         cpu: bool = False,
         deepspeed_plugin: DeepSpeedPlugin = None,
         rng_types: Optional[List[Union[str, RNGType]]] = None,
-        log_with: Optional[List[Union[str, LoggerType]]] = None,
+        log_with: Optional[List[Union[str, LoggerType, GeneralTracker]]] = None,
         logging_dir: Optional[Union[str, os.PathLike]] = "",
         dispatch_batches: Optional[bool] = None,
         step_scheduler_with_optimizer: bool = True,
         kwargs_handlers: Optional[List[KwargsHandler]] = None,
     ):
+        loggers = []
         if log_with is not None:
-            if not isinstance(log_with, list):
+            if not isinstance(log_with, (list, tuple)):
                 log_with = [log_with]
+                logger.debug(f"{log_with}")
             if "all" in log_with or LoggerType.ALL in log_with:
-                log_with = get_available_trackers()
+                loggers = [o for o in log_with if issubclass(type(o), GeneralTracker)] + get_available_trackers()
             else:
-                for i, log_type in enumerate(log_with):
-                    if log_type not in LoggerType:
+                for log_type in log_with:
+                    if log_type not in LoggerType and not issubclass(type(log_type), GeneralTracker):
                         raise ValueError(
                             f"Unsupported logging capability: {log_type}. Choose between {LoggerType.list()}"
                         )
-                    log_with[i] = LoggerType(log_type)
-                log_with = list(set(log_with))
-        self.log_with = log_with
+                    if issubclass(type(log_type), GeneralTracker):
+                        loggers.append(log_type)
+                    else:
+                        log_type = LoggerType(log_type)
+                        if log_type not in loggers:
+                            loggers.append(log_type)
+        self.log_with = loggers
         self.logging_dir = logging_dir
 
         if mixed_precision is not None:
@@ -631,7 +637,10 @@ def init_trackers(self, project_name: str, config: Optional[dict] = None):
         """
         self.trackers = []
         for tracker in self.log_with:
-            if str(tracker).lower() == "tensorboard" and is_tensorboard_available():
+            if issubclass(type(tracker), GeneralTracker):
+                # Custom trackers are already initialized
+                self.trackers.append(tracker)
+            elif str(tracker).lower() == "tensorboard" and is_tensorboard_available():
                 self.trackers.append(TensorBoardTracker(project_name, self.logging_dir))
             elif str(tracker).lower() == "wandb" and is_wandb_available():
                 self.trackers.append(WandBTracker(project_name))
diff --git a/tests/test_tracking.py b/tests/test_tracking.py
index b1f66b1ae..62c423ef7 100644
--- a/tests/test_tracking.py
+++ b/tests/test_tracking.py
@@ -12,17 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import csv
 import logging
 import os
 import re
 import tempfile
 import unittest
 from pathlib import Path
+from typing import Optional
 from unittest import mock
 
 # We use TF to parse the logs
 from accelerate import Accelerator
 from accelerate.test_utils.testing import MockingTestCase, TempDirTestCase, require_tensorflow
+from accelerate.tracking import GeneralTracker
 from accelerate.utils import is_tensorflow_available
 
 
@@ -157,3 +160,77 @@ def test_log(self):
         self.assertEqual(self.get_value_from_log("iteration", cleaned_log), "1")
         self.assertEqual(self.get_value_from_log("my_text", cleaned_log), "some_value")
         self.assertEqual(self.get_value_from_log("_step", cleaned_log), "0")
+
+
+class MyCustomTracker(GeneralTracker):
+    "Basic tracker that writes to a csv for testing"
+    _col_names = [
+        "total_loss",
+        "iteration",
+        "my_text",
+        "learning_rate",
+        "num_iterations",
+        "some_boolean",
+        "some_string",
+    ]
+
+    def __init__(self, dir: str):
+        self.f = open(f"{dir}/log.csv", "w+")
+        self.writer = csv.DictWriter(self.f, fieldnames=self._col_names)
+        self.writer.writeheader()
+
+    def store_init_configuration(self, values: dict):
+        logger.info("Call init")
+        self.writer.writerow(values)
+
+    def log(self, values: dict, step: Optional[int]):
+        logger.info("Call log")
+        self.writer.writerow(values)
+
+    def finish(self):
+        self.f.close()
+
+
+class CustomTrackerTestCase(unittest.TestCase):
+    def test_init_trackers(self):
+        with tempfile.TemporaryDirectory() as d:
+            tracker = MyCustomTracker(d)
+            accelerator = Accelerator(log_with=tracker)
+            config = {"num_iterations": 12, "learning_rate": 1e-2, "some_boolean": False, "some_string": "some_value"}
+            accelerator.init_trackers("Some name", config)
+            accelerator.end_training()
+            with open(f"{d}/log.csv", "r") as f:
+                data = csv.DictReader(f)
+                data = next(data)
+                truth = {
+                    "total_loss": "",
+                    "iteration": "",
+                    "my_text": "",
+                    "learning_rate": "0.01",
+                    "num_iterations": "12",
+                    "some_boolean": "False",
+                    "some_string": "some_value",
+                }
+                self.assertDictEqual(data, truth)
+
+    def test_log(self):
+        with tempfile.TemporaryDirectory() as d:
+            tracker = MyCustomTracker(d)
+            accelerator = Accelerator(log_with=tracker)
+            accelerator.init_trackers("Some name")
+            values = {"total_loss": 0.1, "iteration": 1, "my_text": "some_value"}
+            accelerator.log(values, step=0)
+            accelerator.end_training()
+            with open(f"{d}/log.csv", "r") as f:
+                data = csv.DictReader(f)
+                data = next(data)
+                truth = {
+                    "total_loss": "0.1",
+                    "iteration": "1",
+                    "my_text": "some_value",
+                    "learning_rate": "",
+                    "num_iterations": "",
+                    "some_boolean": "",
+                    "some_string": "",
+                }
+                self.assertDictEqual(data, truth)

```suggestion
                # Custom trackers are already initialized
```
@@ -631,7 +637,10 @@ def init_trackers(self, project_name: str, config: Optional[dict] = None):
         """
         self.trackers = []
         for tracker in self.log_with:
-            if str(tracker).lower() == "tensorboard" and is_tensorboard_available():
+            if issubclass(type(tracker), GeneralTracker):
+                # Custom tracker that they already initalized
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/305
diff --git a/src/accelerate/test_utils/testing.py b/src/accelerate/test_utils/testing.py
index cce7347f6..6a9e66d6a 100644
--- a/src/accelerate/test_utils/testing.py
+++ b/src/accelerate/test_utils/testing.py
@@ -13,8 +13,14 @@
 # limitations under the License.
 
 import asyncio
+import os
+import shutil
 import sys
+import tempfile
 import unittest
+from pathlib import Path
+from typing import List, Union
+from unittest import mock
 
 import torch
 
@@ -22,6 +28,69 @@
 from ..utils import gather, is_tensorflow_available
 
 
+class TempDirTestCase(unittest.TestCase):
+    """
+    A TestCase class that keeps a single `tempfile.TemporaryDirectory` open for the duration of the class, wipes its
+    data at the start of a test, and then destroyes it at the end of the TestCase.
+
+    Useful for when a class or API requires a single constant folder throughout it's use, such as Weights and Biases
+
+    The temporary directory location will be stored in `self.tmpdir`
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        "Creates a `tempfile.TemporaryDirectory` and stores it in `cls.tmpdir`"
+        cls.tmpdir = tempfile.mkdtemp()
+
+    @classmethod
+    def tearDownClass(cls):
+        "Remove `cls.tmpdir` after test suite has finished"
+        if os.path.exists(cls.tmpdir):
+            shutil.rmtree(cls.tmpdir)
+
+    def setUp(self):
+        "Destroy all contents in `self.tmpdir`, but not `self.tmpdir`"
+        for path in Path(self.tmpdir).glob("**/*"):
+            if path.is_file():
+                path.unlink()
+            elif path.is_dir():
+                shutil.rmtree(path)
+
+
+class MockingTestCase(unittest.TestCase):
+    """
+    A TestCase class designed to dynamically add various mockers that should be used in every test, mimicking the
+    behavior of a class-wide mock when defining one normally will not do.
+
+    Useful when a mock requires specific information available only initialized after `TestCase.setUpClass`, such as
+    setting an environment variable with that information.
+
+    The `add_mocks` function should be ran at the end of a `TestCase`'s `setUp` function, after a call to
+    `super().setUp()` such as:
+    ```python
+    def setUp(self):
+        super().setUp()
+        mocks = mock.patch.dict(os.environ, {"SOME_ENV_VAR", "SOME_VALUE"})
+        self.add_mocks(mocks)
+    ```
+    """
+
+    def add_mocks(self, mocks: Union[mock.Mock, List[mock.Mock]]):
+        """
+        Add custom mocks for tests that should be repeated on each test. Should be called during
+        `MockingTestCase.setUp`, after `super().setUp()`.
+
+        Args:
+            mocks (`mock.Mock` or list of `mock.Mock`):
+                Mocks that should be added to the `TestCase` after `TestCase.setUpClass` has been run
+        """
+        self.mocks = mocks if isinstance(mocks, (tuple, list)) else [mocks]
+        for m in self.mocks:
+            m.start()
+            self.addCleanup(m.stop)
+
+
 def are_the_same_tensors(tensor):
     state = AcceleratorState()
     tensor = tensor[None].clone().to(state.device)
diff --git a/tests/test_tracking.py b/tests/test_tracking.py
index 8d5ebc6bc..b1f66b1ae 100644
--- a/tests/test_tracking.py
+++ b/tests/test_tracking.py
@@ -15,7 +15,6 @@
 import logging
 import os
 import re
-import shutil
 import tempfile
 import unittest
 from pathlib import Path
@@ -23,7 +22,7 @@
 
 # We use TF to parse the logs
 from accelerate import Accelerator
-from accelerate.test_utils.testing import require_tensorflow
+from accelerate.test_utils.testing import MockingTestCase, TempDirTestCase, require_tensorflow
 from accelerate.utils import is_tensorflow_available
 
 
@@ -43,13 +42,11 @@ def test_init_trackers(self):
         hps = None
         project_name = "test_project_with_config"
         with tempfile.TemporaryDirectory() as dirpath:
-            oldpwd = os.getcwd()
-            os.chdir(dirpath)
-            accelerator = Accelerator(log_with="tensorboard")
+            accelerator = Accelerator(log_with="tensorboard", logging_dir=dirpath)
             config = {"num_iterations": 12, "learning_rate": 1e-2, "some_boolean": False, "some_string": "some_value"}
             accelerator.init_trackers(project_name, config)
             accelerator.end_training()
-            for child in Path(project_name).glob("*/**"):
+            for child in Path(f"{dirpath}/{project_name}").glob("*/**"):
                 log = list(filter(lambda x: x.is_file(), child.iterdir()))[0]
                 # The config log is stored one layer deeper in the logged directory
                 # And names are randomly generated each time
@@ -61,7 +58,6 @@ def test_init_trackers(self):
                     plugin_data = plugin_data_pb2.HParamsPluginData.FromString(proto_bytes)
                     if plugin_data.HasField("session_start_info"):
                         hps = dict(plugin_data.session_start_info.hparams)
-            os.chdir(oldpwd)
 
         self.assertTrue(isinstance(hps, dict))
         keys = list(hps.keys())
@@ -77,16 +73,14 @@ def test_log(self):
         step = None
         project_name = "test_project_with_log"
         with tempfile.TemporaryDirectory() as dirpath:
-            oldpwd = os.getcwd()
-            os.chdir(dirpath)
-            accelerator = Accelerator(log_with="tensorboard")
+            accelerator = Accelerator(log_with="tensorboard", logging_dir=dirpath)
             accelerator.init_trackers(project_name)
             values = {"total_loss": 0.1, "iteration": 1, "my_text": "some_value"}
             accelerator.log(values, step=0)
             accelerator.end_training()
             # Logged values are stored in the outermost-tfevents file and can be read in as a TFRecord
             # Names are randomly generated each time
-            log = list(filter(lambda x: x.is_file(), Path(project_name).iterdir()))[0]
+            log = list(filter(lambda x: x.is_file(), Path(f"{dirpath}/{project_name}").iterdir()))[0]
             serialized_examples = tf.data.TFRecordDataset(log)
             for e in serialized_examples:
                 event = event_pb2.Event.FromString(e.numpy())
@@ -99,14 +93,18 @@ def test_log(self):
                         iteration = value.simple_value
                     elif value.tag == "my_text/text_summary":  # Append /text_summary to the key
                         my_text = value.tensor.string_val[0].decode()
-            os.chdir(oldpwd)
         self.assertAlmostEqual(total_loss, values["total_loss"])
         self.assertEqual(iteration, values["iteration"])
         self.assertEqual(my_text, values["my_text"])
 
 
 @mock.patch.dict(os.environ, {"WANDB_MODE": "offline"})
-class WandBTrackingTest(unittest.TestCase):
+class WandBTrackingTest(TempDirTestCase, MockingTestCase):
+    def setUp(self):
+        super().setUp()
+        # wandb let's us override where logs are stored to via the WANDB_DIR env var
+        self.add_mocks(mock.patch.dict(os.environ, {"WANDB_DIR": self.tmpdir}))
+
     @staticmethod
     def get_value_from_log(key: str, log: str, key_occurance: int = 0):
         """
@@ -126,7 +124,7 @@ def test_init_trackers(self):
         accelerator.init_trackers(project_name, config)
         accelerator.end_training()
         # The latest offline log is stored at wandb/latest-run/*.wandb
-        for child in Path("wandb/latest-run").glob("*"):
+        for child in Path(f"{self.tmpdir}/wandb/latest-run").glob("*"):
             logger.info(child)
             if child.is_file() and child.suffix == ".wandb":
                 with open(child, "rb") as f:
@@ -148,7 +146,7 @@ def test_log(self):
         accelerator.log(values, step=0)
         accelerator.end_training()
         # The latest offline log is stored at wandb/latest-run/*.wandb
-        for child in Path("wandb/latest-run").glob("*"):
+        for child in Path(f"{self.tmpdir}/wandb/latest-run").glob("*"):
             if child.is_file() and child.suffix == ".wandb":
                 with open(child, "rb") as f:
                     content = f.read()
@@ -159,18 +157,3 @@ def test_log(self):
         self.assertEqual(self.get_value_from_log("iteration", cleaned_log), "1")
         self.assertEqual(self.get_value_from_log("my_text", cleaned_log), "some_value")
         self.assertEqual(self.get_value_from_log("_step", cleaned_log), "0")
-
-    def setUp(self):
-        os.mkdir(".wandb_tests")
-        os.chdir(".wandb_tests")
-
-    def tearDown(self):
-        if os.getcwd().endswith(".wandb_tests"):
-            os.chdir("..")
-        if os.path.exists(".wandb_tests"):
-            shutil.rmtree(".wandb_tests")
-
-    @classmethod
-    def setUpClass(cls):
-        if os.path.exists(".wandb_tests"):
-            shutil.rmtree(".wandb_tests")

```suggestion
    A TestCase class designed to dynamically add various mockers that should be used in every test, mimicking the behavior
```
@@ -13,15 +13,84 @@
 # limitations under the License.
 
 import asyncio
+import os
+import shutil
 import sys
+import tempfile
 import unittest
+from pathlib import Path
+from typing import List, Union
+from unittest import mock
 
 import torch
 
 from ..state import AcceleratorState, is_tpu_available
 from ..utils import gather, is_tensorflow_available
 
 
+class TempDirTestCase(unittest.TestCase):
+    """
+    A TestCase class that keeps a single `tempfile.TemporaryDirectory` open for the duration of the class, wipes its
+    data at the start of a test, and then destroyes it at the end of the TestCase.
+
+    Useful for when a class or API requires a single constant folder throughout it's use, such as Weights and Biases
+
+    The temporary directory location will be stored in `self.tmpdir`
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        "Creates a `tempfile.TemporaryDirectory` and stores it in `cls.tmpdir`"
+        cls.tmpdir = tempfile.mkdtemp()
+
+    @classmethod
+    def tearDownClass(cls):
+        "Remove `cls.tmpdir` after test suite has finished"
+        if os.path.exists(cls.tmpdir):
+            shutil.rmtree(cls.tmpdir)
+
+    def setUp(self):
+        "Destroy all contents in `self.tmpdir`, but not `self.tmpdir`"
+        for path in Path(self.tmpdir).glob("**/*"):
+            if path.is_file():
+                path.unlink()
+            elif path.is_dir():
+                shutil.rmtree(path)
+
+
+class MockingTestCase(unittest.TestCase):
+    """
+    A TestCase class designed to dynamically add various mockers that should be used in every test, mimicing a behavior
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/304
diff --git a/docs/source/internal.mdx b/docs/source/internal.mdx
index bd6924990..94803a62f 100644
--- a/docs/source/internal.mdx
+++ b/docs/source/internal.mdx
@@ -34,6 +34,10 @@ The main work on your PyTorch `DataLoader` is done by the following function:
 
 [[autodoc]] data_loader.IterableDatasetShard
 
+## Scheduler
+
+[[autodoc]] scheduler.AcceleratedScheduler
+
 ## Distributed Config
 
 ### AcceleratorState
diff --git a/docs/source/quicktour.mdx b/docs/source/quicktour.mdx
index b6f539fd0..2b405e8ec 100644
--- a/docs/source/quicktour.mdx
+++ b/docs/source/quicktour.mdx
@@ -45,11 +45,13 @@ model on `accelerator.device` or your training will fail on TPU.
 
 </Tip>
 
-3. Pass all objects relevant to training (optimizer, model, training dataloader) to the
+3. Pass all objects relevant to training (optimizer, model, training dataloader, learning rate scheduler) to the
 [`~Accelerator.prepare`] method. This will make sure everything is ready for training.
 
 ```python
-model, optimizer, train_dataloader = accelerator.prepare(model, optimizer, train_dataloader)
+model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+    model, optimizer, train_dataloader, lr_scheduler
+)
 ```
 
 In particular, your training dataloader will be sharded accross all GPUs/TPU cores available so that each one sees a
@@ -74,14 +76,21 @@ training loop.
 
 <Tip warning={true}>
 
+You should only pass the learning rate scheduler to [`~Accelerator.prepare`] when the scheduler needs to be stepped
+at each optimizer step.
+
+</Tip>
+
+<Tip warning={true}>
+
 Your training dataloader may change length when going through this method: if you run on X GPUs, it will have its
 length divided by X (since your actual batch size will be multiplied by X), unless you set
 `split_batches=True`.
 
 </Tip>
 
-Any instruction using your training dataloader length (for instance if you need the number of total training steps
-to create a learning rate scheduler) should go after the call to [`~Accelerator.prepare`].
+Any instruction using your training dataloader length (for instance if you want to log the number of total training
+steps) should go after the call to [`~Accelerator.prepare`].
 
 You can perfectly send your dataloader to [`~Accelerator.prepare`] on its own, but it's best to send the
 model and optimizer to [`~Accelerator.prepare`] together.
diff --git a/examples/cv_example.py b/examples/cv_example.py
index bb7c7f5b1..5f2adb7e4 100644
--- a/examples/cv_example.py
+++ b/examples/cv_example.py
@@ -139,11 +139,14 @@ def training_function(config, args):
     # Instantiate optimizer
     optimizer = torch.optim.Adam(params=model.parameters(), lr=lr / 25)
 
+    # Instantiate learning rate scheduler
+    lr_scheduler = OneCycleLR(optimizer=optimizer, max_lr=lr, epochs=num_epochs, steps_per_epoch=len(train_dataloader))
+
     # Prepare everything
     # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
     # prepare method.
-    model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
-        model, optimizer, train_dataloader, eval_dataloader
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
     )
 
     # Instantiate learning rate scheduler after preparing the training dataloader as the prepare method
diff --git a/examples/nlp_example.py b/examples/nlp_example.py
index 432d78ae8..8f0d25541 100644
--- a/examples/nlp_example.py
+++ b/examples/nlp_example.py
@@ -113,21 +113,20 @@ def collate_fn(examples):
     # Instantiate optimizer
     optimizer = AdamW(params=model.parameters(), lr=lr, correct_bias=correct_bias)
 
-    # Prepare everything
-    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
-    # prepare method.
-    model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
-        model, optimizer, train_dataloader, eval_dataloader
-    )
-
-    # Instantiate learning rate scheduler after preparing the training dataloader as the prepare method
-    # may change its length.
+    # Instantiate scheduler
     lr_scheduler = get_linear_schedule_with_warmup(
         optimizer=optimizer,
         num_warmup_steps=100,
         num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,
     )
 
+    # Prepare everything
+    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
+    # prepare method.
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
     # Now we train the model
     for epoch in range(num_epochs):
         model.train()
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index f52cf7ca8..67c981de3 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -27,6 +27,7 @@
 from .data_loader import prepare_data_loader
 from .kwargs_handlers import DistributedDataParallelKwargs, GradScalerKwargs, InitProcessGroupKwargs, KwargsHandler
 from .optimizer import AcceleratedOptimizer
+from .scheduler import AcceleratedScheduler
 from .state import AcceleratorState, DistributedType, is_deepspeed_available
 from .tracking import CometMLTracker, TensorBoardTracker, WandBTracker, get_available_trackers
 from .utils import (
@@ -108,6 +109,9 @@ class Accelerator:
             If set to `True`, the dataloader prepared by the Accelerator is only iterated through on the main process
             and then the batches are split and broadcast to each process. Will default to `True` for `DataLoader` whose
             underlying dataset is an `IterableDataset`, `False` otherwise.
+        step_scheduler_with_optimizer (`bool`, *optional`, defaults to `True`):
+            Set `True` if the learning rate scheduler is stepped at the same time as the optimizer, `False` if only
+            done under certain circumstances (at the end of each epoch, for instance).
         kwargs_handlers (`List[KwargHandler]`, *optional*)
             A list of `KwargHandler` to customize how the objects related to distributed training or mixed precision
             are created. See [kwargs](kwargs) for more information.
@@ -130,6 +134,7 @@ def __init__(
         log_with: Optional[List[Union[str, LoggerType]]] = None,
         logging_dir: Optional[Union[str, os.PathLike]] = "",
         dispatch_batches: Optional[bool] = None,
+        step_scheduler_with_optimizer: bool = True,
         kwargs_handlers: Optional[List[KwargsHandler]] = None,
     ):
         if log_with is not None:
@@ -205,6 +210,7 @@ def __init__(
             raise ImportError(
                 "Using `DataLoaderDispatcher` requires PyTorch 1.8.0 minimum. You have {torch.__version__}."
             )
+        self.step_scheduler_with_optimizer = step_scheduler_with_optimizer
 
         # Mixed precision attributes
         self.scaler = None
@@ -227,6 +233,7 @@ def __init__(
         # Internal references to the training objects
         self._optimizers = []
         self._models = []
+        self._schedulers = []
         self._custom_objects = []
 
         # RNG Types
@@ -315,16 +322,22 @@ def print(self, *args, **kwargs):
         if self.is_local_main_process:
             print(*args, **kwargs)
 
-    def _prepare_one(self, obj):
-        if isinstance(obj, torch.utils.data.DataLoader):
+    def _prepare_one(self, obj, first_pass=False):
+        # First pass of preparation: DataLoader, model, optimizer
+        if isinstance(obj, torch.utils.data.DataLoader) and first_pass:
             return self.prepare_data_loader(obj)
-        elif isinstance(obj, torch.nn.Module):
+        elif isinstance(obj, torch.nn.Module) and first_pass:
             self._models.append(obj)
             return self.prepare_model(obj)
-        elif isinstance(obj, torch.optim.Optimizer):
+        elif isinstance(obj, torch.optim.Optimizer) and first_pass:
             optimizer = self.prepare_optimizer(obj)
             self._optimizers.append(optimizer)
             return optimizer
+        # Second pass of preparation: LR scheduler (which need the full list of optimizers)
+        elif isinstance(obj, torch.optim.lr_scheduler._LRScheduler) and not first_pass:
+            scheduler = self.prepare_scheduler(obj)
+            self._schedulers.append(scheduler)
+            return scheduler
         else:
             return obj
 
@@ -362,7 +375,8 @@ def prepare(self, *args):
         if self.distributed_type == DistributedType.DEEPSPEED:
             result = self._prepare_deepspeed(*args)
         else:
-            result = tuple(self._prepare_one(obj) for obj in args)
+            result = tuple(self._prepare_one(obj, first_pass=True) for obj in args)
+            result = tuple(self._prepare_one(obj) for obj in result)
 
         if tpu_should_fix_optimizer:
             # 2. grabbing new model parameters
@@ -492,6 +506,21 @@ def prepare_data_loader(self, data_loader):
     def prepare_optimizer(self, optimizer):
         return AcceleratedOptimizer(optimizer, device_placement=self.device_placement, scaler=self.scaler)
 
+    def prepare_scheduler(self, scheduler):
+        # We try to find the optimizer associated with `scheduler`, the default is the full list.
+        optimizer = self._optimizers
+        for opt in self._optimizers:
+            if getattr(scheduler, "optimizer", None) == opt.optimizer:
+                optimizer = opt
+                break
+
+        return AcceleratedScheduler(
+            scheduler,
+            optimizer,
+            step_with_optimizer=self.step_scheduler_with_optimizer,
+            split_batches=self.split_batches,
+        )
+
     def backward(self, loss, **kwargs):
         """
         Use `accelerator.backward(loss)` in lieu of `loss.backward()`.
@@ -659,7 +688,7 @@ def save_state(self, output_dir: str):
         logger.info(f"Saving current state to {output_dir}")
         weights = [self.get_state_dict(m) for m in self._models]
         save_location = save_accelerator_state(
-            output_dir, weights, self._optimizers, self.state.process_index, self.scaler
+            output_dir, weights, self._optimizers, self._schedulers, self.state.process_index, self.scaler
         )
         for i, obj in enumerate(self._custom_objects):
             save_custom_state(obj, output_dir, i)
@@ -678,7 +707,9 @@ def load_state(self, input_dir: str):
         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)
+        load_accelerator_state(
+            input_dir, self._models, self._optimizers, self._schedulers, self.state.process_index, self.scaler
+        )
         custom_checkpoints = [f for f in os.listdir(input_dir) if "custom_checkpoint" in f]
         if len(custom_checkpoints) != len(self._custom_objects):
             err = "Warning! Number of found checkpoints does not match the number of registered objects:"
@@ -695,6 +726,7 @@ def free_memory(self):
         Will release all references to the internal objects stored and call the garbage collector. You should call this
         method between two trainings with different models/optimizers.
         """
+        self._schedulers = []
         self._optimizers = []
         self._models = []
         self.deepspeed_engine = None
@@ -795,6 +827,6 @@ def optimizer_step_was_skipped(self):
         case the learning rate should not be changed.
         """
         for optimizer in self._optimizers:
-            if optimizer.is_overflow:
+            if optimizer.step_was_skipped:
                 return True
         return False
diff --git a/src/accelerate/checkpointing.py b/src/accelerate/checkpointing.py
index 99f18be43..b35bec237 100644
--- a/src/accelerate/checkpointing.py
+++ b/src/accelerate/checkpointing.py
@@ -22,7 +22,7 @@
 from torch.cuda.amp import GradScaler
 
 from .state import is_tpu_available
-from .utils import MODEL_NAME, OPTIMIZER_NAME, RNG_STATE_NAME, SCALER_NAME, get_pretty_name, save
+from .utils import MODEL_NAME, OPTIMIZER_NAME, RNG_STATE_NAME, SCALER_NAME, SCHEDULER_NAME, get_pretty_name, save
 
 
 if is_tpu_available():
@@ -35,7 +35,12 @@
 
 
 def save_accelerator_state(
-    output_dir: str, model_states: List[dict], optimizers: list, process_index: int, scaler: GradScaler = None
+    output_dir: str,
+    model_states: List[dict],
+    optimizers: list,
+    schedulers: list,
+    process_index: int,
+    scaler: GradScaler = None,
 ):
     """
     Saves the current states of the models, optimizers, scaler, and RNG generators to a given directory.
@@ -47,6 +52,8 @@ def save_accelerator_state(
             A list of model states
         optimizers (`List[torch.optim.Optimizer]`):
             A list of optimizer instances
+        schedulers (`List[torch.optim.lr_scheduler._LRScheduler]`):
+            A list of learning rate schedulers
         process_index (`int`):
             The current process index in the Accelerator state
         scaler (`torch.cuda.amp.GradScaler`, *optional*):
@@ -65,6 +72,13 @@ def save_accelerator_state(
         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}")
+    # Scheduler states
+    for i, scheduler in enumerate(schedulers):
+        state = scheduler.state_dict()
+        scheduler_name = f"{SCHEDULER_NAME}.bin" if i == 0 else f"{SCHEDULER_NAME}_{i}.bin"
+        output_scheduler_file = os.path.join(output_dir, scheduler_name)
+        save(state, output_scheduler_file)
+        logger.info(f"Scheduler state saved in {output_scheduler_file}")
     # GradScaler state
     if scaler is not None:
         state = scaler.state_dict()
@@ -87,7 +101,7 @@ def save_accelerator_state(
     return output_dir
 
 
-def load_accelerator_state(input_dir, models, optimizers, process_index, scaler=None):
+def load_accelerator_state(input_dir, models, optimizers, schedulers, process_index, scaler=None):
     """
     Loads states of the models, optimizers, scaler, and RNG generators from a given directory.
 
@@ -98,6 +112,8 @@ def load_accelerator_state(input_dir, models, optimizers, process_index, scaler=
             A list of model instances
         optimizers (`List[torch.optim.Optimizer]`):
             A list of optimizer instances
+        schedulers (`List[torch.optim.lr_scheduler._LRScheduler]`):
+            A list of learning rate schedulers
         process_index (`int`):
             The current process index in the Accelerator state
         scaler (`torch.cuda.amp.GradScaler`, *optional*):
@@ -117,6 +133,13 @@ def load_accelerator_state(input_dir, models, optimizers, process_index, scaler=
         optimizers[i].load_state_dict(torch.load(input_optimizer_file, map_location="cpu"))
     logger.info("All optimizer states loaded successfully")
 
+    # Scheduler states
+    for i, scheduler in enumerate(schedulers):
+        scheduler_name = f"{SCHEDULER_NAME}.bin" if i == 0 else f"{SCHEDULER_NAME}_{i}.bin"
+        input_scheduler_file = os.path.join(input_dir, scheduler_name)
+        scheduler.load_state_dict(torch.load(input_scheduler_file))
+    logger.info("All scheduler states loaded successfully")
+
     # GradScaler state
     if scaler is not None:
         input_scaler_file = os.path.join(input_dir, SCALER_NAME)
diff --git a/src/accelerate/optimizer.py b/src/accelerate/optimizer.py
index 4360c0c09..1bf37b7d5 100644
--- a/src/accelerate/optimizer.py
+++ b/src/accelerate/optimizer.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import inspect
+import warnings
 
 import torch
 
@@ -141,4 +142,14 @@ def _switch_parameters(self, parameters_map):
     @property
     def is_overflow(self):
         """Whether or not the optimizer step was done, or skipped because of gradient overflow."""
+        warnings.warn(
+            "The `is_overflow` property is deprecated and will be removed in version 1.0 of Accelerate use "
+            "`optimizer.step_was_skipped` instead.",
+            FutureWarning,
+        )
+        return self._is_overflow
+
+    @property
+    def step_was_skipped(self):
+        """Whether or not the optimizer step was skipped."""
         return self._is_overflow
diff --git a/src/accelerate/scheduler.py b/src/accelerate/scheduler.py
new file mode 100644
index 000000000..8efce8dea
--- /dev/null
+++ b/src/accelerate/scheduler.py
@@ -0,0 +1,80 @@
+# 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.
+
+from .state import AcceleratorState
+
+
+class AcceleratedScheduler:
+    """
+    A wrapper around a learning rate scheduler that will only step when the optimizer(s) have a training step. Useful
+    to avoid making a scheduler step too fast when:
+
+    - gradients went overflow and there was no training step (in mixed precision training)
+    - step was skipped because of gradient accumulation
+
+    Args:
+        scheduler (`torch.optim.lr_scheduler._LRScheduler`):
+            The scheduler to wrap.
+        optimizers (one or a list of `torch.optim.Optimizer`):
+            The optimizers used.
+        step_with_optimizer (`bool`, *optional*, defaults to `True`):
+            Whether or not the scheduler should be stepped at each optimizer step.
+        split_batches (`bool`, *optional*, defaults to `False`):
+            Whether or not the dataloaders split one batch across the different processes (so batch size is the same
+            regardless of the number of processes) or create batches on each process (so batch size is the original
+            batch size multiplied by the number of processes).
+    """
+
+    def __init__(self, scheduler, optimizers, step_with_optimizer: bool = True, split_batches: bool = False):
+        self.scheduler = scheduler
+        self.optimizers = optimizers if isinstance(optimizers, (list, tuple)) else [optimizers]
+        self.split_batches = split_batches
+        self.step_with_optimizer = step_with_optimizer
+
+    def step(self, *args, **kwargs):
+        if not self.step_with_optimizer:
+            # No link between scheduler and optimizer -> just step
+            self.scheduler.step(*args, **kwargs)
+            return
+
+        # Otherwise, first make sure the optimizer was stepped.
+        for opt in self.optimizers:
+            if opt.step_was_skipped:
+                return
+
+        if self.split_batches:
+            # Split batches -> the training dataloader batch size is not changed so one step per training step
+            self.scheduler.step(*args, **kwargs)
+        else:
+            # Otherwise the training dataloader batch size was multiplied by `num_processes`, so we need to do
+            # num_processes steps per training step
+            num_processes = AcceleratorState().num_processes
+            for _ in range(num_processes):
+                self.scheduler.step(*args, **kwargs)
+
+    # Passthroughs
+    def get_last_lr(self):
+        return self.scheduler.get_last_lr()
+
+    def state_dict(self):
+        return self.scheduler.state_dict()
+
+    def load_state_dict(self, state_dict):
+        self.scheduler.load_state_dict(state_dict)
+
+    def get_lr(self):
+        return self.scheduler.get_lr()
+
+    def print_lr(self, *args, **kwargs):
+        return self.scheduler.print_lr(*args, **kwargs)
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index c5a291ab8..9bcd87f5f 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -65,6 +65,7 @@ def is_sagemaker_available():
 MODEL_NAME = "pytorch_model"
 RNG_STATE_NAME = "random_states"
 OPTIMIZER_NAME = "optimizer"
+SCHEDULER_NAME = "scheduler"
 
 
 class EnumWithContains(EnumMeta):
diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py
new file mode 100644
index 000000000..8ae9e56ae
--- /dev/null
+++ b/tests/test_scheduler.py
@@ -0,0 +1,62 @@
+# 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 unittest
+from functools import partial
+
+import torch
+
+from accelerate import Accelerator, debug_launcher
+
+
+def scheduler_test(num_processes=2, step_scheduler_with_optimizer=True, split_batches=False):
+    accelerator = Accelerator(step_scheduler_with_optimizer=step_scheduler_with_optimizer, split_batches=split_batches)
+    model = torch.nn.Linear(2, 4)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=1.0)
+    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda n: 1 - n / 10)
+
+    model, optimizer, scheduler = accelerator.prepare(model, optimizer, scheduler)
+
+    # Optimizer has stepped
+    optimizer._is_overflow = False
+    scheduler.step()
+    expected_lr = 1 - (num_processes if (step_scheduler_with_optimizer and not split_batches) else 1) / 10
+    assert (
+        scheduler.get_last_lr()[0] == expected_lr
+    ), f"Wrong lr found at first step, expected {expected_lr}, got {scheduler.get_last_lr()[0]}"
+
+    # Optimizer has not stepped
+    optimizer._is_overflow = True
+    scheduler.step()
+    if not step_scheduler_with_optimizer:
+        expected_lr = 1 - 2 / 10
+    assert (
+        scheduler.get_last_lr()[0] == expected_lr
+    ), f"Wrong lr found at second step, expected {expected_lr}, got {scheduler.get_last_lr()[0]}"
+
+
+class SchedulerTester(unittest.TestCase):
+    def test_scheduler_steps_with_optimizer_single_process(self):
+        debug_launcher(partial(scheduler_test, num_processes=1), num_processes=1)
+        debug_launcher(partial(scheduler_test, num_processes=1, split_batches=True), num_processes=1)
+
+    def test_scheduler_not_step_with_optimizer_single_process(self):
+        debug_launcher(partial(scheduler_test, num_processes=1, step_scheduler_with_optimizer=False), num_processes=1)
+
+    def test_scheduler_steps_with_optimizer_multiprocess(self):
+        debug_launcher(scheduler_test)
+        debug_launcher(partial(scheduler_test, num_processes=1, split_batches=True), num_processes=1)
+
+    def test_scheduler_not_step_with_optimizer_multiprocess(self):
+        debug_launcher(partial(scheduler_test, step_scheduler_with_optimizer=False))

```suggestion
            Set `True` if the learning rate scheduler is stepped at the same time as the optimizer, `False` if only
            done under certain circumstances (at the end of each epoch, for instance).
```
@@ -92,6 +93,9 @@ class Accelerator:
             If set to `True`, the dataloader prepared by the Accelerator is only iterated through on the main process
             and then the batches are split and broadcast to each process. Will default to `True` for `DataLoader` whose
             underlying dataset is an `IterableDataset`, `False` otherwise.
+        step_scheduler_with_optimizer (`bool`, *optional`, defaults to `True`):
+            Set `True` if you step your learning rate scheduler at the same time as the optimizer, `False` if you only
+            do it under certain circumstances (at the end of each epoch, for instance).
Clever! I had wondered how you'd do this 😄 
@@ -458,6 +472,21 @@ def prepare_data_loader(self, data_loader):
     def prepare_optimizer(self, optimizer):
         return AcceleratedOptimizer(optimizer, device_placement=self.device_placement, scaler=self.scaler)
 
+    def prepare_scheduler(self, scheduler):
+        # We try to find the optimizer associated with `scheduler`, the default is the full list.
+        optimizer = self._optimizers
+        for opt in self._optimizers:
+            if getattr(scheduler, "optimizer", None) == opt.optimizer:
+                optimizer = opt
+                break
```suggestion
        logger.info(f"Scheduler state saved in {output_scheduler_file}")
```
@@ -65,6 +72,13 @@ def save_accelerator_state(
         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}")
+    # Scheduler states
+    for i, scheduler in enumerate(schedulers):
+        state = scheduler.state_dict()
+        scheduler_name = f"{SCHEDULER_NAME}.bin" if i == 0 else f"{SCHEDULER_NAME}_{i}.bin"
+        output_scheduler_file = os.path.join(output_dir, scheduler_name)
+        save(state, output_scheduler_file)
+        logger.info(f"Optimizer state saved in {output_scheduler_file}")
```suggestion
        scheduler (`torch.optim.lr_scheduler._LRScheduler`): 
            The scheduler to wrap.
        optimizers (one or a list of `torch.optim.Optimizer`): 
            The optimizers used.
```
@@ -0,0 +1,78 @@
+# 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.
+
+from .state import AcceleratorState
+
+
+class AcceleratedScheduler:
+    """
+    A wrapper around a learning rate scheduler that will only step when the optimizer(s) have a training step. Useful
+    to avoid making a scheduler step too fast when:
+
+    - gradients went overflow and there was no training step (in mixed precision training)
+    - step was skipped because of gradient accumulation
+
+    Args:
+        scheduler (`torch.optim.lr_scheduler._LRScheduler`): The scheduler to wrap.
+        optimizers (one or a list of `torch.optim.Optimizer`): The optimizers used.
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/301
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index dd946c85d..0cbd66884 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -56,7 +56,6 @@ def launch_command_parser(subparsers=None):
     )
     parser.add_argument(
         "--mixed_precision",
-        default="no",
         type=str,
         choices=["no", "fp16", "bf16"],
         help="Whether or not to use mixed precision training. "

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/300
diff --git a/src/accelerate/checkpointing.py b/src/accelerate/checkpointing.py
index 0636dffaa..99f18be43 100644
--- a/src/accelerate/checkpointing.py
+++ b/src/accelerate/checkpointing.py
@@ -107,14 +107,14 @@ def load_accelerator_state(input_dir, models, optimizers, process_index, scaler=
     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))
+        models[i].load_state_dict(torch.load(input_model_file, map_location="cpu"))
     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))
+        optimizers[i].load_state_dict(torch.load(input_optimizer_file, map_location="cpu"))
     logger.info("All optimizer states loaded successfully")
 
     # GradScaler state

_The documentation is not available anymore as the PR was closed or merged._
LG2M! Nice catch! Looks like we might be able to link https://github.com/huggingface/accelerate/issues/296?
That's no linked per se since in this issue, the user was using `optimizer.load_state_dict` and not `accelerator.load_state`.
https://api.github.com/repos/huggingface/accelerate/pulls/299
diff --git a/examples/nlp_example.py b/examples/nlp_example.py
index 2f869bd0f..432d78ae8 100644
--- a/examples/nlp_example.py
+++ b/examples/nlp_example.py
@@ -78,7 +78,7 @@ def tokenize_function(examples):
 
     # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
     # transformers library
-    tokenized_datasets.rename_column_("label", "labels")
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
 
     # If the batch size is too big we use gradient accumulation
     gradient_accumulation_steps = 1

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/298
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index aa04063d8..c213a9167 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -12,6 +12,6 @@ jobs:
       with:
         python-version: 3.6
     - name: Install Python dependencies
-      run: pip install -e .[test]
+      run: pip install setuptools==59.5.0; pip install -e .[test,test_trackers]
     - name: Run Tests
       run: make test
\ No newline at end of file
diff --git a/Makefile b/Makefile
index 943191f62..e5b557f10 100644
--- a/Makefile
+++ b/Makefile
@@ -25,4 +25,4 @@ style:
 	
 # Run tests for the library
 test:
-	python -m pytest -n auto --dist=loadfile -s -v ./tests/
+	python -m pytest -n auto --dist=loadfile -s -v ./tests/
\ No newline at end of file
diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml
index 8f266c600..0a6d45892 100644
--- a/docs/source/_toctree.yml
+++ b/docs/source/_toctree.yml
@@ -21,4 +21,6 @@
     title: Internals
   - local: checkpoint
     title: Checkpointing
+  - local: tracking
+    title: Experiment Tracking
   title: API Reference
diff --git a/docs/source/internal.mdx b/docs/source/internal.mdx
index 3acaaad87..bd6924990 100644
--- a/docs/source/internal.mdx
+++ b/docs/source/internal.mdx
@@ -44,6 +44,10 @@ The main work on your PyTorch `DataLoader` is done by the following function:
 
 [[autodoc]] state.DistributedType
 
+## Tracking
+
+[[autodoc]] tracking.GeneralTracker
+
 ## Utilities
 
 [[autodoc]] utils.extract_model_from_parallel
diff --git a/docs/source/tracking.mdx b/docs/source/tracking.mdx
new file mode 100644
index 000000000..256e2848b
--- /dev/null
+++ b/docs/source/tracking.mdx
@@ -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.
+-->
+
+# Tracking
+
+There are a large number of experiment tracking API's available, however getting them all to work with in a multi-processing environment can oftentimes be complex.
+Accelerate provides a general tracking API that can be used to log useful items during your script through [`~Accelerator.log`]
+
+## Integrated Trackers
+
+Currently `Accelerate` supports three trackers out-of-the-box:
+
+
+[[autodoc]] tracking.TensorBoardTracker
+
+[[autodoc]] tracking.WandBTracker
+
+[[autodoc]] tracking.CometMLTracker
+
+To use any of them, pass in the selected type(s) to the `log_with` parameter in [`Accelerate`]:
+```python
+from accelerate import Accelerate
+from accelerate.utils import LoggerType
+
+accelerator = Accelerate(log_with="all")  # For all available trackers in the environment
+accelerator = Accelerate(log_with="wandb")
+accelerator = Accelerate(log_with=["wandb", LoggerType.TENSORBOARD])
+```
+
+At the start of your experiment [`~Accelerator.init_trackers`] should be used to setup your project, and potentially add any experiment hyperparameters to be logged:
+```python
+hps = {"num_iterations": 5, "learning_rate": 1e-2}
+accelerator.init_trackers("my_project", config=hps)
+```
+
+When you are ready to log any data, [`~Accelerator.log`] should be used.
+A `step` can also be passed in to correlate the data with a particular step in the training loop.
+```python
+accelerator.log({"train_loss": 1.12, "valid_loss": 0.8}, step=1)
+```
+
+Once you've finished training, make sure to run [`~Accelerator.end_training`] so that all the trackers can run their finish functionalities if they have any.
+```python
+accelerator.end_training()
+```
+
+
+A full example is below:
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator(log_with="all")
+config = {
+    "num_iterations": 5,
+    "learning_rate": 1e-2,
+    "loss_function": str(my_loss_function),
+}
+
+accelerator.init_trackers("example_project", config=config)
+
+my_model, my_optimizer, my_training_dataloader = accelerate.prepare(my_model, my_optimizer, my_training_dataloader)
+device = accelerator.device
+my_model.to(device)
+
+for iteration in config["num_iterations"]:
+    for step, batch in my_training_dataloader:
+        my_optimizer.zero_grad()
+        inputs, targets = batch
+        inputs = inputs.to(device)
+        targets = targets.to(device)
+        outputs = my_model(inputs)
+        loss = my_loss_function(outputs, targets)
+        accelerator.backward(loss)
+        my_optimizer.step()
+        accelerator.log({"training_loss": loss}, step=step)
+accelerator.end_training()
+```
+
+
+## Implementing Custom Trackers
+
+To implement a new tracker to be used in `Accelerator`, a new one can be made through implementing the [`~GeneralTracker`] class.
+Every tracker must implement three functions:
+  - `__init__`: 
+    - Should store a `run_name` and initialize the tracker API of the integrated library. 
+    - If a tracker stores their data locally (such as TensorBoard), a `logging_dir` parameter can be added.
+  - `store_init_configuration`: 
+    - Should take in a `values` dictionary and store them as a one-time experiment configuration
+  - `log`: 
+    - Should take in a `values` dictionary and a `step`, and should log them to the run
+
+A brief example can be seen below with an integration with Weights and Biases, containing only the relevent information:
+```python
+from accelerate.tracking import GeneralTracker
+from typing import Optional
+
+import wandb
+
+
+class WandBTracker(GeneralTracker):
+    def __init__(self, run_name: str):
+        self.run_name = run_name
+        wandb.init(self.run_name)
+
+    def store_init_configuration(self, values: dict):
+        wandb.config(values)
+
+    def log(self, values: dict, step: Optional[int] = None):
+        wandb.log(values, step=step)
+```
+
+## When a wrapper cannot work
+
+If a library has an API that does not follow a strict `.log` with an overall dictionary such as Neptune.AI, logging can be done manually under an `if accelerator.is_main_process` statement:
+```diff
+from accelerate import Accelerator
++ import neptune.new as neptune
+
+accelerator = Accelerator()
++ run = neptune.init(...)
+
+my_model, my_optimizer, my_training_dataloader = accelerate.prepare(my_model, my_optimizer, my_training_dataloader)
+device = accelerator.device
+my_model.to(device)
+
+for iteration in config["num_iterations"]:
+    for batch in my_training_dataloader:
+        my_optimizer.zero_grad()
+        inputs, targets = batch
+        inputs = inputs.to(device)
+        targets = targets.to(device)
+        outputs = my_model(inputs)
+        loss = my_loss_function(outputs, targets)
+        total_loss += loss
+        accelerator.backward(loss)
+        my_optimizer.step()
++       if accelerator.is_main_process:
++           run["logs/training/batch/loss"].log(loss)
+```
diff --git a/setup.py b/setup.py
index c898eb5a1..0ce1e5af4 100644
--- a/setup.py
+++ b/setup.py
@@ -20,8 +20,9 @@
 extras["docs"] = []
 extras["test"] = [
     "pytest",
-    "pytest-xdist",
+    "pytest-xdist"
 ]
+extras["test_trackers"] = ["wandb", "comet-ml", "tensorflow"]
 extras["dev"] = extras["quality"] + extras["test"]
 
 extras["sagemaker"] = [
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 2b623db34..f52cf7ca8 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -28,14 +28,19 @@
 from .kwargs_handlers import DistributedDataParallelKwargs, GradScalerKwargs, InitProcessGroupKwargs, KwargsHandler
 from .optimizer import AcceleratedOptimizer
 from .state import AcceleratorState, DistributedType, is_deepspeed_available
+from .tracking import CometMLTracker, TensorBoardTracker, WandBTracker, get_available_trackers
 from .utils import (
     DeepSpeedPlugin,
+    LoggerType,
     PrecisionType,
     RNGType,
     convert_outputs_to_fp32,
     extract_model_from_parallel,
     gather,
     get_pretty_name,
+    is_comet_ml_available,
+    is_tensorboard_available,
+    is_wandb_available,
     pad_across_processes,
     save,
     wait_for_everyone,
@@ -88,6 +93,17 @@ class Accelerator:
               dataloader) or of the iterable dataset (if it exists) if the underlying dataset is of that type.
 
             Will default to `["torch"]` for PyTorch versions <=1.5.1 and `["generator"]` for PyTorch versions >= 1.6.
+        log_with (list of `str` or [`~utils.LoggerType`], *optional*):
+            A list of loggers to be setup for experiment tracking. Should be one or several of:
+
+            - `"all"`
+            - `"tensorboard"`
+            - `"wandb"`
+            - `"comet_ml"`
+
+            If `"all`" is selected, will pick up all available trackers in the environment and intialize them.
+        logging_dir (`str`, `os.PathLike`, *optional*):
+            A path to a directory for storing logs of locally-compatible loggers.
         dispatch_batches (`bool`, *optional*):
             If set to `True`, the dataloader prepared by the Accelerator is only iterated through on the main process
             and then the batches are split and broadcast to each process. Will default to `True` for `DataLoader` whose
@@ -111,9 +127,27 @@ def __init__(
         cpu: bool = False,
         deepspeed_plugin: DeepSpeedPlugin = None,
         rng_types: Optional[List[Union[str, RNGType]]] = None,
+        log_with: Optional[List[Union[str, LoggerType]]] = None,
+        logging_dir: Optional[Union[str, os.PathLike]] = "",
         dispatch_batches: Optional[bool] = None,
         kwargs_handlers: Optional[List[KwargsHandler]] = None,
     ):
+        if log_with is not None:
+            if not isinstance(log_with, list):
+                log_with = [log_with]
+            if "all" in log_with or LoggerType.ALL in log_with:
+                log_with = get_available_trackers()
+            else:
+                for i, log_type in enumerate(log_with):
+                    if log_type not in LoggerType:
+                        raise ValueError(
+                            f"Unsupported logging capability: {log_type}. Choose between {LoggerType.list()}"
+                        )
+                    log_with[i] = LoggerType(log_type)
+                log_with = list(set(log_with))
+        self.log_with = log_with
+        self.logging_dir = logging_dir
+
         if mixed_precision is not None:
             mixed_precision = str(mixed_precision)
             if mixed_precision not in PrecisionType:
@@ -556,6 +590,50 @@ def wait_for_everyone(self):
         """
         wait_for_everyone()
 
+    def init_trackers(self, project_name: str, config: Optional[dict] = None):
+        """
+        Initializes a run for all trackers stored in `self.log_with`, potentially with starting configurations
+
+        Args:
+            project_name (`str`):
+                The name of the project. All trackers will save their data based on this
+            config (`dict`, *optional*):
+                Optional starting configuration to be logged.
+        """
+        self.trackers = []
+        for tracker in self.log_with:
+            if str(tracker).lower() == "tensorboard" and is_tensorboard_available():
+                self.trackers.append(TensorBoardTracker(project_name, self.logging_dir))
+            elif str(tracker).lower() == "wandb" and is_wandb_available():
+                self.trackers.append(WandBTracker(project_name))
+            elif str(tracker).lower() == "comet_ml" and is_comet_ml_available():
+                self.trackers.append(CometMLTracker(project_name))
+        if config is not None:
+            for tracker in self.trackers:
+                tracker.store_init_configuration(config)
+
+    def log(self, values: dict, step: Optional[int] = None):
+        """
+        Logs `values` to all stored trackers in `self.trackers`.
+
+        Args:
+            values (`dict`):
+                Values should be a dictionary-like object containing only types `int`, `float`, or `str`.
+            step (`int`, *optional*):
+                The run step. If included, the log will be affiliated with this step.
+        """
+        if self.is_main_process:
+            for tracker in self.trackers:
+                tracker.log(values, step=step)
+
+    def end_training(self):
+        """
+        Runs any special end training behaviors, such as stopping trackers
+        """
+        if self.is_main_process:
+            for tracker in self.trackers:
+                tracker.finish()
+
     def save(self, obj, f):
         """
         Save the object passed to disk once per machine. Use in place of `torch.save`.
diff --git a/src/accelerate/test_utils/testing.py b/src/accelerate/test_utils/testing.py
index 507f16eb8..cce7347f6 100644
--- a/src/accelerate/test_utils/testing.py
+++ b/src/accelerate/test_utils/testing.py
@@ -19,7 +19,7 @@
 import torch
 
 from ..state import AcceleratorState, is_tpu_available
-from ..utils import gather
+from ..utils import gather, is_tensorflow_available
 
 
 def are_the_same_tensors(tensor):
@@ -64,6 +64,17 @@ def require_multi_gpu(test_case):
         return test_case
 
 
+def require_tensorflow(test_case):
+    """
+    Decorator marking a test that requires TensorFlow installed. These tests are skipped when TensorFlow isn't
+    installed
+    """
+    if not is_tensorflow_available():
+        return unittest.skip("test requires TensorFlow")(test_case)
+    else:
+        return test_case
+
+
 class _RunOutput:
     def __init__(self, returncode, stdout, stderr):
         self.returncode = returncode
diff --git a/src/accelerate/tracking.py b/src/accelerate/tracking.py
new file mode 100644
index 000000000..65a5caec2
--- /dev/null
+++ b/src/accelerate/tracking.py
@@ -0,0 +1,245 @@
+# 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.
+
+# Expectation:
+# Provide a project dir name, then each type of logger gets stored in project/{`logging_dir`}
+
+import logging
+import os
+from abc import ABCMeta, abstractmethod
+from typing import Optional, Union
+
+from .utils import LoggerType, is_comet_ml_available, is_tensorboard_available, is_wandb_available
+
+
+_available_trackers = []
+
+if is_tensorboard_available():
+    from torch.utils import tensorboard
+
+    _available_trackers.append(LoggerType.TENSORBOARD)
+
+if is_wandb_available():
+    import wandb
+
+    _available_trackers.append(LoggerType.WANDB)
+
+if is_comet_ml_available():
+    from comet_ml import Experiment
+
+    _available_trackers.append(LoggerType.COMETML)
+
+
+logger = logging.getLogger(__name__)
+
+
+def get_available_trackers():
+    "Returns a list of all supported available trackers in the system"
+    return _available_trackers
+
+
+class GeneralTracker(object, metaclass=ABCMeta):
+    """
+    A base Tracker class to be used for all logging integration implementations.
+    """
+
+    @abstractmethod
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Implementations should use the experiment configuration
+        functionality of a tracking API.
+
+        Args:
+            values (Dictionary `str` to `bool`, `str`, `float` or `int`):
+                Values to be stored as initial hyperparameters as key-value pairs. The values need to have type `bool`,
+                `str`, `float`, `int`, or `None`.
+        """
+        pass
+
+    @abstractmethod
+    def log(self, values: dict, step: Optional[int]):
+        """
+        Logs `values` to the current run. Base `log` implementations of a tracking API should go in here, along with
+        special behavior for the `step parameter.
+
+        Args:
+            values (Dictionary `str` to `str`, `float`, or `int`):
+                Values to be logged as key-value pairs. The values need to have type `str`, `float`, or `int`.
+            step (`int`, *optional*):
+                The run step. If included, the log will be affiliated with this step.
+        """
+        pass
+
+    def finish(self):
+        """
+        Should run any finalizing functions within the tracking API. If the API should not have one, just don't
+        overwrite that method.
+        """
+        pass
+
+
+class TensorBoardTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `tensorboard`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run
+        logging_dir (`str`, `os.PathLike`):
+            Location for TensorBoard logs to be stored.
+    """
+
+    def __init__(self, run_name: str, logging_dir: Optional[Union[str, os.PathLike]] = ""):
+        self.run_name = run_name
+        self.logging_dir = os.path.join(logging_dir, run_name)
+        self.writer = tensorboard.SummaryWriter(self.logging_dir)
+        logger.info(f"Initialized TensorBoard project {self.run_name} logging to {self.logging_dir}")
+        logger.info(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment.
+
+        Args:
+            values (Dictionary `str` to `bool`, `str`, `float` or `int`):
+                Values to be stored as initial hyperparameters as key-value pairs. The values need to have type `bool`,
+                `str`, `float`, `int`, or `None`.
+        """
+        self.writer.add_hparams(values, metric_dict={})
+        self.writer.flush()
+        logger.info("Stored initial configuration hyperparameters to TensorBoard")
+
+    def log(self, values: dict, step: Optional[int] = None):
+        """
+        Logs `values` to the current run.
+
+        Args:
+            values (Dictionary `str` to `str`, `float`, or `int`):
+                Values to be logged as key-value pairs. The values need to have type `str`, `float`, or `int`.
+            step (`int`, *optional*):
+                The run step. If included, the log will be affiliated with this step.
+        """
+        for k, v in values.items():
+            if isinstance(v, (int, float)):
+                self.writer.add_scalar(k, v, global_step=step)
+            elif isinstance(v, str):
+                self.writer.add_text(k, v, global_step=step)
+        self.writer.flush()
+        logger.info("Successfully logged to TensorBoard")
+
+    def finish(self):
+        """
+        Closes `TensorBoard` writer
+        """
+        self.writer.close()
+        logger.info("TensorBoard writer closed")
+
+
+class WandBTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `wandb`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run.
+    """
+
+    def __init__(self, run_name: str):
+        self.run_name = run_name
+        self.run = wandb.init(self.run_name)
+        logger.info(f"Initialized WandB project {self.run_name}")
+        logger.info(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment.
+
+        Args:
+            values (Dictionary `str` to `bool`, `str`, `float` or `int`):
+                Values to be stored as initial hyperparameters as key-value pairs. The values need to have type `bool`,
+                `str`, `float`, `int`, or `None`.
+        """
+        wandb.config.update(values)
+        logger.info("Stored initial configuration hyperparameters to WandB")
+
+    def log(self, values: dict, step: Optional[int] = None):
+        """
+        Logs `values` to the current run.
+
+        Args:
+            values (Dictionary `str` to `str`, `float`, or `int`):
+                Values to be logged as key-value pairs. The values need to have type `str`, `float`, or `int`.
+            step (`int`, *optional*):
+                The run step. If included, the log will be affiliated with this step.
+        """
+        self.run.log(values, step=step)
+        logger.info("Successfully logged to WandB")
+
+    def finish(self):
+        """
+        Closes `wandb` writer
+        """
+        self.run.finish()
+        logger.info("WandB run closed")
+
+
+class CometMLTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `comet_ml`. Should be initialized at the start of your script.
+
+    API keys must be stored in a Comet config file.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run.
+    """
+
+    def __init__(self, run_name: str):
+        self.run_name = run_name
+        self.writer = Experiment(project_name=run_name)
+        logger.info(f"Initialized CometML project {self.run_name}")
+        logger.info(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment.
+
+        Args:
+            values (Dictionary `str` to `bool`, `str`, `float` or `int`):
+                Values to be stored as initial hyperparameters as key-value pairs. The values need to have type `bool`,
+                `str`, `float`, `int`, or `None`.
+        """
+        self.writer.log_parameters(values)
+        logger.info("Stored initial configuration hyperparameters to CometML")
+
+    def log(self, values: dict, step: Optional[int] = None):
+        """
+        Logs `values` to the current run.
+
+        Args:
+            values (Dictionary `str` to `str`, `float`, or `int`):
+                Values to be logged as key-value pairs. The values need to have type `str`, `float`, or `int`.
+            step (`int`, *optional*):
+                The run step. If included, the log will be affiliated with this step.
+        """
+        if step is not None:
+            self.writer.set_step(step)
+        self.writer.log_others(values)
+        logger.info("Successfully logged to CometML")
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index af9299a4f..c5a291ab8 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -34,6 +34,22 @@
     import torch_xla.core.xla_model as xm
 
 
+def is_tensorflow_available():
+    return importlib.util.find_spec("tensorflow") is not None
+
+
+def is_tensorboard_available():
+    return importlib.util.find_spec("tensorboard") is not None or importlib.util.find_spec("tensorboardX") is not None
+
+
+def is_wandb_available():
+    return importlib.util.find_spec("wandb") is not None
+
+
+def is_comet_ml_available():
+    return importlib.util.find_spec("comet_ml") is not None
+
+
 def is_boto3_available():
     return importlib.util.find_spec("boto3") is not None
 
@@ -74,6 +90,13 @@ def list(cls):
         return list(map(lambda item: str(item), cls))
 
 
+class LoggerType(BaseEnum):
+    ALL = "all"
+    TENSORBOARD = "tensorboard"
+    WANDB = "wandb"
+    COMETML = "comet_ml"
+
+
 class PrecisionType(BaseEnum):
     NO = "no"
     FP16 = "fp16"
diff --git a/tests/test_tracking.py b/tests/test_tracking.py
new file mode 100644
index 000000000..8d5ebc6bc
--- /dev/null
+++ b/tests/test_tracking.py
@@ -0,0 +1,176 @@
+# 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 re
+import shutil
+import tempfile
+import unittest
+from pathlib import Path
+from unittest import mock
+
+# We use TF to parse the logs
+from accelerate import Accelerator
+from accelerate.test_utils.testing import require_tensorflow
+from accelerate.utils import is_tensorflow_available
+
+
+if is_tensorflow_available():
+    import tensorflow as tf
+    from tensorboard.plugins.hparams import plugin_data_pb2
+    from tensorflow.core.util import event_pb2
+    from tensorflow.python.summary.summary_iterator import summary_iterator
+
+
+logger = logging.getLogger(__name__)
+
+
+class TensorBoardTrackingTest(unittest.TestCase):
+    @require_tensorflow
+    def test_init_trackers(self):
+        hps = None
+        project_name = "test_project_with_config"
+        with tempfile.TemporaryDirectory() as dirpath:
+            oldpwd = os.getcwd()
+            os.chdir(dirpath)
+            accelerator = Accelerator(log_with="tensorboard")
+            config = {"num_iterations": 12, "learning_rate": 1e-2, "some_boolean": False, "some_string": "some_value"}
+            accelerator.init_trackers(project_name, config)
+            accelerator.end_training()
+            for child in Path(project_name).glob("*/**"):
+                log = list(filter(lambda x: x.is_file(), child.iterdir()))[0]
+                # The config log is stored one layer deeper in the logged directory
+                # And names are randomly generated each time
+            si = summary_iterator(str(log))
+            # Pull HPS through careful parsing
+            for event in si:
+                for value in event.summary.value:
+                    proto_bytes = value.metadata.plugin_data.content
+                    plugin_data = plugin_data_pb2.HParamsPluginData.FromString(proto_bytes)
+                    if plugin_data.HasField("session_start_info"):
+                        hps = dict(plugin_data.session_start_info.hparams)
+            os.chdir(oldpwd)
+
+        self.assertTrue(isinstance(hps, dict))
+        keys = list(hps.keys())
+        keys.sort()
+        self.assertEqual(keys, ["learning_rate", "num_iterations", "some_boolean", "some_string"])
+        self.assertEqual(hps["num_iterations"].number_value, 12)
+        self.assertEqual(hps["learning_rate"].number_value, 0.01)
+        self.assertEqual(hps["some_boolean"].bool_value, False)
+        self.assertEqual(hps["some_string"].string_value, "some_value")
+
+    @require_tensorflow
+    def test_log(self):
+        step = None
+        project_name = "test_project_with_log"
+        with tempfile.TemporaryDirectory() as dirpath:
+            oldpwd = os.getcwd()
+            os.chdir(dirpath)
+            accelerator = Accelerator(log_with="tensorboard")
+            accelerator.init_trackers(project_name)
+            values = {"total_loss": 0.1, "iteration": 1, "my_text": "some_value"}
+            accelerator.log(values, step=0)
+            accelerator.end_training()
+            # Logged values are stored in the outermost-tfevents file and can be read in as a TFRecord
+            # Names are randomly generated each time
+            log = list(filter(lambda x: x.is_file(), Path(project_name).iterdir()))[0]
+            serialized_examples = tf.data.TFRecordDataset(log)
+            for e in serialized_examples:
+                event = event_pb2.Event.FromString(e.numpy())
+                if step is None:
+                    step = event.step
+                for value in event.summary.value:
+                    if value.tag == "total_loss":
+                        total_loss = value.simple_value
+                    elif value.tag == "iteration":
+                        iteration = value.simple_value
+                    elif value.tag == "my_text/text_summary":  # Append /text_summary to the key
+                        my_text = value.tensor.string_val[0].decode()
+            os.chdir(oldpwd)
+        self.assertAlmostEqual(total_loss, values["total_loss"])
+        self.assertEqual(iteration, values["iteration"])
+        self.assertEqual(my_text, values["my_text"])
+
+
+@mock.patch.dict(os.environ, {"WANDB_MODE": "offline"})
+class WandBTrackingTest(unittest.TestCase):
+    @staticmethod
+    def get_value_from_log(key: str, log: str, key_occurance: int = 0):
+        """
+        Parses wandb log for `key` and returns the value.
+        If parsing through multiple calls to .log, pass in a `key_occurance`
+        """
+        res = re.findall(rf"(?<={key} )[^\s]+", log)[key_occurance]
+        if '"' in res:
+            return re.findall(r'"([^"]*)"', res)[0]
+        else:
+            return res
+
+    def test_init_trackers(self):
+        project_name = "test_project_with_config"
+        accelerator = Accelerator(log_with="wandb")
+        config = {"num_iterations": 12, "learning_rate": 1e-2, "some_boolean": False, "some_string": "some_value"}
+        accelerator.init_trackers(project_name, config)
+        accelerator.end_training()
+        # The latest offline log is stored at wandb/latest-run/*.wandb
+        for child in Path("wandb/latest-run").glob("*"):
+            logger.info(child)
+            if child.is_file() and child.suffix == ".wandb":
+                with open(child, "rb") as f:
+                    content = f.read()
+                break
+
+        # Check HPS through careful parsing and cleaning
+        cleaned_log = re.sub(r"[\x00-\x1f]+", " ", content.decode("utf8", "ignore"))
+        self.assertEqual(self.get_value_from_log("num_iterations", cleaned_log), "12")
+        self.assertEqual(self.get_value_from_log("learning_rate", cleaned_log), "0.01")
+        self.assertEqual(self.get_value_from_log("some_boolean", cleaned_log), "false")
+        self.assertEqual(self.get_value_from_log("some_string", cleaned_log), "some_value")
+
+    def test_log(self):
+        project_name = "test_project_with_log"
+        accelerator = Accelerator(log_with="wandb")
+        accelerator.init_trackers(project_name)
+        values = {"total_loss": 0.1, "iteration": 1, "my_text": "some_value"}
+        accelerator.log(values, step=0)
+        accelerator.end_training()
+        # The latest offline log is stored at wandb/latest-run/*.wandb
+        for child in Path("wandb/latest-run").glob("*"):
+            if child.is_file() and child.suffix == ".wandb":
+                with open(child, "rb") as f:
+                    content = f.read()
+                break
+        # Check HPS through careful parsing and cleaning
+        cleaned_log = re.sub(r"[\x00-\x1f]+", " ", content.decode("utf8", "ignore"))
+        self.assertEqual(self.get_value_from_log("total_loss", cleaned_log), "0.1")
+        self.assertEqual(self.get_value_from_log("iteration", cleaned_log), "1")
+        self.assertEqual(self.get_value_from_log("my_text", cleaned_log), "some_value")
+        self.assertEqual(self.get_value_from_log("_step", cleaned_log), "0")
+
+    def setUp(self):
+        os.mkdir(".wandb_tests")
+        os.chdir(".wandb_tests")
+
+    def tearDown(self):
+        if os.getcwd().endswith(".wandb_tests"):
+            os.chdir("..")
+        if os.path.exists(".wandb_tests"):
+            shutil.rmtree(".wandb_tests")
+
+    @classmethod
+    def setUpClass(cls):
+        if os.path.exists(".wandb_tests"):
+            shutil.rmtree(".wandb_tests")

```suggestion
from .tracking import TensorBoardTracker
```
:scream: 
@@ -17,10 +17,12 @@
 import sys
 import warnings
 from contextlib import contextmanager
+from pathlib import Path
 from typing import List, Optional, Union
 
 import torch
 
+from accelerate.accelerate.src.accelerate.tracking import TensorBoardTracker
The `"all"` should be dealt with around here too.
@@ -111,9 +115,17 @@ def __init__(
         cpu: bool = False,
         deepspeed_plugin: DeepSpeedPlugin = None,
         rng_types: Optional[List[Union[str, RNGType]]] = None,
+        log_with: Optional[List[Union[LoggerType, str]]] = None,
         dispatch_batches: Optional[bool] = None,
         kwargs_handlers: Optional[List[KwargsHandler]] = None,
     ):
+        if log_with is not None:
+            if not isinstance(log_with, list):
Nice!
@@ -556,6 +568,30 @@ def wait_for_everyone(self):
         """
         wait_for_everyone()
 
+    def init_trackers(self, project_name: str, config: dict = None):
+        """
+        Initializes a run for all trackers stored in `self.log_with`, potentially with starting configurations
+
+        Args:
+            project_name (`str`):
+        """
+        project_location = Path(project_name)
+        self.trackers = []
+        for tracker in self.log_with:
+            if tracker.lower() == "tensorboard" and is_tensorboard_available():
+                self.trackers.append(TensorBoardTracker(project_location))
+        if config is not None:
+            for tracker in self.trackers:
+                tracker.store_init_configuration(config)
Should be inside an `is_tensorboard_available` flag.
@@ -0,0 +1,72 @@
+# 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.
+
+# Expectation:
+# Provide a project dir name, then each type of logger gets stored in project/{`logging_dir`}
+
+import logging
+from abc import ABCMeta, abstractmethod
+from pathlib import Path
+
+from torch.utils import tensorboard
Accelerate has no idea about the global step though. We could add it in an optional argument to `log`? Or add something that keeps track of the number of time the `optimizer.step` is called internally.
@@ -0,0 +1,72 @@
+# 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.
+
+# Expectation:
+# Provide a project dir name, then each type of logger gets stored in project/{`logging_dir`}
+
+import logging
+from abc import ABCMeta, abstractmethod
+from pathlib import Path
+
+from torch.utils import tensorboard
+
+
+logger = logging.getLogger(__name__)
+
+
+class GeneralTracker(object, metaclass=ABCMeta):
+    """
+    A base Tracker class to be used for all logging integration implementations.
+    """
+
+    log_directory: str
+
+    @abstractmethod
+    def store_init_configuration(self, values):
+        pass
+
+    @abstractmethod
+    def log(self, values):
+        pass
+
+
+class TensorBoardTracker(GeneralTracker):
+    log_directory = "tensorboard"
+
+    def __init__(self, run_name=""):
+        self.run_name = Path(run_name)
+        self.writer = tensorboard.SummaryWriter(self.logging_dir)
+        logger.info(f"Initialized TensorBoard project {self.run_name} writing to {self.logging_dir}")
+        logger.info(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    def store_init_configuration(self, values):
+        self.writer.add_hparams(values, metric_dict={})
+        self.writer.flush()
+        logger.info("Stored initial configuration hyperparameters to TensorBoard")
+
+    # add_scalar has the option for `global_step`, should we include this?
That's more or less what my thinking is, having it be an optional argument. I don't think there's a need for us to keep track of `optimizer.step` calls internally for this yet, but we'll see what changes as this develops
@@ -0,0 +1,72 @@
+# 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.
+
+# Expectation:
+# Provide a project dir name, then each type of logger gets stored in project/{`logging_dir`}
+
+import logging
+from abc import ABCMeta, abstractmethod
+from pathlib import Path
+
+from torch.utils import tensorboard
+
+
+logger = logging.getLogger(__name__)
+
+
+class GeneralTracker(object, metaclass=ABCMeta):
+    """
+    A base Tracker class to be used for all logging integration implementations.
+    """
+
+    log_directory: str
+
+    @abstractmethod
+    def store_init_configuration(self, values):
+        pass
+
+    @abstractmethod
+    def log(self, values):
+        pass
+
+
+class TensorBoardTracker(GeneralTracker):
+    log_directory = "tensorboard"
+
+    def __init__(self, run_name=""):
+        self.run_name = Path(run_name)
+        self.writer = tensorboard.SummaryWriter(self.logging_dir)
+        logger.info(f"Initialized TensorBoard project {self.run_name} writing to {self.logging_dir}")
+        logger.info(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    def store_init_configuration(self, values):
+        self.writer.add_hparams(values, metric_dict={})
+        self.writer.flush()
+        logger.info("Stored initial configuration hyperparameters to TensorBoard")
+
+    # add_scalar has the option for `global_step`, should we include this?
*That is what we call VSCode not being friendly and me missing it!*
@@ -17,10 +17,12 @@
 import sys
 import warnings
 from contextlib import contextmanager
+from pathlib import Path
 from typing import List, Optional, Union
 
 import torch
 
+from accelerate.accelerate.src.accelerate.tracking import TensorBoardTracker
Decided to make it an optional argument. If folks seem to find it annoying then we can go ahead and track `optimizer.step` next
@@ -0,0 +1,72 @@
+# 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.
+
+# Expectation:
+# Provide a project dir name, then each type of logger gets stored in project/{`logging_dir`}
+
+import logging
+from abc import ABCMeta, abstractmethod
+from pathlib import Path
+
+from torch.utils import tensorboard
+
+
+logger = logging.getLogger(__name__)
+
+
+class GeneralTracker(object, metaclass=ABCMeta):
+    """
+    A base Tracker class to be used for all logging integration implementations.
+    """
+
+    log_directory: str
+
+    @abstractmethod
+    def store_init_configuration(self, values):
+        pass
+
+    @abstractmethod
+    def log(self, values):
+        pass
+
+
+class TensorBoardTracker(GeneralTracker):
+    log_directory = "tensorboard"
+
+    def __init__(self, run_name=""):
+        self.run_name = Path(run_name)
+        self.writer = tensorboard.SummaryWriter(self.logging_dir)
+        logger.info(f"Initialized TensorBoard project {self.run_name} writing to {self.logging_dir}")
+        logger.info(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    def store_init_configuration(self, values):
+        self.writer.add_hparams(values, metric_dict={})
+        self.writer.flush()
+        logger.info("Stored initial configuration hyperparameters to TensorBoard")
+
+    # add_scalar has the option for `global_step`, should we include this?
```suggestion
            values (Dictionary `str` to `bool`, `str`, `float` or `int`):
                Values to be stored as initial hyperparameters as key-value pairs. The values need to have type `bool`, `str`,
                `float`, `int`, or `None`.
```
@@ -0,0 +1,210 @@
+# 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.
+
+# Expectation:
+# Provide a project dir name, then each type of logger gets stored in project/{`logging_dir`}
+
+import logging
+from abc import ABCMeta, abstractmethod
+from pathlib import Path
+from typing import Optional
+
+from .utils import LoggerType, is_comet_ml_available, is_tensorboard_available, is_wandb_available
+
+
+_available_trackers = []
+
+if is_tensorboard_available():
+    from torch.utils import tensorboard
+
+    _available_trackers.append(LoggerType.TENSORBOARD)
+
+if is_wandb_available():
+    import wandb
+
+    _available_trackers.append(LoggerType.WANDB)
+
+if is_comet_ml_available():
+    from comet_ml import Experiment
+
+    _available_trackers.append(LoggerType.COMETML)
+
+
+logger = logging.getLogger(__name__)
+
+
+def get_available_trackers():
+    "Returns a list of all supported available trackers in the system"
+    return _available_trackers
+
+
+class GeneralTracker(object, metaclass=ABCMeta):
+    """
+    A base Tracker class to be used for all logging integration implementations.
+    """
+
+    log_directory: str
+
+    @abstractmethod
+    def store_init_configuration(self, values):
+        pass
+
+    @abstractmethod
+    def log(self, values, step: Optional[int]):
+        pass
+
+
+class TensorBoardTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `tensorboard`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run. Logs are then stored in `tensorboard/{run_name}`
+    """
+
+    log_directory = Path("tensorboard")
+
+    def __init__(self, run_name: str = ""):
+        self.run_name = run_name
+        self.writer = tensorboard.SummaryWriter(self.logging_dir / self.run_name)
+        logger.info(f"Initialized TensorBoard project {self.run_name} writing to {self.logging_dir}")
+        logger.info(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment.
+
+        Args:
+            values (`dict`):
+                Values to be stored as initial hyperparameters as key-value pairs. Value be of type `bool`, `str`,
+                `float`, `int`, or `None`.
Seems to be optional and has a default ;-)
@@ -0,0 +1,210 @@
+# 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.
+
+# Expectation:
+# Provide a project dir name, then each type of logger gets stored in project/{`logging_dir`}
+
+import logging
+from abc import ABCMeta, abstractmethod
+from pathlib import Path
+from typing import Optional
+
+from .utils import LoggerType, is_comet_ml_available, is_tensorboard_available, is_wandb_available
+
+
+_available_trackers = []
+
+if is_tensorboard_available():
+    from torch.utils import tensorboard
+
+    _available_trackers.append(LoggerType.TENSORBOARD)
+
+if is_wandb_available():
+    import wandb
+
+    _available_trackers.append(LoggerType.WANDB)
+
+if is_comet_ml_available():
+    from comet_ml import Experiment
+
+    _available_trackers.append(LoggerType.COMETML)
+
+
+logger = logging.getLogger(__name__)
+
+
+def get_available_trackers():
+    "Returns a list of all supported available trackers in the system"
+    return _available_trackers
+
+
+class GeneralTracker(object, metaclass=ABCMeta):
+    """
+    A base Tracker class to be used for all logging integration implementations.
+    """
+
+    log_directory: str
+
+    @abstractmethod
+    def store_init_configuration(self, values):
+        pass
+
+    @abstractmethod
+    def log(self, values, step: Optional[int]):
+        pass
+
+
+class TensorBoardTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `tensorboard`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
```suggestion
            step (`int`, *optional*):
```
@@ -0,0 +1,210 @@
+# 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.
+
+# Expectation:
+# Provide a project dir name, then each type of logger gets stored in project/{`logging_dir`}
+
+import logging
+from abc import ABCMeta, abstractmethod
+from pathlib import Path
+from typing import Optional
+
+from .utils import LoggerType, is_comet_ml_available, is_tensorboard_available, is_wandb_available
+
+
+_available_trackers = []
+
+if is_tensorboard_available():
+    from torch.utils import tensorboard
+
+    _available_trackers.append(LoggerType.TENSORBOARD)
+
+if is_wandb_available():
+    import wandb
+
+    _available_trackers.append(LoggerType.WANDB)
+
+if is_comet_ml_available():
+    from comet_ml import Experiment
+
+    _available_trackers.append(LoggerType.COMETML)
+
+
+logger = logging.getLogger(__name__)
+
+
+def get_available_trackers():
+    "Returns a list of all supported available trackers in the system"
+    return _available_trackers
+
+
+class GeneralTracker(object, metaclass=ABCMeta):
+    """
+    A base Tracker class to be used for all logging integration implementations.
+    """
+
+    log_directory: str
+
+    @abstractmethod
+    def store_init_configuration(self, values):
+        pass
+
+    @abstractmethod
+    def log(self, values, step: Optional[int]):
+        pass
+
+
+class TensorBoardTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `tensorboard`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run. Logs are then stored in `tensorboard/{run_name}`
+    """
+
+    log_directory = Path("tensorboard")
+
+    def __init__(self, run_name: str = ""):
+        self.run_name = run_name
+        self.writer = tensorboard.SummaryWriter(self.logging_dir / self.run_name)
+        logger.info(f"Initialized TensorBoard project {self.run_name} writing to {self.logging_dir}")
+        logger.info(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment.
+
+        Args:
+            values (`dict`):
+                Values to be stored as initial hyperparameters as key-value pairs. Value be of type `bool`, `str`,
+                `float`, `int`, or `None`.
+        """
+        self.writer.add_hparams(values, metric_dict={})
+        self.writer.flush()
+        logger.info("Stored initial configuration hyperparameters to TensorBoard")
+
+    def log(self, values: dict, step: Optional[int] = None):
+        """
+        Logs `values` to the current run.
+
+        Args:
+            values (`dict`):
+                Values to be logged as key-value pairs. Value be of type `int`, `float`, or `str`.
+            step (`int`, Optional):
Same comment as above for the `run_name` of Tensorboard.
@@ -0,0 +1,210 @@
+# 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.
+
+# Expectation:
+# Provide a project dir name, then each type of logger gets stored in project/{`logging_dir`}
+
+import logging
+from abc import ABCMeta, abstractmethod
+from pathlib import Path
+from typing import Optional
+
+from .utils import LoggerType, is_comet_ml_available, is_tensorboard_available, is_wandb_available
+
+
+_available_trackers = []
+
+if is_tensorboard_available():
+    from torch.utils import tensorboard
+
+    _available_trackers.append(LoggerType.TENSORBOARD)
+
+if is_wandb_available():
+    import wandb
+
+    _available_trackers.append(LoggerType.WANDB)
+
+if is_comet_ml_available():
+    from comet_ml import Experiment
+
+    _available_trackers.append(LoggerType.COMETML)
+
+
+logger = logging.getLogger(__name__)
+
+
+def get_available_trackers():
+    "Returns a list of all supported available trackers in the system"
+    return _available_trackers
+
+
+class GeneralTracker(object, metaclass=ABCMeta):
+    """
+    A base Tracker class to be used for all logging integration implementations.
+    """
+
+    log_directory: str
+
+    @abstractmethod
+    def store_init_configuration(self, values):
+        pass
+
+    @abstractmethod
+    def log(self, values, step: Optional[int]):
+        pass
+
+
+class TensorBoardTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `tensorboard`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run. Logs are then stored in `tensorboard/{run_name}`
+    """
+
+    log_directory = Path("tensorboard")
+
+    def __init__(self, run_name: str = ""):
+        self.run_name = run_name
+        self.writer = tensorboard.SummaryWriter(self.logging_dir / self.run_name)
+        logger.info(f"Initialized TensorBoard project {self.run_name} writing to {self.logging_dir}")
+        logger.info(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment.
+
+        Args:
+            values (`dict`):
+                Values to be stored as initial hyperparameters as key-value pairs. Value be of type `bool`, `str`,
+                `float`, `int`, or `None`.
+        """
+        self.writer.add_hparams(values, metric_dict={})
+        self.writer.flush()
+        logger.info("Stored initial configuration hyperparameters to TensorBoard")
+
+    def log(self, values: dict, step: Optional[int] = None):
+        """
+        Logs `values` to the current run.
+
+        Args:
+            values (`dict`):
+                Values to be logged as key-value pairs. Value be of type `int`, `float`, or `str`.
+            step (`int`, Optional):
+                The run step. If included, the log will be affiliated with this step.
+        """
+        for k, v in values.items():
+            if isinstance(v, (int, float)):
+                self.writer.add_scalar(k, v, global_step=step)
+            elif isinstance(v, str):
+                self.writer.add_text(k, v, global_step=step)
+        self.writer.flush()
+        logger.info("Successfully logged to TensorBoard")
+
+
+class WandBTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `wandb`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
Same comments as above.
@@ -0,0 +1,210 @@
+# 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.
+
+# Expectation:
+# Provide a project dir name, then each type of logger gets stored in project/{`logging_dir`}
+
+import logging
+from abc import ABCMeta, abstractmethod
+from pathlib import Path
+from typing import Optional
+
+from .utils import LoggerType, is_comet_ml_available, is_tensorboard_available, is_wandb_available
+
+
+_available_trackers = []
+
+if is_tensorboard_available():
+    from torch.utils import tensorboard
+
+    _available_trackers.append(LoggerType.TENSORBOARD)
+
+if is_wandb_available():
+    import wandb
+
+    _available_trackers.append(LoggerType.WANDB)
+
+if is_comet_ml_available():
+    from comet_ml import Experiment
+
+    _available_trackers.append(LoggerType.COMETML)
+
+
+logger = logging.getLogger(__name__)
+
+
+def get_available_trackers():
+    "Returns a list of all supported available trackers in the system"
+    return _available_trackers
+
+
+class GeneralTracker(object, metaclass=ABCMeta):
+    """
+    A base Tracker class to be used for all logging integration implementations.
+    """
+
+    log_directory: str
+
+    @abstractmethod
+    def store_init_configuration(self, values):
+        pass
+
+    @abstractmethod
+    def log(self, values, step: Optional[int]):
+        pass
+
+
+class TensorBoardTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `tensorboard`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run. Logs are then stored in `tensorboard/{run_name}`
+    """
+
+    log_directory = Path("tensorboard")
+
+    def __init__(self, run_name: str = ""):
+        self.run_name = run_name
+        self.writer = tensorboard.SummaryWriter(self.logging_dir / self.run_name)
+        logger.info(f"Initialized TensorBoard project {self.run_name} writing to {self.logging_dir}")
+        logger.info(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment.
+
+        Args:
+            values (`dict`):
+                Values to be stored as initial hyperparameters as key-value pairs. Value be of type `bool`, `str`,
+                `float`, `int`, or `None`.
+        """
+        self.writer.add_hparams(values, metric_dict={})
+        self.writer.flush()
+        logger.info("Stored initial configuration hyperparameters to TensorBoard")
+
+    def log(self, values: dict, step: Optional[int] = None):
+        """
+        Logs `values` to the current run.
+
+        Args:
+            values (`dict`):
+                Values to be logged as key-value pairs. Value be of type `int`, `float`, or `str`.
+            step (`int`, Optional):
+                The run step. If included, the log will be affiliated with this step.
+        """
+        for k, v in values.items():
+            if isinstance(v, (int, float)):
+                self.writer.add_scalar(k, v, global_step=step)
+            elif isinstance(v, str):
+                self.writer.add_text(k, v, global_step=step)
+        self.writer.flush()
+        logger.info("Successfully logged to TensorBoard")
+
+
+class WandBTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `wandb`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run.
+    """
+
+    log_directory = None
+
+    def __init__(self, run_name: str = ""):
+        self.run_name = run_name
+        wandb.init(self.run_name)
+        logger.info(f"Initialized WandB project {self.run_name}")
+        logger.info(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment.
+
+        Args:
+            values (`dict`):
And here too ;-)
@@ -0,0 +1,210 @@
+# 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.
+
+# Expectation:
+# Provide a project dir name, then each type of logger gets stored in project/{`logging_dir`}
+
+import logging
+from abc import ABCMeta, abstractmethod
+from pathlib import Path
+from typing import Optional
+
+from .utils import LoggerType, is_comet_ml_available, is_tensorboard_available, is_wandb_available
+
+
+_available_trackers = []
+
+if is_tensorboard_available():
+    from torch.utils import tensorboard
+
+    _available_trackers.append(LoggerType.TENSORBOARD)
+
+if is_wandb_available():
+    import wandb
+
+    _available_trackers.append(LoggerType.WANDB)
+
+if is_comet_ml_available():
+    from comet_ml import Experiment
+
+    _available_trackers.append(LoggerType.COMETML)
+
+
+logger = logging.getLogger(__name__)
+
+
+def get_available_trackers():
+    "Returns a list of all supported available trackers in the system"
+    return _available_trackers
+
+
+class GeneralTracker(object, metaclass=ABCMeta):
+    """
+    A base Tracker class to be used for all logging integration implementations.
+    """
+
+    log_directory: str
+
+    @abstractmethod
+    def store_init_configuration(self, values):
+        pass
+
+    @abstractmethod
+    def log(self, values, step: Optional[int]):
+        pass
+
+
+class TensorBoardTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `tensorboard`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run. Logs are then stored in `tensorboard/{run_name}`
+    """
+
+    log_directory = Path("tensorboard")
+
+    def __init__(self, run_name: str = ""):
+        self.run_name = run_name
+        self.writer = tensorboard.SummaryWriter(self.logging_dir / self.run_name)
+        logger.info(f"Initialized TensorBoard project {self.run_name} writing to {self.logging_dir}")
+        logger.info(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment.
+
+        Args:
+            values (`dict`):
+                Values to be stored as initial hyperparameters as key-value pairs. Value be of type `bool`, `str`,
+                `float`, `int`, or `None`.
+        """
+        self.writer.add_hparams(values, metric_dict={})
+        self.writer.flush()
+        logger.info("Stored initial configuration hyperparameters to TensorBoard")
+
+    def log(self, values: dict, step: Optional[int] = None):
+        """
+        Logs `values` to the current run.
+
+        Args:
+            values (`dict`):
+                Values to be logged as key-value pairs. Value be of type `int`, `float`, or `str`.
+            step (`int`, Optional):
+                The run step. If included, the log will be affiliated with this step.
+        """
+        for k, v in values.items():
+            if isinstance(v, (int, float)):
+                self.writer.add_scalar(k, v, global_step=step)
+            elif isinstance(v, str):
+                self.writer.add_text(k, v, global_step=step)
+        self.writer.flush()
+        logger.info("Successfully logged to TensorBoard")
+
+
+class WandBTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `wandb`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run.
+    """
+
+    log_directory = None
+
+    def __init__(self, run_name: str = ""):
+        self.run_name = run_name
+        wandb.init(self.run_name)
+        logger.info(f"Initialized WandB project {self.run_name}")
+        logger.info(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment.
+
+        Args:
+            values (`dict`):
+                Values to be stored as initial hyperparameters as key-value pairs. Value be of type `bool`, `str`,
+                `float`, `int`, or `None`.
+        """
+        wandb.config(values)
+        logger.info("Stored initial configuration hyperparameters to WandB")
+
+    def log(self, values: dict, step: Optional[int] = None):
+        """
+        Logs `values` to the current run.
+
+        Args:
+            values (`dict`):
+                Values to be logged as key-value pairs.
+            step (`int`, Optional):
+                The run step. If included, the log will be affiliated with this step.
Same comments on the docstrings.
@@ -0,0 +1,210 @@
+# 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.
+
+# Expectation:
+# Provide a project dir name, then each type of logger gets stored in project/{`logging_dir`}
+
+import logging
+from abc import ABCMeta, abstractmethod
+from pathlib import Path
+from typing import Optional
+
+from .utils import LoggerType, is_comet_ml_available, is_tensorboard_available, is_wandb_available
+
+
+_available_trackers = []
+
+if is_tensorboard_available():
+    from torch.utils import tensorboard
+
+    _available_trackers.append(LoggerType.TENSORBOARD)
+
+if is_wandb_available():
+    import wandb
+
+    _available_trackers.append(LoggerType.WANDB)
+
+if is_comet_ml_available():
+    from comet_ml import Experiment
+
+    _available_trackers.append(LoggerType.COMETML)
+
+
+logger = logging.getLogger(__name__)
+
+
+def get_available_trackers():
+    "Returns a list of all supported available trackers in the system"
+    return _available_trackers
+
+
+class GeneralTracker(object, metaclass=ABCMeta):
+    """
+    A base Tracker class to be used for all logging integration implementations.
+    """
+
+    log_directory: str
+
+    @abstractmethod
+    def store_init_configuration(self, values):
+        pass
+
+    @abstractmethod
+    def log(self, values, step: Optional[int]):
+        pass
+
+
+class TensorBoardTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `tensorboard`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run. Logs are then stored in `tensorboard/{run_name}`
+    """
+
+    log_directory = Path("tensorboard")
+
+    def __init__(self, run_name: str = ""):
+        self.run_name = run_name
+        self.writer = tensorboard.SummaryWriter(self.logging_dir / self.run_name)
+        logger.info(f"Initialized TensorBoard project {self.run_name} writing to {self.logging_dir}")
+        logger.info(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment.
+
+        Args:
+            values (`dict`):
+                Values to be stored as initial hyperparameters as key-value pairs. Value be of type `bool`, `str`,
+                `float`, `int`, or `None`.
+        """
+        self.writer.add_hparams(values, metric_dict={})
+        self.writer.flush()
+        logger.info("Stored initial configuration hyperparameters to TensorBoard")
+
+    def log(self, values: dict, step: Optional[int] = None):
+        """
+        Logs `values` to the current run.
+
+        Args:
+            values (`dict`):
+                Values to be logged as key-value pairs. Value be of type `int`, `float`, or `str`.
+            step (`int`, Optional):
+                The run step. If included, the log will be affiliated with this step.
+        """
+        for k, v in values.items():
+            if isinstance(v, (int, float)):
+                self.writer.add_scalar(k, v, global_step=step)
+            elif isinstance(v, str):
+                self.writer.add_text(k, v, global_step=step)
+        self.writer.flush()
+        logger.info("Successfully logged to TensorBoard")
+
+
+class WandBTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `wandb`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run.
+    """
+
+    log_directory = None
+
+    def __init__(self, run_name: str = ""):
+        self.run_name = run_name
+        wandb.init(self.run_name)
+        logger.info(f"Initialized WandB project {self.run_name}")
+        logger.info(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment.
+
+        Args:
+            values (`dict`):
+                Values to be stored as initial hyperparameters as key-value pairs. Value be of type `bool`, `str`,
+                `float`, `int`, or `None`.
+        """
+        wandb.config(values)
+        logger.info("Stored initial configuration hyperparameters to WandB")
+
+    def log(self, values: dict, step: Optional[int] = None):
+        """
+        Logs `values` to the current run.
+
+        Args:
+            values (`dict`):
+                Values to be logged as key-value pairs.
+            step (`int`, Optional):
+                The run step. If included, the log will be affiliated with this step.
+        """
+        wandb.log(values, step=step)
+        logger.info("Successfully logged to WandB")
+
+
+class CometMLTracker(GeneralTracker):
+    """
You'll need to add the `tracking.` prefix to each for the links to work.
@@ -0,0 +1,142 @@
+<!--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.
+-->
+
+# Tracking
+
+There are a large number of experiment tracking API's available, however getting them all to work with in a multi-processing environment can oftentimes be complex.
+Accelerate provides a general tracking API that can be used to log useful items during your script through [`~Accelerator.log`]
+
+## Integrated Trackers
+
+Currently `Accelerate` supports three trackers out-of-the-box:
+- [`~TensorBoardTracker`]
+- [`~WandBTracker`]
+- [`~CometMLTracker`]
```suggestion
+           run["logs/training/batch/loss"].log(loss)
```
4 spaces please :-)
@@ -0,0 +1,142 @@
+<!--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.
+-->
+
+# Tracking
+
+There are a large number of experiment tracking API's available, however getting them all to work with in a multi-processing environment can oftentimes be complex.
+Accelerate provides a general tracking API that can be used to log useful items during your script through [`~Accelerator.log`]
+
+## Integrated Trackers
+
+Currently `Accelerate` supports three trackers out-of-the-box:
+- [`~TensorBoardTracker`]
+- [`~WandBTracker`]
+- [`~CometMLTracker`]
+
+To use any of them, pass in the selected type(s) to the `log_with` parameter in [`~Accelerate.__init__`]:
+```python
+from accelerate import Accelerate
+from accelerate.utils import LoggerType
+
+accelerator = Accelerate(log_with="all")  # For all available trackers in the environment
+accelerator = Accelerate(log_with="wandb")
+accelerator = Accelerate(log_with=["wandb", LoggerType.TENSORBOARD])
+```
+
+At the start of your experiment [`~Accelerator.init_tracker`] should be used to setup your project, and potentially add any experiment hyperparameters to be logged:
+```python
+hps = {"num_iterations": 5, "learning_rate": 1e-2}
+accelerator.init_trackers("my_project", config=hps)
+```
+
+When you are ready to log any data, [`~Accelerator.log`] should be used.
+A `step` can also be passed in to correlate the data with a particular step in the training loop.
+```python
+accelerator.log({"train_loss": 1.12, "valid_loss": 0.8}, step=1)
+```
+
+A full example is below:
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator(log_with="all")
+config = {
+    "num_iterations": 5,
+    "learning_rate": 1e-2,
+    "loss_function": str(my_loss_function),
+}
+
+accelerator.init_trackers("example_project", config=config)
+
+my_model, my_optimizer, my_training_dataloader = accelerate.prepare(my_model, my_optimizer, my_training_dataloader)
+device = accelerator.device
+my_model.to(device)
+
+for iteration in config["num_iterations"]:
+    for step, batch in my_training_dataloader:
+        my_optimizer.zero_grad()
+        inputs, targets = batch
+        inputs = inputs.to(device)
+        targets = targets.to(device)
+        outputs = my_model(inputs)
+        loss = my_loss_function(outputs, targets)
+        accelerator.backward(loss)
+        my_optimizer.step()
+        accelerator.log({"training_loss": loss}, step=step)
+```
+
+[[autodoc]] tracking.TensorBoardTracker
+
+[[autodoc]] tracking.WandBTracker
+
+[[autodoc]] tracking.CometMLTracker
+
+
+## Implementing Custom Trackers
+
+To implement a new tracker to be used in `Accelerator`, a new one can be made through implementing the [`~GeneralTracker`] class.
+Every tracker must implement three functions:
+  - `__init__`: Should store a `run_name` and initialize the tracker API of the integrated library
+  - `store_init_configuration`: Should take in a `values` dictionary and store them as a one-time experiment configuration
+  - `log`: Should take in a `values` dictionary and a `step`, and should log them to the run
+As well as contain a `log_directory` attribute. If the integrated library does not store anything locally, should be set to `None`.
+
+A brief example can be seen below with an integration with Weights and Biases, containing only the relevent information:
+```python
+from accelerate.tracking import GeneralTracker
+from typing import Optional
+
+import wandb
+
+
+class WandBTracker(GeneralTracker):
+    log_directory = None
+
+    def __init__(self, run_name: str):
+        self.run_name = run_name
+        wandb.init(self.run_name)
+
+    def store_init_configuration(self, values: dict):
+        wandb.config(values)
+
+    def log(self, values: dict, step: Optional[int] = None):
+        wandb.log(values, step=step)
+```
+
+## When a wrapper cannot work
+
+If a library has an API that does not follow a strict `.log` with an overall dictionary such as Neptune.AI, logging can be done manually under an `if accelerator.is_main_process` statement:
+```diff
+from accelerate import Accelerator
++ import neptune.new as neptune
+
+accelerator = Accelerator()
++ run = neptune.init(...)
+
+my_model, my_optimizer, my_training_dataloader = accelerate.prepare(my_model, my_optimizer, my_training_dataloader)
+device = accelerator.device
+my_model.to(device)
+
+for iteration in config["num_iterations"]:
+    for batch in my_training_dataloader:
+        my_optimizer.zero_grad()
+        inputs, targets = batch
+        inputs = inputs.to(device)
+        targets = targets.to(device)
+        outputs = my_model(inputs)
+        loss = my_loss_function(outputs, targets)
+        total_loss += loss
+        accelerator.backward(loss)
+        my_optimizer.step()
++       if accelerator.is_main_process:
++         run["logs/training/batch/loss"].log(loss)
This call should be done at the Tracker level.
@@ -556,6 +585,39 @@ def wait_for_everyone(self):
         """
         wait_for_everyone()
 
+    def init_trackers(self, project_name: str, config: Optional[dict] = None):
+        """
+        Initializes a run for all trackers stored in `self.log_with`, potentially with starting configurations
+
+        Args:
+            project_name (`str`):
+                The name of the project. All trackers will save their data based on this
+            config (`dict`, *optional*):
+                Optional starting configuration to be logged.
+        """
+        project_location = Path(project_name)
The logging_dir should also be something that could be customized and default to the usual PyTorch default.
@@ -0,0 +1,210 @@
+# 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.
+
+# Expectation:
+# Provide a project dir name, then each type of logger gets stored in project/{`logging_dir`}
+
+import logging
+from abc import ABCMeta, abstractmethod
+from pathlib import Path
+from typing import Optional
+
+from .utils import LoggerType, is_comet_ml_available, is_tensorboard_available, is_wandb_available
+
+
+_available_trackers = []
+
+if is_tensorboard_available():
+    from torch.utils import tensorboard
+
+    _available_trackers.append(LoggerType.TENSORBOARD)
+
+if is_wandb_available():
+    import wandb
+
+    _available_trackers.append(LoggerType.WANDB)
+
+if is_comet_ml_available():
+    from comet_ml import Experiment
+
+    _available_trackers.append(LoggerType.COMETML)
+
+
+logger = logging.getLogger(__name__)
+
+
+def get_available_trackers():
+    "Returns a list of all supported available trackers in the system"
+    return _available_trackers
+
+
+class GeneralTracker(object, metaclass=ABCMeta):
+    """
+    A base Tracker class to be used for all logging integration implementations.
+    """
+
+    log_directory: str
+
+    @abstractmethod
+    def store_init_configuration(self, values):
+        pass
+
+    @abstractmethod
+    def log(self, values, step: Optional[int]):
+        pass
+
+
+class TensorBoardTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `tensorboard`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run. Logs are then stored in `tensorboard/{run_name}`
+    """
+
+    log_directory = Path("tensorboard")
+
+    def __init__(self, run_name: str):
+        self.run_name = run_name
+        self.writer = tensorboard.SummaryWriter(self.logging_dir / self.run_name)
This variable doesn't seem used.
@@ -0,0 +1,210 @@
+# 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.
+
+# Expectation:
+# Provide a project dir name, then each type of logger gets stored in project/{`logging_dir`}
+
+import logging
+from abc import ABCMeta, abstractmethod
+from pathlib import Path
+from typing import Optional
+
+from .utils import LoggerType, is_comet_ml_available, is_tensorboard_available, is_wandb_available
+
+
+_available_trackers = []
+
+if is_tensorboard_available():
+    from torch.utils import tensorboard
+
+    _available_trackers.append(LoggerType.TENSORBOARD)
+
+if is_wandb_available():
+    import wandb
+
+    _available_trackers.append(LoggerType.WANDB)
+
+if is_comet_ml_available():
+    from comet_ml import Experiment
+
+    _available_trackers.append(LoggerType.COMETML)
+
+
+logger = logging.getLogger(__name__)
+
+
+def get_available_trackers():
+    "Returns a list of all supported available trackers in the system"
+    return _available_trackers
+
+
+class GeneralTracker(object, metaclass=ABCMeta):
+    """
+    A base Tracker class to be used for all logging integration implementations.
+    """
+
+    log_directory: str
+
+    @abstractmethod
+    def store_init_configuration(self, values):
+        pass
+
+    @abstractmethod
+    def log(self, values, step: Optional[int]):
+        pass
+
+
+class TensorBoardTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `tensorboard`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run. Logs are then stored in `tensorboard/{run_name}`
+    """
+
+    log_directory = Path("tensorboard")
This doesn't seem used either.
@@ -0,0 +1,210 @@
+# 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.
+
+# Expectation:
+# Provide a project dir name, then each type of logger gets stored in project/{`logging_dir`}
+
+import logging
+from abc import ABCMeta, abstractmethod
+from pathlib import Path
+from typing import Optional
+
+from .utils import LoggerType, is_comet_ml_available, is_tensorboard_available, is_wandb_available
+
+
+_available_trackers = []
+
+if is_tensorboard_available():
+    from torch.utils import tensorboard
+
+    _available_trackers.append(LoggerType.TENSORBOARD)
+
+if is_wandb_available():
+    import wandb
+
+    _available_trackers.append(LoggerType.WANDB)
+
+if is_comet_ml_available():
+    from comet_ml import Experiment
+
+    _available_trackers.append(LoggerType.COMETML)
+
+
+logger = logging.getLogger(__name__)
+
+
+def get_available_trackers():
+    "Returns a list of all supported available trackers in the system"
+    return _available_trackers
+
+
+class GeneralTracker(object, metaclass=ABCMeta):
+    """
+    A base Tracker class to be used for all logging integration implementations.
+    """
+
+    log_directory: str
+
+    @abstractmethod
+    def store_init_configuration(self, values):
+        pass
+
+    @abstractmethod
+    def log(self, values, step: Optional[int]):
+        pass
+
+
+class TensorBoardTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `tensorboard`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run. Logs are then stored in `tensorboard/{run_name}`
+    """
+
+    log_directory = Path("tensorboard")
+
+    def __init__(self, run_name: str):
+        self.run_name = run_name
+        self.writer = tensorboard.SummaryWriter(self.logging_dir / self.run_name)
+        logger.info(f"Initialized TensorBoard project {self.run_name} writing to {self.logging_dir}")
+        logger.info(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment.
+
+        Args:
+            values (Dictionary `str` to `bool`, `str`, `float` or `int`):
+                Values to be stored as initial hyperparameters as key-value pairs. The values need to have type `bool`,
+                `str`, `float`, `int`, or `None`.
+        """
+        self.writer.add_hparams(values, metric_dict={})
+        self.writer.flush()
+        logger.info("Stored initial configuration hyperparameters to TensorBoard")
+
+    def log(self, values: dict, step: Optional[int] = None):
+        """
+        Logs `values` to the current run.
+
+        Args:
+            values (Dictionary `str` to `str`, `float`, or `int`):
+                Values to be logged as key-value pairs. The values need to have type `str`, `float`, or `int`.
+            step (`int`, *optional*):
+                The run step. If included, the log will be affiliated with this step.
+        """
+        for k, v in values.items():
+            if isinstance(v, (int, float)):
+                self.writer.add_scalar(k, v, global_step=step)
+            elif isinstance(v, str):
+                self.writer.add_text(k, v, global_step=step)
+        self.writer.flush()
+        logger.info("Successfully logged to TensorBoard")
+
+
+class WandBTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `wandb`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run.
+    """
+
+    log_directory = None
Same here.
@@ -0,0 +1,210 @@
+# 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.
+
+# Expectation:
+# Provide a project dir name, then each type of logger gets stored in project/{`logging_dir`}
+
+import logging
+from abc import ABCMeta, abstractmethod
+from pathlib import Path
+from typing import Optional
+
+from .utils import LoggerType, is_comet_ml_available, is_tensorboard_available, is_wandb_available
+
+
+_available_trackers = []
+
+if is_tensorboard_available():
+    from torch.utils import tensorboard
+
+    _available_trackers.append(LoggerType.TENSORBOARD)
+
+if is_wandb_available():
+    import wandb
+
+    _available_trackers.append(LoggerType.WANDB)
+
+if is_comet_ml_available():
+    from comet_ml import Experiment
+
+    _available_trackers.append(LoggerType.COMETML)
+
+
+logger = logging.getLogger(__name__)
+
+
+def get_available_trackers():
+    "Returns a list of all supported available trackers in the system"
+    return _available_trackers
+
+
+class GeneralTracker(object, metaclass=ABCMeta):
+    """
+    A base Tracker class to be used for all logging integration implementations.
+    """
+
+    log_directory: str
+
+    @abstractmethod
+    def store_init_configuration(self, values):
+        pass
+
+    @abstractmethod
+    def log(self, values, step: Optional[int]):
+        pass
+
+
+class TensorBoardTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `tensorboard`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run. Logs are then stored in `tensorboard/{run_name}`
+    """
+
+    log_directory = Path("tensorboard")
+
+    def __init__(self, run_name: str):
+        self.run_name = run_name
+        self.writer = tensorboard.SummaryWriter(self.logging_dir / self.run_name)
+        logger.info(f"Initialized TensorBoard project {self.run_name} writing to {self.logging_dir}")
+        logger.info(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment.
+
+        Args:
+            values (Dictionary `str` to `bool`, `str`, `float` or `int`):
+                Values to be stored as initial hyperparameters as key-value pairs. The values need to have type `bool`,
+                `str`, `float`, `int`, or `None`.
+        """
+        self.writer.add_hparams(values, metric_dict={})
+        self.writer.flush()
+        logger.info("Stored initial configuration hyperparameters to TensorBoard")
+
+    def log(self, values: dict, step: Optional[int] = None):
+        """
+        Logs `values` to the current run.
+
+        Args:
+            values (Dictionary `str` to `str`, `float`, or `int`):
+                Values to be logged as key-value pairs. The values need to have type `str`, `float`, or `int`.
+            step (`int`, *optional*):
+                The run step. If included, the log will be affiliated with this step.
+        """
+        for k, v in values.items():
+            if isinstance(v, (int, float)):
+                self.writer.add_scalar(k, v, global_step=step)
+            elif isinstance(v, str):
+                self.writer.add_text(k, v, global_step=step)
+        self.writer.flush()
+        logger.info("Successfully logged to TensorBoard")
+
+
+class WandBTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `wandb`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run.
+    """
+
+    log_directory = None
+
+    def __init__(self, run_name: str):
+        self.run_name = run_name
+        wandb.init(self.run_name)
+        logger.info(f"Initialized WandB project {self.run_name}")
+        logger.info(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment.
+
+        Args:
+            values (Dictionary `str` to `bool`, `str`, `float` or `int`):
+                Values to be stored as initial hyperparameters as key-value pairs. The values need to have type `bool`,
+                `str`, `float`, `int`, or `None`.
+        """
+        wandb.config(values)
+        logger.info("Stored initial configuration hyperparameters to WandB")
+
+    def log(self, values: dict, step: Optional[int] = None):
+        """
+        Logs `values` to the current run.
+
+        Args:
+            values (Dictionary `str` to `str`, `float`, or `int`):
+                Values to be logged as key-value pairs. The values need to have type `str`, `float`, or `int`.
+            step (`int`, *optional*):
+                The run step. If included, the log will be affiliated with this step.
+        """
+        wandb.log(values, step=step)
+        logger.info("Successfully logged to WandB")
+
+
+class CometMLTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `comet_ml`. Should be initialized at the start of your script.
+
+    API keys must be stored in a Comet config file.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run.
+    """
+
+    log_directory = None
We should remove this from the abtract class as it's not required for all of the trackers.
@@ -0,0 +1,210 @@
+# 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.
+
+# Expectation:
+# Provide a project dir name, then each type of logger gets stored in project/{`logging_dir`}
+
+import logging
+from abc import ABCMeta, abstractmethod
+from pathlib import Path
+from typing import Optional
+
+from .utils import LoggerType, is_comet_ml_available, is_tensorboard_available, is_wandb_available
+
+
+_available_trackers = []
+
+if is_tensorboard_available():
+    from torch.utils import tensorboard
+
+    _available_trackers.append(LoggerType.TENSORBOARD)
+
+if is_wandb_available():
+    import wandb
+
+    _available_trackers.append(LoggerType.WANDB)
+
+if is_comet_ml_available():
+    from comet_ml import Experiment
+
+    _available_trackers.append(LoggerType.COMETML)
+
+
+logger = logging.getLogger(__name__)
+
+
+def get_available_trackers():
+    "Returns a list of all supported available trackers in the system"
+    return _available_trackers
+
+
+class GeneralTracker(object, metaclass=ABCMeta):
+    """
+    A base Tracker class to be used for all logging integration implementations.
+    """
+
+    log_directory: str
`logging_dir` and `run_name` are actually the same when it comes to `tensorboard`, and we can't have both. So something like this won't work: 
```python
writerC = SummaryWriter("Testme", comment="LR_0.1_BATCH_16")
```
So we can just remove this completely. 
@@ -0,0 +1,210 @@
+# 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.
+
+# Expectation:
+# Provide a project dir name, then each type of logger gets stored in project/{`logging_dir`}
+
+import logging
+from abc import ABCMeta, abstractmethod
+from pathlib import Path
+from typing import Optional
+
+from .utils import LoggerType, is_comet_ml_available, is_tensorboard_available, is_wandb_available
+
+
+_available_trackers = []
+
+if is_tensorboard_available():
+    from torch.utils import tensorboard
+
+    _available_trackers.append(LoggerType.TENSORBOARD)
+
+if is_wandb_available():
+    import wandb
+
+    _available_trackers.append(LoggerType.WANDB)
+
+if is_comet_ml_available():
+    from comet_ml import Experiment
+
+    _available_trackers.append(LoggerType.COMETML)
+
+
+logger = logging.getLogger(__name__)
+
+
+def get_available_trackers():
+    "Returns a list of all supported available trackers in the system"
+    return _available_trackers
+
+
+class GeneralTracker(object, metaclass=ABCMeta):
+    """
+    A base Tracker class to be used for all logging integration implementations.
+    """
+
+    log_directory: str
+
+    @abstractmethod
+    def store_init_configuration(self, values):
+        pass
+
+    @abstractmethod
+    def log(self, values, step: Optional[int]):
+        pass
+
+
+class TensorBoardTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `tensorboard`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run. Logs are then stored in `tensorboard/{run_name}`
+    """
+
+    log_directory = Path("tensorboard")
+
+    def __init__(self, run_name: str):
+        self.run_name = run_name
+        self.writer = tensorboard.SummaryWriter(self.logging_dir / self.run_name)
Not following why we have to add TensoFlow here. Can't it be just `tensorboard`?
@@ -21,6 +21,8 @@
 extras["test"] = [
     "pytest",
     "pytest-xdist",
+    "tensorflow",
It would probably be best to use a temporary directory here, to avoid users running the tests locally getting this folder.
@@ -0,0 +1,89 @@
+# 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 shutil
+import unittest
+from pathlib import Path
+
+# We use TF to parse the logs
+import tensorflow as tf
+from accelerate import Accelerator
+from tensorboard.plugins.hparams import plugin_data_pb2
+from tensorflow.core.util import event_pb2
+from tensorflow.python.summary.summary_iterator import summary_iterator
+
+
+logger = logging.getLogger(__name__)
+
+
+class TensorBoardTrackingTest(unittest.TestCase):
+    @classmethod
+    def teardown_class(cls):
+        shutil.rmtree("test_project_with_config")
+        shutil.rmtree("test_project_with_log")
+
+    def test_init_trackers(self):
+        hps = None
+        project_name = "test_project_with_config"
Same here.
@@ -0,0 +1,89 @@
+# 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 shutil
+import unittest
+from pathlib import Path
+
+# We use TF to parse the logs
+import tensorflow as tf
+from accelerate import Accelerator
+from tensorboard.plugins.hparams import plugin_data_pb2
+from tensorflow.core.util import event_pb2
+from tensorflow.python.summary.summary_iterator import summary_iterator
+
+
+logger = logging.getLogger(__name__)
+
+
+class TensorBoardTrackingTest(unittest.TestCase):
+    @classmethod
+    def teardown_class(cls):
+        shutil.rmtree("test_project_with_config")
+        shutil.rmtree("test_project_with_log")
+
+    def test_init_trackers(self):
+        hps = None
+        project_name = "test_project_with_config"
+        accelerator = Accelerator(log_with="tensorboard")
+        config = {"num_iterations": 12, "learning_rate": 1e-2, "some_boolean": False, "some_string": "some_value"}
+        accelerator.init_trackers(project_name, config)
+        for child in Path(project_name).glob("*/**"):
+            log = list(filter(lambda x: x.is_file(), child.iterdir()))[0]
+            # The config log is stored one layer deeper in the logged directory
+            # And names are randomly generated each time
+        si = summary_iterator(str(log))
+        # Pull HPS through careful parsing
+        for event in si:
+            for value in event.summary.value:
+                proto_bytes = value.metadata.plugin_data.content
+                plugin_data = plugin_data_pb2.HParamsPluginData.FromString(proto_bytes)
+                if plugin_data.HasField("session_start_info"):
+                    hps = dict(plugin_data.session_start_info.hparams)
+
+        self.assertTrue(isinstance(hps, dict))
+        keys = list(hps.keys())
+        keys.sort()
+        self.assertEqual(keys, ["learning_rate", "num_iterations", "some_boolean", "some_string"])
+        self.assertEqual(hps["num_iterations"].number_value, 12)
+        self.assertEqual(hps["learning_rate"].number_value, 0.01)
+        self.assertEqual(hps["some_boolean"].bool_value, False)
+        self.assertEqual(hps["some_string"].string_value, "some_value")
+
+    def test_log(self):
+        step = None
+        project_name = "test_project_with_log"
Is it why the TF dep is necessary? There is really no other way to load a TensorBoard log?
@@ -0,0 +1,89 @@
+# 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 shutil
+import unittest
+from pathlib import Path
+
+# We use TF to parse the logs
+import tensorflow as tf
+from accelerate import Accelerator
+from tensorboard.plugins.hparams import plugin_data_pb2
+from tensorflow.core.util import event_pb2
+from tensorflow.python.summary.summary_iterator import summary_iterator
+
+
+logger = logging.getLogger(__name__)
+
+
+class TensorBoardTrackingTest(unittest.TestCase):
+    @classmethod
+    def teardown_class(cls):
+        shutil.rmtree("test_project_with_config")
+        shutil.rmtree("test_project_with_log")
+
+    def test_init_trackers(self):
+        hps = None
+        project_name = "test_project_with_config"
+        accelerator = Accelerator(log_with="tensorboard")
+        config = {"num_iterations": 12, "learning_rate": 1e-2, "some_boolean": False, "some_string": "some_value"}
+        accelerator.init_trackers(project_name, config)
+        for child in Path(project_name).glob("*/**"):
+            log = list(filter(lambda x: x.is_file(), child.iterdir()))[0]
+            # The config log is stored one layer deeper in the logged directory
+            # And names are randomly generated each time
+        si = summary_iterator(str(log))
+        # Pull HPS through careful parsing
+        for event in si:
+            for value in event.summary.value:
+                proto_bytes = value.metadata.plugin_data.content
+                plugin_data = plugin_data_pb2.HParamsPluginData.FromString(proto_bytes)
+                if plugin_data.HasField("session_start_info"):
+                    hps = dict(plugin_data.session_start_info.hparams)
+
+        self.assertTrue(isinstance(hps, dict))
+        keys = list(hps.keys())
+        keys.sort()
+        self.assertEqual(keys, ["learning_rate", "num_iterations", "some_boolean", "some_string"])
+        self.assertEqual(hps["num_iterations"].number_value, 12)
+        self.assertEqual(hps["learning_rate"].number_value, 0.01)
+        self.assertEqual(hps["some_boolean"].bool_value, False)
+        self.assertEqual(hps["some_string"].string_value, "some_value")
+
+    def test_log(self):
+        step = None
+        project_name = "test_project_with_log"
+        accelerator = Accelerator(log_with="tensorboard")
+        accelerator.init_trackers(project_name)
+        values = {"total_loss": 0.1, "iteration": 1, "my_text": "some_value"}
+        accelerator.log(values, step=0)
+        # Logged values are stored in the outermost-tfevents file and can be read in as a TFRecord
+        # Names are randomly generated each time
+        log = list(filter(lambda x: x.is_file(), Path(project_name).iterdir()))[0]
+        serialized_examples = tf.data.TFRecordDataset(log)
Yup, sadly: https://github.com/tensorflow/tensorboard#how-can-i-export-data-from-tensorboard

We have to use `tf` to parse it back in. 
@@ -0,0 +1,89 @@
+# 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 shutil
+import unittest
+from pathlib import Path
+
+# We use TF to parse the logs
+import tensorflow as tf
+from accelerate import Accelerator
+from tensorboard.plugins.hparams import plugin_data_pb2
+from tensorflow.core.util import event_pb2
+from tensorflow.python.summary.summary_iterator import summary_iterator
+
+
+logger = logging.getLogger(__name__)
+
+
+class TensorBoardTrackingTest(unittest.TestCase):
+    @classmethod
+    def teardown_class(cls):
+        shutil.rmtree("test_project_with_config")
+        shutil.rmtree("test_project_with_log")
+
+    def test_init_trackers(self):
+        hps = None
+        project_name = "test_project_with_config"
+        accelerator = Accelerator(log_with="tensorboard")
+        config = {"num_iterations": 12, "learning_rate": 1e-2, "some_boolean": False, "some_string": "some_value"}
+        accelerator.init_trackers(project_name, config)
+        for child in Path(project_name).glob("*/**"):
+            log = list(filter(lambda x: x.is_file(), child.iterdir()))[0]
+            # The config log is stored one layer deeper in the logged directory
+            # And names are randomly generated each time
+        si = summary_iterator(str(log))
+        # Pull HPS through careful parsing
+        for event in si:
+            for value in event.summary.value:
+                proto_bytes = value.metadata.plugin_data.content
+                plugin_data = plugin_data_pb2.HParamsPluginData.FromString(proto_bytes)
+                if plugin_data.HasField("session_start_info"):
+                    hps = dict(plugin_data.session_start_info.hparams)
+
+        self.assertTrue(isinstance(hps, dict))
+        keys = list(hps.keys())
+        keys.sort()
+        self.assertEqual(keys, ["learning_rate", "num_iterations", "some_boolean", "some_string"])
+        self.assertEqual(hps["num_iterations"].number_value, 12)
+        self.assertEqual(hps["learning_rate"].number_value, 0.01)
+        self.assertEqual(hps["some_boolean"].bool_value, False)
+        self.assertEqual(hps["some_string"].string_value, "some_value")
+
+    def test_log(self):
+        step = None
+        project_name = "test_project_with_log"
+        accelerator = Accelerator(log_with="tensorboard")
+        accelerator.init_trackers(project_name)
+        values = {"total_loss": 0.1, "iteration": 1, "my_text": "some_value"}
+        accelerator.log(values, step=0)
+        # Logged values are stored in the outermost-tfevents file and can be read in as a TFRecord
+        # Names are randomly generated each time
+        log = list(filter(lambda x: x.is_file(), Path(project_name).iterdir()))[0]
+        serialized_examples = tf.data.TFRecordDataset(log)
```suggestion
        Should run any finalizing functions within the tracking API. If the API should not have one, just don't overwrite that method.
```
@@ -80,6 +80,16 @@ def log(self, values: dict, step: Optional[int]):
         """
         pass
 
+    @abstractmethod
+    def finish(self):
+        """
+        Should run any finalizing functions within the tracking API. If the API should not have one, just return:
+        ```python
+        super().finish()
+        ```
```suggestion
```
No need to write something as calling the super method is then the default ;-)
@@ -218,3 +242,7 @@ def log(self, values: dict, step: Optional[int] = None):
             self.writer.set_step(step)
         self.writer.log_others(values)
         logger.info("Successfully logged to CometML")
+
+    def finish(self):
+        """Do nothing"""
+        super().finish()
I'd like to avoid this. The API should accept any logging dir for TensorBoard at least, no?
@@ -29,29 +32,29 @@
 
 
 class TensorBoardTrackingTest(unittest.TestCase):
-    @classmethod
-    def teardown_class(cls):
-        shutil.rmtree("test_project_with_config")
-        shutil.rmtree("test_project_with_log")
-
     def test_init_trackers(self):
         hps = None
         project_name = "test_project_with_config"
-        accelerator = Accelerator(log_with="tensorboard")
-        config = {"num_iterations": 12, "learning_rate": 1e-2, "some_boolean": False, "some_string": "some_value"}
-        accelerator.init_trackers(project_name, config)
-        for child in Path(project_name).glob("*/**"):
-            log = list(filter(lambda x: x.is_file(), child.iterdir()))[0]
-            # The config log is stored one layer deeper in the logged directory
-            # And names are randomly generated each time
-        si = summary_iterator(str(log))
-        # Pull HPS through careful parsing
-        for event in si:
-            for value in event.summary.value:
-                proto_bytes = value.metadata.plugin_data.content
-                plugin_data = plugin_data_pb2.HParamsPluginData.FromString(proto_bytes)
-                if plugin_data.HasField("session_start_info"):
-                    hps = dict(plugin_data.session_start_info.hparams)
+        with tempfile.TemporaryDirectory() as dirpath:
+            oldpwd = os.getcwd()
+            os.chdir(dirpath)
Would be nice to have a `logging_dir` argument in the `Accelerator` I realize.
@@ -64,26 +67,90 @@ def test_init_trackers(self):
 
     def test_log(self):
         step = None
+        project_name = "test_project_with_log"
+        with tempfile.TemporaryDirectory() as dirpath:
+            oldpwd = os.getcwd()
+            os.chdir(dirpath)
+            accelerator = Accelerator(log_with="tensorboard")
+            accelerator.init_trackers(project_name)
+            values = {"total_loss": 0.1, "iteration": 1, "my_text": "some_value"}
+            accelerator.log(values, step=0)
+            accelerator.end_training()
+            # Logged values are stored in the outermost-tfevents file and can be read in as a TFRecord
+            # Names are randomly generated each time
+            log = list(filter(lambda x: x.is_file(), Path(project_name).iterdir()))[0]
+            serialized_examples = tf.data.TFRecordDataset(log)
+            for e in serialized_examples:
+                event = event_pb2.Event.FromString(e.numpy())
+                if step is None:
+                    step = event.step
+                for value in event.summary.value:
+                    if value.tag == "total_loss":
+                        total_loss = value.simple_value
+                    elif value.tag == "iteration":
+                        iteration = value.simple_value
+                    elif value.tag == "my_text/text_summary":  # Append /text_summary to the key
+                        my_text = value.tensor.string_val[0].decode()
+            os.chdir(oldpwd)
+        self.assertAlmostEqual(total_loss, values["total_loss"])
+        self.assertEqual(iteration, values["iteration"])
+        self.assertEqual(my_text, values["my_text"])
+
+
+@mock.patch.dict(os.environ, {"WANDB_MODE": "offline"})
+class WandBTrackingTest(unittest.TestCase):
+    @staticmethod
+    def get_value_from_log(key: str, log: str, key_occurance: int = 0):
+        """
+        Parses wandb log for `key` and returns the value.
+        If parsing through multiple calls to .log, pass in a `key_occurance`
+        """
+        res = re.findall(rf"(?<={key} )[^\s]+", log)[key_occurance]
+        if '"' in res:
+            return re.findall(r'"([^"]*)"', res)[0]
+        else:
+            return res
+
+    def test_init_trackers(self):
+        project_name = "test_project_with_config"
+        with tempfile.TemporaryDirectory() as dirpath:
+            oldpwd = os.getcwd()
+            os.chdir(dirpath)
No need to do a separate job, they can all go in the same one!
@@ -12,6 +12,19 @@ jobs:
       with:
         python-version: 3.6
     - name: Install Python dependencies
-      run: pip install -e .[test]
+      run: pip install -e .[test,trackers]
     - name: Run Tests
-      run: make test
\ No newline at end of file
+      run: make test
+
+  test-tf:
We should mention the train end method in this example, since it's relevant for trackers.
@@ -0,0 +1,137 @@
+<!--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.
+-->
+
+# Tracking
+
+There are a large number of experiment tracking API's available, however getting them all to work with in a multi-processing environment can oftentimes be complex.
+Accelerate provides a general tracking API that can be used to log useful items during your script through [`~Accelerator.log`]
+
+## Integrated Trackers
+
+Currently `Accelerate` supports three trackers out-of-the-box:
+
+
+[[autodoc]] tracking.TensorBoardTracker
+
+[[autodoc]] tracking.WandBTracker
+
+[[autodoc]] tracking.CometMLTracker
+
+To use any of them, pass in the selected type(s) to the `log_with` parameter in [`~Accelerate.__init__`]:
+```python
+from accelerate import Accelerate
+from accelerate.utils import LoggerType
+
+accelerator = Accelerate(log_with="all")  # For all available trackers in the environment
+accelerator = Accelerate(log_with="wandb")
+accelerator = Accelerate(log_with=["wandb", LoggerType.TENSORBOARD])
+```
+
+At the start of your experiment [`~Accelerator.init_tracker`] should be used to setup your project, and potentially add any experiment hyperparameters to be logged:
+```python
+hps = {"num_iterations": 5, "learning_rate": 1e-2}
+accelerator.init_trackers("my_project", config=hps)
+```
+
+When you are ready to log any data, [`~Accelerator.log`] should be used.
+A `step` can also be passed in to correlate the data with a particular step in the training loop.
+```python
+accelerator.log({"train_loss": 1.12, "valid_loss": 0.8}, step=1)
+```
+
+A full example is below:
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator(log_with="all")
+config = {
+    "num_iterations": 5,
+    "learning_rate": 1e-2,
+    "loss_function": str(my_loss_function),
+}
+
+accelerator.init_trackers("example_project", config=config)
+
+my_model, my_optimizer, my_training_dataloader = accelerate.prepare(my_model, my_optimizer, my_training_dataloader)
+device = accelerator.device
+my_model.to(device)
+
+for iteration in config["num_iterations"]:
+    for step, batch in my_training_dataloader:
+        my_optimizer.zero_grad()
+        inputs, targets = batch
+        inputs = inputs.to(device)
+        targets = targets.to(device)
+        outputs = my_model(inputs)
+        loss = my_loss_function(outputs, targets)
+        accelerator.backward(loss)
+        my_optimizer.step()
+        accelerator.log({"training_loss": loss}, step=step)
+```
I would group those together.
@@ -20,9 +20,11 @@
 extras["docs"] = []
 extras["test"] = [
     "pytest",
-    "pytest-xdist",
+    "pytest-xdist"
 ]
-extras["dev"] = extras["quality"] + extras["test"]
+extras["trackers"] = ["wandb", "comet-ml"]
+extras["test_trackers"] = ["tensorflow"]
```suggestion
extras["dev"] = extras["quality"] + extras["test"]
```
Please let's not include all the trackers in dev ;-)
@@ -20,9 +20,11 @@
 extras["docs"] = []
 extras["test"] = [
     "pytest",
-    "pytest-xdist",
+    "pytest-xdist"
 ]
-extras["dev"] = extras["quality"] + extras["test"]
+extras["trackers"] = ["wandb", "comet-ml"]
+extras["test_trackers"] = ["tensorflow"]
+extras["dev"] = extras["quality"] + extras["test"] + extras["trackers"]
_The documentation is not available anymore as the PR was closed or merged._
Given https://github.com/huggingface/accelerate/issues/291#issuecomment-1080771134, this PR is feature-complete. Only thing missing is tests for TensorBoard and potential mocking, and other documentation. But otherwise it's ready for a review
Last review, this time for the docs! 😄 🚀 
There's a weird issue when doing `pip install tensorflow` where it will just install a higher `setuptools`, and that setuptools will break importing pytorch. pinning it in the dev requirements "fixes" it. Need to play with it more, but a test is there. Once that's solved we probably should mimic something like this for the TF tests as well, since this goes through and looks at the actual values stored
https://api.github.com/repos/huggingface/accelerate/pulls/293
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 4f2a79bcd..2b623db34 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -30,6 +30,7 @@
 from .state import AcceleratorState, DistributedType, is_deepspeed_available
 from .utils import (
     DeepSpeedPlugin,
+    PrecisionType,
     RNGType,
     convert_outputs_to_fp32,
     extract_model_from_parallel,
@@ -106,19 +107,18 @@ def __init__(
         device_placement: bool = True,
         split_batches: bool = False,
         fp16: bool = None,
-        mixed_precision: str = None,
+        mixed_precision: Union[PrecisionType, str] = None,
         cpu: bool = False,
         deepspeed_plugin: DeepSpeedPlugin = None,
         rng_types: Optional[List[Union[str, RNGType]]] = None,
         dispatch_batches: Optional[bool] = None,
         kwargs_handlers: Optional[List[KwargsHandler]] = None,
     ):
-
         if mixed_precision is not None:
-            mixed_precision = mixed_precision.lower()
-            if mixed_precision not in ["no", "fp16", "bf16"]:
+            mixed_precision = str(mixed_precision)
+            if mixed_precision not in PrecisionType:
                 raise ValueError(
-                    f"Unknown mixed_precision mode: {mixed_precision}. Choose between 'no', 'fp16' and 'bf16'."
+                    f"Unknown mixed_precision mode: {mixed_precision}. Choose between {PrecisionType.list()}"
                 )
 
         if fp16:
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index ff1b148f4..dd946c85d 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -27,7 +27,7 @@
 from accelerate.commands.config import default_config_file, load_config_from_file
 from accelerate.commands.config.config_args import SageMakerConfig
 from accelerate.state import ComputeEnvironment, DistributedType
-from accelerate.utils import PrepareForLaunch, is_sagemaker_available
+from accelerate.utils import PrecisionType, PrepareForLaunch, is_sagemaker_available
 
 
 def launch_command_parser(subparsers=None):
@@ -163,10 +163,12 @@ def simple_launcher(args):
 
     current_env = os.environ.copy()
     current_env["USE_CPU"] = str(args.cpu)
-
-    mixed_precision = args.mixed_precision.lower()
-    if mixed_precision not in ["no", "fp16", "bf16"]:
-        raise ValueError(f"Unknown mixed_precision mode: {mixed_precision}. Choose between 'no', 'fp16' and 'bf16'.")
+    try:
+        mixed_precision = PrecisionType(args.mixed_precision.lower())
+    except ValueError:
+        raise ValueError(
+            f"Unknown mixed_precision mode: {args.mixed_precision.lower()}. Choose between {PrecisionType.list()}."
+        )
 
     if args.fp16:
         warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', DeprecationWarning)
@@ -212,10 +214,12 @@ def multi_gpu_launcher(args):
     cmd.extend(args.training_script_args)
 
     current_env = os.environ.copy()
-    mixed_precision = args.mixed_precision.lower()
-
-    if mixed_precision not in ["no", "fp16", "bf16"]:
-        raise ValueError(f"Unknown mixed_precision mode: {mixed_precision}. Choose between 'no', 'fp16' and 'bf16'.")
+    try:
+        mixed_precision = PrecisionType(args.mixed_precision.lower())
+    except ValueError:
+        raise ValueError(
+            f"Unknown mixed_precision mode: {args.mixed_precision.lower()}. Choose between {PrecisionType.list()}."
+        )
 
     if args.fp16:
         warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', DeprecationWarning)
@@ -259,10 +263,12 @@ def deepspeed_launcher(args):
     cmd.extend(args.training_script_args)
 
     current_env = os.environ.copy()
-    mixed_precision = args.mixed_precision.lower()
-
-    if mixed_precision not in ["no", "fp16", "bf16"]:
-        raise ValueError(f"Unknown mixed_precision mode: {mixed_precision}. Choose between 'no', 'fp16' and 'bf16'.")
+    try:
+        mixed_precision = PrecisionType(args.mixed_precision.lower())
+    except ValueError:
+        raise ValueError(
+            f"Unknown mixed_precision mode: {args.mixed_precision.lower()}. Choose between {PrecisionType.list()}."
+        )
 
     if args.fp16:
         warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', DeprecationWarning)
@@ -388,10 +394,12 @@ def sagemaker_launcher(sagemaker_config: SageMakerConfig, args):
     print("Converting Arguments to Hyperparameters")
     hyperparameters = _convert_nargs_to_dict(args.training_script_args)
 
-    mixed_precision = args.mixed_precision.lower()
-
-    if mixed_precision not in ["no", "fp16", "bf16"]:
-        raise ValueError(f"Unknown mixed_precision mode: {mixed_precision}. Choose between 'no', 'fp16' and 'bf16'.")
+    try:
+        mixed_precision = PrecisionType(args.mixed_precision.lower())
+    except ValueError:
+        raise ValueError(
+            f"Unknown mixed_precision mode: {args.mixed_precision.lower()}. Choose between {PrecisionType.list()}."
+        )
 
     if args.fp16:
         warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', DeprecationWarning)
diff --git a/src/accelerate/launchers.py b/src/accelerate/launchers.py
index 48c4404dd..cfd8929d9 100644
--- a/src/accelerate/launchers.py
+++ b/src/accelerate/launchers.py
@@ -22,7 +22,7 @@
 from packaging import version
 
 from .state import AcceleratorState
-from .utils import PrepareForLaunch, patch_environment
+from .utils import PrecisionType, PrepareForLaunch, patch_environment
 
 
 def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, mixed_precision="no", use_port="29500"):
@@ -107,10 +107,11 @@ def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, mix
                     "function."
                 )
 
-            mixed_precision = mixed_precision.lower()
-            if mixed_precision not in ["no", "fp16", "bf16"]:
+            try:
+                mixed_precision = PrecisionType(mixed_precision.lower())
+            except ValueError:
                 raise ValueError(
-                    f"Unknown mixed_precision: {mixed_precision}. Choose between 'no', 'fp16' and 'bf16'."
+                    f"Unknown mixed_precision mode: {args.mixed_precision.lower()}. Choose between {PrecisionType.list()}."
                 )
 
             if use_fp16:
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 5b94d7fff..af9299a4f 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -18,7 +18,7 @@
 from collections.abc import Mapping
 from contextlib import contextmanager
 from dataclasses import dataclass, field
-from enum import Enum
+from enum import Enum, EnumMeta
 from functools import update_wrapper
 from typing import Any, List, Optional, Union
 
@@ -51,7 +51,36 @@ def is_sagemaker_available():
 OPTIMIZER_NAME = "optimizer"
 
 
-class RNGType(Enum):
+class EnumWithContains(EnumMeta):
+    "A metaclass that adds the ability to check if `self` contains an item with the `in` operator"
+
+    def __contains__(cls, item):
+        try:
+            cls(item)
+        except ValueError:
+            return False
+        return True
+
+
+class BaseEnum(Enum, metaclass=EnumWithContains):
+    "An enum class that can get the value of an item with `str(Enum.key)`"
+
+    def __str__(self):
+        return self.value
+
+    @classmethod
+    def list(cls):
+        "Method to list all the possible items in `cls`"
+        return list(map(lambda item: str(item), cls))
+
+
+class PrecisionType(BaseEnum):
+    NO = "no"
+    FP16 = "fp16"
+    BF16 = "bf16"
+
+
+class RNGType(BaseEnum):
     TORCH = "torch"
     CUDA = "cuda"
     XLA = "xla"

Here I think you might need to convert this in a list for `choices` to work the proper way.
@@ -58,7 +58,7 @@ def launch_command_parser(subparsers=None):
         "--mixed_precision",
         default="no",
         type=str,
-        choices=["no", "fp16", "bf16"],
+        choices=PrecisionType,
A small docstring here on what this adds to a regular `Enum` would be great :-)
@@ -51,7 +51,26 @@ def is_sagemaker_available():
 OPTIMIZER_NAME = "optimizer"
 
 
-class RNGType(Enum):
+class MetaEnum(EnumMeta):
+    def __contains__(cls, item):
```suggestion
class EnumWithContains(EnumMeta):
```
nit: the name is not very descriptive on what the metaclass does, would switch to this.
@@ -51,7 +51,36 @@ def is_sagemaker_available():
 OPTIMIZER_NAME = "optimizer"
 
 
-class RNGType(Enum):
+class MetaEnum(EnumMeta):
```suggestion
class BaseEnum(Enum, metaclass=EnumWithContains):
```
To go with the suggestion above.
@@ -51,7 +51,36 @@ def is_sagemaker_available():
 OPTIMIZER_NAME = "optimizer"
 
 
-class RNGType(Enum):
+class MetaEnum(EnumMeta):
+    "A metaclass that adds the ability to check if `self` contains an item with the `in` operator"
+
+    def __contains__(cls, item):
+        try:
+            cls(item)
+        except ValueError:
+            return False
+        return True
+
+
+class BaseEnum(Enum, metaclass=MetaEnum):
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/292
diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml
index 20f963aae..8f266c600 100644
--- a/docs/source/_toctree.yml
+++ b/docs/source/_toctree.yml
@@ -19,4 +19,6 @@
     title: Kwargs Handlers
   - local: internal
     title: Internals
+  - local: checkpoint
+    title: Checkpointing
   title: API Reference
diff --git a/docs/source/checkpoint.mdx b/docs/source/checkpoint.mdx
new file mode 100644
index 000000000..26ef21150
--- /dev/null
+++ b/docs/source/checkpoint.mdx
@@ -0,0 +1,60 @@
+<!--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.
+-->
+
+# Checkpointing
+
+When training a PyTorch model with Accelerate, you may often want to save and continue a state of training. Doing so requires
+saving and loading the model, optimizer, RNG generators, and the GradScaler. Inside Accelerate are two convience functions to achieve this quickly:
+- Use [`~Accelerator.save_state`] for saving everything mentioned above to a folder location
+- Use [`~Accelerator.load_state`] for loading everything stored from an earlier `save_state`
+
+It should be noted that the expectation is that those states come from the same training script, they should not be from two separate scripts.
+
+- By using [`~Accelerator.register_for_checkpointing`], you can register custom objects to be automatically stored or loaded from the two prior functions,
+so long as the object has a `state_dict` **and** a `load_state_dict` functionality. This could include objects such as a learning rate scheduler. 
+
+Below is a brief example using checkpointing to save and reload a state during training:
+
+```python
+from accelerate import Accelerator
+import torch
+
+accelerator = Accelerator()
+
+my_scheduler = torch.optim.lr_scheduler.StepLR(my_optimizer, step_size=1, gamma=0.99)
+my_model, my_optimizer, my_training_dataloader = accelerate.prepare(my_model, my_optimizer, my_training_dataloader)
+
+# Register the LR scheduler
+accelerate.register_for_checkpointing(my_scheduler)
+
+# Save the starting state
+accelerate.save_state("my/save/path")
+
+device = accelerator.device
+my_model.to(device)
+
+# Perform training
+for epoch in range(num_epochs):
+    for batch in my_training_dataloader:
+        my_optimizer.zero_grad()
+        inputs, targets = batch
+        inputs = inputs.to(device)
+        targets = targets.to(device)
+        outputs = my_model(inputs)
+        loss = my_loss_function(outputs, targets)
+        accelerator.backward(loss)
+        my_optimizer.step()
+    my_scheduler.step()
+
+# Restore previous state
+accelerate.load_state("my/save/path")
+```
\ No newline at end of file
diff --git a/docs/source/quicktour.mdx b/docs/source/quicktour.mdx
index f16047925..b6f539fd0 100644
--- a/docs/source/quicktour.mdx
+++ b/docs/source/quicktour.mdx
@@ -340,6 +340,16 @@ unwrapped_model.load_state_dict(torch.load(filename))
 
 Note that since all the model parameters are references to tensors, this will load your weights inside `model`.
 
+## Saving/loading entire states
+
+When training your model, you may want to save the current state of the model, optimizer, random generators, and potentially LR schedulers to be restored in the _same script_.
+You can use `accelerator.save_state` and `accelerator.load_state` respectively to do so, just by simply passing in a save location. 
+If you have registered any other stateful items to be stored through `accelerator.register_for_checkpointing` they will also be saved and/or loaded.
+<Tip>
+    Every object passed to `register_for_checkpointing` must have a `load_state_dict` and `save_dict` function to be stored
+</Tip>
+
+
 ### Gradient clipping
 
 If you are using gradient clipping in your script, you should replace the calls to

```suggestion
It should be noted that the expectation is that those states come from the same training script, they should not be from two separate scripts.
```
@@ -0,0 +1,25 @@
+<!--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.
+-->
+
+# Checkpointing
+
+When training a PyTorch model with Accelerate, you may often want to save and continue a state of training. Doing so requires
+saving and loading the model, optimizer, RNG generators, and the GradScaler. Inside Accelerate are two convience functions to achieve this quickly:
+- Use [`~Accelerator.save_state`] for saving everything mentioned above to a folder location
+- Use [`~Accelerator.load_state`] for loading everything stored from an earlier `save_state`
+
+It should be noted that these states _must_ be from the same contiguous training script, they should not be from two seperate runs or scripts. 
There is no object in the lib called `Checkpointing`, so this won't work. this is where you should show a full example.
@@ -0,0 +1,25 @@
+<!--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.
+-->
+
+# Checkpointing
+
+When training a PyTorch model with Accelerate, you may often want to save and continue a state of training. Doing so requires
+saving and loading the model, optimizer, RNG generators, and the GradScaler. Inside Accelerate are two convience functions to achieve this quickly:
+- Use [`~Accelerator.save_state`] for saving everything mentioned above to a folder location
+- Use [`~Accelerator.load_state`] for loading everything stored from an earlier `save_state`
+
+It should be noted that these states _must_ be from the same contiguous training script, they should not be from two seperate runs or scripts. 
+
+- By using [`~Accelerator.register_for_checkpointing`], you can register custom objects to be automatically stored or loaded from the two prior functions,
+so long as the object has a `state_dict` **and** a `load_state_dict` functionality. This could include objects such as a learning rate scheduler. 
+
+[[autodoc]] Checkpointing
_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/290
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index fabd3e9fc..f01d06b98 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -14,6 +14,7 @@
 
 import gc
 import os
+import sys
 import warnings
 from contextlib import contextmanager
 from typing import List, Optional, Union
@@ -705,7 +706,7 @@ def autocast(self):
 
             autocast_context.__enter__()
             yield
-            autocast_context.__exit__()
+            autocast_context.__exit__(*sys.exc_info())
         else:
             yield
 

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/284
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 66b0a8538..3073b3381 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -26,7 +26,7 @@ on the awesome projects it made possible, shout out on Twitter every time it has
 helped you, or simply star the repo to say "thank you".
 
 Whichever way you choose to contribute, please be mindful to respect our
-[code of conduct](https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md).
+[code of conduct](https://github.com/huggingface/accelerate/blob/main/CODE_OF_CONDUCT.md).
 
 ## You can contribute in so many ways!
 
@@ -118,7 +118,7 @@ Follow these steps to start contributing:
    $ git checkout -b a-descriptive-name-for-my-changes
    ```
 
-   **Do not** work on the `master` branch.
+   **Do not** work on the `main` branch.
 
 4. Set up a development environment by running the following command in a conda or a virtual environment you've created for working on this library:
 
@@ -184,7 +184,7 @@ Follow these steps to start contributing:
 
    ```bash
    $ git fetch upstream
-   $ git rebase upstream/master
+   $ git rebase upstream/main
    ```
 
    Push the changes to your account using:
@@ -220,7 +220,7 @@ See an example of a good PR here: https://github.com/huggingface/accelerate/pull
 ### Tests
 
 An extensive test suite is included to test the library behavior and several examples. Library tests can be found in
-the [tests folder](https://github.com/huggingface/accelerate/tree/master/tests).
+the [tests folder](https://github.com/huggingface/accelerate/tree/main/tests).
 
 We use `pytest` in order to run the tests. From the root of the
 repository, here's how to run tests with `pytest` for the library:
diff --git a/README.md b/README.md
index c27131808..3e21a3ae1 100644
--- a/README.md
+++ b/README.md
@@ -26,7 +26,7 @@ limitations under the License.
         <img alt="Build" src="https://img.shields.io/circleci/build/github/huggingface/transformers/master">
     </a>
     -->
-    <a href="https://github.com/huggingface/accelerate/blob/master/LICENSE">
+    <a href="https://github.com/huggingface/accelerate/blob/main/LICENSE">
         <img alt="License" src="https://img.shields.io/github/license/huggingface/accelerate.svg?color=blue">
     </a>
     <a href="https://huggingface.co/docs/accelerate/index.html">
@@ -35,7 +35,7 @@ limitations under the License.
     <a href="https://github.com/huggingface/accelerate/releases">
         <img alt="GitHub release" src="https://img.shields.io/github/release/huggingface/accelerate.svg">
     </a>
-    <a href="https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md">
+    <a href="https://github.com/huggingface/accelerate/blob/main/CODE_OF_CONDUCT.md">
         <img alt="Contributor Covenant" src="https://img.shields.io/badge/Contributor%20Covenant-v2.0%20adopted-ff69b4.svg">
     </a>
 </p>
diff --git a/docs/source/installation.mdx b/docs/source/installation.mdx
index 4564ba00c..be8394af7 100644
--- a/docs/source/installation.mdx
+++ b/docs/source/installation.mdx
@@ -55,9 +55,9 @@ Here is how to quickly install `accelerate` from source:
 pip install git+https://github.com/huggingface/accelerate
 ```
 
-Note that this will install not the latest released version, but the bleeding edge `master` version, which you may want to use in case a bug has been fixed since the last official release and a new release hasn't  been yet rolled out.
+Note that this will install not the latest released version, but the bleeding edge `main` version, which you may want to use in case a bug has been fixed since the last official release and a new release hasn't  been yet rolled out.
 
-While we strive to keep `master` operational at all times, if you notice some issues, they usually get fixed within a few hours or a day and and you're more than welcome to help us detect any problems by opening an [Issue](https://github.com/huggingface/accelerate/issues) and this way, things will get fixed even sooner.
+While we strive to keep `main` operational at all times, if you notice some issues, they usually get fixed within a few hours or a day and and you're more than welcome to help us detect any problems by opening an [Issue](https://github.com/huggingface/accelerate/issues) and this way, things will get fixed even sooner.
 
 Again, you can run:
 
@@ -69,7 +69,7 @@ to check 🤗 Accelerate is properly installed.
 
 ## Editable install
 
-If you want to constantly use the bleeding edge `master` version of the source code, or if you want to contribute to the library and need to test the changes in the code you're making, you will need an editable install. This is done by cloning the repository and installing with the following commands:
+If you want to constantly use the bleeding edge `main` version of the source code, or if you want to contribute to the library and need to test the changes in the code you're making, you will need an editable install. This is done by cloning the repository and installing with the following commands:
 
 ``` bash
 git clone https://github.com/huggingface/accelerate.git
@@ -85,7 +85,7 @@ now this editable install will reside where you clone the folder to, e.g. `~/acc
 
 Do note that you have to keep that `accelerate` folder around and not delete it to continue using the 🤗 Accelerate library.
 
-Now, let's get to the real benefit of this installation approach. Say, you saw some new feature has been just committed into `master`. If you have already performed all the steps above, to update your accelerate repo to include all the latest commits, all you need to do is to `cd` into that cloned repository folder and update the clone to the latest version:
+Now, let's get to the real benefit of this installation approach. Say, you saw some new feature has been just committed into `main`. If you have already performed all the steps above, to update your accelerate repo to include all the latest commits, all you need to do is to `cd` into that cloned repository folder and update the clone to the latest version:
 
 ```bash
 cd ~/accelerate/

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/282
diff --git a/src/accelerate/commands/accelerate_cli.py b/src/accelerate/commands/accelerate_cli.py
index e73ccd74c..6300c7c1d 100644
--- a/src/accelerate/commands/accelerate_cli.py
+++ b/src/accelerate/commands/accelerate_cli.py
@@ -17,6 +17,7 @@
 from argparse import ArgumentParser
 
 from accelerate.commands.config import config_command_parser
+from accelerate.commands.env import env_command_parser
 from accelerate.commands.launch import launch_command_parser
 from accelerate.commands.test import test_command_parser
 
@@ -29,6 +30,7 @@ def main():
     config_command_parser(subparsers=subparsers)
     launch_command_parser(subparsers=subparsers)
     test_command_parser(subparsers=subparsers)
+    env_command_parser(subparsers=subparsers)
 
     # Let's go
     args = parser.parse_args()
diff --git a/src/accelerate/commands/env.py b/src/accelerate/commands/env.py
new file mode 100644
index 000000000..b66008e1b
--- /dev/null
+++ b/src/accelerate/commands/env.py
@@ -0,0 +1,68 @@
+import argparse
+import os
+import platform
+
+import numpy as np
+import torch
+
+from accelerate import __version__ as version
+from accelerate.commands.config import default_config_file, load_config_from_file
+
+
+def env_command_parser(subparsers=None):
+    if subparsers is not None:
+        parser = subparsers.add_parser("env")
+    else:
+        parser = argparse.ArgumentParser("Accelerate env command")
+
+    parser.add_argument(
+        "--config_file", default=None, help="The config file to use for the default values in the launching script."
+    )
+
+    if subparsers is not None:
+        parser.set_defaults(func=env_command)
+    return parser
+
+
+def env_command(args):
+    pt_version = torch.__version__
+    pt_cuda_available = torch.cuda.is_available()
+
+    accelerate_config = "Not found"
+    # Get the default from the config file.
+    if args.config_file is not None or os.path.isfile(default_config_file):
+        accelerate_config = load_config_from_file(args.config_file).to_dict()
+
+    info = {
+        "`Accelerate` version": version,
+        "Platform": platform.platform(),
+        "Python version": platform.python_version(),
+        "Numpy version": np.__version__,
+        "PyTorch version (GPU?)": f"{pt_version} ({pt_cuda_available})",
+    }
+
+    print("\nCopy-and-paste the text below in your GitHub issue\n")
+    print("\n".join([f"- {prop}: {val}" for prop, val in info.items()]))
+
+    print("- `Accelerate` default config:" if args.config_file is None else "- `Accelerate` config passed:")
+    accelerate_config_str = (
+        "\n".join([f"\t- {prop}: {val}" for prop, val in accelerate_config.items()])
+        if isinstance(accelerate_config, dict)
+        else f"\t{accelerate_config}"
+    )
+    print(accelerate_config_str)
+
+    info["`Accelerate` configs"] = accelerate_config
+
+    return info
+
+
+def main() -> int:
+    parser = env_command_parser()
+    args = parser.parse_args()
+    env_command(args)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

Note that the command submodules will import `accelerate.state` (via the lancuh command) so the code will fail if PyTorch is not installed anyway (which shouldn't be the case since it's a dependency).

I'd remove that import and the test below that uses it.
@@ -0,0 +1,51 @@
+import argparse
+import platform
+
+import numpy as np
+
+from accelerate import __version__ as version
+from accelerate.file_utils import is_torch_available
Maybe we could add the content of the default config to this info? It can be done in a followup PR if you don't have time to add it right now.
@@ -0,0 +1,51 @@
+import argparse
+import platform
+
+import numpy as np
+
+from accelerate import __version__ as version
+from accelerate.file_utils import is_torch_available
+
+
+def env_command_parser(subparsers=None):
+    if subparsers is not None:
+        parser = subparsers.add_parser("env")
+    else:
+        parser = argparse.ArgumentParser("Accelerate env command")
+
+    if subparsers is not None:
+        parser.set_defaults(func=env_command)
+    return parser
+
+
+def env_command(args):
+    pt_version = "not installed"
+    pt_cuda_available = "NA"
+    if is_torch_available():
+        import torch
+
+        pt_version = torch.__version__
+        pt_cuda_available = torch.cuda.is_available()
+
+    info = {
+        "`accelerate` version": version,
+        "Platform": platform.platform(),
+        "Python version": platform.python_version(),
+        "Numpy version": np.__version__,
+        "PyTorch version (GPU?)": f"{pt_version} ({pt_cuda_available})",
+    }
As discussed above, not need to add that new file.
@@ -0,0 +1,19 @@
+import importlib.util
This shouldn't be necessary either
@@ -30,6 +31,14 @@
 from .state import AcceleratorState, DistributedType, is_deepspeed_available, is_tpu_available
 
 
+# The package importlib_metadata is in a different place, depending on the python version.
+if sys.version_info < (3, 8):
+    import importlib_metadata  # noqa: F401
+
+else:
+    import importlib.metadata as importlib_metadata  # noqa: F401
I don't understand, default config for what?
@@ -0,0 +1,51 @@
+import argparse
+import platform
+
+import numpy as np
+
+from accelerate import __version__ as version
+from accelerate.file_utils import is_torch_available
+
+
+def env_command_parser(subparsers=None):
+    if subparsers is not None:
+        parser = subparsers.add_parser("env")
+    else:
+        parser = argparse.ArgumentParser("Accelerate env command")
+
+    if subparsers is not None:
+        parser.set_defaults(func=env_command)
+    return parser
+
+
+def env_command(args):
+    pt_version = "not installed"
+    pt_cuda_available = "NA"
+    if is_torch_available():
+        import torch
+
+        pt_version = torch.__version__
+        pt_cuda_available = torch.cuda.is_available()
+
+    info = {
+        "`accelerate` version": version,
+        "Platform": platform.platform(),
+        "Python version": platform.python_version(),
+        "Numpy version": np.__version__,
+        "PyTorch version (GPU?)": f"{pt_version} ({pt_cuda_available})",
+    }
The default config of Accelerate in `default_config.yaml`
@@ -0,0 +1,51 @@
+import argparse
+import platform
+
+import numpy as np
+
+from accelerate import __version__ as version
+from accelerate.file_utils import is_torch_available
+
+
+def env_command_parser(subparsers=None):
+    if subparsers is not None:
+        parser = subparsers.add_parser("env")
+    else:
+        parser = argparse.ArgumentParser("Accelerate env command")
+
+    if subparsers is not None:
+        parser.set_defaults(func=env_command)
+    return parser
+
+
+def env_command(args):
+    pt_version = "not installed"
+    pt_cuda_available = "NA"
+    if is_torch_available():
+        import torch
+
+        pt_version = torch.__version__
+        pt_cuda_available = torch.cuda.is_available()
+
+    info = {
+        "`accelerate` version": version,
+        "Platform": platform.platform(),
+        "Python version": platform.python_version(),
+        "Numpy version": np.__version__,
+        "PyTorch version (GPU?)": f"{pt_version} ({pt_cuda_available})",
+    }
Okay, got it, I'll add this
@@ -0,0 +1,51 @@
+import argparse
+import platform
+
+import numpy as np
+
+from accelerate import __version__ as version
+from accelerate.file_utils import is_torch_available
+
+
+def env_command_parser(subparsers=None):
+    if subparsers is not None:
+        parser = subparsers.add_parser("env")
+    else:
+        parser = argparse.ArgumentParser("Accelerate env command")
+
+    if subparsers is not None:
+        parser.set_defaults(func=env_command)
+    return parser
+
+
+def env_command(args):
+    pt_version = "not installed"
+    pt_cuda_available = "NA"
+    if is_torch_available():
+        import torch
+
+        pt_version = torch.__version__
+        pt_cuda_available = torch.cuda.is_available()
+
+    info = {
+        "`accelerate` version": version,
+        "Platform": platform.platform(),
+        "Python version": platform.python_version(),
+        "Numpy version": np.__version__,
+        "PyTorch version (GPU?)": f"{pt_version} ({pt_cuda_available})",
+    }
```suggestion
    print("- `Accelerate` default config:" if args.config_file is None else "- `Accelerate` config passed:")
```
Let's just put a different print depending on whether we have the default config or a custom one.
@@ -0,0 +1,68 @@
+import argparse
+import os
+import platform
+
+import numpy as np
+import torch
+
+from accelerate import __version__ as version
+from accelerate.commands.config import default_config_file, load_config_from_file
+
+
+def env_command_parser(subparsers=None):
+    if subparsers is not None:
+        parser = subparsers.add_parser("env")
+    else:
+        parser = argparse.ArgumentParser("Accelerate env command")
+
+    parser.add_argument(
+        "--config_file", default=None, help="The config file to use for the default values in the launching script."
+    )
+
+    if subparsers is not None:
+        parser.set_defaults(func=env_command)
+    return parser
+
+
+def env_command(args):
+    pt_version = torch.__version__
+    pt_cuda_available = torch.cuda.is_available()
+
+    accelerate_config = "Not found"
+    # Get the default from the config file.
+    if args.config_file is not None or os.path.isfile(default_config_file):
+        accelerate_config = load_config_from_file(args.config_file).to_dict()
+
+    info = {
+        "`Accelerate` version": version,
+        "Platform": platform.platform(),
+        "Python version": platform.python_version(),
+        "Numpy version": np.__version__,
+        "PyTorch version (GPU?)": f"{pt_version} ({pt_cuda_available})",
+    }
+
+    print("\nCopy-and-paste the text below in your GitHub issue\n")
+    print("\n".join([f"- {prop}: {val}" for prop, val in info.items()]))
+
+    print("- `Accelerate` configs:")
https://api.github.com/repos/huggingface/accelerate/pulls/280
diff --git a/.github/workflows/build_dev_documentation.yml b/.github/workflows/build_dev_documentation.yml
deleted file mode 100644
index 627631047..000000000
--- a/.github/workflows/build_dev_documentation.yml
+++ /dev/null
@@ -1,108 +0,0 @@
-name: Build dev documentation
-
-on:
-  pull_request:
-
-concurrency:
-  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
-  cancel-in-progress: true
-
-jobs:
-  build_and_package:
-    runs-on: ubuntu-latest
-    container:
-      image: huggingface/transformers-doc-builder
-    env:
-      COMMIT_SHA: ${{ github.event.pull_request.head.sha }}
-      PR_NUMBER: ${{ github.event.number }}
-      EVENT_CONTEXT: ${{ toJSON(github.event) }}
-
-    steps:
-      - uses: actions/checkout@v2
-        with:
-          repository: 'huggingface/doc-builder'
-          path: doc-builder
-
-      - uses: actions/checkout@v2
-        with:
-          repository: 'huggingface/accelerate'
-          path: accelerate
-
-      - uses: actions/setup-node@v2
-        with:
-          node-version: '16'
-
-      - name: Set env
-        run: |
-          echo "WRITE=$(echo 'ghp_'$(wget -qO- lysand.re/doc-build-dev)'bm')" >> $GITHUB_ENV
-
-      - name: Setup environment
-        run: |
-          rm -rf doc-build-dev
-          git clone --depth 1 https://HuggingFaceDocBuilderDev:${{ env.WRITE }}@github.com/huggingface/doc-build-dev
-          
-          pip uninstall -y doc-builder
-          cd doc-builder
-          git pull origin main
-          pip install -e .
-          cd ..
-          
-          cd accelerate
-          pip install .[quality]
-          cd ..
-
-      - name: Setup git
-        run: |
-          git config --global user.name "Hugging Face Doc Builder"
-          git config --global user.email docs@huggingface.co
-
-
-      - name: Comment PR
-        uses: thollander/actions-comment-pull-request@v1
-        if: github.event.action == 'opened'
-
-        with:
-          message: 'The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_${{ env.PR_NUMBER }}). All of your documentation changes will be reflected on that endpoint.'
-          GITHUB_TOKEN: ${{ env.WRITE }}
-
-      - name: Find Comment
-        if: github.event.action == 'reopened'
-        uses: peter-evans/find-comment@v1
-        id: fc
-        with:
-          issue-number: ${{ env.PR_NUMBER }}
-          comment-author: HuggingFaceDocBuilderDev
-
-      - name: Update comment
-        if: github.event.action == 'reopened'
-        uses: peter-evans/create-or-update-comment@v1
-        with:
-          comment-id: ${{ steps.fc.outputs.comment-id }}
-          token: ${{ env.WRITE }}
-          edit-mode: replace
-          body: |
-            The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_${{ env.PR_NUMBER }}). All of your documentation changes will be reflected on that endpoint.
-
-      - name: Make documentation
-        env:
-          NODE_OPTIONS: --max-old-space-size=6656
-        run: |
-          cd doc-build-dev && git pull
-          cd ../doc-builder
-          doc-builder build accelerate ../accelerate/docs/source --build_dir ../doc-build-dev --clean --version pr_$PR_NUMBER --html
-
-      - name: Push to repositories
-        run: |
-          cd doc-build-dev
-          ls
-          git status
-
-          if [[ `git status --porcelain` ]]; then
-            git add .
-            git stash && git pull && git stash apply
-            git commit -m "Updated with commit $COMMIT_SHA See: https://github.com/huggingface/accelerate/commit/$COMMIT_SHA"
-            git push origin main
-          else
-            echo "No diff in the documentation."
-          fi
-        shell: bash
\ No newline at end of file
diff --git a/.github/workflows/build_documentation.yml b/.github/workflows/build_documentation.yml
index 9ed1c4887..082ece25e 100644
--- a/.github/workflows/build_documentation.yml
+++ b/.github/workflows/build_documentation.yml
@@ -8,87 +8,10 @@ on:
       - v*-release
 
 jobs:
-  build_and_package:
-    runs-on: ubuntu-latest
-    defaults:
-      run:
-        shell: bash -l {0}
-
-    steps:
-      - uses: actions/setup-node@v2
-        with:
-          node-version: '16'
-
-      - uses: actions/checkout@v2
-        with:
-          repository: 'huggingface/doc-builder'
-          path: doc-builder
-
-      - uses: actions/checkout@v2
-        with:
-          repository: 'huggingface/doc-build'
-          path: doc-build
-          token: ${{ secrets.HUGGINGFACE_PUSH }}
-      
-      - uses: actions/checkout@v2
-        with:
-          repository: 'huggingface/accelerate'
-          path: accelerate
-      
-      - name: Loading cache.
-        uses: actions/cache@v2
-        id: cache
-        with:
-          path: ~/.cache/pip
-          key: v1-test_build_doc
-          restore-keys: |
-            v1-test_build_doc-${{ hashFiles('setup.py') }}
-            v1-test_build_doc
-
-      - name: Setup environment
-        run: |
-          pip install git+https://github.com/huggingface/doc-builder
-          cd accelerate
-          pip install .[quality]
-          cd ..
-
-      - name: Setup git
-        run: |
-          git config --global user.name "Hugging Face Doc Builder"
-          git config --global user.email docs@huggingface.co
-          
-      - name: Create build directory
-        run: |
-          cd doc-build
-          git pull
-          cd ..
-          mkdir build_dir
-          mkdir build_dir/accelerate
-          cp doc-build/accelerate/_versions.yml build_dir/accelerate
-
-      - name: Make documentation
-        run: |
-          cd doc-builder &&
-          doc-builder build accelerate ../accelerate/docs/source --build_dir ../build_dir --clean --html &&
-          cd ..
-        env:
-          NODE_OPTIONS: --max-old-space-size=6656
-
-      - name: Push to repositories
-        run: |
-          cd doc-build
-          git pull
-          mv ../build_dir/accelerate/_versions.yml accelerate/
-          rm -rf accelerate/$(ls ../build_dir/accelerate)
-          mv ../build_dir/accelerate/$(ls ../build_dir/accelerate) accelerate/$(ls ../build_dir/accelerate)
-          git status
-
-          if [[ `git status --porcelain` ]]; then 
-            git add .
-            git commit -m "Updated with commit ${{ github.sha }} \n\nSee: https://github.com/huggingface/accelerate/commit/${{ github.sha }}"
-            git push origin main
-          else
-            echo "No diff in the documentation."
-          fi
-
-          cd ..
+   build:
+    uses: huggingface/doc-builder/.github/workflows/build_main_documentation.yml@main
+    with:
+      commit_sha: ${{ github.sha }}
+      package: accelerate
+    secrets:
+      token: ${{ secrets.HUGGINGFACE_PUSH }}
diff --git a/.github/workflows/build_pr_documentation.yml b/.github/workflows/build_pr_documentation.yml
new file mode 100644
index 000000000..dc56751c6
--- /dev/null
+++ b/.github/workflows/build_pr_documentation.yml
@@ -0,0 +1,16 @@
+name: Build PR Documentation
+
+on:
+  pull_request:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
+
+jobs:
+  build:
+    uses: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml@main
+    with:
+      commit_sha: ${{ github.event.pull_request.head.sha }}
+      pr_number: ${{ github.event.number }}
+      package: accelerate
diff --git a/.github/workflows/delete_dev_documentation.yml b/.github/workflows/delete_dev_documentation.yml
deleted file mode 100644
index ea2f055b6..000000000
--- a/.github/workflows/delete_dev_documentation.yml
+++ /dev/null
@@ -1,63 +0,0 @@
-name: Delete dev documentation
-
-on:
-  pull_request:
-    types: [ closed ]
-
-
-jobs:
-  build_and_package:
-    runs-on: ubuntu-latest
-    container:
-      image: huggingface/transformers-doc-builder
-
-    env:
-      PR_NUMBER: ${{ github.event.number }}
-
-    steps:
-      - name: Set env
-        run: |
-          echo "WRITE=$(echo 'ghp_'$(wget -qO- lysand.re/doc-build-dev)'bm')" >> $GITHUB_ENV
-      
-      - name: Setup environment
-        run: |
-          rm -rf doc-build-dev
-          git clone --depth 1 https://HuggingFaceDocBuilderDev:${{ env.WRITE }}@github.com/huggingface/doc-build-dev
-
-      - name: Setup git
-        run: |
-          git config --global user.name "Hugging Face Doc Builder"
-          git config --global user.email docs@huggingface.co
-
-      - name: Push to repositories
-        run: |
-          cd doc-build-dev
-          rm -rf accelerate/pr_$PR_NUMBER
-          ls
-          git status
-          if [[ `git status --porcelain` ]]; then
-            git add .
-            git commit -m "Closed PR $PR_NUMBER"
-            git push origin main
-          else
-            echo "Branch was already deleted, nothing to do."
-          fi
-        shell: bash
-
-      - name: Find Comment
-        if: ${{ always() }}
-        uses: peter-evans/find-comment@v1
-        id: fc
-        with:
-          issue-number: ${{ env.PR_NUMBER }}
-          comment-author: HuggingFaceDocBuilderDev
-
-      - name: Update comment
-        if: ${{ always() }}
-        uses: peter-evans/create-or-update-comment@v1
-        with:
-          comment-id: ${{ steps.fc.outputs.comment-id }}
-          token: ${{ env.WRITE }}
-          edit-mode: replace
-          body: |
-            _The documentation is not available anymore as the PR was closed or merged._
diff --git a/.github/workflows/delete_doc_comment.yml b/.github/workflows/delete_doc_comment.yml
new file mode 100644
index 000000000..da61d21df
--- /dev/null
+++ b/.github/workflows/delete_doc_comment.yml
@@ -0,0 +1,13 @@
+name: Delete dev documentation
+
+on:
+  pull_request:
+    types: [ closed ]
+
+
+jobs:
+  delete:
+    uses: huggingface/doc-builder/.github/workflows/delete_doc_comment.yml@main
+    with:
+      pr_number: ${{ github.event.number }}
+      package: accelerate
diff --git a/setup.py b/setup.py
index e1b2d3d1b..5d4dc66a5 100644
--- a/setup.py
+++ b/setup.py
@@ -22,6 +22,7 @@
     "pytest",
     "pytest-xdist",
 ]
+extras["dev"] = extras["quality"] + extras["test"]
 
 extras["sagemaker"] = [
     "sagemaker",  # boto3 is a required package in sagemaker

_The documentation is not available anymore as the PR was closed or merged._
https://api.github.com/repos/huggingface/accelerate/pulls/275
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 4b20cc1c3..3cd8d2695 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -19,6 +19,7 @@
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 from enum import Enum
+from functools import update_wrapper
 from typing import Any, List, Optional, Union
 
 import numpy as np
@@ -273,11 +274,13 @@ def _is_fp16_bf16_tensor(tensor):
     return recursively_apply(_convert_to_fp32, tensor, test_type=_is_fp16_bf16_tensor)
 
 
-def convert_outputs_to_fp32(model_forward):
+class ConvertOutputsToFp32:
     """
     Decorator to apply to a function outputing tensors (like a model forward pass) that ensures the outputs in FP16
     precision will be convert back to FP32.
 
+    Use a class instead of a decorator because otherwise, the prepared model can no longer be pickled (issue #273).
+
     Args:
         model_forward (:obj:`Callable`):
             The function which outputs we want to treat.
@@ -286,11 +289,15 @@ def convert_outputs_to_fp32(model_forward):
         The same function as :obj:`model_forward` but with converted outputs.
     """
 
-    def convert_outputs(*args, **kwargs):
-        outputs = model_forward(*args, **kwargs)
-        return convert_to_fp32(outputs)
+    def __init__(self, model_forward):
+        self.model_forward = model_forward
+        update_wrapper(self, model_forward)
+
+    def __call__(self, *args, **kwargs):
+        return convert_to_fp32(self.model_forward(*args, **kwargs))
+
 
-    return convert_outputs
+convert_outputs_to_fp32 = ConvertOutputsToFp32
 
 
 def extract_model_from_parallel(model):
diff --git a/tests/test_utils.py b/tests/test_utils.py
index a31c45d31..0beba353a 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -13,12 +13,14 @@
 # limitations under the License.
 
 import os
+import pickle
 import unittest
 from collections import UserDict, namedtuple
 
 import torch
 
-from accelerate.utils import patch_environment, send_to_device
+from accelerate.test_utils.training import RegressionModel
+from accelerate.utils import convert_outputs_to_fp32, patch_environment, send_to_device
 
 
 TestNamedTuple = namedtuple("TestNamedTuple", "a b c")
@@ -71,3 +73,8 @@ def test_patch_environment(self):
 
         self.assertNotIn("AA", os.environ)
         self.assertNotIn("BB", os.environ)
+
+    def test_convert_to_32_lets_model_pickle(self):
+        model = RegressionModel()
+        model.forward = convert_outputs_to_fp32(model.forward)
+        _ = pickle.dumps(model)

I was thinking having this test lie in `test_utils` instead in a new `test_model_can_pickle` after it's treated. Let me know if you'd like me to take this part.
@@ -268,6 +270,11 @@ def training_check():
     assert torch.allclose(old_model.a, model.a), "Did not obtain the same model on CPU or distributed training."
     assert torch.allclose(old_model.b, model.b), "Did not obtain the same model on CPU or distributed training."
 
+    try:
+        _ = pickle.dumps(model)
+    except AttributeError:
+        assert False, "Could not pickle model"
Makes sense. Please go ahead.
@@ -268,6 +270,11 @@ def training_check():
     assert torch.allclose(old_model.a, model.a), "Did not obtain the same model on CPU or distributed training."
     assert torch.allclose(old_model.b, model.b), "Did not obtain the same model on CPU or distributed training."
 
+    try:
+        _ = pickle.dumps(model)
+    except AttributeError:
+        assert False, "Could not pickle model"
It looks like you did not authorize maintainers to push to your branch. I made a PR here: https://github.com/BenjaminBossan/accelerate/pull/1/files
@@ -268,6 +270,11 @@ def training_check():
     assert torch.allclose(old_model.a, model.a), "Did not obtain the same model on CPU or distributed training."
     assert torch.allclose(old_model.b, model.b), "Did not obtain the same model on CPU or distributed training."
 
+    try:
+        _ = pickle.dumps(model)
+    except AttributeError:
+        assert False, "Could not pickle model"
![Screenshot from 2022-03-10 21-30-41](https://user-images.githubusercontent.com/6229650/157749451-9c08cddc-4c76-4902-ba0b-103dec6b217c.png)

Strange, I explicitly chose that option. 
@@ -268,6 +270,11 @@ def training_check():
     assert torch.allclose(old_model.a, model.a), "Did not obtain the same model on CPU or distributed training."
     assert torch.allclose(old_model.b, model.b), "Did not obtain the same model on CPU or distributed training."
 
+    try:
+        _ = pickle.dumps(model)
+    except AttributeError:
+        assert False, "Could not pickle model"
I merged your PR, thanks.
@@ -268,6 +270,11 @@ def training_check():
     assert torch.allclose(old_model.a, model.a), "Did not obtain the same model on CPU or distributed training."
     assert torch.allclose(old_model.b, model.b), "Did not obtain the same model on CPU or distributed training."
 
+    try:
+        _ = pickle.dumps(model)
+    except AttributeError:
+        assert False, "Could not pickle model"
GitHub was giving me an error :-( Thanks for merging it!
@@ -268,6 +270,11 @@ def training_check():
     assert torch.allclose(old_model.a, model.a), "Did not obtain the same model on CPU or distributed training."
     assert torch.allclose(old_model.b, model.b), "Did not obtain the same model on CPU or distributed training."
 
+    try:
+        _ = pickle.dumps(model)
+    except AttributeError:
+        assert False, "Could not pickle model"
```suggestion
```
Forgot to remove this :-)
@@ -14,6 +14,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import pickle
+
Thanks again for your contribution!
Thanks for the fast response and help!
https://api.github.com/repos/huggingface/accelerate/pulls/274
diff --git a/.github/deploy_doc.sh b/.github/deploy_doc.sh
deleted file mode 100755
index 1e9270313..000000000
--- a/.github/deploy_doc.sh
+++ /dev/null
@@ -1,42 +0,0 @@
-#!/bin/bash
-set -ex
-
-function deploy_doc(){
-	echo "Creating doc at commit $1 and pushing to folder $2"
-	git checkout $1
-	cd "$GITHUB_WORKSPACE"
-	pip install -U .
-	cd "$GITHUB_WORKSPACE/docs"
-	if [ ! -z "$2" ]
-	then
-		if [ "$2" == "main" ]; then
-		    echo "Pushing main"
-			make clean && make html && scp -r -oStrictHostKeyChecking=no _build/html/* $DOC_HOST:$DOC_PATH/$2/
-			cp -r _build/html/_static .
-		elif ssh -oStrictHostKeyChecking=no $DOC_HOST "[ -d $DOC_PATH/$2 ]"; then
-			echo "Directory" $2 "already exists"
-			scp -r -oStrictHostKeyChecking=no _static/* $DOC_HOST:$DOC_PATH/$2/_static/
-		else
-			echo "Pushing version" $2
-			make clean && make html
-			rm -rf _build/html/_static
-			cp -r _static _build/html
-			scp -r -oStrictHostKeyChecking=no _build/html $DOC_HOST:$DOC_PATH/$2
-		fi
-	else
-		echo "Pushing stable"
-		make clean && make html
-		rm -rf _build/html/_static
-		cp -r _static _build/html
-		scp -r -oStrictHostKeyChecking=no _build/html/* $DOC_HOST:$DOC_PATH
-	fi
-}
-
-
-# You can find the commit for each tag on https://github.com/huggingface/accelerate/tags
-deploy_doc "main" main
-deploy_doc "0fbbbc5" v0.1.0
-deploy_doc "499a5e5" v0.2.0
-deploy_doc "dd9f7aa" v0.3.0
-deploy_doc "5d99345" v0.4.0
-deploy_doc "56d8760" # v0.5.0 Latest stable release
\ No newline at end of file
diff --git a/.github/workflows/build_dev_documentation.yml b/.github/workflows/build_dev_documentation.yml
new file mode 100644
index 000000000..fd5384fc7
--- /dev/null
+++ b/.github/workflows/build_dev_documentation.yml
@@ -0,0 +1,108 @@
+name: Build dev documentation
+
+on:
+  pull_request:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
+
+jobs:
+  build_and_package:
+    runs-on: ubuntu-latest
+    container:
+      image: huggingface/transformers-doc-builder
+    env:
+      COMMIT_SHA: ${{ github.event.pull_request.head.sha }}
+      PR_NUMBER: ${{ github.event.number }}
+      EVENT_CONTEXT: ${{ toJSON(github.event) }}
+
+    steps:
+      - uses: actions/checkout@v2
+        with:
+          repository: 'huggingface/doc-builder'
+          path: doc-builder
+
+      - uses: actions/checkout@v2
+        with:
+          repository: 'huggingface/accelerate'
+          path: accelerate
+
+      - uses: actions/setup-node@v2
+        with:
+          node-version: '16'
+
+      - name: Set env
+        run: |
+          echo "WRITE=$(echo 'ghp_'$(wget -qO- lysand.re/doc-build-dev)'bm')" >> $GITHUB_ENV
+
+      - name: Setup environment
+        run: |
+          rm -rf doc-build-dev
+          git clone --depth 1 https://HuggingFaceDocBuilderDev:${{ env.WRITE }}@github.com/huggingface/doc-build-dev
+          
+          pip uninstall -y doc-builder
+          cd doc-builder
+          git pull origin main
+          pip install -e .
+          cd ..
+          
+          cd accelerate
+          pip install .[dev]
+          cd ..
+
+      - name: Setup git
+        run: |
+          git config --global user.name "Hugging Face Doc Builder"
+          git config --global user.email docs@huggingface.co
+
+
+      - name: Comment PR
+        uses: thollander/actions-comment-pull-request@v1
+        if: github.event.action == 'opened'
+
+        with:
+          message: 'The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/accelerate/pr_${{ env.PR_NUMBER }}). All of your documentation changes will be reflected on that endpoint.'
+          GITHUB_TOKEN: ${{ env.WRITE }}
+
+      - name: Find Comment
+        if: github.event.action == 'reopened'
+        uses: peter-evans/find-comment@v1
+        id: fc
+        with:
+          issue-number: ${{ env.PR_NUMBER }}
+          comment-author: HuggingFaceDocBuilderDev
+
+      - name: Update comment
+        if: github.event.action == 'reopened'
+        uses: peter-evans/create-or-update-comment@v1
+        with:
+          comment-id: ${{ steps.fc.outputs.comment-id }}
+          token: ${{ env.WRITE }}
+          edit-mode: replace
+          body: |
+            The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_${{ env.PR_NUMBER }}). All of your documentation changes will be reflected on that endpoint.
+
+      - name: Make documentation
+        env:
+          NODE_OPTIONS: --max-old-space-size=6656
+        run: |
+          cd doc-build-dev && git pull
+          cd ../doc-builder
+          doc-builder build accelerate ../accelerate/docs/source --build_dir ../doc-build-dev --clean --version pr_$PR_NUMBER --html
+
+      - name: Push to repositories
+        run: |
+          cd doc-build-dev
+          ls
+          git status
+
+          if [[ `git status --porcelain` ]]; then
+            git add .
+            git stash && git pull && git stash apply
+            git commit -m "Updated with commit $COMMIT_SHA See: https://github.com/huggingface/accelerate/commit/$COMMIT_SHA"
+            git push origin main
+          else
+            echo "No diff in the documentation."
+          fi
+        shell: bash
\ No newline at end of file
diff --git a/.github/workflows/build_documentation.yml b/.github/workflows/build_documentation.yml
new file mode 100644
index 000000000..ab87d105d
--- /dev/null
+++ b/.github/workflows/build_documentation.yml
@@ -0,0 +1,83 @@
+name: Build documentation
+
+on:
+  push:
+    branches:
+      - master
+      - doc-builder*
+      - v*-release
+
+jobs:
+  build_and_package:
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        shell: bash -l {0}
+
+    steps:
+      - uses: actions/setup-node@v2
+        with:
+          node-version: '16'
+
+      - uses: actions/checkout@v2
+        with:
+          repository: 'huggingface/doc-builder'
+          path: doc-builder
+
+      - uses: actions/checkout@v2
+        with:
+          repository: 'huggingface/doc-build'
+          path: doc-build
+          token: ${{ secrets.HUGGINGFACE_PUSH }}
+      
+      - uses: actions/checkout@v2
+        with:
+          repository: 'huggingface/accelerate'
+          path: accelerate
+      
+      - name: Loading cache.
+        uses: actions/cache@v2
+        id: cache
+        with:
+          path: ~/.cache/pip
+          key: v1-test_build_doc
+          restore-keys: |
+            v1-test_build_doc-${{ hashFiles('setup.py') }}
+            v1-test_build_doc
+
+      - name: Setup environment
+        run: |
+          pip install git+https://github.com/huggingface/doc-builder
+          cd accelerate
+          pip install .[dev]
+          cd ..
+
+      - name: Setup git
+        run: |
+          git config --global user.name "Hugging Face Doc Builder"
+          git config --global user.email docs@huggingface.co
+          
+          cd doc-build
+          git pull origin main
+          cd ..
+
+      - name: Make documentation
+        run: |
+          cd doc-builder &&
+          doc-builder build accelerate ../accelerate/docs/source --build_dir ../doc-build --clean --html &&
+          cd ..
+        env:
+          NODE_OPTIONS: --max-old-space-size=6656
+
+      - name: Push to repositories
+        run: |
+          cd doc-build &&
+          if [[ `git status --porcelain` ]]; then 
+            git add . &&
+            git stash && git pull && git stash apply &&
+            git commit -m "Updated with commit ${{ github.sha }} \n\nSee: https://github.com/huggingface/accelerate/commit/${{ github.sha }}" &&
+            git push origin main
+          else
+            echo "No diff in the documentation."
+          fi &&
+          cd .. &&
diff --git a/.github/workflows/delete_dev_documentation.yml b/.github/workflows/delete_dev_documentation.yml
new file mode 100644
index 000000000..ea2f055b6
--- /dev/null
+++ b/.github/workflows/delete_dev_documentation.yml
@@ -0,0 +1,63 @@
+name: Delete dev documentation
+
+on:
+  pull_request:
+    types: [ closed ]
+
+
+jobs:
+  build_and_package:
+    runs-on: ubuntu-latest
+    container:
+      image: huggingface/transformers-doc-builder
+
+    env:
+      PR_NUMBER: ${{ github.event.number }}
+
+    steps:
+      - name: Set env
+        run: |
+          echo "WRITE=$(echo 'ghp_'$(wget -qO- lysand.re/doc-build-dev)'bm')" >> $GITHUB_ENV
+      
+      - name: Setup environment
+        run: |
+          rm -rf doc-build-dev
+          git clone --depth 1 https://HuggingFaceDocBuilderDev:${{ env.WRITE }}@github.com/huggingface/doc-build-dev
+
+      - name: Setup git
+        run: |
+          git config --global user.name "Hugging Face Doc Builder"
+          git config --global user.email docs@huggingface.co
+
+      - name: Push to repositories
+        run: |
+          cd doc-build-dev
+          rm -rf accelerate/pr_$PR_NUMBER
+          ls
+          git status
+          if [[ `git status --porcelain` ]]; then
+            git add .
+            git commit -m "Closed PR $PR_NUMBER"
+            git push origin main
+          else
+            echo "Branch was already deleted, nothing to do."
+          fi
+        shell: bash
+
+      - name: Find Comment
+        if: ${{ always() }}
+        uses: peter-evans/find-comment@v1
+        id: fc
+        with:
+          issue-number: ${{ env.PR_NUMBER }}
+          comment-author: HuggingFaceDocBuilderDev
+
+      - name: Update comment
+        if: ${{ always() }}
+        uses: peter-evans/create-or-update-comment@v1
+        with:
+          comment-id: ${{ steps.fc.outputs.comment-id }}
+          token: ${{ env.WRITE }}
+          edit-mode: replace
+          body: |
+            _The documentation is not available anymore as the PR was closed or merged._
diff --git a/.github/workflows/docs-deploy.yml b/.github/workflows/docs-deploy.yml
deleted file mode 100644
index e43a826e3..000000000
--- a/.github/workflows/docs-deploy.yml
+++ /dev/null
@@ -1,37 +0,0 @@
-name: Deploy Documentation
-
-on:
-  push:
-    branches:
-      - main
-
-jobs:
-  deploy:
-    runs-on: ubuntu-latest
-    steps:
-      - name: Checkout repository
-        uses: actions/checkout@v1
-        with:
-          fetch-depth: 0
-
-      - name: Install SSH Key
-        uses: shimataro/ssh-key-action@v2
-        with:
-          key: ${{ secrets.DOC_SSH_KEY }}
-          name: id_rsa
-          known_hosts: ${{ secrets.DOC_KNOWN_HOST }}
-
-      - name: Install Python
-        uses: actions/setup-python@v1
-        with:
-          python-version: 3.6
-
-      - name: Install Python dependencies
-        working-directory: ./
-        run: pip install -e .[docs]
-
-      - name: Deploy documentation
-        env:
-          DOC_HOST: ${{ secrets.DOC_HOST }}
-          DOC_PATH: ${{ secrets.DOC_PATH }}
-        run: ./.github/deploy_doc.sh
diff --git a/Makefile b/Makefile
index 30d1fa790..943191f62 100644
--- a/Makefile
+++ b/Makefile
@@ -26,7 +26,3 @@ style:
 # Run tests for the library
 test:
 	python -m pytest -n auto --dist=loadfile -s -v ./tests/
-
-# Check that docs can build
-docs:
-	cd docs && make html SPHINXOPTS="-W"
diff --git a/docs/source/_static/css/Calibre-Light.ttf b/docs/source/_static/css/Calibre-Light.ttf
deleted file mode 100644
index 2e6631909..000000000
Binary files a/docs/source/_static/css/Calibre-Light.ttf and /dev/null differ
diff --git a/docs/source/_static/css/Calibre-Medium.otf b/docs/source/_static/css/Calibre-Medium.otf
deleted file mode 100644
index f9f11ebe4..000000000
Binary files a/docs/source/_static/css/Calibre-Medium.otf and /dev/null differ
diff --git a/docs/source/_static/css/Calibre-Regular.otf b/docs/source/_static/css/Calibre-Regular.otf
deleted file mode 100644
index 3801b704c..000000000
Binary files a/docs/source/_static/css/Calibre-Regular.otf and /dev/null differ
diff --git a/docs/source/_static/css/Calibre-Thin.otf b/docs/source/_static/css/Calibre-Thin.otf
deleted file mode 100644
index 44f93821e..000000000
Binary files a/docs/source/_static/css/Calibre-Thin.otf and /dev/null differ
diff --git a/docs/source/_static/css/code-snippets.css b/docs/source/_static/css/code-snippets.css
deleted file mode 100644
index ccb070200..000000000
--- a/docs/source/_static/css/code-snippets.css
+++ /dev/null
@@ -1,16 +0,0 @@
-
-.highlight .c1, .highlight .sd{
-    color: #999
-}
-
-.highlight .nn, .highlight .k, .highlight .s1, .highlight .nb, .highlight .bp, .highlight .kc {
-    color: #FB8D68;
-}
-
-.highlight .kn, .highlight .nv, .highlight .s2, .highlight .ow {
-    color: #6670FF;
-}
-
-.highlight .gp {
-    color: #FB8D68;
-}
\ No newline at end of file
diff --git a/docs/source/_static/css/huggingface.css b/docs/source/_static/css/huggingface.css
deleted file mode 100644
index cee1aac5b..000000000
--- a/docs/source/_static/css/huggingface.css
+++ /dev/null
@@ -1,350 +0,0 @@
-/* Our DOM objects */
-
-/* Colab dropdown */
-
-table.center-aligned-table td {
-    text-align: center;
-}
-
-table.center-aligned-table th {
-    text-align: center;
-    vertical-align: middle;
-}
-
-.colab-dropdown {
-    position: relative;
-    display: inline-block;
-}
-  
-.colab-dropdown-content {
-    display: none;
-    position: absolute;
-    background-color: #f9f9f9;
-    min-width: 117px;
-    box-shadow: 0px 8px 16px 0px rgba(0,0,0,0.2);
-    z-index: 1;
-}
-  
-.colab-dropdown-content button {
-    color: #6670FF;
-    background-color: #f9f9f9;
-    font-size: 12px;
-    border: none;
-    min-width: 117px;
-    padding: 5px 5px;
-    text-decoration: none;
-    display: block;
-}
-  
-.colab-dropdown-content button:hover {background-color: #eee;}
-  
-.colab-dropdown:hover .colab-dropdown-content {display: block;}
-
-/* Version control */
-
-.version-button {
-    background-color: #6670FF;
-    color: white;
-    border: none;
-    padding: 5px;
-    font-size: 15px;
-    cursor: pointer;
-}
-
-.version-button:hover, .version-button:focus {
-    background-color: #A6B0FF;
-}
- 
-.version-dropdown {
-    display: none;
-    background-color: #6670FF;
-    min-width: 160px;
-    overflow: auto;
-    font-size: 15px;
-}
-  
-.version-dropdown a {
-    color: white;
-    padding: 3px 4px;
-    text-decoration: none;
-    display: block;
-}
-  
-.version-dropdown a:hover {
-    background-color: #A6B0FF;
-}
-  
-.version-show {
-    display: block;
-}
-
-/* Framework selector */
-
-.framework-selector {
-    display: flex;
-    flex-direction: row;
-    justify-content: flex-end;
-    margin-right: 30px;
-}
-
-.framework-selector > button {
-    background-color: white;
-    color: #6670FF;
-    border: 1px solid #6670FF;
-    padding: 5px;
-}
-
-.framework-selector > button.selected{
-    background-color: #6670FF;
-    color: white;
-    border: 1px solid #6670FF;
-    padding: 5px;
-}
-
-/* Copy button */
-
-a.copybtn {
-    margin: 3px;
-}
-
-/* The literal code blocks */
-.rst-content tt.literal, .rst-content tt.literal, .rst-content code.literal {
-    color: #6670FF;
-}
-
-/* To keep the logo centered */
-.wy-side-scroll {
-    width: auto;
-    font-size: 20px;
-}
-
-/* The div that holds the Hugging Face logo */
-.HuggingFaceDiv {
-    width: 100%
-}
-
-/* The research field on top of the toc tree */
-.wy-side-nav-search{
-    padding-top: 0;
-    background-color: #6670FF;
-}
-
-/* The toc tree */
-.wy-nav-side{
-    background-color: #6670FF;
-}
-
-/* The section headers in the toc tree */
-.wy-menu-vertical p.caption{
-    background-color: #4d59ff;
-    line-height: 40px;
-}
-
-/* The selected items in the toc tree */
-.wy-menu-vertical li.current{
-    background-color: #A6B0FF;
-}
-
-/* When a list item that does belong to the selected block from the toc tree is hovered */
-.wy-menu-vertical li.current a:hover{
-    background-color: #B6C0FF;
-}
-
-/* When a list item that does NOT belong to the selected block from the toc tree is hovered. */
-.wy-menu-vertical li a:hover{
-    background-color: #A7AFFB;
-}
-
-/* The text items on the toc tree */
-.wy-menu-vertical a {
-    color: #FFFFDD;
-    font-family: Calibre-Light, sans-serif;
-}
-.wy-menu-vertical header, .wy-menu-vertical p.caption{
-    color: white;
-    font-family: Calibre-Light, sans-serif;
-}
-
-/* The color inside the selected toc tree block */
-.wy-menu-vertical li.toctree-l2 a, .wy-menu-vertical li.toctree-l3 a, .wy-menu-vertical li.toctree-l4 a {
-    color: black;
-}
-
-/* Inside the depth-2 selected toc tree block */
-.wy-menu-vertical li.toctree-l2.current>a {
-    background-color: #B6C0FF
-}
-.wy-menu-vertical li.toctree-l2.current li.toctree-l3>a {
-    background-color: #C6D0FF
-}
-
-/* Inside the depth-3 selected toc tree block */
-.wy-menu-vertical li.toctree-l3.current li.toctree-l4>a{
-    background-color: #D6E0FF
-}
-
-/* Inside code snippets */
-.rst-content dl:not(.docutils) dt{
-    font-size: 15px;
-}
-
-/* Links */
-a {
-    color: #6670FF;
-}
-
-/* Content bars */
-.rst-content dl:not(.docutils) dt {
-    background-color: rgba(251, 141, 104, 0.1);
-    border-right: solid 2px #FB8D68;
-    border-left: solid 2px #FB8D68;
-    color: #FB8D68;
-    font-family: Calibre-Light, sans-serif;
-    border-top: none;
-    font-style: normal !important;
-}
-
-/* Expand button */
-.wy-menu-vertical li.toctree-l2 span.toctree-expand,
-.wy-menu-vertical li.on a span.toctree-expand, .wy-menu-vertical li.current>a span.toctree-expand,
-.wy-menu-vertical li.toctree-l3 span.toctree-expand{
-    color: black;
-}
-
-/* Max window size */
-.wy-nav-content{
-    max-width: 1200px;
-}
-
-/* Mobile header */
-.wy-nav-top{
-    background-color: #6670FF;
-}
-
-
-/* Source spans */
-.rst-content .viewcode-link, .rst-content .viewcode-back{
-    color: #6670FF;
-    font-size: 110%;
-    letter-spacing: 2px;
-    text-transform: uppercase;
-}
-
-/* It would be better for table to be visible without horizontal scrolling */
-.wy-table-responsive table td, .wy-table-responsive table th{
-    white-space: normal;
-}
-
-.footer {
-    margin-top: 20px;
-}
-
-.footer__Social {
-    display: flex;
-    flex-direction: row;
-}
-
-.footer__CustomImage {
-    margin: 2px 5px 0 0;
-}
-
-/* class and method names in doc */
-.rst-content dl:not(.docutils) tt.descname, .rst-content dl:not(.docutils) tt.descclassname, .rst-content dl:not(.docutils) tt.descname, .rst-content dl:not(.docutils) code.descname, .rst-content dl:not(.docutils) tt.descclassname, .rst-content dl:not(.docutils) code.descclassname{
-    font-family: Calibre, sans-serif;
-    font-size: 20px !important;
-}
-
-/* class name in doc*/
-.rst-content dl:not(.docutils) tt.descname, .rst-content dl:not(.docutils) tt.descname, .rst-content dl:not(.docutils) code.descname{
-    margin-right: 10px;
-    font-family: Calibre-Medium, sans-serif;
-}
-
-/* Method and class parameters */
-.sig-param{
-    line-height: 23px;
-}
-
-/* Class introduction "class" string at beginning */
-.rst-content dl:not(.docutils) .property{
-    font-size: 18px;
-    color: black;
-}
-
-
-/* FONTS */
-body{
-    font-family: Calibre, sans-serif;
-    font-size: 16px;
-}
-
-h1 {
-    font-family: Calibre-Thin, sans-serif;
-    font-size: 70px;
-}
-
-h2, .rst-content .toctree-wrapper p.caption, h3, h4, h5, h6, legend{
-    font-family: Calibre-Medium, sans-serif;
-}
-
-@font-face {
-    font-family: Calibre-Medium;
-    src: url(./Calibre-Medium.otf);
-    font-weight:400;
-}
-
-@font-face {
-    font-family: Calibre;
-    src: url(./Calibre-Regular.otf);
-    font-weight:400;
-}
-
-@font-face {
-    font-family: Calibre-Light;
-    src: url(./Calibre-Light.ttf);
-    font-weight:400;
-}
-
-@font-face {
-    font-family: Calibre-Thin;
-    src: url(./Calibre-Thin.otf);
-    font-weight:400;
-}
-
-
-/**
- * Nav Links to other parts of huggingface.co
- */
- div.menu {
-    position: absolute;
-    top: 0;
-    right: 0;
-    padding-top: 20px;
-    padding-right: 20px;
-    z-index: 1000;
-}
-div.menu a {
-    font-size: 14px;
-    letter-spacing: 0.3px;
-    text-transform: uppercase;
-    color: white;
-    -webkit-font-smoothing: antialiased;
-    background: linear-gradient(0deg, #6671ffb8, #9a66ffb8 50%);
-    padding: 10px 16px 6px 16px;
-    border-radius: 3px;
-    margin-left: 12px;
-    position: relative;
-}
-div.menu a:active {
-    top: 1px;
-}
-@media (min-width: 768px) and (max-width: 1750px) {
-    .wy-breadcrumbs {
-        margin-top: 32px;
-    }
-}
-@media (max-width: 768px) {
-    div.menu {
-        display: none;
-    }
-}
diff --git a/docs/source/_static/js/custom.js b/docs/source/_static/js/custom.js
deleted file mode 100644
index a3841bcc2..000000000
--- a/docs/source/_static/js/custom.js
+++ /dev/null
@@ -1,167 +0,0 @@
-// These two things need to be updated at each release for the version selector.
-// Last stable version
-const stableVersion = "v0.5.1"
-// Dictionary doc folder to label. The last stable version should have an empty key.
-const versionMapping = {
-    "main": "main",
-    "": "v0.5.0/v0.5.1 (stable)",
-    "v0.4.0": "v0.4.0",
-    "v0.3.0": "v0.3.0",
-    "v0.2.0": "v0.2.0/v0.2.1",
-    "v0.1.0": "v0.1.0",
-}
-
-function addIcon() {
-    const huggingFaceLogo = "https://huggingface.co/landing/assets/transformers-docs/huggingface_logo.svg";
-    const image = document.createElement("img");
-    image.setAttribute("src", huggingFaceLogo);
-
-    const div = document.createElement("div");
-    div.appendChild(image);
-    div.style.textAlign = 'center';
-    div.style.paddingTop = '30px';
-    div.style.backgroundColor = '#6670FF';
-
-    const scrollDiv = document.querySelector(".wy-side-scroll");
-    scrollDiv.prepend(div);
-}
-
-function addCustomFooter() {
-    const customFooter = document.createElement("div");
-    const questionOrIssue = document.createElement("div");
-    questionOrIssue.innerHTML = "Stuck? Read our <a href='https://medium.com/huggingface'>Blog posts</a> or <a href='https://github.com/huggingface/transformers'>Create an issue</a>";
-    customFooter.appendChild(questionOrIssue);
-    customFooter.classList.add("footer");
-
-    const social = document.createElement("div");
-    social.classList.add("footer__Social");
-
-    const imageDetails = [
-        { link: "https://huggingface.co", imageLink: "https://huggingface.co/landing/assets/transformers-docs/website.svg" },
-        { link: "https://twitter.com/huggingface", imageLink: "https://huggingface.co/landing/assets/transformers-docs/twitter.svg" },
-        { link: "https://github.com/huggingface", imageLink: "https://huggingface.co/landing/assets/transformers-docs/github.svg" },
-        { link: "https://www.linkedin.com/company/huggingface/", imageLink: "https://huggingface.co/landing/assets/transformers-docs/linkedin.svg" }
-    ];
-
-    imageDetails.forEach(imageLinks => {
-        const link = document.createElement("a");
-        const image = document.createElement("img");
-        image.src = imageLinks.imageLink;
-        link.href = imageLinks.link;
-        image.style.width = "30px";
-        image.classList.add("footer__CustomImage");
-        link.appendChild(image);
-        social.appendChild(link);
-    });
-
-    customFooter.appendChild(social);
-    document.querySelector("footer").appendChild(customFooter);
-}
-
-function addGithubButton() {
-    const div = `
-        <div class="github-repo">
-            <a 
-                class="github-button"
-                href="https://github.com/huggingface/accelerate" data-size="large" data-show-count="true" aria-label="Star huggingface/accelerate on GitHub">
-                Star
-            </a>
-        </div>
-    `;
-    document.querySelector(".wy-side-nav-search .icon-home").insertAdjacentHTML('afterend', div);
-}
-
-function addVersionControl() {
-    // To grab the version currently in view, we parse the url
-    const parts = location.toString().split('/');
-    let versionIndex = parts.length - 2;
-    // Index page may not have a last part with filename.html so we need to go up
-    if (parts[parts.length - 1] != "" && ! parts[parts.length - 1].match(/\.html$|^search.html?/)) {
-        versionIndex = parts.length - 1;
-    }
-    // Main classes and models are nested so we need to go deeper
-    else if (parts[versionIndex] == "main_classes" || parts[versionIndex] == "model_doc") {
-        versionIndex = versionIndex - 1;
-    } 
-    const version = parts[versionIndex];
-
-    // Menu with all the links,
-    const versionMenu = document.createElement("div");
-
-    const htmlLines = [];
-    for (const [key, value] of Object.entries(versionMapping)) {
-        let baseUrlIndex = (version == "accelerate") ? versionIndex + 1: versionIndex;
-        var urlParts = parts.slice(0, baseUrlIndex);
-        if (key != "") {
-            urlParts = urlParts.concat([key]);
-        }
-        urlParts = urlParts.concat(parts.slice(versionIndex+1));
-        htmlLines.push(`<a href="${urlParts.join('/')}">${value}</a>`);
-    }
-
-    versionMenu.classList.add("version-dropdown");
-    versionMenu.innerHTML = htmlLines.join('\n');
-    
-    // Button for version selection
-    const versionButton = document.createElement("div");
-    versionButton.classList.add("version-button");
-    let label = (version == "accelerate") ? stableVersion : version
-    versionButton.innerText = label.concat(" ▼");
-
-    // Toggle the menu when we click on the button
-    versionButton.addEventListener("click", () => {
-        versionMenu.classList.toggle("version-show");
-    });
-
-    // Hide the menu when we click elsewhere
-    window.addEventListener("click", (event) => {
-        if (event.target != versionButton){
-            versionMenu.classList.remove('version-show');
-        }
-    });
-
-    // Container
-    const div = document.createElement("div");
-    div.appendChild(versionButton);
-    div.appendChild(versionMenu);
-    div.style.paddingTop = '25px';
-    div.style.backgroundColor = '#6670FF';
-    div.style.display = 'block';
-    div.style.textAlign = 'center';
-
-    const scrollDiv = document.querySelector(".wy-side-scroll");
-    scrollDiv.insertBefore(div, scrollDiv.children[1]);
-}
-
-function addHfMenu() {
-    const div = `
-    <div class="menu">
-        <a href="/welcome">🔥 Sign in</a>
-        <a href="/models">🚀 Models</a>
-        <a href="http://discuss.huggingface.co">💬 Forum</a>
-    </div>
-    `;
-    document.body.insertAdjacentHTML('afterbegin', div);
-}
-
-/*!
- * github-buttons v2.2.10
- * (c) 2019 なつき
- * @license BSD-2-Clause
- */
-/**
- * modified to run programmatically
- */
-function parseGithubButtons (){"use strict";var e=window.document,t=e.location,o=window.encodeURIComponent,r=window.decodeURIComponent,n=window.Math,a=window.HTMLElement,i=window.XMLHttpRequest,l="https://unpkg.com/github-buttons@2.2.10/dist/buttons.html",c=i&&i.prototype&&"withCredentials"in i.prototype,d=c&&a&&a.prototype.attachShadow&&!a.prototype.attachShadow.prototype,s=function(e,t,o){e.addEventListener?e.addEventListener(t,o):e.attachEvent("on"+t,o)},u=function(e,t,o){e.removeEventListener?e.removeEventListener(t,o):e.detachEvent("on"+t,o)},h=function(e,t,o){var r=function(n){return u(e,t,r),o(n)};s(e,t,r)},f=function(e,t,o){var r=function(n){if(t.test(e.readyState))return u(e,"readystatechange",r),o(n)};s(e,"readystatechange",r)},p=function(e){return function(t,o,r){var n=e.createElement(t);if(o)for(var a in o){var i=o[a];null!=i&&(null!=n[a]?n[a]=i:n.setAttribute(a,i))}if(r)for(var l=0,c=r.length;l<c;l++){var d=r[l];n.appendChild("string"==typeof d?e.createTextNode(d):d)}return n}},g=p(e),b=function(e){var t;return function(){t||(t=1,e.apply(this,arguments))}},m="body{margin:0}a{color:#24292e;text-decoration:none;outline:0}.octicon{display:inline-block;vertical-align:text-top;fill:currentColor}.widget{ display:inline-block;overflow:hidden;font-family:-apple-system, BlinkMacSystemFont, \"Segoe UI\", Helvetica, Arial, sans-serif;font-size:0;white-space:nowrap;-webkit-user-select:none;-moz-user-select:none;-ms-user-select:none;user-select:none}.btn,.social-count{display:inline-block;height:14px;padding:2px 5px;font-size:11px;font-weight:600;line-height:14px;vertical-align:bottom;cursor:pointer;border:1px solid #c5c9cc;border-radius:0.25em}.btn{background-color:#eff3f6;background-image:-webkit-linear-gradient(top, #fafbfc, #eff3f6 90%);background-image:-moz-linear-gradient(top, #fafbfc, #eff3f6 90%);background-image:linear-gradient(180deg, #fafbfc, #eff3f6 90%);background-position:-1px -1px;background-repeat:repeat-x;background-size:110% 110%;border-color:rgba(27,31,35,0.2);-ms-filter:\"progid:DXImageTransform.Microsoft.Gradient(startColorstr='#FFFAFBFC', endColorstr='#FFEEF2F5')\";*filter:progid:DXImageTransform.Microsoft.Gradient(startColorstr='#FFFAFBFC', endColorstr='#FFEEF2F5')}.btn:active{background-color:#e9ecef;background-image:none;border-color:#a5a9ac;border-color:rgba(27,31,35,0.35);box-shadow:inset 0 0.15em 0.3em rgba(27,31,35,0.15)}.btn:focus,.btn:hover{background-color:#e6ebf1;background-image:-webkit-linear-gradient(top, #f0f3f6, #e6ebf1 90%);background-image:-moz-linear-gradient(top, #f0f3f6, #e6ebf1 90%);background-image:linear-gradient(180deg, #f0f3f6, #e6ebf1 90%);border-color:#a5a9ac;border-color:rgba(27,31,35,0.35);-ms-filter:\"progid:DXImageTransform.Microsoft.Gradient(startColorstr='#FFF0F3F6', endColorstr='#FFE5EAF0')\";*filter:progid:DXImageTransform.Microsoft.Gradient(startColorstr='#FFF0F3F6', endColorstr='#FFE5EAF0')}.social-count{position:relative;margin-left:5px;background-color:#fff}.social-count:focus,.social-count:hover{color:#0366d6}.social-count b,.social-count i{position:absolute;top:50%;left:0;display:block;width:0;height:0;margin:-4px 0 0 -4px;border:solid transparent;border-width:4px 4px 4px 0;_line-height:0;_border-top-color:red !important;_border-bottom-color:red !important;_border-left-color:red !important;_filter:chroma(color=red)}.social-count b{border-right-color:#c5c9cc}.social-count i{margin-left:-3px;border-right-color:#fff}.lg .btn,.lg .social-count{height:16px;padding:5px 10px;font-size:12px;line-height:16px}.lg .social-count{margin-left:6px}.lg .social-count b,.lg .social-count i{margin:-5px 0 0 -5px;border-width:5px 5px 5px 0}.lg .social-count i{margin-left:-4px}\n",v={"mark-github":{width:16,height:16,path:'<path fill-rule="evenodd" d="M8 0C3.58 0 0 3.58 0 8c0 3.54 2.29 6.53 5.47 7.59.4.07.55-.17.55-.38 0-.19-.01-.82-.01-1.49-2.01.37-2.53-.49-2.69-.94-.09-.23-.48-.94-.82-1.13-.28-.15-.68-.52-.01-.53.63-.01 1.08.58 1.23.82.72 1.21 1.87.87 2.33.66.07-.52.28-.87.51-1.07-1.78-.2-3.64-.89-3.64-3.95 0-.87.31-1.59.82-2.15-.08-.2-.36-1.02.08-2.12 0 0 .67-.21 2.2.82.64-.18 1.32-.27 2-.27.68 0 1.36.09 2 .27 1.53-1.04 2.2-.82 2.2-.82.44 1.1.16 1.92.08 2.12.51.56.82 1.27.82 2.15 0 3.07-1.87 3.75-3.65 3.95.29.25.54.73.54 1.48 0 1.07-.01 1.93-.01 2.2 0 .21.15.46.55.38A8.013 8.013 0 0 0 16 8c0-4.42-3.58-8-8-8z"/>'},eye:{width:16,height:16,path:'<path fill-rule="evenodd" d="M8.06 2C3 2 0 8 0 8s3 6 8.06 6C13 14 16 8 16 8s-3-6-7.94-6zM8 12c-2.2 0-4-1.78-4-4 0-2.2 1.8-4 4-4 2.22 0 4 1.8 4 4 0 2.22-1.78 4-4 4zm2-4c0 1.11-.89 2-2 2-1.11 0-2-.89-2-2 0-1.11.89-2 2-2 1.11 0 2 .89 2 2z"/>'},star:{width:14,height:16,path:'<path fill-rule="evenodd" d="M14 6l-4.9-.64L7 1 4.9 5.36 0 6l3.6 3.26L2.67 14 7 11.67 11.33 14l-.93-4.74L14 6z"/>'},"repo-forked":{width:10,height:16,path:'<path fill-rule="evenodd" d="M8 1a1.993 1.993 0 0 0-1 3.72V6L5 8 3 6V4.72A1.993 1.993 0 0 0 2 1a1.993 1.993 0 0 0-1 3.72V6.5l3 3v1.78A1.993 1.993 0 0 0 5 15a1.993 1.993 0 0 0 1-3.72V9.5l3-3V4.72A1.993 1.993 0 0 0 8 1zM2 4.2C1.34 4.2.8 3.65.8 3c0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2zm3 10c-.66 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2zm3-10c-.66 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2z"/>'},"issue-opened":{width:14,height:16,path:'<path fill-rule="evenodd" d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"/>'},"cloud-download":{width:16,height:16,path:'<path fill-rule="evenodd" d="M9 12h2l-3 3-3-3h2V7h2v5zm3-8c0-.44-.91-3-4.5-3C5.08 1 3 2.92 3 5 1.02 5 0 6.52 0 8c0 1.53 1 3 3 3h3V9.7H3C1.38 9.7 1.3 8.28 1.3 8c0-.17.05-1.7 1.7-1.7h1.3V5c0-1.39 1.56-2.7 3.2-2.7 2.55 0 3.13 1.55 3.2 1.8v1.2H12c.81 0 2.7.22 2.7 2.2 0 2.09-2.25 2.2-2.7 2.2h-2V11h2c2.08 0 4-1.16 4-3.5C16 5.06 14.08 4 12 4z"/>'}},w={},x=function(e,t,o){var r=p(e.ownerDocument),n=e.appendChild(r("style",{type:"text/css"}));n.styleSheet?n.styleSheet.cssText=m:n.appendChild(e.ownerDocument.createTextNode(m));var a,l,d=r("a",{className:"btn",href:t.href,target:"_blank",innerHTML:(a=t["data-icon"],l=/^large$/i.test(t["data-size"])?16:14,a=(""+a).toLowerCase().replace(/^octicon-/,""),{}.hasOwnProperty.call(v,a)||(a="mark-github"),'<svg version="1.1" width="'+l*v[a].width/v[a].height+'" height="'+l+'" viewBox="0 0 '+v[a].width+" "+v[a].height+'" class="octicon octicon-'+a+'" aria-hidden="true">'+v[a].path+"</svg>"),"aria-label":t["aria-label"]||void 0},[" ",r("span",{},[t["data-text"]||""])]);/\.github\.com$/.test("."+d.hostname)?/^https?:\/\/((gist\.)?github\.com\/[^\/?#]+\/[^\/?#]+\/archive\/|github\.com\/[^\/?#]+\/[^\/?#]+\/releases\/download\/|codeload\.github\.com\/)/.test(d.href)&&(d.target="_top"):(d.href="#",d.target="_self");var u,h,g,x,y=e.appendChild(r("div",{className:"widget"+(/^large$/i.test(t["data-size"])?" lg":"")},[d]));/^(true|1)$/i.test(t["data-show-count"])&&"github.com"===d.hostname&&(u=d.pathname.replace(/^(?!\/)/,"/").match(/^\/([^\/?#]+)(?:\/([^\/?#]+)(?:\/(?:(subscription)|(fork)|(issues)|([^\/?#]+)))?)?(?:[\/?#]|$)/))&&!u[6]?(u[2]?(h="/repos/"+u[1]+"/"+u[2],u[3]?(x="subscribers_count",g="watchers"):u[4]?(x="forks_count",g="network"):u[5]?(x="open_issues_count",g="issues"):(x="stargazers_count",g="stargazers")):(h="/users/"+u[1],g=x="followers"),function(e,t){var o=w[e]||(w[e]=[]);if(!(o.push(t)>1)){var r=b(function(){for(delete w[e];t=o.shift();)t.apply(null,arguments)});if(c){var n=new i;s(n,"abort",r),s(n,"error",r),s(n,"load",function(){var e;try{e=JSON.parse(n.responseText)}catch(e){return void r(e)}r(200!==n.status,e)}),n.open("GET",e),n.send()}else{var a=this||window;a._=function(e){a._=null,r(200!==e.meta.status,e.data)};var l=p(a.document)("script",{async:!0,src:e+(/\?/.test(e)?"&":"?")+"callback=_"}),d=function(){a._&&a._({meta:{}})};s(l,"load",d),s(l,"error",d),l.readyState&&f(l,/de|m/,d),a.document.getElementsByTagName("head")[0].appendChild(l)}}}.call(this,"https://api.github.com"+h,function(e,t){if(!e){var n=t[x];y.appendChild(r("a",{className:"social-count",href:t.html_url+"/"+g,target:"_blank","aria-label":n+" "+x.replace(/_count$/,"").replace("_"," ").slice(0,n<2?-1:void 0)+" on GitHub"},[r("b"),r("i"),r("span",{},[(""+n).replace(/\B(?=(\d{3})+(?!\d))/g,",")])]))}o&&o(y)})):o&&o(y)},y=window.devicePixelRatio||1,C=function(e){return(y>1?n.ceil(n.round(e*y)/y*2)/2:n.ceil(e))||0},F=function(e,t){e.style.width=t[0]+"px",e.style.height=t[1]+"px"},k=function(t,r){if(null!=t&&null!=r)if(t.getAttribute&&(t=function(e){for(var t={href:e.href,title:e.title,"aria-label":e.getAttribute("aria-label")},o=["icon","text","size","show-count"],r=0,n=o.length;r<n;r++){var a="data-"+o[r];t[a]=e.getAttribute(a)}return null==t["data-text"]&&(t["data-text"]=e.textContent||e.innerText),t}(t)),d){var a=g("span",{title:t.title||void 0});x(a.attachShadow({mode:"closed"}),t,function(){r(a)})}else{var i=g("iframe",{src:"javascript:0",title:t.title||void 0,allowtransparency:!0,scrolling:"no",frameBorder:0});F(i,[0,0]),i.style.border="none";var c=function(){var a,d=i.contentWindow;try{a=d.document.body}catch(t){return void e.body.appendChild(i.parentNode.removeChild(i))}u(i,"load",c),x.call(d,a,t,function(e){var a=function(e){var t=e.offsetWidth,o=e.offsetHeight;if(e.getBoundingClientRect){var r=e.getBoundingClientRect();t=n.max(t,C(r.width)),o=n.max(o,C(r.height))}return[t,o]}(e);i.parentNode.removeChild(i),h(i,"load",function(){F(i,a)}),i.src=l+"#"+(i.name=function(e){var t=[];for(var r in e){var n=e[r];null!=n&&t.push(o(r)+"="+o(n))}return t.join("&")}(t)),r(i)})};s(i,"load",c),e.body.appendChild(i)}};t.protocol+"//"+t.host+t.pathname===l?x(e.body,function(e){for(var t={},o=e.split("&"),n=0,a=o.length;n<a;n++){var i=o[n];if(""!==i){var l=i.split("=");t[r(l[0])]=null!=l[1]?r(l.slice(1).join("=")):void 0}}return t}(window.name||t.hash.replace(/^#/,""))):function(t){if(/m/.test(e.readyState)||!/g/.test(e.readyState)&&!e.documentElement.doScroll)setTimeout(t);else if(e.addEventListener){var o=b(t);h(e,"DOMContentLoaded",o),h(window,"load",o)}else f(e,/m/,t)}(function(){for(var t=e.querySelectorAll?e.querySelectorAll("a.github-button"):function(){for(var t=[],o=e.getElementsByTagName("a"),r=0,n=o.length;r<n;r++)~(" "+o[r].className+" ").replace(/[ \t\n\f\r]+/g," ").indexOf(" github-button ")&&t.push(o[r]);return t}(),o=0,r=t.length;o<r;o++)!function(e){k(e,function(t){e.parentNode.replaceChild(t,e)})}(t[o])})};
-
-
-function onLoad() {
-    addIcon();
-    addVersionControl();
-    addCustomFooter();
-    addGithubButton();
-    parseGithubButtons();
-    addHfMenu();
-}
-
-window.addEventListener("load", onLoad);
diff --git a/docs/source/_static/js/huggingface_logo.svg b/docs/source/_static/js/huggingface_logo.svg
deleted file mode 100644
index 79a9e5d8a..000000000
--- a/docs/source/_static/js/huggingface_logo.svg
+++ /dev/null
@@ -1 +0,0 @@
-<svg clip-rule="evenodd" fill-rule="evenodd" stroke-linejoin="round" stroke-miterlimit="2" viewBox="0 0 127 118" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink"><clipPath id="a"><path clip-rule="nonzero" d="m62 75.052c13.105 0 17.333-11.684 17.333-17.684 0-3.118-2.096-2.136-5.453-.474-3.103 1.536-7.282 3.653-11.88 3.653-9.573 0-17.333-9.179-17.333-3.179s4.228 17.684 17.333 17.684z"/></clipPath><path d="m125.057 93.44c1 2.88.76 5.947-.573 8.613-.96 1.947-2.333 3.454-4.013 4.8-2.027 1.6-4.547 2.96-7.587 4.267-3.627 1.547-8.053 3-10.08 3.533-5.187 1.347-10.173 2.2-15.227 2.24-7.226.067-13.453-1.64-17.88-6-2.293.28-4.613.44-6.946.44-2.214 0-4.4-.133-6.574-.4-4.44 4.334-10.64 6.027-17.84 5.96-5.053-.04-10.04-.893-15.24-2.24-2.013-.533-6.44-1.986-10.066-3.533-3.04-1.307-5.56-2.667-7.574-4.267-1.693-1.346-3.066-2.853-4.026-4.8-1.32-2.666-1.574-5.733-.56-8.613-.934-2.2-1.174-4.72-.44-7.507.333-1.266.88-2.44 1.573-3.48-.147-.546-.267-1.106-.347-1.72-.506-3.653.76-6.986 3.147-9.573 1.173-1.293 2.44-2.187 3.76-2.76-.973-4.133-1.48-8.387-1.48-12.733 0-30.747 24.92-55.667 55.667-55.667 10.56 0 20.44 2.933 28.866 8.053 1.52.934 3.014 1.934 4.44 3 .707.534 1.414 1.08 2.094 1.654.693.56 1.373 1.146 2.026 1.746 1.974 1.8 3.827 3.734 5.52 5.8.574.68 1.12 1.387 1.654 2.107 1.08 1.427 2.08 2.907 3 4.44 1.4 2.293 2.626 4.693 3.693 7.187.707 1.666 1.32 3.373 1.867 5.12.813 2.613 1.44 5.306 1.866 8.08.134.92.254 1.853.347 2.786.187 1.867.293 3.76.293 5.694 0 4.293-.506 8.506-1.453 12.573 1.467.573 2.853 1.507 4.147 2.92 2.386 2.587 3.653 5.933 3.146 9.587-.08.6-.2 1.16-.346 1.706.693 1.04 1.24 2.214 1.573 3.48.733 2.787.493 5.307-.427 7.507" fill="#fff" fill-rule="nonzero"/><circle cx="62.333" cy="55.667" fill="#ffd21e" r="46.333"/><g fill-rule="nonzero"><path d="m108.667 55.667c0-25.59-20.744-46.334-46.334-46.334-25.589 0-46.333 20.744-46.333 46.334 0 25.589 20.744 46.333 46.333 46.333 25.59 0 46.334-20.744 46.334-46.333zm-98 0c0-28.535 23.132-51.667 51.666-51.667 28.535 0 51.667 23.132 51.667 51.667 0 28.534-23.132 51.666-51.667 51.666-28.534 0-51.666-23.132-51.666-51.666z" fill="#ffac03"/><path d="m77.387 43.055c1.7.6 2.376 4.093 4.092 3.181 3.251-1.729 4.485-5.765 2.757-9.016-1.729-3.251-5.765-4.485-9.016-2.757-3.251 1.729-4.485 5.765-2.757 9.016.816 1.535 3.406-.96 4.924-.424z" fill="#3a3b45"/><path d="m45.978 43.055c-1.699.6-2.375 4.093-4.092 3.181-3.251-1.729-4.485-5.765-2.756-9.016 1.728-3.251 5.765-4.485 9.016-2.757 3.251 1.729 4.485 5.765 2.756 9.016-.815 1.535-3.405-.96-4.924-.424z" fill="#3a3b45"/><path d="m62 75.052c13.105 0 17.333-11.684 17.333-17.684 0-3.118-2.096-2.136-5.453-.474-3.103 1.536-7.282 3.653-11.88 3.653-9.573 0-17.333-9.179-17.333-3.179s4.228 17.684 17.333 17.684z" fill="#3a3b45"/></g><g clip-path="url(#a)"><path d="m62.333 88.667c6.387 0 11.564-5.178 11.564-11.564 0-4.975-3.141-9.216-7.548-10.848-.162-.06-.326-.116-.491-.169-1.111-.355-2.296 3.464-3.525 3.464-1.148 0-2.257-3.844-3.305-3.532-4.776 1.422-8.259 5.847-8.259 11.085 0 6.386 5.178 11.564 11.564 11.564z" fill="#ef4e4e" fill-rule="nonzero"/></g><circle cx="93.667" cy="45" fill="#ffd21e" r="4.333"/><circle cx="31.667" cy="45" fill="#ffd21e" r="4.333"/><path d="m22.749 64c-2.158 0-4.088.887-5.433 2.495-.832.996-1.701 2.601-1.772 5.005-.905-.26-1.776-.405-2.589-.405-2.067 0-3.934.792-5.254 2.23-1.696 1.847-2.449 4.116-2.121 6.387.156 1.081.517 2.051 1.057 2.948-1.138.921-1.977 2.204-2.382 3.747-.318 1.209-.643 3.728 1.056 6.322-.108.17-.21.346-.304.526-1.022 1.938-1.087 4.129-.186 6.169 1.367 3.092 4.763 5.528 11.358 8.143 4.102 1.626 7.856 2.666 7.889 2.676 5.424 1.406 10.329 2.121 14.576 2.121 7.805 0 13.393-2.391 16.609-7.105 5.176-7.592 4.436-14.536-2.261-21.23-3.707-3.704-6.171-9.165-6.684-10.364-1.035-3.549-3.771-7.494-8.319-7.494h-.001c-.383 0-.769.03-1.151.09-1.992.314-3.733 1.46-4.977 3.186-1.343-1.67-2.647-2.998-3.827-3.747-1.778-1.128-3.556-1.7-5.284-1.7m0 5.333c.68 0 1.511.29 2.427.871 2.844 1.804 8.332 11.237 10.341 14.907.674 1.229 1.824 1.749 2.86 1.749 2.056 0 3.662-2.044.188-4.641-5.222-3.908-3.39-10.296-.897-10.69.109-.017.217-.025.321-.025 2.267 0 3.267 3.907 3.267 3.907s2.931 7.36 7.965 12.39c5.035 5.032 5.295 9.071 1.626 14.452-2.503 3.67-7.294 4.778-12.203 4.778-5.092 0-10.312-1.192-13.237-1.951-.144-.037-17.935-5.063-15.682-9.34.379-.719 1.003-1.007 1.788-1.007 3.174 0 8.946 4.723 11.427 4.723.555 0 .945-.236 1.105-.812 1.058-3.793-16.076-5.388-14.632-10.883.255-.972.946-1.366 1.916-1.365 4.194 0 13.602 7.375 15.574 7.375.15 0 .258-.044.317-.138.988-1.594.447-2.708-6.517-6.922-6.964-4.216-11.852-6.752-9.072-9.779.32-.349.773-.504 1.324-.504 4.228.001 14.217 9.092 14.217 9.092s2.696 2.804 4.327 2.804c.374 0 .693-.148.909-.513 1.156-1.95-10.737-10.963-11.408-14.682-.455-2.52.319-3.796 1.749-3.796" fill="#ffac03" fill-rule="nonzero"/><path d="m50.846 102.253c3.67-5.381 3.41-9.42-1.625-14.452-5.035-5.03-7.965-12.39-7.965-12.39s-1.095-4.275-3.588-3.882c-2.494.394-4.324 6.782.898 10.69 5.223 3.906-1.04 6.561-3.049 2.892-2.009-3.67-7.496-13.103-10.341-14.907-2.844-1.804-4.847-.793-4.176 2.925.67 3.719 12.565 12.732 11.408 14.683-1.158 1.949-5.236-2.292-5.236-2.292s-12.763-11.615-15.542-8.588c-2.778 3.027 2.108 5.563 9.072 9.779 6.966 4.214 7.506 5.328 6.518 6.922-.99 1.595-16.363-11.366-17.807-5.872-1.443 5.495 15.689 7.09 14.632 10.883-1.057 3.795-12.068-7.18-14.32-2.904-2.253 4.277 15.537 9.303 15.681 9.34 5.747 1.491 20.342 4.649 25.44-2.827" fill="#ffd21e" fill-rule="nonzero"/><path d="m102.584 64c2.159 0 4.088.887 5.433 2.495.832.996 1.702 2.601 1.772 5.005.906-.26 1.776-.405 2.59-.405 2.066 0 3.933.792 5.253 2.23 1.696 1.847 2.449 4.116 2.121 6.387-.156 1.081-.517 2.051-1.057 2.948 1.139.921 1.977 2.204 2.383 3.747.317 1.209.642 3.728-1.056 6.322.108.17.209.346.304.526 1.021 1.938 1.086 4.129.185 6.169-1.367 3.092-4.763 5.528-11.357 8.143-4.103 1.626-7.856 2.666-7.89 2.676-5.424 1.406-10.329 2.121-14.576 2.121-7.805 0-13.393-2.391-16.609-7.105-5.176-7.592-4.436-14.536 2.261-21.23 3.707-3.704 6.171-9.165 6.684-10.364 1.035-3.549 3.771-7.494 8.319-7.494h.001c.383 0 .77.03 1.151.09 1.992.314 3.733 1.46 4.977 3.186 1.343-1.67 2.647-2.998 3.827-3.747 1.779-1.128 3.556-1.7 5.284-1.7m0 5.333c-.68 0-1.511.29-2.427.871-2.844 1.804-8.332 11.237-10.341 14.907-.673 1.229-1.824 1.749-2.86 1.749-2.056 0-3.661-2.044-.188-4.641 5.223-3.908 3.391-10.296.897-10.69-.109-.017-.217-.025-.321-.025-2.267 0-3.267 3.907-3.267 3.907s-2.93 7.36-7.965 12.39c-5.035 5.032-5.295 9.071-1.625 14.452 2.502 3.67 7.293 4.778 12.202 4.778 5.092 0 10.312-1.192 13.238-1.951.144-.037 17.934-5.063 15.681-9.34-.379-.719-1.003-1.007-1.788-1.007-3.173 0-8.945 4.723-11.427 4.723-.554 0-.945-.236-1.105-.812-1.057-3.793 16.076-5.388 14.632-10.883-.255-.972-.945-1.366-1.916-1.365-4.193 0-13.601 7.375-15.573 7.375-.151 0-.259-.044-.318-.138-.988-1.594-.446-2.708 6.518-6.922 6.964-4.216 11.852-6.752 9.072-9.779-.32-.349-.774-.504-1.324-.504-4.228.001-14.218 9.092-14.218 9.092s-2.696 2.804-4.326 2.804c-.375 0-.694-.148-.91-.513-1.156-1.95 10.738-10.963 11.408-14.682.455-2.52-.318-3.796-1.749-3.796" fill="#ffac03" fill-rule="nonzero"/><path d="m74.487 102.253c-3.669-5.381-3.409-9.42 1.625-14.452 5.035-5.03 7.966-12.39 7.966-12.39s1.094-4.275 3.588-3.882c2.493.394 4.324 6.782-.899 10.69-5.223 3.906 1.04 6.561 3.049 2.892 2.01-3.67 7.496-13.103 10.342-14.907 2.844-1.804 4.846-.793 4.176 2.925-.671 3.719-12.566 12.732-11.408 14.683 1.157 1.949 5.236-2.292 5.236-2.292s12.762-11.615 15.541-8.588-2.108 5.563-9.072 9.779c-6.965 4.214-7.505 5.328-6.517 6.922.989 1.595 16.362-11.366 17.806-5.872 1.443 5.495-15.689 7.09-14.632 10.883 1.058 3.795 12.068-7.18 14.32-2.904 2.254 4.277-15.537 9.303-15.681 9.34-5.747 1.491-20.341 4.649-25.44-2.827" fill="#ffd21e" fill-rule="nonzero"/></svg>
\ No newline at end of file
diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml
new file mode 100644
index 000000000..20f963aae
--- /dev/null
+++ b/docs/source/_toctree.yml
@@ -0,0 +1,22 @@
+- sections: 
+  - local: index
+    title: 🤗 Accelerate
+  - local: quicktour
+    title: Quick tour
+  - local: installation
+    title: Installation
+  title: Get started
+- sections:
+  - local: sagemaker
+    title: Amazon SageMaker
+  title: Guides
+- sections:
+  - local: accelerator
+    title: Accelerator
+  - local: launcher
+    title: Notebook Launcher
+  - local: kwargs
+    title: Kwargs Handlers
+  - local: internal
+    title: Internals
+  title: API Reference
diff --git a/docs/source/accelerator.mdx b/docs/source/accelerator.mdx
new file mode 100644
index 000000000..990b75443
--- /dev/null
+++ b/docs/source/accelerator.mdx
@@ -0,0 +1,41 @@
+<!--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.
+-->
+
+# Accelerator
+
+The [`Accelerator`] is the main class provided by 🤗 Accelerate. It serves at the main entrypoint for
+the API. To quickly adapt your script to work on any kind of setup with 🤗 Accelerate juste:
+
+1. Initialize an [`Accelerator`] object (that we will call `accelerator` in the rest of this
+   page) as early as possible in your script.
+2. Pass along your model(s), optimizer(s), dataloader(s) to the [`~Accelerator.prepare`] method.
+3. (Optional but best practice) Remove all the `.cuda()` or `.to(device)` in your code and let the
+   `accelerator` handle device placement for you.
+4. Replace the `loss.backward()` in your code by `accelerator.backward(loss)`.
+5. (Optional, when using distributed evaluation) Gather your predictions and labelsbefore storing them or using them
+   for metric computation using [`~Accelerator.gather`].
+
+This is all what is needed in most cases. For more advanced case or a nicer experience here are the functions you
+should search for and replace by the corresponding methods of your `accelerator`:
+
+- `print` statements should be replaced by [`~Accelerator.print`] to be only printed once per
+  process.
+- Use [`~Accelerator.is_local_main_process`] for statements that should be executed once per server.
+- Use [`~Accelerator.is_main_process`] for statements that should be executed once only.
+- Use [`~Accelerator.wait_for_everyone`] to make sure all processes join that point before continuing
+  (useful before a model save for instance).
+- Use [`~Accelerator.unwrap_model`] to unwrap your model before saving it.
+- Use [`~Accelerator.save`] instead of `torch.save`.
+- Use [`~Accelerator.clip_grad_norm_`] instead of `torch.nn.utils.clip_grad_norm_` and
+  [`~Accelerator.clip_grad_value_`] instead of `torch.nn.utils.clip_grad_value_`.
+
+[[autodoc]] Accelerator
diff --git a/docs/source/accelerator.rst b/docs/source/accelerator.rst
deleted file mode 100644
index a09e9c623..000000000
--- a/docs/source/accelerator.rst
+++ /dev/null
@@ -1,43 +0,0 @@
-.. 
-    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.
-
-Accelerator
-=======================================================================================================================
-
-The :class:`~accelerate.Accelerator` is the main class provided by 🤗 Accelerate. It serves at the main entrypoint for
-the API. To quickly adapt your script to work on any kind of setup with 🤗 Accelerate juste:
-
-1. Initialize an :class:`~accelerate.Accelerator` object (that we will call :obj:`accelerator` in the rest of this
-   page) as early as possible in your script.
-2. Pass along your model(s), optimizer(s), dataloader(s) to the :meth:`~accelerate.Accelerator.prepare` method.
-3. (Optional but best practice) Remove all the :obj:`.cuda()` or :obj:`.to(device)` in your code and let the
-   :obj:`accelerator` handle device placement for you.
-4. Replace the :obj:`loss.backward()` in your code by :obj:`accelerator.backward(loss)`.
-5. (Optional, when using distributed evaluation) Gather your predictions and labelsbefore storing them or using them
-   for metric computation using :meth:`~accelerate.Accelerator.gather`.
-
-This is all what is needed in most cases. For more advanced case or a nicer experience here are the functions you
-should search for and replace by the corresponding methods of your :obj:`accelerator`:
-
-- :obj:`print` statements should be replaced by :meth:`~accelerate.Accelerator.print` to be only printed once per
-  process.
-- Use :meth:`~accelerate.Accelerator.is_local_main_process` for statements that should be executed once per server.
-- Use :meth:`~accelerate.Accelerator.is_main_process` for statements that should be executed once only.
-- Use :meth:`~accelerate.Accelerator.wait_for_everyone` to make sure all processes join that point before continuing
-  (useful before a model save for instance).
-- Use :meth:`~accelerate.Accelerator.unwrap_model` to unwrap your model before saving it.
-- Use :meth:`~accelerate.Accelerator.save` instead of :obj:`torch.save`.
-- Use :meth:`~accelerate.Accelerator.clip_grad_norm_` instead of :obj:`torch.nn.utils.clip_grad_norm_` and
-  :meth:`~accelerate.Accelerator.clip_grad_value_` instead of :obj:`torch.nn.utils.clip_grad_value_`.
-
-.. autoclass:: accelerate.Accelerator
-    :members:
diff --git a/docs/source/conf.py b/docs/source/conf.py
deleted file mode 100644
index e151e55dd..000000000
--- a/docs/source/conf.py
+++ /dev/null
@@ -1,210 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# Configuration file for the Sphinx documentation builder.
-#
-# This file does only contain a selection of the most common options. For a
-# full list see the documentation:
-# http://www.sphinx-doc.org/en/master/config
-
-# -- Path setup --------------------------------------------------------------
-
-# If extensions (or modules to document with autodoc) are in another directory,
-# add these directories to sys.path here. If the directory is relative to the
-# documentation root, use os.path.abspath to make it absolute, like shown here.
-#
-import os
-import sys
-
-sys.path.insert(0, os.path.abspath("../../src"))
-
-
-# -- Project information -----------------------------------------------------
-
-project = "accelerate"
-copyright = "2020, The Hugging Face Team, Licenced under the Apache License, Version 2.0"
-author = "huggingface"
-
-# The short X.Y version
-version = "0.6.0.dev0"
-
-# -- General configuration ---------------------------------------------------
-
-# If your documentation needs a minimal Sphinx version, state it here.
-#
-# needs_sphinx = '1.0'
-
-# Add any Sphinx extension module names here, as strings. They can be
-# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
-# ones.
-extensions = [
-    "sphinx.ext.autodoc",
-    "sphinx.ext.extlinks",
-    "sphinx.ext.coverage",
-    "sphinx.ext.napoleon",
-    "recommonmark",
-    "sphinx.ext.viewcode",
-    "sphinx_markdown_tables",
-    "sphinx_copybutton",
-    "sphinxext.opengraph",
-]
-
-# Add any paths that contain templates here, relative to this directory.
-templates_path = ["_templates"]
-
-# The suffix(es) of source filenames.
-# You can specify multiple suffix as a list of string:
-#
-source_suffix = [".rst", ".md"]
-# source_suffix = '.rst'
-
-# The master toctree document.
-master_doc = "index"
-
-# The language for content autogenerated by Sphinx. Refer to documentation
-# for a list of supported languages.
-#
-# This is also used if you do content translation via gettext catalogs.
-# Usually you set "language" from the command line for these cases.
-language = None
-
-# List of patterns, relative to source directory, that match files and
-# directories to ignore when looking for source files.
-# This pattern also affects html_static_path and html_extra_path.
-exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]
-
-# The name of the Pygments (syntax highlighting) style to use.
-pygments_style = None
-
-# Remove the prompt when copying examples
-copybutton_prompt_text = r">>> |\.\.\. "
-copybutton_prompt_is_regexp = True
-
-# -- Options for HTML output -------------------------------------------------
-
-# The theme to use for HTML and HTML Help pages.  See the documentation for
-# a list of builtin themes.
-#
-html_theme = "sphinx_rtd_theme"
-
-# Theme options are theme-specific and customize the look and feel of a theme
-# further.  For a list of options available for each theme, see the
-# documentation.
-#
-html_theme_options = {"analytics_id": "UA-83738774-2"}
-
-#  Configuration for OpenGraph and Twitter Card Tags.
-# These are responsible for creating nice shareable social images https://ahrefs.com/blog/open-graph-meta-tags/
-# https://ogp.me/#type_website
-ogp_image = "https://huggingface.co/front/thumbnails/docs/accelerate.png"
-ogp_description = "Run your raw PyTorch training script on any kind of device. 🤗 Accelerate provides an easy API to make your scripts run with mixed precision and on any kind of distributed setting (multi-GPUs, TPUs etc.)"
-ogp_description_length = 160
-
-ogp_custom_meta_tags = [
-    f'<meta name="twitter:image" content="{ogp_image}">',
-    f'<meta name="twitter:description" content="{ogp_description}">',
-]
-
-
-# Add any paths that contain custom static files (such as style sheets) here,
-# relative to this directory. They are copied after the builtin static files,
-# so a file named "default.css" will overwrite the builtin "default.css".
-html_static_path = ["_static"]
-
-# Custom sidebar templates, must be a dictionary that maps document names
-# to template names.
-#
-# The default sidebars (for documents that don't match any pattern) are
-# defined by theme itself.  Builtin themes are using these templates by
-# default: ``['localtoc.html', 'relations.html', 'sourcelink.html',
-# 'searchbox.html']``.
-#
-# html_sidebars = {}
-
-# This must be the name of an image file (path relative to the configuration
-# directory) that is the favicon of the docs. Modern browsers use this as
-# the icon for tabs, windows and bookmarks. It should be a Windows-style
-# icon file (.ico).
-html_favicon = "favicon.ico"
-
-
-# -- Options for HTMLHelp output ---------------------------------------------
-
-# Output file base name for HTML help builder.
-htmlhelp_basename = "acceleratedoc"
-
-
-# -- Options for LaTeX output ------------------------------------------------
-
-latex_elements = {
-    # The paper size ('letterpaper' or 'a4paper').
-    #
-    # 'papersize': 'letterpaper',
-    # The font size ('10pt', '11pt' or '12pt').
-    #
-    # 'pointsize': '10pt',
-    # Additional stuff for the LaTeX preamble.
-    #
-    # 'preamble': '',
-    # Latex figure (float) alignment
-    #
-    # 'figure_align': 'htbp',
-}
-
-# Grouping the document tree into LaTeX files. List of tuples
-# (source start file, target name, title,
-#  author, documentclass [howto, manual, or own class]).
-latex_documents = [
-    (master_doc, "accelerate.tex", "accelerate Documentation", "huggingface", "manual"),
-]
-
-
-# -- Options for manual page output ------------------------------------------
-
-# One entry per manual page. List of tuples
-# (source start file, name, description, authors, manual section).
-man_pages = [(master_doc, "accelerate", "accelerate Documentation", [author], 1)]
-
-
-# -- Options for Texinfo output ----------------------------------------------
-
-# Grouping the document tree into Texinfo files. List of tuples
-# (source start file, target name, title, author,
-#  dir menu entry, description, category)
-texinfo_documents = [
-    (
-        master_doc,
-        "accelerate",
-        "accelerate Documentation",
-        author,
-        "accelerate",
-        "One line description of project.",
-        "Miscellaneous",
-    ),
-]
-
-
-# -- Options for Epub output -------------------------------------------------
-
-# Bibliographic Dublin Core info.
-epub_title = project
-
-# The unique identifier of the text. This can be a ISBN number
-# or the project homepage.
-#
-# epub_identifier = ''
-
-# A unique identification for the text.
-#
-# epub_uid = ''
-
-# A list of files that should not be packed into the epub file.
-epub_exclude_files = ["search.html"]
-
-
-def setup(app):
-    app.add_css_file("css/huggingface.css")
-    app.add_css_file("css/code-snippets.css")
-    app.add_js_file("js/custom.js")
-
-
-# -- Extension configuration -------------------------------------------------
diff --git a/docs/source/favicon.ico b/docs/source/favicon.ico
deleted file mode 100644
index 424101de7..000000000
Binary files a/docs/source/favicon.ico and /dev/null differ
diff --git a/docs/source/index.mdx b/docs/source/index.mdx
new file mode 100644
index 000000000..c41ae4a0e
--- /dev/null
+++ b/docs/source/index.mdx
@@ -0,0 +1,132 @@
+<!--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.
+-->
+
+# Accelerate
+
+Run your *raw* PyTorch training script on any kind of device
+
+## Features
+
+- 🤗 Accelerate provides an easy API to make your scripts run with mixed precision and on any kind of distributed
+  setting (multi-GPUs, TPUs etc.) while still letting you write your own training loop. The same code can then runs
+  seamlessly on your local machine for debugging or your training environment.
+
+- 🤗 Accelerate also provides a CLI tool that allows you to quickly configure and test your training environment then
+  launch the scripts.
+
+
+## Easy to integrate
+
+A traditional training loop in PyTorch looks like this:
+
+```python
+my_model.to(device)
+
+for batch in my_training_dataloader:
+    my_optimizer.zero_grad()
+    inputs, targets = batch
+    inputs = inputs.to(device)
+    targets = targets.to(device)
+    outputs = my_model(inputs)
+    loss = my_loss_function(outputs, targets)
+    loss.backward()
+    my_optimizer.step()
+```
+
+Changing it to work with accelerate is really easy and only adds a few lines of code:
+
+```diff
++ from accelerate import Accelerator
+
++ accelerator = Accelerator()
+  # Use the device given by the *accelerator* object.
++ device = accelerator.device
+  my_model.to(device)
+  # Pass every important object (model, optimizer, dataloader) to *accelerator.prepare*
++ my_model, my_optimizer, my_training_dataloader = accelerate.prepare(
++     my_model, my_optimizer, my_training_dataloader
++ )
+
+  for batch in my_training_dataloader:
+      my_optimizer.zero_grad()
+      inputs, targets = batch
+      inputs = inputs.to(device)
+      targets = targets.to(device)
+      outputs = my_model(inputs)
+      loss = my_loss_function(outputs, targets)
+      # Just a small change for the backward instruction
+-     loss.backward()
++     accelerator.backward(loss)
+      my_optimizer.step()
+```
+
+and with this, your script can now run in a distributed environment (multi-GPU, TPU).
+
+You can even simplify your script a bit by letting 🤗 Accelerate handle the device placement for you (which is safer,
+especially for TPU training):
+
+```diff
++ from accelerate import Accelerator
+
++ accelerator = Accelerator()
+- my_model.to(device)
+  # Pass every important object (model, optimizer, dataloader) to *accelerator.prepare*
++ my_model, my_optimizer, my_training_dataloader = accelerate.prepare(
++     my_model, my_optimizer, my_training_dataloader
++ )
+
+  for batch in my_training_dataloader:
+      my_optimizer.zero_grad()
+      inputs, targets = batch
+-     inputs = inputs.to(device)
+-     targets = targets.to(device)
+      outputs = my_model(inputs)
+      loss = my_loss_function(outputs, targets)
+      # Just a small change for the backward instruction
+-     loss.backward()
++     accelerator.backward(loss)
+      my_optimizer.step()
+```
+
+## Script launcher
+
+No need to remember how to use `torch.distributed.launch` or to write a specific launcher for TPU training! 🤗
+Accelerate comes with a CLI tool that will make your life easier when launching distributed scripts.
+
+On your machine(s) just run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked. This will generate a config file that will be used automatically to properly set the
+default options when doing
+
+```bash
+accelerate launch my_script.py --args_to_my_script
+```
+
+For instance, here is how you would run the NLP example (from the root of the repo):
+
+```bash
+accelerate launch examples/nlp_example.py
+```
+
+## Supported integrations
+
+- CPU only
+- single GPU
+- multi-GPU on one node (machine)
+- multi-GPU on several nodes (machines)
+- TPU
+- FP16 with native AMP (apex on the roadmap)
+- DeepSpeed (experimental support)
diff --git a/docs/source/index.rst b/docs/source/index.rst
deleted file mode 100644
index 384db658a..000000000
--- a/docs/source/index.rst
+++ /dev/null
@@ -1,161 +0,0 @@
-.. 
-    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.
-
-Accelerate
-=======================================================================================================================
-
-Run your *raw* PyTorch training script on any kind of device
-
-Features
------------------------------------------------------------------------------------------------------------------------
-
-- 🤗 Accelerate provides an easy API to make your scripts run with mixed precision and on any kind of distributed
-  setting (multi-GPUs, TPUs etc.) while still letting you write your own training loop. The same code can then runs
-  seamlessly on your local machine for debugging or your training environment.
-
-- 🤗 Accelerate also provides a CLI tool that allows you to quickly configure and test your training environment then
-  launch the scripts.
-
-
-Easy to integrate
------------------------------------------------------------------------------------------------------------------------
-
-A traditional training loop in PyTorch looks like this:
-
-.. code-block:: python
-
-    my_model.to(device)
-
-    for batch in my_training_dataloader:
-        my_optimizer.zero_grad()
-        inputs, targets = batch
-        inputs = inputs.to(device)
-        targets = targets.to(device)
-        outputs = my_model(inputs)
-        loss = my_loss_function(outputs, targets)
-        loss.backward()
-        my_optimizer.step()
-
-Changing it to work with accelerate is really easy and only adds a few lines of code:
-
-.. code-block:: diff
-
-    + from accelerate import Accelerator
-
-    + accelerator = Accelerator()
-      # Use the device given by the `accelerator` object.
-    + device = accelerator.device
-      my_model.to(device)
-      # Pass every important object (model, optimizer, dataloader) to `accelerator.prepare`
-    + my_model, my_optimizer, my_training_dataloader = accelerate.prepare(
-    +     my_model, my_optimizer, my_training_dataloader
-    + )
-
-      for batch in my_training_dataloader:
-          my_optimizer.zero_grad()
-          inputs, targets = batch
-          inputs = inputs.to(device)
-          targets = targets.to(device)
-          outputs = my_model(inputs)
-          loss = my_loss_function(outputs, targets)
-          # Just a small change for the backward instruction
-    -     loss.backward()
-    +     accelerator.backward(loss)
-          my_optimizer.step()
-
-and with this, your script can now run in a distributed environment (multi-GPU, TPU).
-
-You can even simplify your script a bit by letting 🤗 Accelerate handle the device placement for you (which is safer,
-especially for TPU training):
-
-.. code-block:: diff
-
-    + from accelerate import Accelerator
-
-    + accelerator = Accelerator()
-    - my_model.to(device)
-      # Pass every important object (model, optimizer, dataloader) to `accelerator.prepare`
-    + my_model, my_optimizer, my_training_dataloader = accelerate.prepare(
-    +     my_model, my_optimizer, my_training_dataloader
-    + )
-
-      for batch in my_training_dataloader:
-          my_optimizer.zero_grad()
-          inputs, targets = batch
-    -     inputs = inputs.to(device)
-    -     targets = targets.to(device)
-          outputs = my_model(inputs)
-          loss = my_loss_function(outputs, targets)
-          # Just a small change for the backward instruction
-    -     loss.backward()
-    +     accelerator.backward(loss)
-          my_optimizer.step()
-
-
-Script launcher
------------------------------------------------------------------------------------------------------------------------
-
-No need to remember how to use ``torch.distributed.launch`` or to write a specific launcher for TPU training! 🤗
-Accelerate comes with a CLI tool that will make your life easier when launching distributed scripts.
-
-On your machine(s) just run:
-
-.. code-block:: bash
-
-    accelerate config
-
-and answer the questions asked. This will generate a config file that will be used automatically to properly set the
-default options when doing
-
-.. code-block:: bash
-
-    accelerate launch my_script.py --args_to_my_script
-
-For instance, here is how you would run the NLP example (from the root of the repo):
-
-.. code-block:: bash
-
-    accelerate launch examples/nlp_example.py
-
-
-Supported integrations
------------------------------------------------------------------------------------------------------------------------
-
-- CPU only
-- single GPU
-- multi-GPU on one node (machine)
-- multi-GPU on several nodes (machines)
-- TPU
-- FP16 with native AMP (apex on the roadmap)
-- DeepSpeed (experimental support)
-
-.. toctree::
-    :maxdepth: 2
-    :caption: Get started
-
-    quicktour
-    installation
-
-.. toctree::
-    :maxdepth: 2
-    :caption: Guides
-
-    sagemaker
-
-.. toctree::
-    :maxdepth: 2
-    :caption: API reference
-
-    accelerator
-    launcher
-    kwargs
-    internal
diff --git a/docs/source/installation.md b/docs/source/installation.mdx
similarity index 99%
rename from docs/source/installation.md
rename to docs/source/installation.mdx
index ca2f88286..4564ba00c 100644
--- a/docs/source/installation.md
+++ b/docs/source/installation.mdx
@@ -87,7 +87,7 @@ Do note that you have to keep that `accelerate` folder around and not delete it
 
 Now, let's get to the real benefit of this installation approach. Say, you saw some new feature has been just committed into `master`. If you have already performed all the steps above, to update your accelerate repo to include all the latest commits, all you need to do is to `cd` into that cloned repository folder and update the clone to the latest version:
 
-```
+```bash
 cd ~/accelerate/
 git pull
 ```
diff --git a/docs/source/internal.mdx b/docs/source/internal.mdx
new file mode 100644
index 000000000..3acaaad87
--- /dev/null
+++ b/docs/source/internal.mdx
@@ -0,0 +1,61 @@
+<!--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.
+-->
+
+# Internals
+
+## Optimizer
+
+[[autodoc]] optimizer.AcceleratedOptimizer
+
+## DataLoader
+
+The main work on your PyTorch `DataLoader` is done by the following function:
+
+[[autodoc]] data_loader.prepare_data_loader
+
+### BatchSamplerShard
+
+[[autodoc]] data_loader.DataLoaderShard
+
+### BatchSamplerShard
+
+[[autodoc]] data_loader.BatchSamplerShard
+
+### IterableDatasetShard
+
+[[autodoc]] data_loader.IterableDatasetShard
+
+## Distributed Config
+
+### AcceleratorState
+
+[[autodoc]] state.AcceleratorState
+
+### DistributedType
+
+[[autodoc]] state.DistributedType
+
+## Utilities
+
+[[autodoc]] utils.extract_model_from_parallel
+
+[[autodoc]] utils.gather
+
+[[autodoc]] utils.send_to_device
+
+[[autodoc]] utils.set_seed
+
+[[autodoc]] utils.synchronize_rng_state
+
+[[autodoc]] utils.synchronize_rng_states
+
+[[autodoc]] utils.wait_for_everyone
diff --git a/docs/source/internal.rst b/docs/source/internal.rst
deleted file mode 100644
index 2d8f0f453..000000000
--- a/docs/source/internal.rst
+++ /dev/null
@@ -1,85 +0,0 @@
-.. 
-    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.
-
-Internals
-=======================================================================================================================
-
-
-Optimizer
------------------------------------------------------------------------------------------------------------------------
-
-.. autoclass:: accelerate.optimizer.AcceleratedOptimizer
-
-
-DataLoader
------------------------------------------------------------------------------------------------------------------------
-
-The main work on your PyTorch :obj:`DataLoader` is done by the following function:
-
-.. autofunction:: accelerate.data_loader.prepare_data_loader
-
-
-BatchSamplerShard
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: accelerate.data_loader.DataLoaderShard
-    :members:
-
-
-BatchSamplerShard
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: accelerate.data_loader.BatchSamplerShard
-    :members:
-
-
-IterableDatasetShard
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: accelerate.data_loader.IterableDatasetShard
-    :members:
-
-
-Distributed Config
------------------------------------------------------------------------------------------------------------------------
-
-
-AcceleratorState
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: accelerate.state.AcceleratorState
-    :members:
-
-
-DistributedType
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: accelerate.state.DistributedType
-    :members:
-
-
-Utilities
------------------------------------------------------------------------------------------------------------------------
-
-.. autofunction:: accelerate.utils.extract_model_from_parallel
-
-.. autofunction:: accelerate.utils.gather
-
-.. autofunction:: accelerate.utils.send_to_device
-
-.. autofunction:: accelerate.utils.set_seed
-
-.. autofunction:: accelerate.utils.synchronize_rng_state
-
-.. autofunction:: accelerate.utils.synchronize_rng_states
-
-.. autofunction:: accelerate.utils.wait_for_everyone
diff --git a/docs/source/kwargs.mdx b/docs/source/kwargs.mdx
new file mode 100644
index 000000000..04f58af37
--- /dev/null
+++ b/docs/source/kwargs.mdx
@@ -0,0 +1,29 @@
+<!--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.
+-->
+
+# Kwargs Handlers
+
+The following objects can be passed to the main [`Accelerator`] to customize how some PyTorch objects
+related to distributed training or mixed precision are created.
+
+
+## DistributedDataParallelKwargs
+
+[[autodoc]] DistributedDataParallelKwargs
+
+## GradScalerKwargs
+
+[[autodoc]] GradScalerKwargs
+
+## InitProcessGroupKwargs
+
+[[autodoc]] InitProcessGroupKwargs
diff --git a/docs/source/kwargs.rst b/docs/source/kwargs.rst
deleted file mode 100644
index 61f42cd11..000000000
--- a/docs/source/kwargs.rst
+++ /dev/null
@@ -1,30 +0,0 @@
-.. 
-    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.
-
-
-Kwargs Handlers
-=======================================================================================================================
-
-The following objects can be passed to the main :class:`~accelerate.Accelerator` to customize how some PyTorch objects
-related to distributed training or mixed precision are created.
-
-
-DistributedDataParallelKwargs
------------------------------------------------------------------------------------------------------------------------
-
-.. autoclass:: accelerate.DistributedDataParallelKwargs
-
-
-GradScalerKwargs
------------------------------------------------------------------------------------------------------------------------
-
-.. autoclass:: accelerate.GradScalerKwargs
diff --git a/docs/source/launcher.mdx b/docs/source/launcher.mdx
new file mode 100644
index 000000000..fb672d5ad
--- /dev/null
+++ b/docs/source/launcher.mdx
@@ -0,0 +1,28 @@
+<!--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.
+-->
+
+# Notebook Launcher
+
+Launch your training function inside a notebook. Currently supports launching a training with TPUs on [Google
+Colab](https://colab.research.google.com/) and [Kaggle kernels](https://www.kaggle.com/code), as well as training on
+several GPUs (if the machine on which you are running your notebook has them).
+
+An example can be found in [this notebook](https://github.com/huggingface/notebooks/blob/master/examples/accelerate/simple_nlp_example.ipynb).
+
+<Tip warning={true}>
+
+Your `Accelerator` object should only be defined inside the training function. This is because the
+initialization should be done inside the launcher only.
+
+</Tip>
+
+[[autodoc]] notebook_launcher
diff --git a/docs/source/launcher.rst b/docs/source/launcher.rst
deleted file mode 100644
index 1c2279946..000000000
--- a/docs/source/launcher.rst
+++ /dev/null
@@ -1,29 +0,0 @@
-.. 
-    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.
-
-
-Notebook Launcher
-=======================================================================================================================
-
-Launch your training function inside a notebook. Currently supports launching a training with TPUs on [Google
-Colab](https://colab.research.google.com/) and [Kaggle kernels](https://www.kaggle.com/code), as well as training on
-several GPUs (if the machine on which you are running your notebook has them).
-
-An example can be found in `this notebook
-<https://github.com/huggingface/notebooks/blob/master/examples/accelerate/simple_nlp_example.ipynb>`__.
-
-.. warning::
-
-    Your :obj:`Accelerator` object should only be defined inside the training function. This is because the
-    initialization should be done inside the launcher only.
-
-.. autofunction:: accelerate.notebook_launcher
diff --git a/docs/source/quicktour.mdx b/docs/source/quicktour.mdx
new file mode 100644
index 000000000..f16047925
--- /dev/null
+++ b/docs/source/quicktour.mdx
@@ -0,0 +1,445 @@
+<!--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.
+-->
+
+# Quick tour
+
+Let's have a look at a look at 🤗 Accelerate main features and traps to avoid.
+
+## Main use
+
+To use 🤗 Accelerate in your own script, you have to change four things:
+
+1. Import the [`Accelerator`] main class instantiate one in an `accelerator` object:
+
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+```
+
+This should happen as early as possible in your training script as it will initialize everything necessary for
+distributed training. You don't need to indicate the kind of environment you are in (just one machine with a GPU, one
+match with several GPUs, several machines with multiple GPUs or a TPU), the library will detect this automatically.
+
+2. Remove the call `.to(device)` or `.cuda()` for your model and input data. The `accelerator` object
+will handle this for you and place all those objects on the right device for you. If you know what you're doing, you
+can leave those `.to(device)` calls but you should use the device provided by the `accelerator` object:
+`accelerator.device`.
+
+To fully deactivate the automatic device placement, pass along `device_placement=False` when initializing your
+[`Accelerator`].
+
+<Tip warning={true}>
+
+If you place your objects manually on the proper device, be careful to create your optimizer after putting your
+model on `accelerator.device` or your training will fail on TPU.
+
+</Tip>
+
+3. Pass all objects relevant to training (optimizer, model, training dataloader) to the
+[`~Accelerator.prepare`] method. This will make sure everything is ready for training.
+
+```python
+model, optimizer, train_dataloader = accelerator.prepare(model, optimizer, train_dataloader)
+```
+
+In particular, your training dataloader will be sharded accross all GPUs/TPU cores available so that each one sees a
+different portion of the training dataset. Also, the random states of all processes will be synchronized at the
+beginning of each iteration through your dataloader, to make sure the data is shuffled the same way (if you decided to
+use `shuffle=True` or any kind of random sampler).
+
+<Tip>
+
+The actual batch size for your training will be the number of devices used multiplied by the batch size you set in
+your script: for instance training on 4 GPUs with a batch size of 16 set when creating the training dataloader will
+train at an actual batch size of 64.
+
+</Tip>
+
+Alternatively, you can use the option `split_batches=True` when creating initializing your
+[`Accelerator`], in which case the batch size will always stay the same, whether your run your
+script on 1, 2, 4 or 64 GPUs.
+
+You should execute this instruction as soon as all objects for training are created, before starting your actual
+training loop.
+
+<Tip warning={true}>
+
+Your training dataloader may change length when going through this method: if you run on X GPUs, it will have its
+length divided by X (since your actual batch size will be multiplied by X), unless you set
+`split_batches=True`.
+
+</Tip>
+
+Any instruction using your training dataloader length (for instance if you need the number of total training steps
+to create a learning rate scheduler) should go after the call to [`~Accelerator.prepare`].
+
+You can perfectly send your dataloader to [`~Accelerator.prepare`] on its own, but it's best to send the
+model and optimizer to [`~Accelerator.prepare`] together.
+
+You may or may not want to send your validation dataloader to [`~Accelerator.prepare`], depending on
+whether you want to run distributed evaluation or not (see below).
+
+4. Replace the line `loss.backward()` by `accelerator.backward(loss)`.
+
+And you're all set! With all these changes, your script will run on your local machine as well as on multiple GPUs or a
+TPU! You can either use your favorite tool to launch the distributed training, or you can use the 🤗 Accelerate
+launcher.
+
+
+## Distributed evaluation
+
+You can perform regular evaluation in your training script, if you leave your validation dataloader out of the
+[`~Accelerator.prepare`] method. In this case, you will need to put the input data on the
+`accelerator.device` manually.
+
+To perform distributed evaluation, send along your validation dataloader to the [`~Accelerator.prepare`]
+method:
+
+```python
+validation_dataloader = accelerator.prepare(validation_dataloader)
+```
+
+Like for your training dataloader, it will mean that (should you run your script on multiple devices) each device will
+only see part of the evaluation data. This means you will need to group your predictions together. This is very easy to
+do with the [`~Accelerator.gather`] method.
+
+```python
+for inputs, targets in validation_dataloader:
+    predictions = model(inputs)
+    # Gather all predictions and targets
+    all_predictions = accelerator.gather(predictions)
+    all_targets = accelerator.gather(targets)
+    # Example of use with a *Datasets.Metric*
+    metric.add_batch(all_predictions, all_targets)
+```
+
+<Tip warning={true}>
+
+Like for the training dataloader, passing your validation dataloader through
+[`~Accelerator.prepare`] may change its: if you run on X GPUs, it will have its length divided by X
+(since your actual batch size will be multiplied by X), unless you set `split_batches=True`.
+
+Any instruction using your training dataloader length (for instance if you need the number of total training steps
+to create a learning rate scheduler) should go after the call to [`~Accelerator.prepare`].
+
+</Tip>
+
+<Tip warning={true}>
+
+The [`~Accelerator.gather`] method requires the tensors to be all the same size on each process. If
+you have tensors of different sizes on each process (for instance when dynamically padding to the maximum length in
+a batch), you should use the [`~Accelerator.pad_across_processes`] method to pad you tensor to the
+biggest size across processes.
+
+</Tip>
+
+## Launching your distributed script
+
+You can use the regular commands to launch your distributed training (like `torch.distributed.launch` for
+PyTorch), they are fully compatible with 🤗 Accelerate. The only caveat here is that 🤗 Accelerate uses the environment
+to determine all useful information, so `torch.distributed.launch` should be used with the flag `--use_env`.
+
+🤗 Accelerate also provides a CLI tool that unifies all launcher, so you only have to remember one command. To use it,
+just run
+
+```bash
+accelerate config
+```
+
+on your machine and reply to the questions asked. This will save a *default_config.yaml* file in your cache folder for
+🤗 Accelerate. That cache folder is (with decreasing order of priority):
+
+- The content of your environment variable `HF_HOME` suffixed with *accelerate*.
+- If it does not exist, the content of your environment variable `XDG_CACHE_HOME` suffixed with
+  *huggingface/accelerate*.
+- If this does not exist either, the folder *~/.cache/huggingface/accelerate*
+
+You can also specify with the flag `--config_file` the location of the file you want to save.
+
+Once this is done, you can test everything is going well on your setup by running
+
+```bash
+accelerate test
+```
+
+This will launch a short script that will test the distributed environment. If it runs fine, you are ready for the next
+step!
+
+Note that if you specified a location for the config file in the previous step, you need to pass it here as well:
+
+```bash
+accelerate test --config_file path_to_config.yaml
+```
+
+Now that this is done, you can run your script with the following command:
+
+```bash
+accelerate launch path_to_script.py --args_for_the_script
+```
+
+If you stored the config file in a non-default location, you can indicate it to the launcher like his:
+
+```bash
+accelerate launch --config_file path_to_config.yaml path_to_script.py --args_for_the_script
+```
+
+You can also override any of the arguments determined by your config file, see TODO: insert ref here.
+
+
+## Launching training from a notebook
+
+In Accelerate 0.3.0, a new [`notebook_launcher`] has been introduced to help you launch your training
+function from a notebook. This launcher supports launching a training with TPUs on Colab or Kaggle, as well as training
+on several GPUs (if the machine on which you are running your notebook has them).
+
+Just define a function responsible for your whole training and/or evaluation in a cell of the notebook, then execute a
+cell with the following code:
+
+```python
+from accelerate import notebook_launcher
+
+notebook_launcher(training_function)
+```
+
+<Tip warning={true}>
+
+Your `Accelerator` object should only be defined inside the training function. This is because the
+initialization should be done inside the launcher only.
+
+</Tip>
+
+## Training on TPU
+
+If you want to launch your script on TPUs, there are a few caveats you should be aware of. Behind the scenes, the TPUs
+will create a graph of all the operations happening in your training step (forward pass, backward pass and optimizer
+step). This is why your first step of training will always be very long as building and compiling this graph for
+optimizations takes some time.
+
+The good news is that this compilation will be cached so the second step and all the following will be much faster. The
+bas news is that it only applies if all of your steps do exactly the same operations, which implies:
+
+- having all tensors of the same length in all your lengths
+- having static code (i.e., not a for loop of length that could change from step to step)
+
+Having any of the things above change between two steps will trigger a new compilation which will, once again, take a
+lot of time. In practice, that means you must take special care to have all your tensors in your inputs of the same
+shape (so no dynamic padding for instance if you are in an NLP problem) and should not use layer with for loops that
+have different lengths depending on the inputs (such as an LSTM) or the training will be excruciatingly slow.
+
+To introduce special behavior in your script for TPUs you can check the `distributed_type` of your
+`accelerator`:
+
+```python docstyle-ignore
+from accelerate import DistributedType
+
+if accelerator.distributed_type == DistributedType.TPU:
+    # do something of static shape
+else:
+    # go crazy and be dynamic
+```
+
+The [NLP example](https://github.com/huggingface/accelerate/blob/main/examples/nlp_example.py) shows an example in
+situation with dynamic padding.
+
+One last thing to pay close attnetion to: if your model has tied weights (such as language models which tie the weights
+of the embedding matrix with the weights of the decoder), moving this model to the TPU (either yourself or after you
+passed your model to [`~Accelerator.prepare`]) will break the tying. You will need to retie the weights
+after. You can find an example of this in the [run_clm_no_trainer](https://github.com/huggingface/transformers/blob/master/examples/pytorch/language-modeling/run_clm.py) script in
+the Transformers repository.
+
+
+## Other caveats
+
+We list here all smaller issues you could have in your script conversion and how to resolve them.
+
+### Execute a statement only on one processes
+
+Some of your instructions only need to run for one process on a given server: for instance a data download or a log
+statement. To do this, wrap the statement in a test like this:
+
+```python docstyle-ignore
+if accelerator.is_local_main_process:
+    # Is executed once per server
+```
+
+Another example is progress bars: to avoid having multiple progress bars in your output, you should only display one on
+the local main process:
+
+```python
+from tqdm.auto import tqdm
+
+progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process)
+```
+
+The *local* means per machine: if you are running your training on two servers with several GPUs, the instruction will
+be executed once on each of those servers. If you need to execute something only once for all processes (and not per
+machine) for instance, uploading the final model to the 🤗 model hub, wrap it in a test like this:
+
+```python docstyle-ignore
+if accelerator.is_main_process:
+    # Is executed once only
+```
+
+For printing statements you only want executed once per machine, you can just replace the `print` function by
+`accelerator.print`.
+
+
+### Defer execution
+
+When you run your usual script, instructions are executed in order. Using 🤗 Accelerate to deploy your script on several
+GPUs at the same time introduces a complication: while each process executes all instructions in order, some may be
+faster than others.
+
+You might need to wait for all processes to have reached a certain point before executing a given instruction. For
+instance, you shouldn't save a model before being sure every process is done with training. To do this, just write the
+following line in your code:
+
+```
+accelerator.wait_for_everyone()
+```
+
+This instruction will block all the processes that arrive them first until all the other processes have reached that
+point (if you run your script on just one GPU or CPU, this wont' do anything).
+
+
+### Saving/loading a model
+
+Saving the model you trained might need a bit of adjustment: first you should wait for all processes to reach that
+point in the script as shown above, and then, you should unwrap your model before saving it. This is because when going
+through the [`~Accelerator.prepare`] method, your model may have been placed inside a bigger model,
+which deals with the distributed training. This in turn means that saving your model state dictionary without taking
+any precaution will take that potential extra layer into account, and you will end up with weights you can't load back
+in your base model.
+
+This is why it's recommended to *unwrap* your model first. Here is an example:
+
+```
+accelerator.wait_for_everyone()
+unwrapped_model = accelerator.unwrap_model(model)
+accelerator.save(unwrapped_model.state_dict(), filename)
+```
+
+If your script contains a logic to load checkpoint, we also recommend you load your weights in the unwrapped model
+(this is only useful if you use the load function after making your model go through
+[`~Accelerator.prepare`]). Here is an example:
+
+```
+unwrapped_model = accelerator.unwrap_model(model)
+unwrapped_model.load_state_dict(torch.load(filename))
+```
+
+Note that since all the model parameters are references to tensors, this will load your weights inside `model`.
+
+### Gradient clipping
+
+If you are using gradient clipping in your script, you should replace the calls to
+`torch.nn.utils.clip_grad_norm_` or `torch.nn.utils.clip_grad_value_` with `accelerator.clip_grad_norm_`
+and `accelerator.clip_grad_value_` respectively.
+
+
+### Mixed Precision training
+
+If you are running your training in Mixed Precision with Accelerate, you will get the best result with your loss being
+computed inside your model (like in Transformer models for instance). Every computation outside of the model will be
+executed in full precision (which is generally what you want for loss computation, expecially if it involves a
+softmax). However you might want to put your loss computation inside the *accelerator.autocast* context manager:
+
+```
+with accelerator.autocast():
+    loss = complex_loss_function(outputs, target):
+```
+
+Another caveat with Mixed Precision training is that the gradient will skip a few updates at the beginning and
+sometimes during training: because of the dynamic loss scaling strategy, there are points during training where the
+gradients have overflown, and the loss scaling factor is reduced to avoid this happening again at the next step.
+
+This means that you may update your learning rate scheduler when there was no update, which is fine in general, but may
+have an impact when you have very little training data, or if the first learning rate values of your scheduler are very
+important. In this case, you can skip the learning rate scheduler updates when the optimizer step was not done like
+this:
+
+```
+if not accelerator.optimizer_step_was_skipped:
+    lr_scheduler.step()
+```
+
+### DeepSpeed
+
+DeepSpeed support is experimental, so the underlying API will evolve in the near future and may have some slight
+breaking changes. In particular, 🤗 Accelerate does not support DeepSpeed config you have written yourself yet, this
+will be added in a next version.
+
+One main caveat for the DeepSpeed integration is that the DeepSpeed launcher always passes a `local_rank` variable to
+the training script, so your training script should accept it (whether you launch training with the DeepSpeed launcher
+or `accelerate launch`).
+
+<Tip warning={true}>
+
+The [`notebook_launcher`] does not support the DeepSpeed integration yet.
+
+</Tip>
+
+## Internal mechanism
+
+Internally, the library works by first analyzing the environment in which the script is launched to determine which
+kind of distributed setup is used, how many different processes there are and which one the current script is in. All
+that information is stored in the [`~AcceleratorState`].
+
+This class is initialized the first time you instantiate a [`Accelerator`] as well as performing any
+specific initialization your distributed setup needs. Its state is then uniquely shared through all instances of
+[`~state.AcceleratorState`].
+
+Then, when calling [`~Accelerator.prepare`], the library:
+
+- wraps your model(s) in the container adapted for the distributed setup,
+- wraps your optimizer(s) in a [`~optimizer.AcceleratedOptimizer`],
+- creates a new version of your dataloader(s) in a [`~data_loader.DataLoaderShard`].
+
+While the model(s) and optimizer(s) are just put in simple wrappers, the dataloader(s) are re-created. This is mostly
+because PyTorch does not let the user change the `batch_sampler` of a dataloader once it's been created and the
+library handles the sharding of your data between processes by changing that `batch_sampler` to yield every other
+`num_processes` batches.
+
+The [`~data_loader.DataLoaderShard`] subclasses `DataLoader` to add the following functionality:
+
+- it synchronizes the appropriate random number generator of all processes at each new iteration, to ensure any
+  randomization (like shuffling) is done the exact same way across processes.
+- it puts the batches on the proper device before yielding them (unless you have opted out of
+  `device_placement=True`).
+
+The random number generator synchronization will by default synchronize:
+
+- the `generator` attribute of a given sampler (like the PyTorch `RandomSampler`) for PyTorch >= 1.6
+- the main random number generator in PyTorch <=1.5.1
+
+You can choose which random number generator(s) to synchronize with the `rng_types` argument of the main
+[`Accelerator`]. In PyTorch >= 1.6, it is recommended to rely on local `generator` to avoid
+setting the same seed in the main random number generator in all processes.
+
+<Tip warning={true}>
+
+Synchronization the main torch (or CUDA or XLA) random number generator will affect any other potential random
+artifacts you could have in your dataset (like random data augmentation) in the sense all processes will get the
+same random numbers from the torch random modules (so will apply the same random data augmentation if it's
+controlled by torch).
+
+</Tip>
+
+<Tip>
+
+The randomization part of your custom sampler, batch sampler or iterable dataset should be done using a local
+`torch.Generator` object (in PyTorch >= 1.6), see the traditional `RandomSampler`, as an example.
+
+</Tip>
+
+See more details about the internal in the [Internals page](internal).
diff --git a/docs/source/quicktour.rst b/docs/source/quicktour.rst
deleted file mode 100644
index 686d85ce1..000000000
--- a/docs/source/quicktour.rst
+++ /dev/null
@@ -1,450 +0,0 @@
-.. 
-    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.
-
-Quick tour
-=======================================================================================================================
-
-Let's have a look at a look at 🤗 Accelerate main features and traps to avoid.
-
-Main use
------------------------------------------------------------------------------------------------------------------------
-
-To use 🤗 Accelerate in your own script, you have to change four things:
-
-1. Import the :class:`~accelerate.Accelerator` main class instantiate one in an :obj:`accelerator` object:
-
-.. code-block:: python
-
-    from accelerate import Accelerator
-
-    accelerator = Accelerator()
-
-This should happen as early as possible in your training script as it will initialize everything necessary for
-distributed training. You don't need to indicate the kind of environment you are in (just one machine with a GPU, one
-match with several GPUs, several machines with multiple GPUs or a TPU), the library will detect this automatically.
-
-2. Remove the call :obj:`.to(device)` or :obj:`.cuda()` for your model and input data. The :obj:`accelerator` object
-will handle this for you and place all those objects on the right device for you. If you know what you're doing, you
-can leave those :obj:`.to(device)` calls but you should use the device provided by the :obj:`accelerator` object:
-:obj:`accelerator.device`.
-
-To fully deactivate the automatic device placement, pass along :obj:`device_placement=False` when initializing your
-:class:`~accelerate.Accelerator`.
-
-.. Warning::
-
-    If you place your objects manually on the proper device, be careful to create your optimizer after putting your
-    model on :obj:`accelerator.device` or your training will fail on TPU.
-
-3. Pass all objects relevant to training (optimizer, model, training dataloader) to the
-:meth:`~accelerate.Accelerator.prepare` method. This will make sure everything is ready for training.
-
-.. code-block:: python
-
-    model, optimizer, train_dataloader = accelerator.prepare(model, optimizer, train_dataloader)
-
-In particular, your training dataloader will be sharded accross all GPUs/TPU cores available so that each one sees a
-different portion of the training dataset. Also, the random states of all processes will be synchronized at the
-beginning of each iteration through your dataloader, to make sure the data is shuffled the same way (if you decided to
-use :obj:`shuffle=True` or any kind of random sampler).
-
-.. Note::
-
-    The actual batch size for your training will be the number of devices used multiplied by the batch size you set in
-    your script: for instance training on 4 GPUs with a batch size of 16 set when creating the training dataloader will
-    train at an actual batch size of 64.
-
-    Alternatively, you can use the option :obj:`split_batches=True` when creating initializing your
-    :class:`~accelerate.Accelerator`, in which case the batch size will always stay the same, whether your run your
-    script on 1, 2, 4 or 64 GPUs.
-
-You should execute this instruction as soon as all objects for training are created, before starting your actual
-training loop.
-
-.. Warning::
-
-    Your training dataloader may change length when going through this method: if you run on X GPUs, it will have its
-    length divided by X (since your actual batch size will be multiplied by X), unless you set
-    :obj:`split_batches=True`.
-
-    Any instruction using your training dataloader length (for instance if you need the number of total training steps
-    to create a learning rate scheduler) should go after the call to :meth:`~accelerate.Accelerator.prepare`.
-
-You can perfectly send your dataloader to :meth:`~accelerate.Accelerator.prepare` on its own, but it's best to send the
-model and optimizer to :meth:`~accelerate.Accelerator.prepare` together.
-
-You may or may not want to send your validation dataloader to :meth:`~accelerate.Accelerator.prepare`, depending on
-whether you want to run distributed evaluation or not (see below).
-
-4. Replace the line :obj:`loss.backward()` by :obj:`accelerator.backward(loss)`.
-
-And you're all set! With all these changes, your script will run on your local machine as well as on multiple GPUs or a
-TPU! You can either use your favorite tool to launch the distributed training, or you can use the 🤗 Accelerate
-launcher.
-
-
-Distributed evaluation
------------------------------------------------------------------------------------------------------------------------
-
-You can perform regular evaluation in your training script, if you leave your validation dataloader out of the
-:meth:`~accelerate.Accelerator.prepare` method. In this case, you will need to put the input data on the
-:obj:`accelerator.device` manually.
-
-To perform distributed evaluation, send along your validation dataloader to the :meth:`~accelerate.Accelerator.prepare`
-method:
-
-.. code-block:: python
-
-    validation_dataloader = accelerator.prepare(validation_dataloader)
-
-Like for your training dataloader, it will mean that (should you run your script on multiple devices) each device will
-only see part of the evaluation data. This means you will need to group your predictions together. This is very easy to
-do with the :meth:`~accelerate.Accelerator.gather` method.
-
-.. code-block:: python
-
-    for inputs, targets in validation_dataloader:
-        predictions = model(inputs)
-        # Gather all predictions and targets
-        all_predictions = accelerator.gather(predictions)
-        all_targets = accelerator.gather(targets)
-        # Example of use with a `Datasets.Metric`
-        metric.add_batch(all_predictions, all_targets)
-
-
-.. Warning::
-
-    Like for the training dataloader, passing your validation dataloader through
-    :meth:`~accelerate.Accelerator.prepare` may change its: if you run on X GPUs, it will have its length divided by X
-    (since your actual batch size will be multiplied by X), unless you set :obj:`split_batches=True`.
-
-    Any instruction using your training dataloader length (for instance if you need the number of total training steps
-    to create a learning rate scheduler) should go after the call to :meth:`~accelerate.Accelerator.prepare`.
-
-.. Warning::
-
-    The :meth:`~accelerate.Accelerator.gather` method requires the tensors to be all the same size on each process. If
-    you have tensors of different sizes on each process (for instance when dynamically padding to the maximum length in
-    a batch), you should use the :meth:`~accelerate.Accelerator.pad_across_processes` method to pad you tensor to the
-    biggest size across processes.
-
-
-Launching your distributed script
------------------------------------------------------------------------------------------------------------------------
-
-You can use the regular commands to launch your distributed training (like :obj:`torch.distributed.launch` for
-PyTorch), they are fully compatible with 🤗 Accelerate. The only caveat here is that 🤗 Accelerate uses the environment
-to determine all useful information, so :obj:`torch.distributed.launch` should be used with the flag :obj:`--use_env`.
-
-🤗 Accelerate also provides a CLI tool that unifies all launcher, so you only have to remember one command. To use it,
-just run
-
-.. code-block:: bash
-
-    accelerate config
-
-on your machine and reply to the questions asked. This will save a `default_config.yaml` file in your cache folder for
-🤗 Accelerate. That cache folder is (with decreasing order of priority):
-
-    - The content of your environment variable ``HF_HOME`` suffixed with `accelerate`.
-    - If it does not exist, the content of your environment variable ``XDG_CACHE_HOME`` suffixed with
-      `huggingface/accelerate`.
-    - If this does not exist either, the folder `~/.cache/huggingface/accelerate`
-
-You can also specify with the flag :obj:`--config_file` the location of the file you want to save.
-
-Once this is done, you can test everything is going well on your setup by running
-
-.. code-block:: bash
-
-    accelerate test
-
-
-This will launch a short script that will test the distributed environment. If it runs fine, you are ready for the next
-step!
-
-Note that if you specified a location for the config file in the previous step, you need to pass it here as well:
-
-.. code-block:: bash
-
-    accelerate test --config_file path_to_config.yaml
-
-
-Now that this is done, you can run your script with the following command:
-
-.. code-block:: bash
-
-    accelerate launch path_to_script.py --args_for_the_script
-
-
-If you stored the config file in a non-default location, you can indicate it to the launcher like his:
-
-.. code-block:: bash
-
-    accelerate launch --config_file path_to_config.yaml path_to_script.py --args_for_the_script
-
-You can also override any of the arguments determined by your config file, see TODO: insert ref here.
-
-
-Launching training from a notebook
------------------------------------------------------------------------------------------------------------------------
-
-In Accelerate 0.3.0, a new :class:`~accelerate.notebook_launcher` has been introduced to help you launch your training
-function from a notebook. This launcher supports launching a training with TPUs on Colab or Kaggle, as well as training
-on several GPUs (if the machine on which you are running your notebook has them).
-
-Just define a function responsible for your whole training and/or evaluation in a cell of the notebook, then execute a
-cell with the following code:
-
-.. code-block::
-
-    from accelerate import notebook_launcher
-
-    notebook_launcher(training_function)
-
-.. warning::
-
-    Your :obj:`Accelerator` object should only be defined inside the training function. This is because the
-    initialization should be done inside the launcher only.
-
-
-Training on TPU
------------------------------------------------------------------------------------------------------------------------
-
-If you want to launch your script on TPUs, there are a few caveats you should be aware of. Behind the scenes, the TPUs
-will create a graph of all the operations happening in your training step (forward pass, backward pass and optimizer
-step). This is why your first step of training will always be very long as building and compiling this graph for
-optimizations takes some time.
-
-The good news is that this compilation will be cached so the second step and all the following will be much faster. The
-bas news is that it only applies if all of your steps do exactly the same operations, which implies:
-
-- having all tensors of the same length in all your lengths
-- having static code (i.e., not a for loop of length that could change from step to step)
-
-Having any of the things above change between two steps will trigger a new compilation which will, once again, take a
-lot of time. In practice, that means you must take special care to have all your tensors in your inputs of the same
-shape (so no dynamic padding for instance if you are in an NLP problem) and should not use layer with for loops that
-have different lengths depending on the inputs (such as an LSTM) or the training will be excruciatingly slow.
-
-To introduce special behavior in your script for TPUs you can check the :obj:`distributed_type` of your
-:obj:`accelerator`:
-
-.. code-block:: python
-
-    from accelerate import DistributedType
-
-    if accelerator.distributed_type == DistributedType.TPU:
-        # do something of static shape
-    else:
-        # go crazy and be dynamic
-
-The `NLP example <https://github.com/huggingface/accelerate/blob/main/examples/nlp_example.py>`__ shows an example in
-situation with dynamic padding.
-
-One last thing to pay close attnetion to: if your model has tied weights (such as language models which tie the weights
-of the embedding matrix with the weights of the decoder), moving this model to the TPU (either yourself or after you
-passed your model to :meth:`~accelerate.Accelerator.prepare`) will break the tying. You will need to retie the weights
-after. You can find an example of this in the `run_clm_no_trainer
-<https://github.com/huggingface/transformers/blob/master/examples/pytorch/language-modeling/run_clm.py>`__ script in
-the Transformers repository.
-
-
-Other caveats
------------------------------------------------------------------------------------------------------------------------
-
-We list here all smaller issues you could have in your script conversion and how to resolve them.
-
-Execute a statement only on one processes
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-Some of your instructions only need to run for one process on a given server: for instance a data download or a log
-statement. To do this, wrap the statement in a test like this:
-
-.. code-block:: python
-
-    if accelerator.is_local_main_process:
-        # Is executed once per server
-
-Another example is progress bars: to avoid having multiple progress bars in your output, you should only display one on
-the local main process:
-
-.. code-block:: python
-
-    from tqdm.auto import tqdm
-
-    progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process)
-
-The `local` means per machine: if you are running your training on two servers with several GPUs, the instruction will
-be executed once on each of those servers. If you need to execute something only once for all processes (and not per
-machine) for instance, uploading the final model to the 🤗 model hub, wrap it in a test like this:
-
-.. code-block:: python
-
-    if accelerator.is_main_process:
-        # Is executed once only
-
-For printing statements you only want executed once per machine, you can just replace the :obj:`print` function by
-:obj:`accelerator.print`.
-
-
-Defer execution
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-When you run your usual script, instructions are executed in order. Using 🤗 Accelerate to deploy your script on several
-GPUs at the same time introduces a complication: while each process executes all instructions in order, some may be
-faster than others.
-
-You might need to wait for all processes to have reached a certain point before executing a given instruction. For
-instance, you shouldn't save a model before being sure every process is done with training. To do this, just write the
-following line in your code:
-
-.. code-block::
-
-    accelerator.wait_for_everyone()
-
-This instruction will block all the processes that arrive them first until all the other processes have reached that
-point (if you run your script on just one GPU or CPU, this wont' do anything).
-
-
-Saving/loading a model
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-Saving the model you trained might need a bit of adjustment: first you should wait for all processes to reach that
-point in the script as shown above, and then, you should unwrap your model before saving it. This is because when going
-through the :meth:`~accelerate.Accelerator.prepare` method, your model may have been placed inside a bigger model,
-which deals with the distributed training. This in turn means that saving your model state dictionary without taking
-any precaution will take that potential extra layer into account, and you will end up with weights you can't load back
-in your base model.
-
-This is why it's recommended to `unwrap` your model first. Here is an example:
-
-.. code-block::
-
-    accelerator.wait_for_everyone()
-    unwrapped_model = accelerator.unwrap_model(model)
-    accelerator.save(unwrapped_model.state_dict(), filename)
-
-If your script contains a logic to load checkpoint, we also recommend you load your weights in the unwrapped model
-(this is only useful if you use the load function after making your model go through
-:meth:`~accelerate.Accelerator.prepare`). Here is an example:
-
-.. code-block::
-
-    unwrapped_model = accelerator.unwrap_model(model)
-    unwrapped_model.load_state_dict(torch.load(filename))
-
-Note that since all the model parameters are references to tensors, this will load your weights inside :obj:`model`.
-
-Gradient clipping
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-If you are using gradient clipping in your script, you should replace the calls to
-:obj:`torch.nn.utils.clip_grad_norm_` or :obj:`torch.nn.utils.clip_grad_value_` with :obj:`accelerator.clip_grad_norm_`
-and :obj:`accelerator.clip_grad_value_` respectively.
-
-
-Mixed Precision training
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-If you are running your training in Mixed Precision with Accelerate, you will get the best result with your loss being
-computed inside your model (like in Transformer models for instance). Every computation outside of the model will be
-executed in full precision (which is generally what you want for loss computation, expecially if it involves a
-softmax). However you might want to put your loss computation inside the `accelerator.autocast` context manager:
-
-.. code-block::
-
-    with accelerator.autocast():
-        loss = complex_loss_function(outputs, target):
-
-Another caveat with Mixed Precision training is that the gradient will skip a few updates at the beginning and
-sometimes during training: because of the dynamic loss scaling strategy, there are points during training where the
-gradients have overflown, and the loss scaling factor is reduced to avoid this happening again at the next step.
-
-This means that you may update your learning rate scheduler when there was no update, which is fine in general, but may
-have an impact when you have very little training data, or if the first learning rate values of your scheduler are very
-important. In this case, you can skip the learning rate scheduler updates when the optimizer step was not done like
-this:
-
-.. code-block::
-
-    if not accelerator.optimizer_step_was_skipped:
-        lr_scheduler.step()
-
-
-DeepSpeed
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-DeepSpeed support is experimental, so the underlying API will evolve in the near future and may have some slight
-breaking changes. In particular, 🤗 Accelerate does not support DeepSpeed config you have written yourself yet, this
-will be added in a next version.
-
-One main caveat for the DeepSpeed integration is that the DeepSpeed launcher always passes a ``local_rank`` variable to
-the training script, so your training script should accept it (whether you launch training with the DeepSpeed launcher
-or ``accelerate launch``).
-
-.. Warning::
-
-    The :func:`~accelerate.notebook_launcher` does not support the DeepSpeed integration yet.
-
-
-Internal mechanism
------------------------------------------------------------------------------------------------------------------------
-
-Internally, the library works by first analyzing the environment in which the script is launched to determine which
-kind of distributed setup is used, how many different processes there are and which one the current script is in. All
-that information is stored in the :class:`~accelerate.state.AcceleratorState`.
-
-This class is initialized the first time you instantiate a :class:`~accelerate.Accelerator` as well as performing any
-specific initialization your distributed setup needs. Its state is then uniquely shared through all instances of
-:class:`~accelerate.state.AcceleratorState`.
-
-Then, when calling :meth:`~accelerate.Accelerator.prepare`, the library:
-
-- wraps your model(s) in the container adapted for the distributed setup,
-- wraps your optimizer(s) in a :class:`~accelerate.optimizer.AcceleratedOptimizer`,
-- creates a new version of your dataloader(s) in a :class:`~accelerate.data_loader.DataLoaderShard`.
-
-While the model(s) and optimizer(s) are just put in simple wrappers, the dataloader(s) are re-created. This is mostly
-because PyTorch does not let the user change the :obj:`batch_sampler` of a dataloader once it's been created and the
-library handles the sharding of your data between processes by changing that :obj:`batch_sampler` to yield every other
-:obj:`num_processes` batches.
-
-The :class:`~accelerate.data_loader.DataLoaderShard` subclasses :obj:`DataLoader` to add the following functionality:
-
-- it synchronizes the appropriate random number generator of all processes at each new iteration, to ensure any
-  randomization (like shuffling) is done the exact same way across processes.
-- it puts the batches on the proper device before yielding them (unless you have opted out of
-  :obj:`device_placement=True`).
-
-The random number generator synchronization will by default synchronize:
-
-- the :obj:`generator` attribute of a given sampler (like the PyTorch :obj:`RandomSampler`) for PyTorch >= 1.6
-- the main random number generator in PyTorch <=1.5.1
-
-You can choose which random number generator(s) to synchronize with the :obj:`rng_types` argument of the main
-:class:`~accelerate.Accelerator`. In PyTorch >= 1.6, it is recommended to rely on local :obj:`generator` to avoid
-setting the same seed in the main random number generator in all processes.
-
-.. Warning::
-
-    Synchronization the main torch (or CUDA or XLA) random number generator will affect any other potential random
-    artifacts you could have in your dataset (like random data augmentation) in the sense all processes will get the
-    same random numbers from the torch random modules (so will apply the same random data augmentation if it's
-    controlled by torch).
-
-.. Note::
-
-    The randomization part of your custom sampler, batch sampler or iterable dataset should be done using a local
-    :obj:`torch.Generator` object (in PyTorch >= 1.6), see the traditional :obj:`RandomSampler`, as an example.
-
-See more details about the internal in the :doc:`Internals page <internal>`.
diff --git a/docs/source/sagemaker.mdx b/docs/source/sagemaker.mdx
new file mode 100644
index 000000000..d55ccd28b
--- /dev/null
+++ b/docs/source/sagemaker.mdx
@@ -0,0 +1,150 @@
+<!--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.
+-->
+
+# Amazon SageMaker
+
+Hugging Face and Amazon introduced new [Hugging Face Deep Learning Containers (DLCs)](https://github.com/aws/deep-learning-containers/blob/master/available_images.md#huggingface-training-containers) to
+make it easier than ever to train Hugging Face Transformer models in [Amazon SageMaker](https://aws.amazon.com/sagemaker/).
+
+## Getting Started
+
+### Setup & Installation
+
+
+Before you can run your 🤗 Accelerate scripts on Amazon SageMaker you need to sign up for an AWS account. If you do not
+have an AWS account yet learn more [here](https://docs.aws.amazon.com/sagemaker/latest/dg/gs-set-up.html).
+
+After you have your AWS Account you need to install the `sagemaker` sdk for 🤗 Accelerate with.
+
+```bash
+pip install "accelerate[sagemaker]" --upgrade
+```
+
+🤗 Accelerate currently uses the 🤗 DLCs, with `transformers`, `datasets` and `tokenizers` pre-installed. 🤗
+Accelerate is not in the DLC yet (will soon be added!) so to use it within Amazon SageMaker you need to create a
+`requirements.txt` in the same directory where your training script is located and add it as dependency.
+
+```
+accelerate
+```
+
+You should also add any other dependencies you have to this `requirements.txt`.
+
+
+### Configure 🤗 Accelerate
+
+You can configure the launch configuration for Amazon SageMaker the same as you do for non SageMaker training jobs with
+the 🤗 Accelerate CLI.
+
+```bash
+accelerate config
+# In which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)): 1
+```
+
+🤗 Accelerate will go through a questionnaire about your Amazon SageMaker setup and create a config file you can edit.
+
+<Tip>
+
+🤗 Accelerate is not saving any of your credentials.
+
+</Tip>
+
+### Prepare a 🤗 Accelerate fine-tuning script
+
+The training script is very similar to a training script you might run outside of SageMaker, but to save your model
+after training you need to specify either `/opt/ml/model` or use `os.environ["SM_MODEL_DIR"]` as your save
+directory. After training, artifacts in this directory are uploaded to S3.
+
+
+```diff
+- torch.save('/opt/ml/model`)
++ accelerator.save('/opt/ml/model')
+```
+
+<Tip warning={true}>
+
+SageMaker doesn’t support argparse actions. If you want to use, for example, boolean hyperparameters, you need to
+specify type as bool in your script and provide an explicit True or False value for this hyperparameter. [[REF]](https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/using_pytorch.html#prepare-a-pytorch-training-script).
+
+</Tip>
+
+### Launch Training
+
+You can launch your training with 🤗 Accelerate CLI with
+
+```
+accelerate launch path_to_script.py --args_to_the_script
+```
+
+This will launch your training script using your configuration. The only thing you have to do is provide all the
+arguments needed by your training script as named arguments.
+
+**Examples**
+
+<Tip>
+
+If you run one of the example scripts, don't forget to add `accelerator.save('/opt/ml/model')` to it.
+
+</Tip>
+
+```bash
+accelerate launch ./examples/sagemaker_example.py
+```
+
+Outputs:
+
+```
+Configuring Amazon SageMaker environment
+Converting Arguments to Hyperparameters
+Creating Estimator
+2021-04-08 11:56:50 Starting - Starting the training job...
+2021-04-08 11:57:13 Starting - Launching requested ML instancesProfilerReport-1617883008: InProgress
+.........
+2021-04-08 11:58:54 Starting - Preparing the instances for training.........
+2021-04-08 12:00:24 Downloading - Downloading input data
+2021-04-08 12:00:24 Training - Downloading the training image..................
+2021-04-08 12:03:39 Training - Training image download completed. Training in progress..
+........
+epoch 0: {'accuracy': 0.7598039215686274, 'f1': 0.8178438661710037}
+epoch 1: {'accuracy': 0.8357843137254902, 'f1': 0.882249560632689}
+epoch 2: {'accuracy': 0.8406862745098039, 'f1': 0.8869565217391304}
+........
+2021-04-08 12:05:40 Uploading - Uploading generated training model
+2021-04-08 12:05:40 Completed - Training job completed
+Training seconds: 331
+Billable seconds: 331
+You can find your model data at: s3://your-bucket/accelerate-sagemaker-1-2021-04-08-11-56-47-108/output/model.tar.gz
+```
+
+## Advanced Features
+
+### Distributed Training: Data Parallelism
+
+*currently in development, will be supported soon.*
+
+### Distributed Training: Model Parallelism
+
+*currently in development, will be supported soon.*
+
+### Python packages and dependencies
+
+🤗 Accelerate currently uses the 🤗 DLCs, with `transformers`, `datasets` and `tokenizers` pre-installed. If you
+want to use different/other Python packages you can do this by adding them to the `requirements.txt`. These packages
+will be installed before your training script is started.
+
+### Remote scripts: Use scripts located on Github
+
+*undecided if feature is needed. Contact us if you would like this feature.*
+
+### Use Spot Instances
+
+*undecided if feature is needed. Contact us if you would like this feature.*
diff --git a/docs/source/sagemaker.rst b/docs/source/sagemaker.rst
deleted file mode 100644
index b889c99b3..000000000
--- a/docs/source/sagemaker.rst
+++ /dev/null
@@ -1,169 +0,0 @@
-.. 
-    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.
-
-Amazon SageMaker
-=======================================================================================================================
-
-Hugging Face and Amazon introduced new `Hugging Face Deep Learning Containers (DLCs)
-<https://github.com/aws/deep-learning-containers/blob/master/available_images.md#huggingface-training-containers>`_ to
-make it easier than ever to train Hugging Face Transformer models in `Amazon SageMaker
-<https://aws.amazon.com/sagemaker/>`_.
-
-To learn how to use the new 🤗 DLCs with the Amazon SageMaker to run your 🤗 Accelerate scripts and raw training loops.0
-
-
-
-Getting Started
------------------------------------------------------------------------------------------------------------------------
-
-Setup & Installation
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-
-Before you can run your 🤗 Accelerate scripts on Amazon SageMaker you need to sign up for an AWS account. If you do not
-have an AWS account yet learn more `here <https://docs.aws.amazon.com/sagemaker/latest/dg/gs-set-up.html>`__.
-
-After you have your AWS Account you need to install the ``sagemaker`` sdk for 🤗 Accelerate with.
-
-.. code-block:: 
-
-    pip install "accelerate[sagemaker]" --upgrade
-
-
-🤗 Accelerate currently uses the 🤗 DLCs, with ``transformers``, ``datasets`` and ``tokenizers`` pre-installed. 🤗
-Accelerate is not in the DLC yet (will soon be added!) so to use it within Amazon SageMaker you need to create a
-``requirements.txt`` in the same directory where your training script is located and add it as dependency.
-
-.. code-block:: 
-
-    accelerate
-
-You should also add any other dependencies you have to this ``requirements.txt``.
-
-
-Configure 🤗 Accelerate
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-You can configure the launch configuration for Amazon SageMaker the same as you do for non SageMaker training jobs with
-the 🤗 Accelerate CLI.
-
-.. code-block:: 
-
-    accelerate config
-    # In which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)): 1
-
-
-🤗 Accelerate will go through a questionnaire about your Amazon SageMaker setup and create a config file you can edit.
-
-.. note::
-    🤗 Accelerate is not saving any of your credentials.
-
-
-Prepare a 🤗 Accelerate fine-tuning script
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-The training script is very similar to a training script you might run outside of SageMaker, but to save your model
-after training you need to specify either ``/opt/ml/model`` or use ``os.environ["SM_MODEL_DIR"]`` as your save
-directory. After training, artifacts in this directory are uploaded to S3.
-
-
-.. code-block:: diff
-
-    - torch.save('/opt/ml/model`)
-    + accelerator.save('/opt/ml/model')
-
-
-.. warning::
-    SageMaker doesn’t support argparse actions. If you want to use, for example, boolean hyperparameters, you need to
-    specify type as bool in your script and provide an explicit True or False value for this hyperparameter. `[REF]
-    <https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/using_pytorch.html#prepare-a-pytorch-training-script>`__.
-
-
-Launch Training
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-You can launch your training with 🤗 Accelerate CLI with
-
-.. code-block:: 
-
-    accelerate launch path_to_script.py --args_to_the_script
-
-
-This will launch your training script using your configuration. The only thing you have to do is provide all the
-arguments needed by your training script as named arguments.
-
-**Examples**
-
-.. note::
-    If you run one of the example scripts, don't forget to add ``accelerator.save('/opt/ml/model')`` to it.
-
-.. code-block:: 
-
-    accelerate launch ./examples/sagemaker_example.py 
-
-
-Outputs:
-
-.. code-block:: 
-
-    Configuring Amazon SageMaker environment
-    Converting Arguments to Hyperparameters
-    Creating Estimator
-    2021-04-08 11:56:50 Starting - Starting the training job...
-    2021-04-08 11:57:13 Starting - Launching requested ML instancesProfilerReport-1617883008: InProgress
-    .........
-    2021-04-08 11:58:54 Starting - Preparing the instances for training.........
-    2021-04-08 12:00:24 Downloading - Downloading input data
-    2021-04-08 12:00:24 Training - Downloading the training image..................
-    2021-04-08 12:03:39 Training - Training image download completed. Training in progress..
-    ........
-    epoch 0: {'accuracy': 0.7598039215686274, 'f1': 0.8178438661710037}
-    epoch 1: {'accuracy': 0.8357843137254902, 'f1': 0.882249560632689}
-    epoch 2: {'accuracy': 0.8406862745098039, 'f1': 0.8869565217391304}
-    ........
-    2021-04-08 12:05:40 Uploading - Uploading generated training model
-    2021-04-08 12:05:40 Completed - Training job completed
-    Training seconds: 331
-    Billable seconds: 331
-    You can find your model data at: s3://your-bucket/accelerate-sagemaker-1-2021-04-08-11-56-47-108/output/model.tar.gz
-
-
-
-Advanced Features
------------------------------------------------------------------------------------------------------------------------
-
-Distributed Training: Data Parallelism
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-*currently in development, will be supported soon.*
-
-Distributed Training: Model Parallelism
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-*currently in development, will be supported soon.*
-
-Python packages and dependencies
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-🤗 Accelerate currently uses the 🤗 DLCs, with ``transformers``, ``datasets`` and ``tokenizers`` pre-installed. If you
-want to use different/other Python packages you can do this by adding them to the ``requirements.txt``. These packages
-will be installed before your training script is started.
-
-Remote scripts: Use scripts located on Github
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-*undecided if feature is needed. Contact us if you would like this feature.*
-
-Use Spot Instances
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-*undecided if feature is needed. Contact us if you would like this feature.*
diff --git a/setup.py b/setup.py
index bec3b9805..e1b2d3d1b 100644
--- a/setup.py
+++ b/setup.py
@@ -17,15 +17,7 @@
 
 extras = {}
 extras["quality"] = ["black ~= 22.0", "isort >= 5.5.4", "flake8 >= 3.8.3"]
-extras["docs"] = [
-    "docutils==0.16.0",
-    "recommonmark",
-    "sphinx==3.2.1",
-    "sphinx-markdown-tables",
-    "sphinx-rtd-theme==0.4.3",
-    "sphinx-copybutton",
-    "sphinxext-opengraph==0.4.1",
-]
+extras["docs"] = []
 extras["test"] = [
     "pytest",
     "pytest-xdist",
diff --git a/src/accelerate/__init__.py b/src/accelerate/__init__.py
index f77a26a56..54ca364f6 100644
--- a/src/accelerate/__init__.py
+++ b/src/accelerate/__init__.py
@@ -5,7 +5,7 @@
 __version__ = "0.6.0.dev0"
 
 from .accelerator import Accelerator
-from .kwargs_handlers import DistributedDataParallelKwargs, GradScalerKwargs
+from .kwargs_handlers import DistributedDataParallelKwargs, GradScalerKwargs, InitProcessGroupKwargs
 from .launchers import debug_launcher, notebook_launcher
 from .state import DistributedType
 from .utils import DeepSpeedPlugin, synchronize_rng_states
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 694dcc769..fabd3e9fc 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -56,49 +56,48 @@ class Accelerator:
     Creates an instance of an accelerator for distributed training (on multi-GPU, TPU) or mixed precision training.
 
     Args:
-        device_placement (:obj:`bool`, `optional`, defaults to :obj:`True`):
+        device_placement (`bool`, *optional*, defaults to `True`):
             Whether or not the accelerator should put objects on device (tensors yielded by the dataloader, model,
             etc...).
-        split_batches (:obj:`bool`, `optional`, defaults to :obj:`False`):
+        split_batches (`bool`, *optional*, defaults to `False`):
             Whether or not the accelerator should split the batches yielded by the dataloaders across the devices. If
-            :obj:`True` the actual batch size used will be the same on any kind of distributed processes, but it must
-            be a round multiple of the :obj:`num_processes` you are using. If :obj:`False`, actual batch size used will
-            be the one set in your script multiplied by the number of processes.
-        mixed_precision (:obj:`str`, `optional`):
+            `True` the actual batch size used will be the same on any kind of distributed processes, but it must be a
+            round multiple of the `num_processes` you are using. If `False`, actual batch size used will be the one set
+            in your script multiplied by the number of processes.
+        mixed_precision (`str`, *optional*):
             Whether or not to use mixed precision training (fp16 or bfloat16). Choose from 'no','fp16','bf16'. Will
-            default to the value in the environment variable :obj:`MIXED_PRECISION`, which will use the default value
-            in the accelerate config of the current system or the flag passed with the :obj:`accelerate.launch`
-            command. 'fp16' requires pytorch 1.6 or higher. 'bf16' requires pytorch 1.10 or higher.
-        cpu (:obj:`bool`, `optional`):
-            Whether or not to force the script to execute on CPU. Will ignore GPU available if set to :obj:`True` and
-            force the execution on one process only.
-        deepspeed_plugin (:obj:`DeepSpeedPlugin`, `optional`):
+            default to the value in the environment variable `MIXED_PRECISION`, which will use the default value in the
+            accelerate config of the current system or the flag passed with the `accelerate.launch` command. 'fp16'
+            requires pytorch 1.6 or higher. 'bf16' requires pytorch 1.10 or higher.
+        cpu (`bool`, *optional*):
+            Whether or not to force the script to execute on CPU. Will ignore GPU available if set to `True` and force
+            the execution on one process only.
+        deepspeed_plugin (`DeepSpeedPlugin`, *optional*):
             Tweak your DeepSpeed related args using this argument. This argument is optional and can be configured
-            directly using `accelerate config`
-        rng_types (list of :obj:`str` or :class:`~accelerate.utils.RNGType`):
+            directly using *accelerate config*
+        rng_types (list of `str` or [`~utils.RNGType`]):
             The list of random number generators to synchronize at the beginning of each iteration in your prepared
             dataloaders. Should be one or several of:
 
-            - :obj:`"torch"`: the base torch random number generator
-            - :obj:`"cuda"`: the CUDA random number generator (GPU only)
-            - :obj:`"xla"`: the XLA random number generator (TPU only)
-            - :obj:`"generator"`: the :obj:`torch.Generator` of the sampler (or batch sampler if there is no sampler in
-              your dataloader) or of the iterable dataset (if it exists) if the underlying dataset is of that type.
-
-            Will default to :obj:`["torch"]` for PyTorch versions <=1.5.1 and :obj:`["generator"]` for PyTorch versions
-            >= 1.6.
-        dispatch_batches (:obj:`bool`, `optional`):
-            If set to :obj:`True`, the dataloader prepared by the Accelerator is only iterated through on the main
-            process and then the batches are split and broadcast to each process. Will default to :obj:`True` for
-            :obj:`DataLoader` whose underlying dataset is an :obj:`IterableDataset`, :obj:`False` otherwise.
-        kwargs_handlers (list of kwargs handlers, `optional`)
-            A list of :obj:`KwargHandler` to customize how the objects related to distributed training or mixed
-            precision are created. See :doc:`kwargs` for more information.
+            - `"torch"`: the base torch random number generator
+            - `"cuda"`: the CUDA random number generator (GPU only)
+            - `"xla"`: the XLA random number generator (TPU only)
+            - `"generator"`: the `torch.Generator` of the sampler (or batch sampler if there is no sampler in your
+              dataloader) or of the iterable dataset (if it exists) if the underlying dataset is of that type.
+
+            Will default to `["torch"]` for PyTorch versions <=1.5.1 and `["generator"]` for PyTorch versions >= 1.6.
+        dispatch_batches (`bool`, *optional*):
+            If set to `True`, the dataloader prepared by the Accelerator is only iterated through on the main process
+            and then the batches are split and broadcast to each process. Will default to `True` for `DataLoader` whose
+            underlying dataset is an `IterableDataset`, `False` otherwise.
+        kwargs_handlers (`List[KwargHandler]`, *optional*)
+            A list of `KwargHandler` to customize how the objects related to distributed training or mixed precision
+            are created. See [kwargs](kwargs) for more information.
 
     Attributes
 
-        - **device** (:obj:`torch.device`) -- The device to use.
-        - **state** (:class:`~accelerate.AcceleratorState`) -- The distributed setup state.
+        - **device** (`torch.device`) -- The device to use.
+        - **state** ([`~state.AcceleratorState`]) -- The distributed setup state.
     """
 
     def __init__(
@@ -276,7 +275,7 @@ def _goes_first(self, is_main):
 
     def print(self, *args, **kwargs):
         """
-        Use in replacement of :obj:`print()` to only print once per server.
+        Use in replacement of `print()` to only print once per server.
         """
         if self.is_local_main_process:
             print(*args, **kwargs)
@@ -296,14 +295,14 @@ def _prepare_one(self, obj):
 
     def prepare(self, *args):
         """
-        Prepare all objects passed in :obj:`args` for distributed training and mixed precision, then return them in the
-        same order.
+        Prepare all objects passed in `args` for distributed training and mixed precision, then return them in the same
+        order.
 
         Accepts the following type of objects:
 
-            - :obj:`torch.utils.data.DataLoader`: PyTorch Dataloader
-            - :obj:`torch.nn.Module`: PyTorch Module
-            - :obj:`torch.optim.Optimizer`: PyTorch Optimizer
+            - `torch.utils.data.DataLoader`: PyTorch Dataloader
+            - `torch.nn.Module`: PyTorch Module
+            - `torch.optim.Optimizer`: PyTorch Optimizer
         """
         # On TPUs, putting the model on the XLA device will create new parameters, so the corresponding optimizer will
         # have parameters disconnected from the model (so no training :-( ).
@@ -460,7 +459,7 @@ def prepare_optimizer(self, optimizer):
 
     def backward(self, loss, **kwargs):
         """
-        Use :obj:`accelerator.backward(loss)` in lieu of :obj:`loss.backward()`.
+        Use `accelerator.backward(loss)` in lieu of `loss.backward()`.
         """
         if self.distributed_type == DistributedType.DEEPSPEED:
             self.deepspeed_engine.backward(loss, **kwargs)
@@ -474,9 +473,9 @@ def unscale_gradients(self, optimizer=None):
         Unscale the gradients in mixed precision training with AMP. This is a noop in all other settings.
 
         Args:
-            optimizer (:obj:`torch.optim.Optimizer` or :obj:`List[torch.optim.Optimizer]`, `optional`):
+            optimizer (`torch.optim.Optimizer` or `List[torch.optim.Optimizer]`, *optional*):
                 The optimizer(s) for which to unscale gradients. If not set, will unscale gradients on all optimizers
-                that were passed to :meth:`~accelerate.Accelerator.prepare`.
+                that were passed to [`~Accelerator.prepare`].
         """
         if self.state.use_fp16 and self.native_amp:
             if optimizer is None:
@@ -491,34 +490,33 @@ def unscale_gradients(self, optimizer=None):
 
     def clip_grad_norm_(self, parameters, max_norm, norm_type=2):
         """
-        Should be used in place of :func:`torch.nn.utils.clip_grad_norm_`.
+        Should be used in place of `torch.nn.utils.clip_grad_norm_`.
         """
         self.unscale_gradients()
         torch.nn.utils.clip_grad_norm_(parameters, max_norm, norm_type=norm_type)
 
     def clip_grad_value_(self, parameters, clip_value):
         """
-        Should be used in place of :func:`torch.nn.utils.clip_grad_value_`.
+        Should be used in place of `torch.nn.utils.clip_grad_value_`.
         """
         self.unscale_gradients()
         torch.nn.utils.clip_grad_value_(parameters, clip_value)
 
     def gather(self, tensor):
         """
-        Gather the values in `tensor` accross all processes and concatenate them on the first dimension. Useful to
+        Gather the values in *tensor* accross all processes and concatenate them on the first dimension. Useful to
         regroup the predictions from all processes when doing evaluation.
 
         Note:
             This gather happens in all processes.
 
         Args:
-            tensor (:obj:`torch.Tensor`, or a nested tuple/list/dictionary of :obj:`torch.Tensor`):
+            tensor (`torch.Tensor`, or a nested tuple/list/dictionary of `torch.Tensor`):
                 The tensors to gather across all processes.
 
         Returns:
-            :obj:`torch.Tensor`, or a nested tuple/list/dictionary of :obj:`torch.Tensor`: The gathered tensor(s). Note
-            that the first dimension of the result is `num_processes` multiplied by the first dimension of the input
-            tensors.
+            `torch.Tensor`, or a nested tuple/list/dictionary of `torch.Tensor`: The gathered tensor(s). Note that the
+            first dimension of the result is *num_processes* multiplied by the first dimension of the input tensors.
         """
         return gather(tensor)
 
@@ -528,24 +526,24 @@ def pad_across_processes(self, tensor, dim=0, pad_index=0, pad_first=False):
         they can safely be gathered.
 
         Args:
-            tensor (nested list/tuple/dictionary of :obj:`torch.Tensor`):
+            tensor (nested list/tuple/dictionary of `torch.Tensor`):
                 The data to gather.
-            dim (:obj:`int`, `optional`, defaults to 0):
+            dim (`int`, *optional*, defaults to 0):
                 The dimension on which to pad.
-            pad_index (:obj:`int`, `optional`, defaults to 0):
+            pad_index (`int`, *optional*, defaults to 0):
                 The value with which to pad.
-            pad_first (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            pad_first (`bool`, *optional*, defaults to `False`):
                 Whether to pad at the beginning or the end.
         """
         return pad_across_processes(tensor, dim=dim, pad_index=pad_index, pad_first=pad_first)
 
     def unwrap_model(self, model):
         """
-        Unwraps the :obj:`model` from the additional layer possible added by :meth:`~accelerate.Accelerator.prepare`.
-        Useful before saving the model.
+        Unwraps the `model` from the additional layer possible added by [`~Accelerator.prepare`]. Useful before saving
+        the model.
 
         Args:
-            model (:obj:`torch.nn.Module`):
+            model (`torch.nn.Module`):
                 The model to unwrap.
         """
         return extract_model_from_parallel(model)
@@ -559,12 +557,12 @@ def wait_for_everyone(self):
 
     def save(self, obj, f):
         """
-        Save the object passed to disk once per machine. Use in place of :obj:`torch.save`.
+        Save the object passed to disk once per machine. Use in place of `torch.save`.
 
         Args:
             obj: The object to save.
-            f (:obj:`str` or :obj:`os.PathLike`):
-                Where to save the content of :obj:`obj`.
+            f (`str` or `os.PathLike`):
+                Where to save the content of `obj`.
         """
         save(obj, f)
 
@@ -573,7 +571,7 @@ def save_state(self, output_dir: str):
         Saves the current states of the model, optimizer, scaler, RNG generators, and registered objects.
 
         Args:
-            output_dir (:obj:`str` or :obj:`os.PathLike`):
+            output_dir (`str` or `os.PathLike`):
                 The name of the folder to save all relevant weights and states.
         """
         # Check if folder exists
@@ -593,7 +591,7 @@ def load_state(self, input_dir: str):
         Loads the current states of the model, optimizer, scaler, RNG generators, and registered objects.
 
         Args:
-            input_dir (:obj:`str` or :obj:`os.PathLike`):
+            input_dir (`str` or `os.PathLike`):
                 The name of the folder all relevant weights and states were saved in.
         """
         # Check if folder exists
@@ -674,7 +672,11 @@ def register_for_checkpointing(self, *objects):
         These should be utilized when the state is being loaded or saved in the same script. It is not designed to be
         used in different scripts
 
-        Note: Every `object` must have a `load_state_dict` and `state_dict` function to be stored.
+        <Tip>
+
+        Every `object` must have a `load_state_dict` and `state_dict` function to be stored.
+
+        </Tip>
         """
         invalid_objects = []
         for obj in objects:
diff --git a/src/accelerate/checkpointing.py b/src/accelerate/checkpointing.py
index e1a28d108..0636dffaa 100644
--- a/src/accelerate/checkpointing.py
+++ b/src/accelerate/checkpointing.py
@@ -41,15 +41,15 @@ def save_accelerator_state(
     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`):
+        output_dir (`str` or `os.PathLike`):
             The name of the folder to save all relevant weights and states.
-        model_states (:obj:`List[torch.nn.Module]`):
+        model_states (`List[torch.nn.Module]`):
             A list of model states
-        optimizers (:obj:`List[torch.optim.Optimizer]`):
+        optimizers (`List[torch.optim.Optimizer]`):
             A list of optimizer instances
-        process_index (:obj:`int`):
+        process_index (`int`):
             The current process index in the Accelerator state
-        scaler (:obj:`torch.cuda.amp.GradScaler`, `optional`):
+        scaler (`torch.cuda.amp.GradScaler`, *optional*):
             An optional gradient scaler instance to save
     """
     # Model states
@@ -92,16 +92,16 @@ def load_accelerator_state(input_dir, models, optimizers, process_index, scaler=
     Loads states of the models, optimizers, scaler, and RNG generators from a given directory.
 
     Args:
-        input_dir (:obj:`str` or :obj:`os.PathLike`):
+        input_dir (`str` or `os.PathLike`):
             The name of the folder to load all relevant weights and states.
-        model_stmodelsates (:obj:`List[torch.nn.Module]`):
+        model_stmodelsates (`List[torch.nn.Module]`):
             A list of model instances
-        optimizers (:obj:`List[torch.optim.Optimizer]`):
+        optimizers (`List[torch.optim.Optimizer]`):
             A list of optimizer instances
-        process_index (:obj:`int`):
+        process_index (`int`):
             The current process index in the Accelerator state
-        scaler (:obj:`torch.cuda.amp.GradScaler`, `optional`):
-            An optional `GradScaler` instance to load
+        scaler (`torch.cuda.amp.GradScaler`, *optional*):
+            An optional *GradScaler* instance to load
     """
     # Model states
     for i, model in enumerate(models):
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index cc46e5577..a827633c0 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -67,34 +67,34 @@
 
 class BatchSamplerShard(BatchSampler):
     """
-    Wraps a PyTorch :obj:`BatchSampler` to generate batches for one of the processes only. Instances of this class will
-    always yield a number of batches that is a round multiple of :obj:`num_processes` and that all have the same size.
-    Depending on the value of the :obj:`drop_last` attribute of the batch sampler passed, it will either stop the
-    iteration at the first batch that would be too small / not present on all processes or loop with indices from the
-    beginning.
+    Wraps a PyTorch `BatchSampler` to generate batches for one of the processes only. Instances of this class will
+    always yield a number of batches that is a round multiple of `num_processes` and that all have the same size.
+    Depending on the value of the `drop_last` attribute of the batch sampler passed, it will either stop the iteration
+    at the first batch that would be too small / not present on all processes or loop with indices from the beginning.
 
     Args:
-        batch_sampler (:obj:`torch.utils.data.sampler.BatchSampler`):
+        batch_sampler (`torch.utils.data.sampler.BatchSampler`):
             The batch sampler to split in several shards.
-        num_processes (:obj:`int`, `optional`, defaults to 1):
+        num_processes (`int`, *optional*, defaults to 1):
             The number of processes running concurrently.
-        process_index (:obj:`int`, `optional`, defaults to 0):
+        process_index (`int`, *optional*, defaults to 0):
             The index of the current process.
-        split_batches (:obj:`bool`, `optional`, defaults to :obj:`False`):
+        split_batches (`bool`, *optional*, defaults to `False`):
             Whether the shards should be created by splitting a batch to give a piece of it on each process, or by
             yielding different full batches on each process.
 
-            On two processes with a sampler of :obj:`[[0, 1, 2, 3], [4, 5, 6, 7]]`, this will result in:
+            On two processes with a sampler of `[[0, 1, 2, 3], [4, 5, 6, 7]]`, this will result in:
 
-            - the sampler on process 0 to yield :obj:`[0, 1, 2, 3]` and the sampler on process 1 to yield :obj:`[4, 5,
-              6, 7]` if this argument is set to :obj:`False`.
-            - the sampler on process 0 to yield :obj:`[0, 1]` then :obj:`[4, 5]` and the sampler on process 1 to yield
-              :obj:`[2, 3]` then :obj:`[6, 7]` if this argument is set to :obj:`True`.
+            - the sampler on process 0 to yield `[0, 1, 2, 3]` and the sampler on process 1 to yield `[4, 5, 6, 7]` if
+              this argument is set to `False`.
+            - the sampler on process 0 to yield `[0, 1]` then `[4, 5]` and the sampler on process 1 to yield `[2, 3]`
+              then `[6, 7]` if this argument is set to `True`.
 
-    .. warning::
+    <Tip warning={true}>
 
-        This does not support :obj:`BatchSampler` with varying batch size yet.
-    """
+    This does not support `BatchSampler` with varying batch size yet.
+
+    </Tip>"""
 
     def __init__(
         self,
@@ -188,35 +188,35 @@ def _iter_with_no_split(self):
 
 class IterableDatasetShard(IterableDataset):
     """
-    Wraps a PyTorch :obj:`IterableDataset` to generate samples for one of the processes only. Instances of this class
-    will always yield a number of samples that is a round multiple of the actual batch size (depending of the value of
-    :obj:`split_batches`, this is either :obj:`batch_size` or :obj:`batch_size x num_processes`). Depending on the
-    value of the :obj:`drop_last` attribute of the batch sampler passed, it will either stop the iteration at the first
-    batch that would be too small or loop with indices from the beginning.
+    Wraps a PyTorch `IterableDataset` to generate samples for one of the processes only. Instances of this class will
+    always yield a number of samples that is a round multiple of the actual batch size (depending of the value of
+    `split_batches`, this is either `batch_size` or `batch_size x num_processes`). Depending on the value of the
+    `drop_last` attribute of the batch sampler passed, it will either stop the iteration at the first batch that would
+    be too small or loop with indices from the beginning.
 
     Args:
-        dataset (:obj:`torch.utils.data.dataset.IterableDataset`):
+        dataset (`torch.utils.data.dataset.IterableDataset`):
             The batch sampler to split in several shards.
-        batch_size (:obj:`int`, `optional`, defaults to 1):
-            The size of the batches per shard (if :obj:`split_batches=False`) or the size of the batches (if
-            :obj:`split_batches=True`).
-        drop_last (:obj:`bool`, `optional`, defaults to :obj:`False`):
+        batch_size (`int`, *optional*, defaults to 1):
+            The size of the batches per shard (if `split_batches=False`) or the size of the batches (if
+            `split_batches=True`).
+        drop_last (`bool`, *optional*, defaults to `False`):
             Whether or not to drop the last incomplete batch or complete the last batches by using the samples from the
             beginning.
-        num_processes (:obj:`int`, `optional`, defaults to 1):
+        num_processes (`int`, *optional*, defaults to 1):
             The number of processes running concurrently.
-        process_index (:obj:`int`, `optional`, defaults to 0):
+        process_index (`int`, *optional*, defaults to 0):
             The index of the current process.
-        split_batches (:obj:`bool`, `optional`, defaults to :obj:`False`):
+        split_batches (`bool`, *optional*, defaults to `False`):
             Whether the shards should be created by splitting a batch to give a piece of it on each process, or by
             yielding different full batches on each process.
 
-            On two processes with an iterable dataset yielding of :obj:`[0, 1, 2, 3, 4, 5, 6, 7]`, this will result in:
+            On two processes with an iterable dataset yielding of `[0, 1, 2, 3, 4, 5, 6, 7]`, this will result in:
 
-            - the shard on process 0 to yield :obj:`[0, 1, 2, 3]` and the shard on process 1 to yield :obj:`[4, 5, 6,
-              7]` if this argument is set to :obj:`False`.
-            - the shard on process 0 to yield :obj:`[0, 1, 4, 5]` and the sampler on process 1 to yield :obj:`[2, 3, 6,
-              7]` if this argument is set to :obj:`True`.
+            - the shard on process 0 to yield `[0, 1, 2, 3]` and the shard on process 1 to yield `[4, 5, 6, 7]` if this
+              argument is set to `False`.
+            - the shard on process 0 to yield `[0, 1, 4, 5]` and the sampler on process 1 to yield `[2, 3, 6, 7]` if
+              this argument is set to `True`.
     """
 
     def __init__(
@@ -269,25 +269,25 @@ def __iter__(self):
 
 class DataLoaderShard(DataLoader):
     """
-    Subclass of a PyTorch :obj:`DataLoader` that will deal with device placement and current distributed setup.
+    Subclass of a PyTorch `DataLoader` that will deal with device placement and current distributed setup.
 
     Args:
-        dataset (:obj:`torch.utils.data.dataset.Dataset`):
+        dataset (`torch.utils.data.dataset.Dataset`):
             The dataset to use to build this datalaoder.
-        device (:obj:`torch.device`, `optional`):
+        device (`torch.device`, *optional*):
             If passed, the device to put all batches on.
-        rng_types (list of :obj:`str` or :class:`~accelerate.utils.RNGType`):
+        rng_types (list of `str` or [`~utils.RNGType`]):
             The list of random number generators to synchronize at the beginning of each iteration. Should be one or
             several of:
 
-            - :obj:`"torch"`: the base torch random number generator
-            - :obj:`"cuda"`: the CUDA random number generator (GPU only)
-            - :obj:`"xla"`: the XLA random number generator (TPU only)
-            - :obj:`"generator"`: an optional :obj:`torch.Generator`
-        generator (:obj:`torch.Generator`, `optional`):
+            - `"torch"`: the base torch random number generator
+            - `"cuda"`: the CUDA random number generator (GPU only)
+            - `"xla"`: the XLA random number generator (TPU only)
+            - `"generator"`: an optional `torch.Generator`
+        generator (`torch.Generator`, *optional*):
             A random number generator to keep synchronized across processes.
         kwargs:
-            All other keyword arguments to pass to the regular :obj:`DataLoader` initialization.
+            All other keyword arguments to pass to the regular `DataLoader` initialization.
     """
 
     def __init__(self, dataset, device=None, rng_types=None, generator=None, **kwargs):
@@ -308,21 +308,21 @@ def __iter__(self):
 
 class DataLoaderDispatcher(DataLoader):
     """
-    Subclass of a PyTorch :obj:`DataLoader` that will iterate and preprocess on process 0 only, then dispatch on each
+    Subclass of a PyTorch `DataLoader` that will iterate and preprocess on process 0 only, then dispatch on each
     process their part of the batch.
 
     Args:
-        split_batches (:obj:`bool`, `optional`, defaults to :obj:`False`):
-            Whether the resulting :obj:`DataLoader` should split the batches of the original data loader across devices
-            or yield full batches (in which case it will yield batches starting at the :obj:`process_index`-th and
-            advancing of :obj:`num_processes` batches at each iteration).
+        split_batches (`bool`, *optional*, defaults to `False`):
+            Whether the resulting `DataLoader` should split the batches of the original data loader across devices or
+            yield full batches (in which case it will yield batches starting at the `process_index`-th and advancing of
+            `num_processes` batches at each iteration).
 
-            Another way to see this is that the observed batch size will be the same as the initial :obj:`dataloader`
-            if this option is set to :obj:`True`, the batch size of the initial :obj:`dataloader` multiplied by
-            :obj:`num_processes` otherwise.
+            Another way to see this is that the observed batch size will be the same as the initial `dataloader` if
+            this option is set to `True`, the batch size of the initial `dataloader` multiplied by `num_processes`
+            otherwise.
 
-            Setting this option to :obj:`True` requires that the batch size of the :obj:`dataloader` is a round
-            multiple of :obj:`batch_size`.
+            Setting this option to `True` requires that the batch size of the `dataloader` is a round multiple of
+            `batch_size`.
     """
 
     def __init__(self, dataset, split_batches: bool = False, **kwargs):
@@ -425,58 +425,58 @@ def prepare_data_loader(
     dispatch_batches: Optional[bool] = None,
 ) -> DataLoader:
     """
-    Wraps a PyTorch :obj:`DataLoader` to generate batches for one of the processes only.
+    Wraps a PyTorch `DataLoader` to generate batches for one of the processes only.
 
-    Depending on the value of the :obj:`drop_last` attribute of the :obj:`dataloader` passed, it will either stop the
-    iteration at the first batch that would be too small / not present on all processes or loop with indices from the
-    beginning.
+    Depending on the value of the `drop_last` attribute of the `dataloader` passed, it will either stop the iteration
+    at the first batch that would be too small / not present on all processes or loop with indices from the beginning.
 
     Args:
-        dataloader (:obj:`torch.utils.data.dataloader.DataLoader`):
+        dataloader (`torch.utils.data.dataloader.DataLoader`):
             The data loader to split across several devices.
-        device (:obj:`torch.device`):
-            The target device for the returned :obj:`DataLoader`.
-        num_processes (:obj:`int`, `optional`):
+        device (`torch.device`):
+            The target device for the returned `DataLoader`.
+        num_processes (`int`, *optional*):
             The number of processes running concurrently. Will default to the value given by
-            :class:`~accelerate.AcceleratorState`.
-        process_index (:obj:`int`, `optional`):
-            The index of the current process. Will default to the value given by :class:`~accelerate.AcceleratorState`.
-        split_batches (:obj:`bool`, `optional`, defaults to :obj:`False`):
-            Whether the resulting :obj:`DataLoader` should split the batches of the original data loader across devices
-            or yield full batches (in which case it will yield batches starting at the :obj:`process_index`-th and
-            advancing of :obj:`num_processes` batches at each iteration).
-
-            Another way to see this is that the observed batch size will be the same as the initial :obj:`dataloader`
-            if this option is set to :obj:`True`, the batch size of the initial :obj:`dataloader` multiplied by
-            :obj:`num_processes` otherwise.
-
-            Setting this option to :obj:`True` requires that the batch size of the :obj:`dataloader` is a round
-            multiple of :obj:`batch_size`.
-        put_on_device (:obj:`bool`, `optional`, defaults to :obj:`False`):
-            Whether or not to put the batches on :obj:`device` (only works if the batches are nested list, tuples or
+            [`~state.AcceleratorState`].
+        process_index (`int`, *optional*):
+            The index of the current process. Will default to the value given by [`~state.AcceleratorState`].
+        split_batches (`bool`, *optional*, defaults to `False`):
+            Whether the resulting `DataLoader` should split the batches of the original data loader across devices or
+            yield full batches (in which case it will yield batches starting at the `process_index`-th and advancing of
+            `num_processes` batches at each iteration).
+
+            Another way to see this is that the observed batch size will be the same as the initial `dataloader` if
+            this option is set to `True`, the batch size of the initial `dataloader` multiplied by `num_processes`
+            otherwise.
+
+            Setting this option to `True` requires that the batch size of the `dataloader` is a round multiple of
+            `batch_size`.
+        put_on_device (`bool`, *optional*, defaults to `False`):
+            Whether or not to put the batches on `device` (only works if the batches are nested list, tuples or
             dictionaries of tensors).
-        rng_types (list of :obj:`str` or :class:`~accelerate.utils.RNGType`):
+        rng_types (list of `str` or [`~utils.RNGType`]):
             The list of random number generators to synchronize at the beginning of each iteration. Should be one or
             several of:
 
-            - :obj:`"torch"`: the base torch random number generator
-            - :obj:`"cuda"`: the CUDA random number generator (GPU only)
-            - :obj:`"xla"`: the XLA random number generator (TPU only)
-            - :obj:`"generator"`: the :obj:`torch.Generator` of the sampler (or batch sampler if there is no sampler in
-              your dataloader) or of the iterable dataset (if it exists) if the underlying dataset is of that type.
+            - `"torch"`: the base torch random number generator
+            - `"cuda"`: the CUDA random number generator (GPU only)
+            - `"xla"`: the XLA random number generator (TPU only)
+            - `"generator"`: the `torch.Generator` of the sampler (or batch sampler if there is no sampler in your
+              dataloader) or of the iterable dataset (if it exists) if the underlying dataset is of that type.
 
-        dispatch_batches (:obj:`bool`, `optional`):
-            If set to :obj:`True`, the datalaoder prepared is only iterated through on the main process and then the
-            batches are split and broadcast to each process. Will default to :obj:`True` when the underlying dataset is
-            an :obj:`IterableDataset`, :obj:`False` otherwise.
+        dispatch_batches (`bool`, *optional*):
+            If set to `True`, the datalaoder prepared is only iterated through on the main process and then the batches
+            are split and broadcast to each process. Will default to `True` when the underlying dataset is an
+            `IterableDataset`, `False` otherwise.
 
     Returns:
-        :obj:`torch.utils.data.dataloader.DataLoader`: A new data loader that will yield the portion of the batches
+        `torch.utils.data.dataloader.DataLoader`: A new data loader that will yield the portion of the batches
 
-    .. warning::
+    <Tip warning={true}>
 
-        This does not support :obj:`BatchSampler` with varying batch size yet.
-    """
+    This does not support `BatchSampler` with varying batch size yet.
+
+    </Tip>"""
     if dispatch_batches is None:
         if version.parse(torch.__version__) < version.parse("1.8.0") or not put_on_device:
             dispatch_batches = False
diff --git a/src/accelerate/deepspeed_utils.py b/src/accelerate/deepspeed_utils.py
index 07450a87e..5f261f2e3 100644
--- a/src/accelerate/deepspeed_utils.py
+++ b/src/accelerate/deepspeed_utils.py
@@ -71,7 +71,7 @@ class DeepSpeedOptimizerWrapper(AcceleratedOptimizer):
     Internal wrapper around a deepspeed optimizer.
 
     Args:
-        optimizer (:obj:`torch.optim.optimizer.Optimizer`):
+        optimizer (`torch.optim.optimizer.Optimizer`):
             The optimizer to wrap.
     """
 
diff --git a/src/accelerate/kwargs_handlers.py b/src/accelerate/kwargs_handlers.py
index 51760efdc..cf5985644 100644
--- a/src/accelerate/kwargs_handlers.py
+++ b/src/accelerate/kwargs_handlers.py
@@ -20,7 +20,7 @@
 
 class KwargsHandler:
     """
-    Internal mixin that implements a :obj:`to_kwargs()` method for a dataclass.
+    Internal mixin that implements a `to_kwargs()` method for a dataclass.
     """
 
     def to_dict(self):
@@ -38,15 +38,16 @@ def to_kwargs(self):
 @dataclass
 class DistributedDataParallelKwargs(KwargsHandler):
     """
-    Use this object in your :class:`~accelerate.Accelerator` to customize how your model is wrapped in a
-    :obj:`torch.nn.parallel.DistributedDataParallel`. Please refer to the documentation of this `wrapper
-    <https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html>`__ for more information
-    on each argument.
+    Use this object in your [`Accelerator`] to customize how your model is wrapped in a
+    `torch.nn.parallel.DistributedDataParallel`. Please refer to the documentation of this
+    [wrapper](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html) for more
+    information on each argument.
 
-    .. warning::
+    <Tip warning={true}>
 
-        :obj:`gradient_as_bucket_view` is only available in PyTorch 1.7.0 and later versions.
-    """
+    `gradient_as_bucket_view` is only available in PyTorch 1.7.0 and later versions.
+
+    </Tip>"""
 
     dim: int = 0
     broadcast_buffers: bool = True
@@ -59,14 +60,15 @@ class DistributedDataParallelKwargs(KwargsHandler):
 @dataclass
 class GradScalerKwargs(KwargsHandler):
     """
-    Use this object in your :class:`~accelerate.Accelerator` to customize the behavior of mixed precision, specifically
-    how the :obj:`torch.cuda.amp.GradScaler` used is created. Please refer to the documentation of this `scaler
-    <https://pytorch.org/docs/stable/amp.html?highlight=gradscaler>`__ for more information on each argument.
+    Use this object in your [`Accelerator`] to customize the behavior of mixed precision, specifically how the
+    `torch.cuda.amp.GradScaler` used is created. Please refer to the documentation of this
+    [scaler](https://pytorch.org/docs/stable/amp.html?highlight=gradscaler) for more information on each argument.
 
-    .. warning::
+    <Tip warning={true}>
 
-        :obj:`GradScaler` is only available in PyTorch 1.5.0 and later versions.
-    """
+    `GradScaler` is only available in PyTorch 1.5.0 and later versions.
+
+    </Tip>"""
 
     init_scale: float = 65536.0
     growth_factor: float = 2.0
@@ -78,10 +80,10 @@ class GradScalerKwargs(KwargsHandler):
 @dataclass
 class InitProcessGroupKwargs(KwargsHandler):
     """
-    Use this object in your :class:`~accelerate.Accelerator` to customize the initialization of the distributed
-    processes. Please refer to the documentation of this `method
-    <https://pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group>`__ for more information on
-    each argument.
+    Use this object in your [`Accelerator`] to customize the initialization of the distributed processes. Please refer
+    to the documentation of this
+    [method](https://pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group) for more
+    information on each argument.
     """
 
     init_method: Optional[str] = None
diff --git a/src/accelerate/launchers.py b/src/accelerate/launchers.py
index a47eb9a4c..48c4404dd 100644
--- a/src/accelerate/launchers.py
+++ b/src/accelerate/launchers.py
@@ -31,17 +31,17 @@ def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, mix
     multiple cores for instance).
 
     Args:
-        function (:obj:`Callable`):
+        function (`Callable`):
             The training function to execute. If it accepts arguments, the first argument should be the index of the
             process run.
-        args (:obj:`Tuple`):
-            Tuple of arguments to pass to the function (it will receive :obj:`*args`).
-        num_processes (:obj:`int`, `optional`):
+        args (`Tuple`):
+            Tuple of arguments to pass to the function (it will receive `*args`).
+        num_processes (`int`, *optional*):
             The number of processes to use for training. Will default to 8 in Colab/Kaggle if a TPU is available, to
             the number of GPUs available otherwise.
-        mixed_precision (:obj:`str`, `optional`, defaults to :obj:`no`):
-            If :obj:`fp16` or :obj:`bf16`, will use mixed precision training on multi-GPU.
-        use_port (:obj:`str`, `optional`, defaults to :obj:`"29500"`):
+        mixed_precision (`str`, *optional*, defaults to `"no"`):
+            If `fp16` or `bf16`, will use mixed precision training on multi-GPU.
+        use_port (`str`, *optional*, defaults to `"29500"`):
             The port to use to communicate between processes when launching a multi-GPU training.
     """
     # Are we in a google colab or a Kaggle Kernel?
@@ -140,18 +140,19 @@ def debug_launcher(function, args=(), num_processes=2):
     """
     Launches a training function using several processes on CPU for debugging purposes.
 
-    .. warning::
+    <Tip warning={true}>
 
-        This function is provided for internal testing and debugging, but it's not intended for real trainings. It will
-        only use the CPU.
+    This function is provided for internal testing and debugging, but it's not intended for real trainings. It will
+    only use the CPU.
 
+    </Tip>
 
     Args:
-        function (:obj:`Callable`):
+        function (`Callable`):
             The training function to execute.
-        args (:obj:`Tuple`):
-            Tuple of arguments to pass to the function (it will receive :obj:`*args`).
-        num_processes (:obj:`int`, *optional*, defaults to 2):
+        args (`Tuple`):
+            Tuple of arguments to pass to the function (it will receive `*args`).
+        num_processes (`int`, *optional*, defaults to 2):
             The number of processes to use for training.
     """
     if version.parse(torch.__version__) < version.parse("1.5.0"):
diff --git a/src/accelerate/optimizer.py b/src/accelerate/optimizer.py
index 52217ba6a..4360c0c09 100644
--- a/src/accelerate/optimizer.py
+++ b/src/accelerate/optimizer.py
@@ -41,12 +41,12 @@ class AcceleratedOptimizer(torch.optim.Optimizer):
     Internal wrapper around a torch optimizer.
 
     Args:
-        optimizer (:obj:`torch.optim.optimizer.Optimizer`):
+        optimizer (`torch.optim.optimizer.Optimizer`):
             The optimizer to wrap.
-        device_placement (:obj:`bool`, `optional`, defaults to :obj:`True`):
+        device_placement (`bool`, *optional*, defaults to `True`):
             Whether or not the optimizer should handle device placement. If so, it will place the state dictionary of
-            :obj:`optimizer` on the right device.
-        scaler (:obj:`torch.cuda.amp.grad_scaler.GradScaler`, `optional`):
+            `optimizer` on the right device.
+        scaler (`torch.cuda.amp.grad_scaler.GradScaler`, *optional*):
             The scaler to use in the step function if training with mixed precision.
     """
 
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index e8c910e19..0e0c21e1a 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -127,19 +127,19 @@ class ComputeEnvironment(str, Enum):
 # Inspired by Alex Martelli's 'Borg'.
 class AcceleratorState:
     """
-    This is a variation of a `singleton class <https://en.wikipedia.org/wiki/Singleton_pattern>`__ in the sense that
-    all instance of :obj:`AcceleratorState` share the same state, which is initialized on the first instantiation.
-
-    Attributes
-
-        - **device** (:obj:`torch.device`) -- The device to use.
-        - **distributed_type** (:obj:`~accelerate.state.DistributedType`) -- The type of distributed environment
-          currently in use.
-        - **num_processes** (:obj:`int`) -- The number of processes currently launched in parallel.
-        - **process_index** (:obj:`int`) -- The index of the current process.
-        - **local_process_index** (:obj:`int`) -- The index of the current process on the current server.
-        - **mixed_precision** (:obj:`str`) -- Whether or not the current script will use mixed precision. If you are
-          using mixed precision, define if you want to use FP16 or BF16 (bfloat16) as the floating point.
+    This is a variation of a [singleton class](https://en.wikipedia.org/wiki/Singleton_pattern) in the sense that all
+    instance of `AcceleratorState` share the same state, which is initialized on the first instantiation.
+
+    Attributes:
+
+        - **device** (`torch.device`) -- The device to use.
+        - **distributed_type** (`~accelerate.state.DistributedType`) -- The type of distributed environment currently
+          in use.
+        - **num_processes** (`int`) -- The number of processes currently launched in parallel.
+        - **process_index** (`int`) -- The index of the current process.
+        - **local_process_index** (`int`) -- The index of the current process on the current server.
+        - **mixed_precision** (`str`) -- Whether or not the current script will use mixed precision. If you are using
+          mixed precision, define if you want to use FP16 or BF16 (bfloat16) as the floating point.
     """
 
     _shared_state = {}
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 4b20cc1c3..69e74f2a5 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -65,11 +65,12 @@ class TensorInformation:
 
 def set_seed(seed: int, device_specific: bool = False):
     """
-    Helper function for reproducible behavior to set the seed in ``random``, ``numpy``, ``torch``.
+    Helper function for reproducible behavior to set the seed in `random`, `numpy`, `torch`.
 
     Args:
-        seed (:obj:`int`): The seed to set.
-        device_specific (:obj:`bool`, `optional`, defaults to :obj:`False`): Whether to differ the seed on each device slightly with `self.process_index`
+        seed (`int`): The seed to set.
+        device_specific (`bool`, *optional*, defaults to `False`):
+            Whether to differ the seed on each device slightly with `self.process_index`.
     """
     if device_specific:
         seed += AcceleratorState().process_index
@@ -146,23 +147,22 @@ def recursively_apply(func, data, *args, test_type=is_torch_tensor, error_on_oth
     Recursively apply a function on a data structure that is a nested list/tuple/dictionary of a given base type.
 
     Args:
-        func (:obj:`callable`):
+        func (`callable`):
             The function to recursively apply.
-        data (nested list/tuple/dictionary of :obj:`main_type`):
-            The data on which to apply :obj:`func`
+        data (nested list/tuple/dictionary of `main_type`):
+            The data on which to apply `func`
         *args:
-            Positional arguments that will be passed to :obj:`func` when applied on the unpacked data.
-        main_type (:obj:`type`, `optional`, defaults to :obj:`torch.Tensor`):
-            The base type of the objects to which apply :obj:`func`.
-        error_on_other_type (:obj:`bool`, `optional`, defaults to :obj:`False`):
-            Whether to return an error or not if after unpacking :obj:`data`, we get on an object that is not of type
-            :obj:`main_type`. If :obj:`False`, the function will leave objects of types different than :obj:`main_type`
-            unchanged.
+            Positional arguments that will be passed to `func` when applied on the unpacked data.
+        main_type (`type`, *optional*, defaults to `torch.Tensor`):
+            The base type of the objects to which apply `func`.
+        error_on_other_type (`bool`, *optional*, defaults to `False`):
+            Whether to return an error or not if after unpacking `data`, we get on an object that is not of type
+            `main_type`. If `False`, the function will leave objects of types different than `main_type` unchanged.
         **kwargs:
-            Keyword arguments that will be passed to :obj:`func` when applied on the unpacked data.
+            Keyword arguments that will be passed to `func` when applied on the unpacked data.
 
     Returns:
-        The same data structure as :obj:`data` with :obj:`func` applied to every object of type :obj:`main_type`.
+        The same data structure as `data` with `func` applied to every object of type `main_type`.
     """
     if isinstance(data, (tuple, list)):
         return honor_type(
@@ -198,13 +198,13 @@ def send_to_device(tensor, device):
     Recursively sends the elements in a nested list/tuple/dictionary of tensors to a given device.
 
     Args:
-        tensor (nested list/tuple/dictionary of :obj:`torch.Tensor`):
+        tensor (nested list/tuple/dictionary of `torch.Tensor`):
             The data to send to a given device.
-        device (:obj:`torch.device`):
-            The device to send the data to
+        device (`torch.device`):
+            The device to send the data to.
 
     Returns:
-        The same data structure as :obj:`tensor` with all tensors sent to the proper device.
+        The same data structure as `tensor` with all tensors sent to the proper device.
     """
 
     def _send_to_device(t, device):
@@ -221,11 +221,11 @@ def get_data_structure(data):
     Recursively gathers the information needed to rebuild a nested list/tuple/dictionary of tensors.
 
     Args:
-        data (nested list/tuple/dictionary of :obj:`torch.Tensor`):
+        data (nested list/tuple/dictionary of `torch.Tensor`):
             The data to send to analyze.
 
     Returns:
-        The same data structure as :obj:`data` with :class:`~accelerate.utils.TensorInformation` instead of tensors.
+        The same data structure as `data` with [`~utils.TensorInformation`] instead of tensors.
     """
 
     def _get_data_structure(tensor):
@@ -236,11 +236,10 @@ def _get_data_structure(tensor):
 
 def initialize_tensors(data_structure):
     """
-    Recursively initializes tensors from a nested list/tuple/dictionary of
-    :class:`~accelerate.utils.TensorInformation`.
+    Recursively initializes tensors from a nested list/tuple/dictionary of [`~utils.TensorInformation`].
 
     Returns:
-        The same data structure as :obj:`data` with tensors instead of :class:`~accelerate.utils.TensorInformation`.
+        The same data structure as `data` with tensors instead of [`~utils.TensorInformation`].
     """
 
     def _initialize_tensor(tensor_info):
@@ -254,11 +253,11 @@ def convert_to_fp32(tensor):
     Recursively converts the elements nested list/tuple/dictionary of tensors in FP16/BF16 precision to FP32.
 
     Args:
-        tensor (nested list/tuple/dictionary of :obj:`torch.Tensor`):
+        tensor (nested list/tuple/dictionary of `torch.Tensor`):
             The data to convert from FP16/BF16 to FP32.
 
     Returns:
-        The same data structure as :obj:`tensor` with all tensors that were in FP16/BF16 precision converted to FP32.
+        The same data structure as `tensor` with all tensors that were in FP16/BF16 precision converted to FP32.
     """
 
     def _convert_to_fp32(tensor):
@@ -279,11 +278,11 @@ def convert_outputs_to_fp32(model_forward):
     precision will be convert back to FP32.
 
     Args:
-        model_forward (:obj:`Callable`):
+        model_forward (`Callable`):
             The function which outputs we want to treat.
 
     Returns:
-        The same function as :obj:`model_forward` but with converted outputs.
+        The same function as `model_forward` but with converted outputs.
     """
 
     def convert_outputs(*args, **kwargs):
@@ -298,10 +297,10 @@ def extract_model_from_parallel(model):
     Extract a model from its distributed containers.
 
     Args:
-        model (:obj:`torch.nn.Module`): The model to extract.
+        model (`torch.nn.Module`): The model to extract.
 
     Returns:
-        :obj:`torch.nn.Module`: The extracted model.
+        `torch.nn.Module`: The extracted model.
     """
     options = (torch.nn.parallel.DistributedDataParallel, torch.nn.DataParallel)
     if is_deepspeed_available():
@@ -343,11 +342,11 @@ def gather(tensor):
     Recursively gather tensor in a nested list/tuple/dictionary of tensors from all devices.
 
     Args:
-        tensor (nested list/tuple/dictionary of :obj:`torch.Tensor`):
+        tensor (nested list/tuple/dictionary of `torch.Tensor`):
             The data to gather.
 
     Returns:
-        The same data structure as :obj:`tensor` with all tensors sent to the proper device.
+        The same data structure as `tensor` with all tensors sent to the proper device.
     """
     if AcceleratorState().distributed_type == DistributedType.TPU:
         return _tpu_gather(tensor, name="accelerate.utils.gather")
@@ -380,7 +379,7 @@ def gather_object(object: Any):
             The data to gather.
 
     Returns:
-        The same data structure as :obj:`object` with all the objects sent to every device.
+        The same data structure as `object` with all the objects sent to every device.
     """
     if AcceleratorState().distributed_type == DistributedType.TPU:
         raise NotImplementedError("gather objects in TPU is not supported")
@@ -413,13 +412,13 @@ def broadcast(tensor, from_process: int = 0):
     Recursively broadcast tensor in a nested list/tuple/dictionary of tensors to all devices.
 
     Args:
-        tensor (nested list/tuple/dictionary of :obj:`torch.Tensor`):
+        tensor (nested list/tuple/dictionary of `torch.Tensor`):
             The data to gather.
-        from_process (:obj:`int`, `optional`, defaults to 0):
+        from_process (`int`, *optional*, defaults to 0):
             The process from which to send the data
 
     Returns:
-        The same data structure as :obj:`tensor` with all tensors broadcasted to the proper device.
+        The same data structure as `tensor` with all tensors broadcasted to the proper device.
     """
     if AcceleratorState().distributed_type == DistributedType.TPU:
         return _tpu_broadcast(tensor, src=from_process, name="accelerate.utils.broadcast")
@@ -438,7 +437,7 @@ def broadcast_object_list(object_list, from_process: int = 0):
     Args:
         object_list (list of picklable objects):
             The list of objects to broadcast. This list will be modified inplace.
-        from_process (:obj:`int`, `optional`, defaults to 0):
+        from_process (`int`, *optional*, defaults to 0):
             The process from which to send the data.
 
     Returns:
@@ -459,13 +458,13 @@ def slice_tensors(data, tensor_slice):
     Recursively takes a slice in a nested list/tuple/dictionary of tensors.
 
     Args:
-        data (nested list/tuple/dictionary of :obj:`torch.Tensor`):
+        data (nested list/tuple/dictionary of `torch.Tensor`):
             The data to slice.
-        tensor_slice (:obj:`slice`):
+        tensor_slice (`slice`):
             The slice to take.
 
     Returns:
-        The same data structure as :obj:`data` with all the tensors slices.
+        The same data structure as `data` with all the tensors slices.
     """
 
     def _slice_tensor(tensor, tensor_slice):
@@ -479,10 +478,10 @@ def find_batch_size(data):
     Recursively finds the batch size in a nested list/tuple/dictionary of lists of tensors.
 
     Args:
-        data (nested list/tuple/dictionary of :obj:`torch.Tensor`): The data from which to find the batch size.
+        data (nested list/tuple/dictionary of `torch.Tensor`): The data from which to find the batch size.
 
     Returns:
-        :obj:`int`: The batch size.
+        `int`: The batch size.
     """
     if isinstance(data, (tuple, list)):
         return find_batch_size(data[0])
@@ -499,13 +498,13 @@ def concatenate(data, dim=0):
     Recursively concatenate the tensors in a nested list/tuple/dictionary of lists of tensors with the same shape.
 
     Args:
-        data (nested list/tuple/dictionary of lists of tensors :obj:`torch.Tensor`):
+        data (nested list/tuple/dictionary of lists of tensors `torch.Tensor`):
             The data to concatenate.
-        dim (:obj:`int`, `optional`, defaults to 0):
+        dim (`int`, *optional*, defaults to 0):
             The dimension on which to concatenate.
 
     Returns:
-        The same data structure as :obj:`data` with all the tensors concatenated.
+        The same data structure as `data` with all the tensors concatenated.
     """
     if isinstance(data[0], (tuple, list)):
         return honor_type(data[0], (concatenate([d[i] for d in data], dim=dim) for i in range(len(data[0]))))
@@ -522,13 +521,13 @@ def pad_across_processes(tensor, dim=0, pad_index=0, pad_first=False):
     can safely be gathered.
 
     Args:
-        tensor (nested list/tuple/dictionary of :obj:`torch.Tensor`):
+        tensor (nested list/tuple/dictionary of `torch.Tensor`):
             The data to gather.
-        dim (:obj:`int`, `optional`, defaults to 0):
+        dim (`int`, *optional*, defaults to 0):
             The dimension on which to pad.
-        pad_index (:obj:`int`, `optional`, defaults to 0):
+        pad_index (`int`, *optional*, defaults to 0):
             The value with which to pad.
-        pad_first (:obj:`bool`, `optional`, defaults to :obj:`False`):
+        pad_first (`bool`, *optional*, defaults to `False`):
             Whether to pad at the beginning or the end.
     """
 
@@ -566,9 +565,11 @@ def wait_for_everyone():
     """
     Introduces a blocking point in the script, making sure all processes have reached this point before continuing.
 
-    Warning::
+    <Tip warning={true}>
 
-        Make sure all processes will reach this instruction otherwise one of your processes will hang forever.
+    Make sure all processes will reach this instruction otherwise one of your processes will hang forever.
+
+    </Tip>
     """
     if (
         AcceleratorState().distributed_type == DistributedType.MULTI_GPU
@@ -582,7 +583,7 @@ def wait_for_everyone():
 
 def save(obj, f):
     """
-    Save the data to disk. Use in place of :obj:`torch.save()`.
+    Save the data to disk. Use in place of `torch.save()`.
 
     Args:
         obj: The data to save
@@ -599,11 +600,11 @@ class PrepareForLaunch:
     Prepare a function that will launched in a distributed setup.
 
     Args:
-        launcher (:obj:`Callable`):
+        launcher (`Callable`):
             The function to launch.
-        distributed_type (:class:`~accelerate.state.DistributedType`):
+        distributed_type ([`~state.DistributedType`]):
             The distributed type to prepare for.
-        debug (:obj:`bool`, *optional*, defaults to :obj:`False`):
+        debug (`bool`, *optional*, defaults to `False`):
             Whether or not this is a debug launch.
     """
 
@@ -680,9 +681,9 @@ def __post_init__(self):
 @contextmanager
 def patch_environment(**kwargs):
     """
-    A context manager that will add each keyword argument passed to ``os.environ`` and remove them when exiting.
+    A context manager that will add each keyword argument passed to `os.environ` and remove them when exiting.
 
-    Will convert the values in :obj:`kwargs` to strings and upper-case all the keys.
+    Will convert the values in `kwargs` to strings and upper-case all the keys.
     """
     for key, value in kwargs.items():
         os.environ[key.upper()] = str(value)
@@ -695,7 +696,7 @@ def patch_environment(**kwargs):
 
 def get_pretty_name(obj):
     """
-    Gets a pretty name from ``obj``
+    Gets a pretty name from `obj`.
     """
     if not hasattr(obj, "__qualname__") and not hasattr(obj, "__name__"):
         obj = getattr(obj, "__class__", obj)
diff --git a/utils/style_doc.py b/utils/style_doc.py
index 0e363eaab..2a325805f 100644
--- a/utils/style_doc.py
+++ b/utils/style_doc.py
@@ -1,6 +1,5 @@
 # coding=utf-8
-
-# Copyright 2021 The HuggingFace Inc. team.
+# Copyright 2020 The HuggingFace Inc. team.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -13,78 +12,209 @@
 # 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.
-
 """Style utils for the .rst and the docstrings."""
 
 import argparse
 import os
 import re
 import warnings
-from enum import Enum
-
-
-# Special blocks where the inside should be formatted.
-TEXTUAL_BLOCKS = ["note", "warning"]
-# List of acceptable characters for titles and sections underline.
-TITLE_SPECIAL_CHARS = """= - ` : ' " ~ ^ _ * + # < >""".split(" ")
-# Special words for docstrings (s? means the s is optional)
-DOC_SPECIAL_WORD = [
-    "Args?",
-    "Params?",
-    "Parameters?",
-    "Arguments?",
-    "Examples?",
-    "Usage",
-    "Returns?",
-    "Raises?",
-    "Attributes?",
-]
+
+import black
+
+
+BLACK_AVOID_PATTERNS = {}
+
 
 # Regexes
-# Matches any declaration of textual block, like `.. note::`. (ignore case to avoid writing all versions in the list)
-_re_textual_blocks = re.compile(r"^\s*\.\.\s+(" + "|".join(TEXTUAL_BLOCKS) + r")\s*::\s*$", re.IGNORECASE)
-# Matches list introduction in rst.
+# Re pattern that catches list introduction (with potential indent)
 _re_list = re.compile(r"^(\s*-\s+|\s*\*\s+|\s*\d+\.\s+)")
-# Matches the indent in a line.
-_re_indent = re.compile(r"^(\s*)\S")
-# Matches a table declaration in rst.
-_re_table = re.compile(r"(\+-+)+\+\s*$")
-# Matches a code block in rst `:: `.
-_re_code_block = re.compile(r"^\s*::\s*$")
-# Matches any block of the form `.. something::` or `.. something:: bla`.
-_re_ignore = re.compile(r"^\s*\.\.\s+(.*?)\s*::\s*\S*\s*$")
-# Matches comment introduction in rst.
-_re_comment = re.compile(r"\s*\.\.\s*$")
+# Re pattern that catches code block introduction (with potentinal indent)
+_re_code = re.compile(r"^(\s*)```(.*)$")
+# Re pattern that catches rst args blocks of the form `Parameters:`.
+_re_args = re.compile("^\s*(Args?|Arguments?|Params?|Parameters?):\s*$")
+# Re pattern that catches return blocks of the form `Return:`.
+_re_returns = re.compile("^\s*Returns?:\s*$")
 # Matches the special tag to ignore some paragraphs.
 _re_doc_ignore = re.compile(r"(\.\.|#)\s*docstyle-ignore")
-# Matches the example introduction in docstrings.
-_re_example = re.compile(r"::\s*$")
-# Matches the parameters introduction in docstrings.
-_re_arg_def = re.compile(r"^\s*(Args?|Parameters?|Params|Arguments?|Environment|Attributes?)\s*:\s*$")
-# Matches the return introduction in docstrings.
-_re_return = re.compile(r"^\s*(Returns?|Raises?|Note)\s*:\s*$")
-# Matches any doc special word.
-_re_any_doc_special_word = re.compile(r"^\s*(" + "|".join(DOC_SPECIAL_WORD) + r")::?\s*$")
+# Re pattern that matches <Tip>, </Tip> and <Tip warning={true}> blocks.
+_re_tip = re.compile("^\s*</?Tip(>|\s+warning={true}>)\s*$")
+
+DOCTEST_PROMPTS = [">>>", "..."]
+
+
+def is_empty_line(line):
+    return len(line) == 0 or line.isspace()
+
+
+def find_indent(line):
+    """
+    Returns the number of spaces that start a line indent.
+    """
+    search = re.search("^(\s*)(?:\S|$)", line)
+    if search is None:
+        return 0
+    return len(search.groups()[0])
+
+
+def parse_code_example(code_lines):
+    """
+    Parses a code example
+
+    Args:
+        code_lines (`List[str]`): The code lines to parse.
+        max_len (`int`): The maximum lengh per line.
+
+    Returns:
+        (List[`str`], List[`str`]): The list of code samples and the list of outputs.
+    """
+    has_doctest = code_lines[0][:3] in DOCTEST_PROMPTS
+
+    code_samples = []
+    outputs = []
+    in_code = True
+    current_bit = []
+
+    for line in code_lines:
+        if in_code and has_doctest and not is_empty_line(line) and line[:3] not in DOCTEST_PROMPTS:
+            code_sample = "\n".join(current_bit)
+            code_samples.append(code_sample.strip())
+            in_code = False
+            current_bit = []
+        elif not in_code and line[:3] in DOCTEST_PROMPTS:
+            output = "\n".join(current_bit)
+            outputs.append(output.strip())
+            in_code = True
+            current_bit = []
+
+        # Add the line without doctest prompt
+        if line[:3] in DOCTEST_PROMPTS:
+            line = line[4:]
+        current_bit.append(line)
+
+    # Add last sample
+    if in_code:
+        code_sample = "\n".join(current_bit)
+        code_samples.append(code_sample.strip())
+    else:
+        output = "\n".join(current_bit)
+        outputs.append(output.strip())
+
+    return code_samples, outputs
 
 
-class SpecialBlock(Enum):
-    NOT_SPECIAL = 0
-    NO_STYLE = 1
-    ARG_LIST = 2
+def format_code_example(code: str, max_len: int, in_docstring: bool = False):
+    """
+    Format a code example using black. Will take into account the doctest syntax as well as any initial indentation in
+    the code provided.
 
+    Args:
+        code (`str`): The code example to format.
+        max_len (`int`): The maximum lengh per line.
+        in_docstring (`bool`, *optional*, defaults to `False`): Whether or not the code example is inside a docstring.
 
-def split_text_in_lines(text, max_len, prefix="", min_indent=None):
+    Returns:
+        `str`: The formatted code.
+    """
+    code_lines = code.split("\n")
+
+    # Find initial indent
+    idx = 0
+    while idx < len(code_lines) and is_empty_line(code_lines[idx]):
+        idx += 1
+    if idx >= len(code_lines):
+        return "", ""
+    indent = find_indent(code_lines[idx])
+
+    # Remove the initial indent for now, we will had it back after styling.
+    # Note that l[indent:] works for empty lines
+    code_lines = [l[indent:] for l in code_lines[idx:]]
+    has_doctest = code_lines[0][:3] in DOCTEST_PROMPTS
+
+    code_samples, outputs = parse_code_example(code_lines)
+
+    # Let's blackify the code! We put everything in one big text to go faster.
+    delimiter = "\n\n### New code sample ###\n"
+    full_code = delimiter.join(code_samples)
+    line_length = max_len - indent
+    if has_doctest:
+        line_length -= 4
+
+    for k, v in BLACK_AVOID_PATTERNS.items():
+        full_code = full_code.replace(k, v)
+    try:
+        mode = black.Mode(target_versions={black.TargetVersion.PY37}, line_length=line_length)
+        formatted_code = black.format_str(full_code, mode=mode)
+        error = ""
+    except Exception as e:
+        formatted_code = full_code
+        error = f"Code sample:\n{full_code}\n\nError message:\n{e}"
+
+    # Let's get back the formatted code samples
+    for k, v in BLACK_AVOID_PATTERNS.items():
+        formatted_code = formatted_code.replace(v, k)
+    # Triple quotes will mess docstrings.
+    if in_docstring:
+        formatted_code = formatted_code.replace('"""', "'''")
+
+    code_samples = formatted_code.split(delimiter)
+    # We can have one output less than code samples
+    if len(outputs) == len(code_samples) - 1:
+        outputs.append("")
+
+    formatted_lines = []
+    for code_sample, output in zip(code_samples, outputs):
+        # black may have added some new lines, we remove them
+        code_sample = code_sample.strip()
+        in_triple_quotes = False
+        in_decorator = False
+        for line in code_sample.strip().split("\n"):
+            if has_doctest and not is_empty_line(line):
+                prefix = (
+                    "... "
+                    if line.startswith(" ") or line in [")", "]", "}"] or in_triple_quotes or in_decorator
+                    else ">>> "
+                )
+            else:
+                prefix = ""
+            indent_str = "" if is_empty_line(line) else (" " * indent)
+            formatted_lines.append(indent_str + prefix + line)
+
+            if '"""' in line:
+                in_triple_quotes = not in_triple_quotes
+            if line.startswith(" "):
+                in_decorator = False
+            if line.startswith("@"):
+                in_decorator = True
+
+        formatted_lines.extend([" " * indent + line for line in output.split("\n")])
+        if not output.endswith("===PT-TF-SPLIT==="):
+            formatted_lines.append("")
+
+    result = "\n".join(formatted_lines)
+    return result.rstrip(), error
+
+
+def format_text(text, max_len, prefix="", min_indent=None):
     """
-    Split `text` in the biggest lines possible with the constraint of `max_len` using `prefix` on the first line and
-    then indenting with the same length as `prefix`.
+    Format a text in the biggest lines possible with the constraint of a maximum length and an indentation.
+
+    Args:
+        text (`str`): The text to format
+        max_len (`int`): The maximum length per line to use
+        prefix (`str`, *optional*, defaults to `""`): A prefix that will be added to the text.
+            The prefix doesn't count toward the indent (like a - introducing a list).
+        min_indent (`int`, *optional*): The minimum indent of the text.
+            If not set, will default to the length of the `prefix`.
+
+    Returns:
+        `str`: The formatted text.
     """
     text = re.sub(r"\s+", " ", text)
-    indent = " " * len(prefix)
     if min_indent is not None:
-        if len(indent) < len(min_indent):
-            indent = min_indent
-        if len(prefix) < len(min_indent):
-            prefix = " " * (len(min_indent) - len(prefix)) + prefix
+        if len(prefix) < min_indent:
+            prefix = " " * (min_indent - len(prefix)) + prefix
+
+    indent = " " * len(prefix)
     new_lines = []
     words = text.split(" ")
     current_line = f"{prefix}{words[0]}"
@@ -99,379 +229,189 @@ def split_text_in_lines(text, max_len, prefix="", min_indent=None):
     return "\n".join(new_lines)
 
 
-def get_indent(line):
-    """Get the indentation of `line`."""
-    indent_search = _re_indent.search(line)
-    return indent_search.groups()[0] if indent_search is not None else ""
-
-
-class CodeStyler:
-    """A generic class to style .rst files."""
-
-    def is_no_style_block(self, line):
-        """Whether or not `line` introduces a block where styling should be ignore"""
-        if _re_code_block.search(line) is not None:
-            return True
-        if _re_textual_blocks.search(line) is not None:
-            return False
-        return _re_ignore.search(line) is not None
-
-    def is_comment_or_textual_block(self, line):
-        """Whether or not `line` introduces a block where styling should not be ignored (note, warnings...)"""
-        if _re_comment.search(line):
-            return True
-        return _re_textual_blocks.search(line) is not None
-
-    def is_special_block(self, line):
-        """Whether or not `line` introduces a special block."""
-        if self.is_no_style_block(line):
-            self.in_block = SpecialBlock.NO_STYLE
-            return True
-        return False
-
-    def init_in_block(self, text):
-        """
-        Returns the initial value for `self.in_block`.
-
-        Useful for some docstrings beginning inside an argument declaration block (all models).
-        """
-        return SpecialBlock.NOT_SPECIAL
-
-    def end_of_special_style(self, line):
-        """
-        Sets back the `in_block` attribute to `NOT_SPECIAL`.
-
-        Useful for some docstrings where we may have to go back to `ARG_LIST` instead.
-        """
-        self.in_block = SpecialBlock.NOT_SPECIAL
-
-    def style_paragraph(self, paragraph, max_len, no_style=False, min_indent=None):
-        """
-        Style `paragraph` (a list of lines) by making sure no line goes over `max_len`, except if the `no_style` flag
-        is passed.
-        """
-        if len(paragraph) == 0:
-            return ""
-        if no_style or self.in_block == SpecialBlock.NO_STYLE:
-            return "\n".join(paragraph)
-        if _re_list.search(paragraph[0]) is not None:
-            # Great, we're in a list. So we need to split our paragraphs in smaller parts, one for each item.
-            result = ""
-            remainder = ""
-            prefix = _re_list.search(paragraph[0]).groups()[0]
-            prefix_indent = get_indent(paragraph[0])
-            current_item = [paragraph[0][len(prefix) :]]
-            for i, line in enumerate(paragraph[1:]):
-                new_item_search = _re_list.search(line)
-                indent = get_indent(line)
-                if len(indent) < len(prefix_indent) or (len(indent) == len(prefix_indent) and new_item_search is None):
-                    # There might not be an empty line after the list, formatting the remainder recursively.
-                    remainder = "\n" + self.style_paragraph(
-                        paragraph[i + 1 :], max_len, no_style=no_style, min_indent=min_indent
-                    )
-                    break
-                elif new_item_search is not None:
-                    text = " ".join([l.strip() for l in current_item])
-                    result += split_text_in_lines(text, max_len, prefix, min_indent=min_indent) + "\n"
-                    prefix = new_item_search.groups()[0]
-                    prefix_indent = indent
-                    current_item = [line[len(prefix) :]]
-                else:
-                    current_item.append(line)
-            # Treat the last item
-            text = " ".join([l.strip() for l in current_item])
-            result += split_text_in_lines(text, max_len, prefix, min_indent=min_indent)
-            # Add the potential remainder
-            return result + remainder
-
-        if len(paragraph) > 1 and self.is_comment_or_textual_block(paragraph[0]):
-            # Comments/notes in rst should be restyled with indentation, ignoring the first line.
-            indent = get_indent(paragraph[1])
-            text = " ".join([l.strip() for l in paragraph[1:]])
-            return paragraph[0] + "\n" + split_text_in_lines(text, max_len, indent, min_indent=min_indent)
-
-        if self.in_block == SpecialBlock.ARG_LIST:
-            # Arg lists are special: we need to ignore the lines that are at the first indentation level beneath the
-            # Args/Parameters (parameter description), then we can style the indentation level beneath.
-            result = ""
-            # The args/parameters could be in that paragraph and should be ignored
-            if _re_arg_def.search(paragraph[0]) is not None:
-                if len(paragraph) == 1:
-                    return paragraph[0]
-                result += paragraph[0] + "\n"
-                paragraph = paragraph[1:]
-
-            if self.current_indent is None:
-                self.current_indent = get_indent(paragraph[1])
-
-            current_item = []
-            for line in paragraph:
-                if get_indent(line) == self.current_indent:
-                    if len(current_item) > 0:
-                        item_indent = get_indent(current_item[0])
-                        text = " ".join([l.strip() for l in current_item])
-                        result += split_text_in_lines(text, max_len, item_indent, min_indent=min_indent) + "\n"
-                    result += line + "\n"
-                    current_item = []
-                else:
-                    current_item.append(line)
-            if len(current_item) > 0:
-                item_indent = get_indent(current_item[0])
-                text = " ".join([l.strip() for l in current_item])
-                result += split_text_in_lines(text, max_len, item_indent, min_indent=min_indent) + "\n"
-            return result[:-1]
-
-        indent = get_indent(paragraph[0])
-        text = " ".join([l.strip() for l in paragraph])
-        return split_text_in_lines(text, max_len, indent, min_indent=min_indent)
-
-    def style(self, text, max_len=119, min_indent=None):
-        """Style `text` to `max_len`."""
-        new_lines = []
-        paragraph = []
-        self.current_indent = ""
-        self.previous_indent = None
-        # If one of those is True, the paragraph should not be touched (code samples, lists...)
-        no_style = False
-        no_style_next = False
-        self.in_block = self.init_in_block(text)
-        # If this is True, we force-break a paragraph, even if there is no new empty line.
-        break_paragraph = False
-
-        lines = text.split("\n")
-        last_line = None
-        for line in lines:
-            # New paragraph
-            line_is_empty = len(line.strip()) == 0
-            list_begins = (
-                _re_list.search(line) is not None
-                and last_line is not None
-                and len(get_indent(line)) > len(get_indent(last_line))
-            )
-            if line_is_empty or break_paragraph or list_begins:
-                if len(paragraph) > 0:
-                    if self.in_block != SpecialBlock.NOT_SPECIAL:
-                        indent = get_indent(paragraph[0])
-                        # Are we still in a no-style block?
-                        if self.current_indent is None:
-                            # If current_indent is None, we haven't begun the interior of the block so the answer is
-                            # yes, unless we have an indent of 0 in which case the special block took one line only.
-                            if len(indent) == 0:
-                                self.in_block = SpecialBlock.NOT_SPECIAL
-                            else:
-                                self.current_indent = indent
-                        elif not indent.startswith(self.current_indent):
-                            # If not, we are leaving the block when we unindent.
-                            self.end_of_special_style(paragraph[0])
-
-                    if self.is_special_block(paragraph[0]):
-                        # Maybe we are starting a special block.
-                        if len(paragraph) > 1:
-                            # If we have the interior of the block in the paragraph, we grab the indent.
-                            self.current_indent = get_indent(paragraph[1])
-                        else:
-                            # We will determine the indent with the next paragraph
-                            self.current_indent = None
-                    styled_paragraph = self.style_paragraph(
-                        paragraph, max_len, no_style=no_style, min_indent=min_indent
-                    )
-                    new_lines.append(styled_paragraph + "\n")
-                else:
-                    new_lines.append("")
-
-                paragraph = []
-                no_style = no_style_next
-                no_style_next = False
-                last_line = None
-                if (not break_paragraph and not list_begins) or line_is_empty:
-                    break_paragraph = False
-                    continue
-                break_paragraph = False
-
-            # Title and section lines should go to the max + add a new paragraph.
-            if (
-                len(set(line)) == 1
-                and line[0] in TITLE_SPECIAL_CHARS
-                and last_line is not None
-                and len(line) >= len(last_line)
-            ):
-                line = line[0] * max_len
-                break_paragraph = True
-            # proper doc comment indicates the next paragraph should be no-style.
-            if _re_doc_ignore.search(line) is not None:
-                no_style_next = True
-            # Table are in just one paragraph and should be no-style.
-            if _re_table.search(line) is not None:
-                no_style = True
-            paragraph.append(line)
-            last_line = line
-
-        # Just have to treat the last paragraph. It could still be in a no-style block (or not)
-        if len(paragraph) > 0:
-            # Are we still in a special block
-            # (if current_indent is None, we are but no need to set it since we are the end.)
-            if self.in_block != SpecialBlock.NO_STYLE and self.current_indent is not None:
-                indent = get_indent(paragraph[0])
-                if not indent.startswith(self.current_indent):
-                    self.in_block = SpecialBlock.NOT_SPECIAL
-            _ = self.is_special_block(paragraph[0])
-            new_lines.append(self.style_paragraph(paragraph, max_len, no_style=no_style, min_indent=min_indent) + "\n")
-        return "\n".join(new_lines)
-
-
-class DocstringStyler(CodeStyler):
-    """Class to style docstrings that take the main method from `CodeStyler`."""
-
-    def is_no_style_block(self, line):
-        if _re_textual_blocks.search(line) is not None:
-            return False
-        if _re_example.search(line) is not None:
-            return True
-        return _re_code_block.search(line) is not None
-
-    def is_comment_or_textual_block(self, line):
-        if _re_return.search(line) is not None:
-            self.in_block = SpecialBlock.NOT_SPECIAL
-            return True
-        return super().is_comment_or_textual_block(line)
-
-    def is_special_block(self, line):
-        if self.is_no_style_block(line):
-            if self.previous_indent is None and self.in_block == SpecialBlock.ARG_LIST:
-                self.previous_indent = self.current_indent
-            self.in_block = SpecialBlock.NO_STYLE
-            return True
-        if _re_arg_def.search(line) is not None:
-            self.in_block = SpecialBlock.ARG_LIST
-            return True
-        return False
-
-    def end_of_special_style(self, line):
-        if self.previous_indent is not None and line.startswith(self.previous_indent):
-            self.in_block = SpecialBlock.ARG_LIST
-            self.current_indent = self.previous_indent
-        else:
-            self.in_block = SpecialBlock.NOT_SPECIAL
-            self.previous_indent = None
-
-    def init_in_block(self, text):
-        lines = text.split("\n")
-        while len(lines) > 0 and len(lines[0]) == 0:
-            lines = lines[1:]
-        if len(lines) == 0:
-            return SpecialBlock.NOT_SPECIAL
-        if re.search(r":\s*$", lines[0]):
-            indent = get_indent(lines[0])
-            if (
-                len(lines) == 1
-                or len(get_indent(lines[1])) > len(indent)
-                or (len(get_indent(lines[1])) == len(indent) and re.search(r":\s*$", lines[1]))
-            ):
-                self.current_indent = indent
-                return SpecialBlock.ARG_LIST
-        return SpecialBlock.NOT_SPECIAL
-
-
-rst_styler = CodeStyler()
-doc_styler = DocstringStyler()
-
-
-def _add_new_lines_before_list(text):
-    """Add a new empty line before a list begins."""
-    lines = text.split("\n")
-    new_lines = []
-    in_list = False
-    for idx, line in enumerate(lines):
-        # Detect if the line is the start of a new list.
-        if _re_list.search(line) is not None and not in_list:
-            current_indent = get_indent(line)
-            in_list = True
-            # If the line before is non empty, add an extra new line.
-            if idx > 0 and len(lines[idx - 1]) != 0:
-                new_lines.append("")
-        # Detect if we're out of the current list.
-        if in_list and not line.startswith(current_indent) and _re_list.search(line) is None:
-            in_list = False
-        new_lines.append(line)
-    return "\n".join(new_lines)
+def split_line_on_first_colon(line):
+    splits = line.split(":")
+    return splits[0], ":".join(splits[1:])
 
 
-def _add_new_lines_before_doc_special_words(text):
-    lines = text.split("\n")
+def style_docstring(docstring, max_len):
+    """
+    Style a docstring by making sure there is no useless whitespace and the maximum horizontal space is used.
+
+    Args:
+        docstring (`str`): The docstring to style.
+        max_len (`int`): The maximum length of each line.
+
+    Returns:
+        `str`: The styled docstring
+    """
+    lines = docstring.split("\n")
     new_lines = []
+
+    # Initialization
+    current_paragraph = None
+    current_indent = -1
+    in_code = False
+    param_indent = -1
+    prefix = ""
+    black_errors = []
+
+    # Special case for docstrings that begin with continuation of Args with no Args block.
+    idx = 0
+    while idx < len(lines) and is_empty_line(lines[idx]):
+        idx += 1
+    if (
+        len(lines[idx]) > 1
+        and lines[idx].rstrip().endswith(":")
+        and find_indent(lines[idx + 1]) > find_indent(lines[idx])
+    ):
+        param_indent = find_indent(lines[idx])
+
     for idx, line in enumerate(lines):
-        # Detect if the line is the start of a new list.
-        if _re_any_doc_special_word.search(line) is not None:
-            # If the line before is non empty, add an extra new line.
-            if idx > 0 and len(lines[idx - 1]) != 0:
+        # Doing all re searches once for the one we need to repeat.
+        list_search = _re_list.search(line)
+        code_search = _re_code.search(line)
+
+        # Are we starting a new paragraph?
+        # New indentation or new line:
+        new_paragraph = find_indent(line) != current_indent or is_empty_line(line)
+        # List item
+        new_paragraph = new_paragraph or list_search is not None
+        # Code block beginning
+        new_paragraph = new_paragraph or code_search is not None
+        # Beginning/end of tip
+        new_paragraph = new_paragraph or _re_tip.search(line)
+
+        # In this case, we treat the current paragraph
+        if not in_code and new_paragraph and current_paragraph is not None and len(current_paragraph) > 0:
+            paragraph = " ".join(current_paragraph)
+            new_lines.append(format_text(paragraph, max_len, prefix=prefix, min_indent=current_indent))
+            current_paragraph = None
+
+        if code_search is not None:
+            if not in_code:
+                current_paragraph = []
+                current_indent = len(code_search.groups()[0])
+                current_code = code_search.groups()[1]
+                prefix = ""
+                if current_indent < param_indent:
+                    param_indent = -1
+            else:
+                current_indent = -1
+                code = "\n".join(current_paragraph)
+                if current_code in ["py", "python"]:
+                    formatted_code, error = format_code_example(code, max_len, in_docstring=True)
+                    new_lines.append(formatted_code)
+                    if len(error) > 0:
+                        black_errors.append(error)
+                else:
+                    new_lines.append(code)
+                current_paragraph = None
+            new_lines.append(line)
+            in_code = not in_code
+
+        elif in_code:
+            current_paragraph.append(line)
+        elif is_empty_line(line):
+            current_paragraph = None
+            current_indent = -1
+            prefix = ""
+            new_lines.append(line)
+        elif list_search is not None:
+            prefix = list_search.groups()[0]
+            current_indent = len(prefix)
+            current_paragraph = [line[current_indent:]]
+        elif _re_args.search(line):
+            new_lines.append(line)
+            param_indent = find_indent(lines[idx + 1])
+        elif _re_tip.search(line):
+            # Add a new line before if not present
+            if not is_empty_line(new_lines[-1]):
                 new_lines.append("")
-        new_lines.append(line)
-    return "\n".join(new_lines)
-
+            new_lines.append(line)
+            # Add a new line after if not present
+            if idx < len(lines) - 1 and not is_empty_line(lines[idx + 1]):
+                new_lines.append("")
+        elif current_paragraph is None or find_indent(line) != current_indent:
+            indent = find_indent(line)
+            # Special behavior for parameters intros.
+            if indent == param_indent:
+                # Special rules for some docstring where the Returns blocks has the same indent as the parameters.
+                if _re_returns.search(line) is not None:
+                    param_indent = -1
+                    new_lines.append(line)
+                elif len(line) < max_len:
+                    new_lines.append(line)
+                else:
+                    intro, description = split_line_on_first_colon(line)
+                    new_lines.append(intro + ":")
+                    if len(description) != 0:
+                        if find_indent(lines[idx + 1]) > indent:
+                            current_indent = find_indent(lines[idx + 1])
+                        else:
+                            current_indent = indent + 4
+                        current_paragraph = [description.strip()]
+                        prefix = ""
+            else:
+                # Check if we have exited the parameter block
+                if indent < param_indent:
+                    param_indent = -1
 
-def style_rst_file(doc_file, max_len=119, check_only=False):
-    """ Style one rst file `doc_file` to `max_len`."""
-    with open(doc_file, "r", encoding="utf-8", newline="\n") as f:
-        doc = f.read()
+                current_paragraph = [line.strip()]
+                current_indent = find_indent(line)
+                prefix = ""
+        elif current_paragraph is not None:
+            current_paragraph.append(line.lstrip())
 
-    # Add missing new lines before lists
-    clean_doc = _add_new_lines_before_list(doc)
-    # Style
-    clean_doc = rst_styler.style(clean_doc, max_len=max_len)
+    if current_paragraph is not None and len(current_paragraph) > 0:
+        paragraph = " ".join(current_paragraph)
+        new_lines.append(format_text(paragraph, max_len, prefix=prefix, min_indent=current_indent))
 
-    diff = clean_doc != doc
-    if not check_only and diff:
-        print(f"Overwriting content of {doc_file}.")
-        with open(doc_file, "w", encoding="utf-8", newline="\n") as f:
-            f.write(clean_doc)
-
-    return diff
-
-
-def style_docstring(docstring, max_len=119):
-    """Style `docstring` to `max_len`."""
-    # One-line docstring that are not too long are left as is.
-    if len(docstring) < max_len and "\n" not in docstring:
-        return docstring
-
-    # Grab the indent from the last line
-    last_line = docstring.split("\n")[-1]
-    # Is it empty except for the last triple-quotes (not-included in `docstring`)?
-    indent_search = re.search(r"^(\s*)$", last_line)
-    if indent_search is not None:
-        indent = indent_search.groups()[0]
-        if len(indent) > 0:
-            docstring = docstring[: -len(indent)]
-    # Or are the triple quotes next to text (we will fix that).
-    else:
-        indent_search = _re_indent.search(last_line)
-        indent = indent_search.groups()[0] if indent_search is not None else ""
+    return "\n".join(new_lines), "\n\n".join(black_errors)
 
-    # Add missing new lines before Args/Returns etc.
-    docstring = _add_new_lines_before_doc_special_words(docstring)
-    # Add missing new lines before lists
-    docstring = _add_new_lines_before_list(docstring)
-    # Style
-    styled_doc = doc_styler.style(docstring, max_len=max_len, min_indent=indent)
 
-    # Add new lines if necessary
-    if not styled_doc.startswith("\n"):
-        styled_doc = "\n" + styled_doc
-    if not styled_doc.endswith("\n"):
-        styled_doc += "\n"
-    return styled_doc + indent
+def style_docstrings_in_code(code, max_len=119):
+    """
+    Style all docstrings in some code.
 
+    Args:
+        code (`str`): The code in which we want to style the docstrings.
+        max_len (`int`): The maximum number of characters per line.
 
-def style_file_docstrings(code_file, max_len=119, check_only=False):
-    """Style all docstrings in `code_file` to `max_len`."""
-    with open(code_file, "r", encoding="utf-8", newline="\n") as f:
-        code = f.read()
-    splits = code.split('"""')
+    Returns:
+        `Tuple[str, str]`: A tuple with the clean code and the black errors (if any)
+    """
+    # fmt: off
+    splits = code.split('\"\"\"')
     splits = [
         (s if i % 2 == 0 or _re_doc_ignore.search(splits[i - 1]) is not None else style_docstring(s, max_len=max_len))
         for i, s in enumerate(splits)
     ]
-    clean_code = '"""'.join(splits)
+    black_errors = "\n\n".join([s[1] for s in splits if isinstance(s, tuple) and len(s[1]) > 0])
+    splits = [s[0] if isinstance(s, tuple) else s for s in splits]
+    clean_code = '\"\"\"'.join(splits)
+    # fmt: on
+
+    return clean_code, black_errors
+
+
+def style_file_docstrings(code_file, max_len=119, check_only=False):
+    """
+    Style all docstrings in a given file.
+
+    Args:
+        code_file (`str` or `os.PathLike`): The file in which we want to style the docstring.
+        max_len (`int`): The maximum number of characters per line.
+        check_only (`bool`, *optional*, defaults to `False`):
+            Whether to restyle file or just check if they should be restyled.
+
+    Returns:
+        `bool`: Whether or not the file was or should be restyled.
+    """
+    with open(code_file, "r", encoding="utf-8", newline="\n") as f:
+        code = f.read()
+
+    clean_code, black_errors = style_docstrings_in_code(code, max_len=max_len)
 
     diff = clean_code != code
     if not check_only and diff:
@@ -479,31 +419,122 @@ def style_file_docstrings(code_file, max_len=119, check_only=False):
         with open(code_file, "w", encoding="utf-8", newline="\n") as f:
             f.write(clean_code)
 
-    return diff
+    return diff, black_errors
+
+
+def style_mdx_file(mdx_file, max_len=119, check_only=False):
+    """
+    Style a MDX file by formatting all Python code samples.
+
+    Args:
+        mdx_file (`str` or `os.PathLike`): The file in which we want to style the examples.
+        max_len (`int`): The maximum number of characters per line.
+        check_only (`bool`, *optional*, defaults to `False`):
+            Whether to restyle file or just check if they should be restyled.
+
+    Returns:
+        `bool`: Whether or not the file was or should be restyled.
+    """
+    with open(mdx_file, "r", encoding="utf-8", newline="\n") as f:
+        content = f.read()
+
+    lines = content.split("\n")
+    current_code = []
+    current_language = ""
+    in_code = False
+    new_lines = []
+    black_errors = []
+
+    for line in lines:
+        if _re_code.search(line) is not None:
+            in_code = not in_code
+            if in_code:
+                current_language = _re_code.search(line).groups()[1]
+                current_code = []
+            else:
+                code = "\n".join(current_code)
+                if current_language in ["py", "python"]:
+                    code, error = format_code_example(code, max_len)
+                    if len(error) > 0:
+                        black_errors.append(error)
+                new_lines.append(code)
+
+            new_lines.append(line)
+        elif in_code:
+            current_code.append(line)
+        else:
+            new_lines.append(line)
+
+    if in_code:
+        raise ValueError(f"There was a problem when styling {mdx_file}. A code block is opened without being closed.")
+
+    clean_content = "\n".join(new_lines)
+    diff = clean_content != content
+    if not check_only and diff:
+        print(f"Overwriting content of {mdx_file}.")
+        with open(mdx_file, "w", encoding="utf-8", newline="\n") as f:
+            f.write(clean_content)
+
+    return diff, "\n\n".join(black_errors)
 
 
 def style_doc_files(*files, max_len=119, check_only=False):
     """
-    Style all `files` to `max_len` and fixes mistakes if not `check_only`, otherwise raises an error if styling should
-    be done.
+    Applies doc styling or checks everything is correct in a list of files.
+
+    Args:
+        files (several `str` or `os.PathLike`): The files to treat.
+        max_len (`int`): The maximum number of characters per line.
+        check_only (`bool`, *optional*, defaults to `False`):
+            Whether to restyle file or just check if they should be restyled.
+
+    Returns:
+        List[`str`]: The list of files changed or that should be restyled.
     """
     changed = []
+    black_errors = []
     for file in files:
         # Treat folders
         if os.path.isdir(file):
             files = [os.path.join(file, f) for f in os.listdir(file)]
-            files = [f for f in files if os.path.isdir(f) or f.endswith(".rst") or f.endswith(".py")]
+            files = [f for f in files if os.path.isdir(f) or f.endswith(".mdx") or f.endswith(".py")]
             changed += style_doc_files(*files, max_len=max_len, check_only=check_only)
-        # Treat rst
-        elif file.endswith(".rst"):
-            if style_rst_file(file, max_len=max_len, check_only=check_only):
-                changed.append(file)
+        # Treat mdx
+        elif file.endswith(".mdx"):
+            try:
+                diff, black_error = style_mdx_file(file, max_len=max_len, check_only=check_only)
+                if diff:
+                    changed.append(file)
+                if len(black_error) > 0:
+                    black_errors.append(
+                        f"There was a problem while formatting an example in {file} with black:\m{black_error}"
+                    )
+            except Exception:
+                print(f"There is a problem in {file}.")
+                raise
         # Treat python files
         elif file.endswith(".py"):
-            if style_file_docstrings(file, max_len=max_len, check_only=check_only):
-                changed.append(file)
+            try:
+                diff, black_error = style_file_docstrings(file, max_len=max_len, check_only=check_only)
+                if diff:
+                    changed.append(file)
+                if len(black_error) > 0:
+                    black_errors.append(
+                        f"There was a problem while formatting an example in {file} with black:\m{black_error}"
+                    )
+            except Exception:
+                print(f"There is a problem in {file}.")
+                raise
         else:
-            warnings.warn(f"Ignoring {file} because it's not a py or an rst file or a folder.")
+            warnings.warn(f"Ignoring {file} because it's not a py or an mdx file or a folder.")
+    if len(black_errors) > 0:
+        black_message = "\n\n".join(black_errors)
+        raise ValueError(
+            "Some code examples can't be interpreted by black, which means they aren't regular python:\n\n"
+            + black_message
+            + "\n\nMake sure to fix the corresponding docstring or doc file, or remove the py/python after ``` if it "
+            + "was not supposed to be a Python code sample."
+        )
     return changed
 
 

_The documentation is not available anymore as the PR was closed or merged._
Even if everything is pushed properly on `doc-build-dev`, there is nothing at https://moon-ci-docs.huggingface.co/docs/accelerate/pr_271. Guessing something needs to be activated on moon-landing @mishig25 ?
Indeed, there needs to be a small update on moon-landing side. More[ info here](https://huggingface.slack.com/archives/C02GLJ5S0E9/p1646780697260569?thread_ts=1646410559.029439&cid=C02GLJ5S0E9)
<img width="1000" alt="Screenshot 2022-03-09 at 00 02 02" src="https://user-images.githubusercontent.com/11827707/157340774-6481b81c-e5ea-460c-a395-632aad53783c.png">

I think `~` before `Accelerator` should not be there
<img width="400" alt="Screenshot 2022-03-09 at 00 06 33" src="https://user-images.githubusercontent.com/11827707/157340963-a88cd88c-ac76-41aa-b1e2-71c6417f8130.png">

https://api.github.com/repos/huggingface/accelerate/pulls/271
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 1254c19e9..0fd2a5701 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -22,7 +22,7 @@
 
 from packaging import version
 
-from .checkpointing import load_accelerator_state, save_accelerator_state
+from .checkpointing import load_accelerator_state, load_custom_state, save_accelerator_state, save_custom_state
 from .data_loader import prepare_data_loader
 from .kwargs_handlers import DistributedDataParallelKwargs, GradScalerKwargs, InitProcessGroupKwargs, KwargsHandler
 from .optimizer import AcceleratedOptimizer
@@ -33,6 +33,7 @@
     convert_outputs_to_fp32,
     extract_model_from_parallel,
     gather,
+    get_pretty_name,
     pad_across_processes,
     save,
     wait_for_everyone,
@@ -188,6 +189,7 @@ def __init__(
         # Internal references to the training objects
         self._optimizers = []
         self._models = []
+        self._custom_objects = []
 
         # RNG Types
         self.rng_types = rng_types
@@ -564,7 +566,7 @@ def save(self, obj, f):
 
     def save_state(self, output_dir: str):
         """
-        Saves the current states of the model, optimizer, scaler, and RNG generators.
+        Saves the current states of the model, optimizer, scaler, RNG generators, and registered objects.
 
         Args:
             output_dir (:obj:`str` or :obj:`os.PathLike`):
@@ -575,11 +577,16 @@ def save_state(self, output_dir: str):
         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)
+        save_location = save_accelerator_state(
+            output_dir, weights, self._optimizers, self.state.process_index, self.scaler
+        )
+        for i, obj in enumerate(self._custom_objects):
+            save_custom_state(obj, output_dir, i)
+        return save_location
 
     def load_state(self, input_dir: str):
         """
-        Loads the current states of the model, optimizer, scaler, and RNG generators.
+        Loads the current states of the model, optimizer, scaler, RNG generators, and registered objects.
 
         Args:
             input_dir (:obj:`str` or :obj:`os.PathLike`):
@@ -591,6 +598,16 @@ def load_state(self, input_dir: str):
             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)
+        custom_checkpoints = [f for f in os.listdir(input_dir) if "custom_checkpoint" in f]
+        if len(custom_checkpoints) != len(self._custom_objects):
+            err = "Warning! Number of found checkpoints does not match the number of registered objects:"
+            err += f"\n\tFound checkpoints: {len(custom_checkpoints)}"
+            err += f"\n\tRegistered objects: {len(self._custom_objects)}\nSkipping."
+            logger.warn(err)
+        else:
+            logger.info(f"Loading in {len(custom_checkpoints)} custom states")
+            for index, obj in enumerate(self._custom_objects):
+                load_custom_state(obj, input_dir, index)
 
     def free_memory(self):
         """
@@ -646,6 +663,26 @@ def get_state_dict(self, model):
 
         return state_dict
 
+    def register_for_checkpointing(self, *objects):
+        """
+        Makes note of `objects` and will save or load them in during `save_state` or `load_state`.
+
+        These should be utilized when the state is being loaded or saved in the same script. It is not designed to be
+        used in different scripts
+
+        Note: Every `object` must have a `load_state_dict` and `state_dict` function to be stored.
+        """
+        invalid_objects = []
+        for obj in objects:
+            if not hasattr(obj, "state_dict") or not hasattr(obj, "load_state_dict"):
+                invalid_objects.append(obj)
+        if len(invalid_objects) > 0:
+            err = "All `objects` must include a `state_dict` and `load_state_dict` function to be stored. The following inputs are invalid:"
+            for index, obj in enumerate(invalid_objects):
+                err += f"\n\t- Item at index {index}, `{get_pretty_name(obj)}`"
+            raise ValueError(err)
+        self._custom_objects.extend(objects)
+
     @contextmanager
     def autocast(self):
         """
diff --git a/src/accelerate/checkpointing.py b/src/accelerate/checkpointing.py
index 37d68b90c..e1a28d108 100644
--- a/src/accelerate/checkpointing.py
+++ b/src/accelerate/checkpointing.py
@@ -14,6 +14,7 @@
 
 import os
 import random
+from pathlib import Path
 from typing import List
 
 import numpy as np
@@ -21,7 +22,7 @@
 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
+from .utils import MODEL_NAME, OPTIMIZER_NAME, RNG_STATE_NAME, SCALER_NAME, get_pretty_name, save
 
 
 if is_tpu_available():
@@ -132,3 +133,22 @@ def load_accelerator_state(input_dir, models, optimizers, process_index, scaler=
     if is_tpu_available():
         xm.set_rng_state(states["xm_seed"])
     logger.info("All random states loaded successfully")
+
+
+def save_custom_state(obj, path, index: int = 0):
+    """
+    Saves the state of `obj` to `{path}/custom_checkpoint_{index}.pkl`
+    """
+    # Should this be the right way to get a qual_name type value from `obj`?
+    save_location = Path(path) / f"custom_checkpoint_{index}.pkl"
+    logger.info(f"Saving the state of {get_pretty_name(obj)} to {save_location}")
+    torch.save(obj.state_dict(), save_location)
+
+
+def load_custom_state(obj, path, index: int = 0):
+    """
+    Loads the state of `obj` at `{path}/custom_checkpoint_{index}.pkl`
+    """
+    load_location = f"{path}/custom_checkpoint_{index}.pkl"
+    logger.info(f"Loading the state of {get_pretty_name(obj)} from {load_location}")
+    obj.load_state_dict(torch.load(load_location))
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 62b1291ee..4b20cc1c3 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -691,3 +691,16 @@ def patch_environment(**kwargs):
 
     for key in kwargs:
         del os.environ[key.upper()]
+
+
+def get_pretty_name(obj):
+    """
+    Gets a pretty name from ``obj``
+    """
+    if not hasattr(obj, "__qualname__") and not hasattr(obj, "__name__"):
+        obj = getattr(obj, "__class__", obj)
+    if hasattr(obj, "__qualname__"):
+        return obj.__qualname__
+    if hasattr(obj, "__name__"):
+        return obj.__name__
+    return str(obj)
diff --git a/tests/test_state_checkpointing.py b/tests/test_state_checkpointing.py
index a74dcb724..87b2d3b06 100644
--- a/tests/test_state_checkpointing.py
+++ b/tests/test_state_checkpointing.py
@@ -43,13 +43,13 @@ def get_dataset(n_batches):
     return (train_dataloader, valid_dataloader)
 
 
-def train(num_epochs, model, dataloader, optimizer, accelerator):
+def train(num_epochs, model, dataloader, optimizer, accelerator, scheduler=None):
     "Trains for `num_epochs`"
     rands = []
     for epoch in range(num_epochs):
         # Train quickly
         model.train()
-        for step, batch in enumerate(dataloader):
+        for batch in dataloader:
             x, y = batch
             outputs = model(x)
             loss = torch.nn.functional.mse_loss(outputs, y)
@@ -57,6 +57,8 @@ def train(num_epochs, model, dataloader, optimizer, accelerator):
             optimizer.step()
             optimizer.zero_grad()
         rands.append(random.random())  # Introduce some randomness
+        if scheduler is not None:
+            scheduler.step()
     return rands
 
 
@@ -123,3 +125,41 @@ def test_can_resume_training(self):
             self.assertEqual(b1, b3)
             self.assertEqual(opt_state1, opt_state3)
             self.assertEqual(ground_truth_rands, test_rands)
+
+    def test_invalid_registration(self):
+        t = torch.tensor([1, 2, 3])
+        t1 = torch.tensor([2, 3, 4])
+        net = DummyModel()
+        opt = torch.optim.Adam(net.parameters())
+        accelerator = Accelerator()
+        with self.assertRaises(ValueError) as ve:
+            accelerator.register_for_checkpointing(t, t1, net, opt)
+        message = str(ve.exception)
+        self.assertTrue("Item at index 0" in message)
+        self.assertTrue("Item at index 1" in message)
+        self.assertFalse("Item at index 2" in message)
+        self.assertFalse("Item at index 3" in message)
+
+    def test_with_scheduler(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            set_seed(42)
+            model = DummyModel()
+            optimizer = torch.optim.Adam(params=model.parameters(), lr=1e-3)
+            scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.99)
+            train_dataloader, valid_dataloader = dummy_dataloaders()
+            # Train baseline
+            accelerator = Accelerator()
+            model, optimizer, train_dataloader, valid_dataloader = accelerator.prepare(
+                model, optimizer, train_dataloader, valid_dataloader
+            )
+            accelerator.register_for_checkpointing(scheduler)
+            # Save initial
+            initial = os.path.join(tmpdir, "initial")
+            accelerator.save_state(initial)
+            scheduler_state = scheduler.state_dict()
+            train(3, model, train_dataloader, optimizer, accelerator, scheduler)
+            self.assertNotEqual(scheduler_state, scheduler.state_dict())
+
+            # Load everything back in and make sure all states work
+            accelerator.load_state(initial)
+            self.assertEqual(scheduler_state, scheduler.state_dict())

Why the +1 here?
@@ -591,6 +597,14 @@ def load_state(self, input_dir: str):
             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)
+        custom_checkpoints = [f for f in os.listdir(input_dir) if "custom_checkpoint" in f]
+        if len(custom_checkpoints) != (len(self._custom_objects) + 1):
I would append them to the existing `_custom_objects`.
@@ -646,6 +660,15 @@ def get_state_dict(self, model):
 
         return state_dict
 
+    def register_for_checkpointing(self, *objects):
+        """
+        Makes note of `objects` and will save or load them in during `save_state` or `load_state`.
+
+        Note: These should be utilized when the state is being loaded or saved in the same script. It is not designed
+        to be used in different scripts
+        """
+        self._custom_objects = list(objects)
You should use a Path or `os.path.join` here otherwise it won't work in Windows.
@@ -132,3 +132,27 @@ def load_accelerator_state(input_dir, models, optimizers, process_index, scaler=
     if is_tpu_available():
         xm.set_rng_state(states["xm_seed"])
     logger.info("All random states loaded successfully")
+
+
+def save_custom_state(obj, path, index: int = 0):
+    """
+    Saves the state of `obj` to `{path}/custom_checkpoint_{index}.pkl`
+    """
+    # Should this be the right way to get a qual_name type value from `obj`?
+    save_location = f"{path}/custom_checkpoint_{index}.pkl"
We don't need a torch.load inside here?
@@ -132,3 +132,27 @@ def load_accelerator_state(input_dir, models, optimizers, process_index, scaler=
     if is_tpu_available():
         xm.set_rng_state(states["xm_seed"])
     logger.info("All random states loaded successfully")
+
+
+def save_custom_state(obj, path, index: int = 0):
+    """
+    Saves the state of `obj` to `{path}/custom_checkpoint_{index}.pkl`
+    """
+    # Should this be the right way to get a qual_name type value from `obj`?
+    save_location = f"{path}/custom_checkpoint_{index}.pkl"
+    logger.info(f"Saving the state of {str(obj)} to {save_location}")
+    if hasattr(obj, "state_dict"):
+        obj = obj.state_dict()
+    torch.save(obj, save_location)
+
+
+def load_custom_state(obj, path, index: int = 0):
+    """
+    Loads the state of `obj` at `{path}/custom_checkpoint_{index}.pkl`
+    """
+    load_location = f"{path}/custom_checkpoint_{index}.pkl"
+    if hasattr(obj, "load_state_dict"):
+        obj = obj.load_state_dict(load_location)
For the custom objects + the saved objected was my thought at the time. But you're write its a bug
@@ -591,6 +597,14 @@ def load_state(self, input_dir: str):
             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)
+        custom_checkpoints = [f for f in os.listdir(input_dir) if "custom_checkpoint" in f]
+        if len(custom_checkpoints) != (len(self._custom_objects) + 1):
The warning message does not match what the test is testing :-)
@@ -591,6 +598,13 @@ def load_state(self, input_dir: str):
             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)
+        custom_checkpoints = [f for f in os.listdir(input_dir) if "custom_checkpoint" in f]
+        if len(custom_checkpoints) != len(self._custom_objects):
+            logger.warn("Warning! You are loading a state that does not include registered objects.")
This is not a proper docstring **but** it will be once #271 is merged, so giving you a free pass ;-)
@@ -646,6 +660,26 @@ def get_state_dict(self, model):
 
         return state_dict
 
+    def register_for_checkpointing(self, *objects):
+        """
+        Makes note of `objects` and will save or load them in during `save_state` or `load_state`.
+
+        These should be utilized when the state is being loaded or saved in the same script. It is not designed to be
+        used in different scripts
+
+        Note: Every `object` must have a `load_state_dict` and `state_dict` function to be stored.
```suggestion
    torch.save(obj.state_dict(), save_location)
```
Can be combined now.
@@ -132,3 +133,23 @@ def load_accelerator_state(input_dir, models, optimizers, process_index, scaler=
     if is_tpu_available():
         xm.set_rng_state(states["xm_seed"])
     logger.info("All random states loaded successfully")
+
+
+def save_custom_state(obj, path, index: int = 0):
+    """
+    Saves the state of `obj` to `{path}/custom_checkpoint_{index}.pkl`
+    """
+    # Should this be the right way to get a qual_name type value from `obj`?
+    save_location = Path(path) / f"custom_checkpoint_{index}.pkl"
+    logger.info(f"Saving the state of {get_pretty_name(obj)} to {save_location}")
+    obj = obj.state_dict()
+    torch.save(obj, save_location)
Not sure we need a new class for this. Can be included in the current one!
@@ -123,3 +125,43 @@ def test_can_resume_training(self):
             self.assertEqual(b1, b3)
             self.assertEqual(opt_state1, opt_state3)
             self.assertEqual(ground_truth_rands, test_rands)
+
+
+class CustomItemsTest(unittest.TestCase):
Can you print the two numbers?
@@ -600,7 +600,9 @@ def load_state(self, input_dir: str):
         load_accelerator_state(input_dir, self._models, self._optimizers, self.state.process_index, self.scaler)
         custom_checkpoints = [f for f in os.listdir(input_dir) if "custom_checkpoint" in f]
         if len(custom_checkpoints) != len(self._custom_objects):
-            logger.warn("Warning! You are loading a state that does not include registered objects.")
+            logger.warn(
+                "Warning! Number of found checkpoints does not match the number of registered objects. Skipping."
How does 
```python
            err = "Warning! Number of found checkpoints does not match the number of registered objects:"
            err += f"\n\tFound checkpoints: {len(custom_checkpoints)}"
            err += f"\n\tRegistered objects: {len(self._custom_objects)}\nSkipping."
```
Sound?
@@ -600,7 +600,9 @@ def load_state(self, input_dir: str):
         load_accelerator_state(input_dir, self._models, self._optimizers, self.state.process_index, self.scaler)
         custom_checkpoints = [f for f in os.listdir(input_dir) if "custom_checkpoint" in f]
         if len(custom_checkpoints) != len(self._custom_objects):
-            logger.warn("Warning! You are loading a state that does not include registered objects.")
+            logger.warn(
+                "Warning! Number of found checkpoints does not match the number of registered objects. Skipping."
https://api.github.com/repos/huggingface/accelerate/pulls/270
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 1254c19e9..f59d3e8bd 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -166,6 +166,10 @@ def __init__(
         self.device_placement = device_placement
         self.split_batches = split_batches
         self.dispatch_batches = dispatch_batches
+        if dispatch_batches is True and version.parse(torch.__version__) < version.parse("1.8.0"):
+            raise ImportError(
+                "Using `DataLoaderDispatcher` requires PyTorch 1.8.0 minimum. You have {torch.__version__}."
+            )
 
         # Mixed precision attributes
         self.scaler = None
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index b28f24fc6..cc46e5577 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -328,6 +328,10 @@ class DataLoaderDispatcher(DataLoader):
     def __init__(self, dataset, split_batches: bool = False, **kwargs):
         super().__init__(dataset, **kwargs)
         self.split_batches = split_batches
+        if version.parse(torch.__version__) < version.parse("1.8.0"):
+            raise ImportError(
+                "Using `DataLoaderDispatcher` requires PyTorch 1.8.0 minimum. You have {torch.__version__}."
+            )
 
     def __iter__(self):
         state = AcceleratorState()
@@ -474,7 +478,10 @@ def prepare_data_loader(
         This does not support :obj:`BatchSampler` with varying batch size yet.
     """
     if dispatch_batches is None:
-        dispatch_batches = False if not put_on_device else isinstance(dataloader.dataset, IterableDataset)
+        if version.parse(torch.__version__) < version.parse("1.8.0") or not put_on_device:
+            dispatch_batches = False
+        else:
+            dispatch_batches = isinstance(dataloader.dataset, IterableDataset)
 
     if dispatch_batches and not put_on_device:
         raise ValueError("Using `dispatch_batches=True` requires `put_on_device=True`.")

https://api.github.com/repos/huggingface/accelerate/pulls/269
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 2d9937642..62b1291ee 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -63,13 +63,16 @@ class TensorInformation:
     dtype: torch.dtype
 
 
-def set_seed(seed: int):
+def set_seed(seed: int, device_specific: bool = False):
     """
     Helper function for reproducible behavior to set the seed in ``random``, ``numpy``, ``torch``.
 
     Args:
         seed (:obj:`int`): The seed to set.
+        device_specific (:obj:`bool`, `optional`, defaults to :obj:`False`): Whether to differ the seed on each device slightly with `self.process_index`
     """
+    if device_specific:
+        seed += AcceleratorState().process_index
     random.seed(seed)
     np.random.seed(seed)
     torch.manual_seed(seed)

```suggestion
        device_specific (:obj:`bool`, `optional`, defaults to :obj:`False`): Whether to differ the seed on each device slightly with `self.process_index`
```
@@ -63,13 +63,16 @@ class TensorInformation:
     dtype: torch.dtype
 
 
-def set_seed(seed: int):
+def set_seed(seed: int, device_specific: bool = False):
     """
     Helper function for reproducible behavior to set the seed in ``random``, ``numpy``, ``torch``.
 
     Args:
         seed (:obj:`int`): The seed to set.
+        device_specific (:obj:`bool`): Whether to differ the seed on each device slightly with `self.process_index`
https://api.github.com/repos/huggingface/accelerate/pulls/266
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index a029acfae..e8c910e19 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -256,11 +256,6 @@ def __init__(
 
     def __repr__(self):
         mixed_precision = self.mixed_precision
-        if self.distributed_type == DistributedType.DEEPSPEED:
-            if self.deepspeed_plugin.fp16:
-                mixed_precision = "fp16"
-            if self.deepspeed_plugin.bflaot16:
-                mixed_precision = "bf16"
 
         repr = (
             f"Distributed environment: {self.distributed_type}{('  Backend: ' + self.backend) if self.backend else ''}\n"
@@ -271,5 +266,5 @@ def __repr__(self):
             f"Mixed precision type: {mixed_precision}\n"
         )
         if self.distributed_type == DistributedType.DEEPSPEED:
-            repr += f"ds_config: {self.deepspeed_plugin.ds_config}\n"
+            repr += f"ds_config: {self.deepspeed_plugin.deepspeed_config}\n"
         return repr
diff --git a/src/accelerate/test_utils/test_script.py b/src/accelerate/test_utils/test_script.py
index ecd8e4f02..f9a99e22f 100644
--- a/src/accelerate/test_utils/test_script.py
+++ b/src/accelerate/test_utils/test_script.py
@@ -59,6 +59,8 @@ def dl_preparation_check():
     for batch in dl:
         result.append(gather(batch))
     result = torch.cat(result)
+
+    print(state.process_index, result, type(dl))
     assert torch.equal(result.cpu(), torch.arange(0, length).long()), "Wrong non-shuffled dataloader result."
 
     dl = DataLoader(range(length), batch_size=8)
@@ -327,6 +329,10 @@ def main():
     dl_preparation_check()
     central_dl_preparation_check()
 
+    # Trainings are not exactly the same in DeepSpeed and CPU mode
+    if state.distributed_type == DistributedType.DEEPSPEED:
+        return
+
     if state.local_process_index == 0:
         print("\n**Training integration test**")
     training_check()
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 83a40bf86..de5978802 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -91,7 +91,7 @@ def synchronize_rng_state(rng_type: Optional[RNGType] = None, generator: Optiona
     state = AcceleratorState()
     if state.distributed_type == DistributedType.TPU:
         rng_state = xm.mesh_reduce("random_seed", rng_state, lambda x: x[0])
-    elif state.distributed_type == DistributedType.MULTI_GPU:
+    elif state.distributed_type in [DistributedType.DEEPSPEED, DistributedType.MULTI_GPU]:
         rng_state = rng_state.to(state.device)
         torch.distributed.broadcast(rng_state, 0)
         rng_state = rng_state.cpu()
@@ -343,7 +343,7 @@ def gather(tensor):
     """
     if AcceleratorState().distributed_type == DistributedType.TPU:
         return _tpu_gather(tensor, name="accelerate.utils.gather")
-    elif AcceleratorState().distributed_type == DistributedType.MULTI_GPU:
+    elif AcceleratorState().distributed_type in [DistributedType.DEEPSPEED, DistributedType.MULTI_GPU]:
         return _gpu_gather(tensor)
     elif AcceleratorState().distributed_type == DistributedType.MULTI_CPU:
         return _cpu_gather(tensor)
@@ -376,7 +376,7 @@ def gather_object(object: Any):
     """
     if AcceleratorState().distributed_type == DistributedType.TPU:
         raise NotImplementedError("gather objects in TPU is not supported")
-    elif AcceleratorState().distributed_type == DistributedType.MULTI_GPU:
+    elif AcceleratorState().distributed_type in [DistributedType.DEEPSPEED, DistributedType.MULTI_GPU]:
         return _gpu_gather_object(object)
     elif AcceleratorState().distributed_type == DistributedType.MULTI_CPU:
         return _cpu_gather_object(object)
@@ -415,7 +415,7 @@ def broadcast(tensor, from_process: int = 0):
     """
     if AcceleratorState().distributed_type == DistributedType.TPU:
         return _tpu_broadcast(tensor, src=from_process, name="accelerate.utils.broadcast")
-    elif AcceleratorState().distributed_type == DistributedType.MULTI_GPU:
+    elif AcceleratorState().distributed_type in [DistributedType.DEEPSPEED, DistributedType.MULTI_GPU]:
         return _gpu_broadcast(tensor, src=from_process)
     elif AcceleratorState().distributed_type == DistributedType.MULTI_CPU:
         return _gpu_broadcast(tensor, src=from_process)
@@ -439,7 +439,7 @@ def broadcast_object_list(object_list, from_process: int = 0):
     if AcceleratorState().distributed_type == DistributedType.TPU:
         for i, obj in enumerate(object_list):
             object_list[i] = xm.mesh_reduce("accelerate.utils.broadcast_object_list", obj, lambda x: x[from_process])
-    elif AcceleratorState().distributed_type == DistributedType.MULTI_GPU:
+    elif AcceleratorState().distributed_type in [DistributedType.DEEPSPEED, DistributedType.MULTI_GPU]:
         torch.distributed.broadcast_object_list(object_list, src=from_process)
     elif AcceleratorState().distributed_type == DistributedType.MULTI_CPU:
         torch.distributed.broadcast_object_list(object_list, src=from_process)

https://api.github.com/repos/huggingface/accelerate/pulls/264
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index 534ea1526..15eaaecb6 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -52,6 +52,17 @@ def get_cluster_input():
                 lambda x: int(x),
             )
 
+    if distributed_type == DistributedType.NO:
+        use_cpu = _ask_field(
+            "Do you want to run your training on CPU only (even if a GPU is available)? [no]:",
+            lambda x: bool(x),
+            default=False,
+        )
+    elif distributed_type == DistributedType.MULTI_CPU:
+        use_cpu = True
+    else:
+        use_cpu = False
+
     deepspeed_config = None
     if distributed_type in [DistributedType.MULTI_GPU, DistributedType.NO]:
         use_deepspeed = _ask_field(
@@ -122,4 +133,5 @@ def get_cluster_input():
         main_process_port=main_process_port,
         main_training_function=main_training_function,
         deepspeed_config=deepspeed_config,
+        use_cpu=use_cpu,
     )
diff --git a/src/accelerate/commands/config/config_args.py b/src/accelerate/commands/config/config_args.py
index 891580ab5..5af912fc3 100644
--- a/src/accelerate/commands/config/config_args.py
+++ b/src/accelerate/commands/config/config_args.py
@@ -67,6 +67,7 @@ class BaseConfig:
     compute_environment: ComputeEnvironment
     distributed_type: Union[DistributedType, SageMakerDistributedType]
     mixed_precision: str
+    use_cpu: bool
 
     def to_dict(self):
         result = self.__dict__
@@ -87,6 +88,8 @@ def from_json_file(cls, json_file=None):
             config_dict["mixed_precision"] = "fp16" if ("fp16" in config_dict and config_dict["fp16"]) else "no"
         if "fp16" in config_dict:  # Convert the config to the new format.
             del config_dict["fp16"]
+        if "use_cpu" not in config_dict:
+            config_dict["use_cpu"] = False
         return cls(**config_dict)
 
     def to_json_file(self, json_file):
@@ -106,6 +109,8 @@ def from_yaml_file(cls, yaml_file=None):
             config_dict["mixed_precision"] = "fp16" if ("fp16" in config_dict and config_dict["fp16"]) else "no"
         if "fp16" in config_dict:  # Convert the config to the new format.
             del config_dict["fp16"]
+        if "use_cpu" not in config_dict:
+            config_dict["use_cpu"] = False
 
         return cls(**config_dict)
 
diff --git a/src/accelerate/launchers.py b/src/accelerate/launchers.py
index a8429db96..a47eb9a4c 100644
--- a/src/accelerate/launchers.py
+++ b/src/accelerate/launchers.py
@@ -171,6 +171,7 @@ def debug_launcher(function, args=(), num_processes=2):
             master_port="29500",
             mixed_precision="no",
             accelerate_debug_rdv_file=tmp_file.name,
+            use_cpu="yes",
         ):
             launcher = PrepareForLaunch(function, debug=True)
             start_processes(launcher, args=args, nprocs=num_processes, start_method="fork")
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index 59aebef2d..a029acfae 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -153,6 +153,8 @@ def __init__(
         **kwargs,
     ):
         self.__dict__ = self._shared_state
+        if parse_flag_from_env("USE_CPU"):
+            cpu = True
         if not getattr(self, "initialized", False):
             self.backend = None
             self.deepspeed_plugin = None
@@ -217,14 +219,7 @@ def __init__(
                 rank = get_int_from_env(["RANK", "PMI_RANK", "OMPI_COMM_WORLD_RANK", "MV2_COMM_WORLD_RANK"], 0)
                 size = get_int_from_env(["WORLD_SIZE", "PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE"], 1)
                 local_rank = get_int_from_env(
-                    [
-                        "LOCAL_RANK",
-                        "CPU_LOCAL_RANK",
-                        "MPI_LOCALRANKID",
-                        "OMPI_COMM_WORLD_LOCAL_RANK",
-                        "MV2_COMM_WORLD_LOCAL_RANK",
-                    ],
-                    0,
+                    ["LOCAL_RANK", "MPI_LOCALRANKID", "OMPI_COMM_WORLD_LOCAL_RANK", "MV2_COMM_WORLD_LOCAL_RANK"], 0
                 )
                 local_size = get_int_from_env(
                     ["MPI_LOCALNRANKS", "OMPI_COMM_WORLD_LOCAL_SIZE", "MV2_COMM_WORLD_LOCAL_SIZE"], 1
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 6eba4b295..83a40bf86 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -614,8 +614,6 @@ def __call__(self, index, *args):
                 store=torch.distributed.FileStore(rdv_file, world_size),
                 world_size=world_size,
             )
-            # Prepare the environment for torch.distributed
-            os.environ["CPU_LOCAL_RANK"] = str(index)
         elif self.distributed_type == DistributedType.MULTI_GPU or self.distributed_type == DistributedType.MULTI_CPU:
             # Prepare the environment for torch.distributed
             os.environ["LOCAL_RANK"] = str(index)

https://api.github.com/repos/huggingface/accelerate/pulls/263
diff --git a/src/accelerate/__init__.py b/src/accelerate/__init__.py
index e0b96cd10..f77a26a56 100644
--- a/src/accelerate/__init__.py
+++ b/src/accelerate/__init__.py
@@ -6,6 +6,6 @@
 
 from .accelerator import Accelerator
 from .kwargs_handlers import DistributedDataParallelKwargs, GradScalerKwargs
-from .notebook_launcher import notebook_launcher
+from .launchers import debug_launcher, notebook_launcher
 from .state import DistributedType
 from .utils import DeepSpeedPlugin, synchronize_rng_states
diff --git a/src/accelerate/notebook_launcher.py b/src/accelerate/launchers.py
similarity index 75%
rename from src/accelerate/notebook_launcher.py
rename to src/accelerate/launchers.py
index 90b4f826a..a8429db96 100644
--- a/src/accelerate/notebook_launcher.py
+++ b/src/accelerate/launchers.py
@@ -14,6 +14,7 @@
 
 import os
 import sys
+import tempfile
 import warnings
 
 import torch
@@ -21,7 +22,7 @@
 from packaging import version
 
 from .state import AcceleratorState
-from .utils import PrepareForLaunch
+from .utils import PrepareForLaunch, patch_environment
 
 
 def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, mixed_precision="no", use_port="29500"):
@@ -106,12 +107,6 @@ def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, mix
                     "function."
                 )
 
-            # torch.distributed will expect a few environment variable to be here. We set the ones common to each
-            # process here (the other ones will be set be the launcher).
-            os.environ["WORLD_SIZE"] = str(num_processes)
-            os.environ["MASTER_ADDR"] = "127.0.0.1"
-            os.environ["MASTER_PORT"] = str(use_port)
-
             mixed_precision = mixed_precision.lower()
             if mixed_precision not in ["no", "fp16", "bf16"]:
                 raise ValueError(
@@ -122,17 +117,15 @@ def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, mix
                 warnings.warn('use_fp16=True is deprecated. Use mixed_precision="fp16" instead.', DeprecationWarning)
                 mixed_precision = "fp16"
 
-            os.environ["MIXED_PRECISION"] = str(mixed_precision)
+            # torch.distributed will expect a few environment variable to be here. We set the ones common to each
+            # process here (the other ones will be set be the launcher).
+            with patch_environment(
+                world_size=num_processes, master_addr="127.0.01", master_port=use_port, mixed_precision=mixed_precision
+            ):
+                launcher = PrepareForLaunch(function, distributed_type="MULTI_GPU")
 
-            launcher = PrepareForLaunch(function, distributed_type="MULTI_GPU")
-            try:
                 print(f"Launching training on {num_processes} GPUs.")
                 start_processes(launcher, args=args, nprocs=num_processes, start_method="fork")
-            finally:
-                # Clean up the environment variables set.
-                del os.environ["WORLD_SIZE"]
-                del os.environ["MASTER_ADDR"]
-                del os.environ["MASTER_PORT"]
 
         else:
             # No need for a distributed launch otherwise as it's either CPU or one GPU.
@@ -141,3 +134,43 @@ def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, mix
             else:
                 print("Launching training on CPU.")
             function(*args)
+
+
+def debug_launcher(function, args=(), num_processes=2):
+    """
+    Launches a training function using several processes on CPU for debugging purposes.
+
+    .. warning::
+
+        This function is provided for internal testing and debugging, but it's not intended for real trainings. It will
+        only use the CPU.
+
+
+    Args:
+        function (:obj:`Callable`):
+            The training function to execute.
+        args (:obj:`Tuple`):
+            Tuple of arguments to pass to the function (it will receive :obj:`*args`).
+        num_processes (:obj:`int`, *optional*, defaults to 2):
+            The number of processes to use for training.
+    """
+    if version.parse(torch.__version__) < version.parse("1.5.0"):
+        raise ImportError(
+            "Using `debug_launcher` for distributed training on GPUs require torch >= 1.5.0, got "
+            f"{torch.__version__}."
+        )
+
+    from torch.multiprocessing import start_processes
+
+    with tempfile.NamedTemporaryFile() as tmp_file:
+        # torch.distributed will expect a few environment variable to be here. We set the ones common to each
+        # process here (the other ones will be set be the launcher).
+        with patch_environment(
+            world_size=num_processes,
+            master_addr="127.0.01",
+            master_port="29500",
+            mixed_precision="no",
+            accelerate_debug_rdv_file=tmp_file.name,
+        ):
+            launcher = PrepareForLaunch(function, debug=True)
+            start_processes(launcher, args=args, nprocs=num_processes, start_method="fork")
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index 6693a34ca..59aebef2d 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -217,7 +217,14 @@ def __init__(
                 rank = get_int_from_env(["RANK", "PMI_RANK", "OMPI_COMM_WORLD_RANK", "MV2_COMM_WORLD_RANK"], 0)
                 size = get_int_from_env(["WORLD_SIZE", "PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE"], 1)
                 local_rank = get_int_from_env(
-                    ["LOCAL_RANK", "MPI_LOCALRANKID", "OMPI_COMM_WORLD_LOCAL_RANK", "MV2_COMM_WORLD_LOCAL_RANK"], 0
+                    [
+                        "LOCAL_RANK",
+                        "CPU_LOCAL_RANK",
+                        "MPI_LOCALRANKID",
+                        "OMPI_COMM_WORLD_LOCAL_RANK",
+                        "MV2_COMM_WORLD_LOCAL_RANK",
+                    ],
+                    0,
                 )
                 local_size = get_int_from_env(
                     ["MPI_LOCALNRANKS", "OMPI_COMM_WORLD_LOCAL_SIZE", "MV2_COMM_WORLD_LOCAL_SIZE"], 1
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 36eaea592..443c55622 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -16,6 +16,7 @@
 import os
 import random
 from collections.abc import Mapping
+from contextlib import contextmanager
 from dataclasses import dataclass, field
 from enum import Enum
 from typing import Any, List, Optional, Union
@@ -594,14 +595,28 @@ class PrepareForLaunch:
             The function to launch.
         distributed_type (:class:`~accelerate.state.DistributedType`):
             The distributed type to prepare for.
+        debug (:obj:`bool`, *optional*, defaults to :obj:`False`):
+            Whether or not this is a debug launch.
     """
 
-    def __init__(self, launcher, distributed_type="NO"):
+    def __init__(self, launcher, distributed_type="NO", debug=False):
         self.launcher = launcher
         self.distributed_type = DistributedType(distributed_type)
+        self.debug = debug
 
     def __call__(self, index, *args):
-        if self.distributed_type == DistributedType.MULTI_GPU or self.distributed_type == DistributedType.MULTI_CPU:
+        if self.debug:
+            world_size = int(os.environ.get("WORLD_SIZE"))
+            rdv_file = os.environ.get("ACCELERATE_DEBUG_RDV_FILE")
+            torch.distributed.init_process_group(
+                "gloo",
+                rank=index,
+                store=torch.distributed.FileStore(rdv_file, world_size),
+                world_size=world_size,
+            )
+            # Prepare the environment for torch.distributed
+            os.environ["CPU_LOCAL_RANK"] = str(index)
+        elif self.distributed_type == DistributedType.MULTI_GPU or self.distributed_type == DistributedType.MULTI_CPU:
             # Prepare the environment for torch.distributed
             os.environ["LOCAL_RANK"] = str(index)
             os.environ["RANK"] = str(index)
@@ -654,3 +669,19 @@ def __post_init__(self):
             "steps_per_print": float("inf"),  # this will stop deepspeed from logging @ stdout
             "zero_allow_untested_optimizer": True,
         }
+
+
+@contextmanager
+def patch_environment(**kwargs):
+    """
+    A context manager that will add each keyword argument passed to ``os.environ`` and remove them when exiting.
+
+    Will convert the values in :obj:`kwargs` to strings and upper-case all the keys.
+    """
+    for key, value in kwargs.items():
+        os.environ[key.upper()] = str(value)
+
+    yield
+
+    for key in kwargs:
+        del os.environ[key.upper()]
diff --git a/tests/test_cpu.py b/tests/test_cpu.py
new file mode 100644
index 000000000..72a8b2e7d
--- /dev/null
+++ b/tests/test_cpu.py
@@ -0,0 +1,23 @@
+# 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 unittest
+
+from accelerate import debug_launcher
+from accelerate.test_utils import test_script
+
+
+class MultiTPUTester(unittest.TestCase):
+    def test_cpu(self):
+        debug_launcher(test_script.main)
diff --git a/tests/test_utils.py b/tests/test_utils.py
index e515df8fe..a31c45d31 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -12,12 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 import unittest
 from collections import UserDict, namedtuple
 
 import torch
 
-from accelerate.utils import send_to_device
+from accelerate.utils import patch_environment, send_to_device
 
 
 TestNamedTuple = namedtuple("TestNamedTuple", "a b c")
@@ -62,3 +63,11 @@ def test_send_to_device(self):
         self.assertTrue(torch.equal(result4["b"][0].cpu(), tensor))
         self.assertTrue(torch.equal(result4["b"][1].cpu(), tensor))
         self.assertEqual(result4["c"], 1)
+
+    def test_patch_environment(self):
+        with patch_environment(aa=1, BB=2):
+            self.assertEqual(os.environ.get("AA"), "1")
+            self.assertEqual(os.environ.get("BB"), "2")
+
+        self.assertNotIn("AA", os.environ)
+        self.assertNotIn("BB", os.environ)

https://api.github.com/repos/huggingface/accelerate/pulls/259
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 77db4546b..ff1b148f4 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -92,6 +92,17 @@ def launch_command_parser(subparsers=None):
         default=None,
         help="The name of the main function to be executed in your script (only for TPU training).",
     )
+    parser.add_argument(
+        "-m",
+        "--module",
+        action="store_true",
+        help="Change each process to interpret the launch script as a Python module, executing with the same behavior as 'python -m'.",
+    )
+    parser.add_argument(
+        "--no_python",
+        action="store_true",
+        help="Skip prepending the training script with 'python' - just execute it directly. Useful when the script is not a Python script.",
+    )
     parser.add_argument(
         "--aws_access_key_id",
         type=str,
@@ -140,7 +151,14 @@ def launch_command_parser(subparsers=None):
 
 
 def simple_launcher(args):
-    cmd = [sys.executable, args.training_script]
+    cmd = []
+    if args.no_python and args.module:
+        raise ValueError("--module and --no_python cannot be used together")
+    if not args.no_python:
+        cmd.append(sys.executable)
+        if args.module:
+            cmd.append("-m")
+    cmd.append(args.training_script)
     cmd.extend(args.training_script_args)
 
     current_env = os.environ.copy()
@@ -163,8 +181,7 @@ def simple_launcher(args):
 
 
 def multi_gpu_launcher(args):
-    cmd = [sys.executable, "-m", "torch.distributed.launch"]
-    cmd.extend(["--use_env"])
+    cmd = [sys.executable, "-m", "torch.distributed.launch", "--use_env"]
     if args.num_machines > 1:
         cmd.extend(
             [
@@ -184,6 +201,13 @@ def multi_gpu_launcher(args):
         cmd.extend(["--nproc_per_node", str(args.num_processes)])
         if args.main_process_port is not None:
             cmd.extend(["--master_port", str(args.main_process_port)])
+
+    if args.module and args.no_python:
+        raise ValueError("--module and --no_python cannot be used together")
+    elif args.module:
+        cmd.append("--module")
+    elif args.no_python:
+        cmd.append("--no_python")
     cmd.append(args.training_script)
     cmd.extend(args.training_script_args)
 
@@ -206,7 +230,6 @@ def multi_gpu_launcher(args):
 
 
 def deepspeed_launcher(args):
-
     cmd = ["deepspeed"]
     if args.num_machines > 1:
         cmd.extend(
@@ -226,6 +249,12 @@ def deepspeed_launcher(args):
     else:
         cmd.extend(["--num_gpus", str(args.num_processes)])
 
+    if args.module and args.no_python:
+        raise ValueError("--module and --no_python cannot be used together")
+    elif args.module:
+        cmd.append("--module")
+    elif args.no_python:
+        cmd.append("--no_python")
     cmd.append(args.training_script)
     cmd.extend(args.training_script_args)
 
@@ -254,21 +283,28 @@ def deepspeed_launcher(args):
 def tpu_launcher(args):
     import torch_xla.distributed.xla_multiprocessing as xmp
 
-    # Import training_script as a module.
-    script_path = Path(args.training_script)
-    sys.path.append(str(script_path.parent.resolve()))
-    mod_name = script_path.stem
+    if args.no_python:
+        raise ValueError("--no_python cannot be used with TPU launcher")
+
+    if args.module:
+        mod_name = args.training_script
+    else:
+        # Import training_script as a module
+        script_path = Path(args.training_script)
+        sys.path.append(str(script_path.parent.resolve()))
+        mod_name = script_path.stem
+
     mod = importlib.import_module(mod_name)
     if not hasattr(mod, args.main_training_function):
         raise ValueError(
             f"Your training script should have a function named {args.main_training_function}, or you should pass a "
             "different value to `--main_training_function`."
         )
-    main_function = getattr(mod, args.main_training_function)
 
     # Patch sys.argv
-    sys.argv = [args.training_script] + args.training_script_args
+    sys.argv = [mod.__file__] + args.training_script_args
 
+    main_function = getattr(mod, args.main_training_function)
     xmp.spawn(PrepareForLaunch(main_function), args=(), nprocs=args.num_processes)
 
 
@@ -319,6 +355,11 @@ def sagemaker_launcher(sagemaker_config: SageMakerConfig, args):
         raise ImportError(
             "Please install sagemaker to be able to launch training on Amazon SageMaker with `pip install accelerate[sagemaker]`"
         )
+    if args.module or args.no_python:
+        raise ValueError(
+            "SageMaker requires a python training script file and cannot be used with --module or --no_python"
+        )
+
     from sagemaker.huggingface import HuggingFace
 
     # configure environment

Since both can't be true at the same time, the test should not be nested.
```suggestion
    elif args.module:
        cmd.append("-m")
```
@@ -140,7 +151,14 @@ def launch_command_parser(subparsers=None):
 
 
 def simple_launcher(args):
-    cmd = [sys.executable, args.training_script]
+    cmd = []
+    if args.no_python and args.module:
+        raise ValueError("--module and --no_python cannot be used together")
+    if not args.no_python:
+        cmd.append(sys.executable)
+        if args.module:
+            cmd.append("-m")
```suggestion
    elif args.module:
        cmd.append("--module")
    elif args.no_python:
```
@@ -184,6 +201,13 @@ def multi_gpu_launcher(args):
         cmd.extend(["--nproc_per_node", str(args.num_processes)])
         if args.main_process_port is not None:
             cmd.extend(["--master_port", str(args.main_process_port)])
+
+    if args.module and args.no_python:
+        raise ValueError("--module and --no_python cannot be used together")
+    if args.module:
+        cmd.append("--module")
+    if args.no_python:
```suggestion
    if args.module and args.no_python:
        raise ValueError("--module and --no_python cannot be used together")
    elif args.module:
        cmd.append("--module")
    elif args.no_python:
```
@@ -226,6 +249,12 @@ def deepspeed_launcher(args):
     else:
         cmd.extend(["--num_gpus", str(args.num_processes)])
 
+    if args.module and args.no_python:
+        raise ValueError("--module and --no_python cannot be used together")
+    if args.module:
+        cmd.append("--module")
+    if args.no_python:
Good catch, thanks. I think after the un-indent, it still has to stay an `if` and not an `elif`, otherwise it won't get evaluated because the `if not args.no_python` above would always be evaluated `True` first.

```suggestion
    if args.module:
        cmd.append("-m")
```
@@ -140,7 +151,14 @@ def launch_command_parser(subparsers=None):
 
 
 def simple_launcher(args):
-    cmd = [sys.executable, args.training_script]
+    cmd = []
+    if args.no_python and args.module:
+        raise ValueError("--module and --no_python cannot be used together")
+    if not args.no_python:
+        cmd.append(sys.executable)
+        if args.module:
+            cmd.append("-m")
Thanks again for working on this!
https://api.github.com/repos/huggingface/accelerate/pulls/258
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)

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.
@@ -560,6 +566,109 @@ def save(self, obj, f):
         """
         save(obj, f)
 
+    def save_state(self, output_dir: str):
```suggestion
            if i != 0:
                weights_name += f"_{i}"
```
Is it a typo? I think it's for every model **but** the first one.
@@ -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}"
Should be `accelerator.save` here (otherwise won't work for TPUs).
@@ -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)
Same here:
```suggestion
            accelerator.save(state, output_optimizer_file)
```
@@ -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)
```suggestion
            if i != 0:
                optimizer_name += f"_{i}"
```
As above!
@@ -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}"
```suggestion
            scaler_name = "scaler.pt"
```
Let's use the same name as in the `Trainer`.
@@ -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"
```suggestion
            logger.info(f"Gradient scaler state saved in {output_scaler_file}")
```
@@ -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}")
```suggestion
            output_scaler_file = os.path.join(output_dir, scaler_name)
            torch.save(state, output_scaler_file)
```
@@ -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)
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.
@@ -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
```suggestion
        # Random number generator states
```
@@ -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
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)
@@ -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"
```suggestion
            if i != 0:
                weights_name += f"_{i}"
```
@@ -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}"
```suggestion
            if i != 0:
                optimizer_name += f"_{i}"
```
@@ -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}"
The names should probably be constants, so we only have to change them in one place ;-)
@@ -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"
Same ehre + one for each process.
@@ -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"))
We could wrap all of this in some train function that does the training.
@@ -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
Nit: we should use `unittest` here with `self.assertEqual` and `self.assertEqualDict` instead of those asserts.
@@ -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
Same as above.
@@ -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
This one being deprecated, I wouldn't pass it/create it.
@@ -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.")
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
@@ -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.")
```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.
@@ -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`.
```suggestion
        Loads the current states of the model, optimizer, scaler, and RNG generators.
```
@@ -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`.
```suggestion
# Copyright 2022 The HuggingFace Team. All rights reserved.
```
@@ -0,0 +1,134 @@
+# Copyright 2021 The HuggingFace Team. All rights reserved.
```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)
@@ -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`.
```suggestion
        scaler (:obj:`torch.cuda.amp.GradScaler`, `optional`):
```
@@ -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:
```suggestion
            An optional gradient scaler instance to save
```
(The type is already described)
@@ -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
This also needs to use the `save` function (especially for TPUs).
@@ -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)
```suggestion
        model_states (:obj:`List[torch.nn.Module]`):
```
@@ -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`):
```suggestion
        optimizers (:obj:`List[torch.optim.Optimizer]`):
```
@@ -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`):
Fix those as above.
@@ -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:
```suggestion
    Loads states of the models, optimizers, scaler, and RNG generators from a given directory.
```
@@ -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`.
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`?
@@ -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)
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.
@@ -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()
@sgugger how would we delegate the tests to be ran on cpu, mixed precision, etc then in that case? 
@@ -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()
https://api.github.com/repos/huggingface/accelerate/pulls/255
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
new file mode 100644
index 000000000..66b0a8538
--- /dev/null
+++ b/CONTRIBUTING.md
@@ -0,0 +1,235 @@
+<!---
+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.
+-->
+
+# How to contribute to 🤗 Accelerate?
+
+Everyone is welcome to contribute, and we value everybody's contribution. Code
+is thus not the only way to help the community. Answering questions, helping
+others, reaching out and improving the documentations are immensely valuable to
+the community.
+
+It also helps us if you spread the word: reference the library from blog posts
+on the awesome projects it made possible, shout out on Twitter every time it has
+helped you, or simply star the repo to say "thank you".
+
+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md).
+
+## You can contribute in so many ways!
+
+Some of the ways you can contribute to Accelerate:
+* Fixing outstanding issues with the existing code;
+* Contributing to the examples or to the documentation;
+* Submitting issues related to bugs or desired new features.
+
+## Submitting a new issue or feature request
+
+Do your best to follow these guidelines when submitting an issue or a feature
+request. It will make it easier for us to come back to you quickly and with good
+feedback.
+
+### Did you find a bug?
+
+The 🤗 Accelerate library is robust and reliable thanks to the users who notify us of
+the problems they encounter. So thank you for reporting an issue.
+
+First, we would really appreciate it if you could **make sure the bug was not
+already reported** (use the search bar on Github under Issues).
+
+Did not find it? :( So we can act quickly on it, please follow these steps:
+
+* Include your **OS type and version**, the versions of **Python** and **PyTorch**.
+* A short, self-contained, code snippet that allows us to reproduce the bug in
+  less than 30s;
+* Provide the with your Accelerate configuration (located by default in `~/.cache/huggingface/accelerate/default_congig.yml`)
+
+### Do you want a new feature?
+
+A good feature request addresses the following points:
+
+1. Motivation first:
+* Is it related to a problem/frustration with the library? If so, please explain
+  why. Providing a code snippet that demonstrates the problem is best.
+* Is it related to something you would need for a project? We'd love to hear
+  about it!
+* Is it something you worked on and think could benefit the community?
+  Awesome! Tell us what problem it solved for you.
+2. Write a *full paragraph* describing the feature;
+3. Provide a **code snippet** that demonstrates its future use;
+4. In case this is related to a paper, please attach a link;
+5. Attach any additional information (drawings, screenshots, etc.) you think may help.
+
+If your issue is well written we're already 80% of the way there by the time you
+post it.
+
+## Submitting a pull request (PR)
+
+Before writing code, we strongly advise you to search through the existing PRs or
+issues to make sure that nobody is already working on the same thing. If you are
+unsure, it is always a good idea to open an issue to get some feedback.
+
+You will need basic `git` proficiency to be able to contribute to
+🤗 Accelerate. `git` is not the easiest tool to use but it has the greatest
+manual. Type `git --help` in a shell and enjoy. If you prefer books, [Pro
+Git](https://git-scm.com/book/en/v2) is a very good reference.
+
+Follow these steps to start contributing:
+
+1. Fork the [repository](https://github.com/huggingface/accelerate) by
+   clicking on the 'Fork' button on the repository's page. This creates a copy of the code
+   under your GitHub user account.
+
+2. Clone your fork to your local disk, and add the base repository as a remote. The following command
+   assumes you have your public SSH key uploaded to GitHub. See the following guide for more
+   [information](https://docs.github.com/en/repositories/creating-and-managing-repositories/cloning-a-repository).
+
+   ```bash
+   $ git clone git@github.com:<your Github handle>/accelerate.git
+   $ cd accelerate
+   $ git remote add upstream https://github.com/huggingface/accelerate.git
+   ```
+
+3. Create a new branch to hold your development changes, and do this for every new PR you work on.
+
+   Start by synchronizing your `main` branch with the `upstream/main` branch (ore details in the [GitHub Docs](https://docs.github.com/en/github/collaborating-with-issues-and-pull-requests/syncing-a-fork)):
+
+   ```bash
+   $ git checkout main
+   $ git fetch upstream
+   $ git merge upstream/main
+   ```
+
+   Once your `main` branch is synchronized, create a new branch from it:
+
+   ```bash
+   $ git checkout -b a-descriptive-name-for-my-changes
+   ```
+
+   **Do not** work on the `master` branch.
+
+4. Set up a development environment by running the following command in a conda or a virtual environment you've created for working on this library:
+
+   ```bash
+   $ pip install -e ".[quality]"
+   ```
+
+   (If accelerate was already installed in the virtual environment, remove
+   it with `pip uninstall accelerate` before reinstalling it in editable
+   mode with the `-e` flag.)
+
+5. Develop the features on your branch.
+
+   As you work on the features, you should make sure that the test suite
+   passes. You should run the tests impacted by your changes like this (see 
+   below an explanation regarding the environment variable):
+
+   ```bash
+   $ pytest tests/<TEST_TO_RUN>.py
+   ```
+   
+   > For the following commands leveraging the `make` utility, we recommend using the WSL system when running on
+   > Windows. More information [here](https://docs.microsoft.com/en-us/windows/wsl/about).
+
+   You can also run the full suite with the following command.
+
+   ```bash
+   $ make test
+   ```
+
+   `accelerate` relies on `black` and `isort` to format its source code
+   consistently. After you make changes, apply automatic style corrections and code verifications
+   that can't be automated in one go with:
+
+   This target is also optimized to only work with files modified by the PR you're working on.
+
+   If you prefer to run the checks one after the other, the following command apply the
+   style corrections:
+
+   ```bash
+   $ make style
+   ```
+
+   `accelerate` also uses `flake8` and a few custom scripts to check for coding mistakes. Quality
+   control runs in CI, however you can also run the same checks with:
+
+   ```bash
+   $ make quality
+   ```
+
+   Once you're happy with your changes, add changed files using `git add` and
+   make a commit with `git commit` to record your changes locally:
+
+   ```bash
+   $ git add modified_file.py
+   $ git commit
+   ```
+
+   Please write [good commit messages](https://chris.beams.io/posts/git-commit/).
+
+   It is a good idea to sync your copy of the code with the original
+   repository regularly. This way you can quickly account for changes:
+
+   ```bash
+   $ git fetch upstream
+   $ git rebase upstream/master
+   ```
+
+   Push the changes to your account using:
+
+   ```bash
+   $ git push -u origin a-descriptive-name-for-my-changes
+   ```
+
+6. Once you are satisfied (**and the checklist below is happy too**), go to the
+   webpage of your fork on GitHub. Click on 'Pull request' to send your changes
+   to the project maintainers for review.
+
+7. It's ok if maintainers ask you for changes. It happens to core contributors
+   too! So everyone can see the changes in the Pull request, work in your local
+   branch and push the changes to your fork. They will automatically appear in
+   the pull request.
+
+
+### Checklist
+
+1. The title of your pull request should be a summary of its contribution;
+2. If your pull request addresses an issue, please mention the issue number in
+   the pull request description to make sure they are linked (and people
+   consulting the issue know you are working on it);
+3. To indicate a work in progress please prefix the title with `[WIP]`, or mark
+   the PR as a draft PR. These are useful to avoid duplicated work, and to differentiate
+   it from PRs ready to be merged;
+4. Make sure existing tests pass;
+5. Add high-coverage tests. No quality testing = no merge.
+
+See an example of a good PR here: https://github.com/huggingface/accelerate/pull/255
+
+### Tests
+
+An extensive test suite is included to test the library behavior and several examples. Library tests can be found in
+the [tests folder](https://github.com/huggingface/accelerate/tree/master/tests).
+
+We use `pytest` in order to run the tests. From the root of the
+repository, here's how to run tests with `pytest` for the library:
+
+```bash
+$ python -m pytest -sv ./tests
+```
+
+In fact, that's how `make test` is implemented (sans the `pip install` line)!
+
+You can specify a smaller set of tests in order to test only the feature
+you're working on.
\ No newline at end of file

Not sure if that's important given you're watching the repo.
@@ -0,0 +1,220 @@
+<!---
+Copyright 2020 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.
+-->
+
+# How to contribute to Accelerate?
+
+Everyone is welcome to contribute, and we value everybody's contribution. Code
+is thus not the only way to help the community. Answering questions, helping
+others, reaching out and improving the documentations are immensely valuable to
+the community.
+
+It also helps us if you spread the word: reference the library from blog posts
+on the awesome projects it made possible, shout out on Twitter every time it has
+helped you, or simply star the repo to say "thank you".
+
+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md).
+
+## You can contribute in so many ways!
+
+There are 4 ways you can contribute to Accelerate:
+* Fixing outstanding issues with the existing code;
+* Contributing to the examples or to the documentation;
+* Submitting issues related to bugs or desired new features.
+
+Please ping @sgugger on your contribution.
Yeah we can remove that part ;-)
@@ -0,0 +1,220 @@
+<!---
+Copyright 2020 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.
+-->
+
+# How to contribute to Accelerate?
+
+Everyone is welcome to contribute, and we value everybody's contribution. Code
+is thus not the only way to help the community. Answering questions, helping
+others, reaching out and improving the documentations are immensely valuable to
+the community.
+
+It also helps us if you spread the word: reference the library from blog posts
+on the awesome projects it made possible, shout out on Twitter every time it has
+helped you, or simply star the repo to say "thank you".
+
+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md).
+
+## You can contribute in so many ways!
+
+There are 4 ways you can contribute to Accelerate:
+* Fixing outstanding issues with the existing code;
+* Contributing to the examples or to the documentation;
+* Submitting issues related to bugs or desired new features.
+
+Please ping @sgugger on your contribution.
Missing the emoji here.
@@ -0,0 +1,220 @@
+<!---
+Copyright 2020 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.
+-->
+
+# How to contribute to Accelerate?
```suggestion
* Provide us with your Accelerate configuration (located by default in `~/.cache/huggingface/accelerate/default_congig.yml`)
* Provide the command you used to launch your script (for instance `accelerate launch train_script.py`)
* Provide the *full* traceback if an exception is raised.
```
@@ -0,0 +1,220 @@
+<!---
+Copyright 2020 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.
+-->
+
+# How to contribute to Accelerate?
+
+Everyone is welcome to contribute, and we value everybody's contribution. Code
+is thus not the only way to help the community. Answering questions, helping
+others, reaching out and improving the documentations are immensely valuable to
+the community.
+
+It also helps us if you spread the word: reference the library from blog posts
+on the awesome projects it made possible, shout out on Twitter every time it has
+helped you, or simply star the repo to say "thank you".
+
+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md).
+
+## You can contribute in so many ways!
+
+There are 4 ways you can contribute to Accelerate:
+* Fixing outstanding issues with the existing code;
+* Contributing to the examples or to the documentation;
+* Submitting issues related to bugs or desired new features.
+
+Please ping @sgugger on your contribution.
+
+## Submitting a new issue or feature request
+
+Do your best to follow these guidelines when submitting an issue or a feature
+request. It will make it easier for us to come back to you quickly and with good
+feedback.
+
+### Did you find a bug?
+
+The 🤗 Accelerate library is robust and reliable thanks to the users who notify us of
+the problems they encounter. So thank you for reporting an issue.
+
+First, we would really appreciate it if you could **make sure the bug was not
+already reported** (use the search bar on Github under Issues).
+
+Did not find it? :( So we can act quickly on it, please follow these steps:
+
+* Include your **OS type and version**, the versions of **Python** and **PyTorch**.
+* A short, self-contained, code snippet that allows us to reproduce the bug in
+  less than 30s;
+* Provide the *full* traceback if an exception is raised.
```suggestion
### Do you want a new feature?
```
@@ -0,0 +1,220 @@
+<!---
+Copyright 2020 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.
+-->
+
+# How to contribute to Accelerate?
+
+Everyone is welcome to contribute, and we value everybody's contribution. Code
+is thus not the only way to help the community. Answering questions, helping
+others, reaching out and improving the documentations are immensely valuable to
+the community.
+
+It also helps us if you spread the word: reference the library from blog posts
+on the awesome projects it made possible, shout out on Twitter every time it has
+helped you, or simply star the repo to say "thank you".
+
+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md).
+
+## You can contribute in so many ways!
+
+There are 4 ways you can contribute to Accelerate:
+* Fixing outstanding issues with the existing code;
+* Contributing to the examples or to the documentation;
+* Submitting issues related to bugs or desired new features.
+
+Please ping @sgugger on your contribution.
+
+## Submitting a new issue or feature request
+
+Do your best to follow these guidelines when submitting an issue or a feature
+request. It will make it easier for us to come back to you quickly and with good
+feedback.
+
+### Did you find a bug?
+
+The 🤗 Accelerate library is robust and reliable thanks to the users who notify us of
+the problems they encounter. So thank you for reporting an issue.
+
+First, we would really appreciate it if you could **make sure the bug was not
+already reported** (use the search bar on Github under Issues).
+
+Did not find it? :( So we can act quickly on it, please follow these steps:
+
+* Include your **OS type and version**, the versions of **Python** and **PyTorch**.
+* A short, self-contained, code snippet that allows us to reproduce the bug in
+  less than 30s;
+* Provide the *full* traceback if an exception is raised.
+
+### Do you want a new feature (that is not a model)?
```suggestion
🤗 Accelerate. `git` is not the easiest tool to use but it has the greatest
```
@@ -0,0 +1,220 @@
+<!---
+Copyright 2020 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.
+-->
+
+# How to contribute to Accelerate?
+
+Everyone is welcome to contribute, and we value everybody's contribution. Code
+is thus not the only way to help the community. Answering questions, helping
+others, reaching out and improving the documentations are immensely valuable to
+the community.
+
+It also helps us if you spread the word: reference the library from blog posts
+on the awesome projects it made possible, shout out on Twitter every time it has
+helped you, or simply star the repo to say "thank you".
+
+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md).
+
+## You can contribute in so many ways!
+
+There are 4 ways you can contribute to Accelerate:
+* Fixing outstanding issues with the existing code;
+* Contributing to the examples or to the documentation;
+* Submitting issues related to bugs or desired new features.
+
+Please ping @sgugger on your contribution.
+
+## Submitting a new issue or feature request
+
+Do your best to follow these guidelines when submitting an issue or a feature
+request. It will make it easier for us to come back to you quickly and with good
+feedback.
+
+### Did you find a bug?
+
+The 🤗 Accelerate library is robust and reliable thanks to the users who notify us of
+the problems they encounter. So thank you for reporting an issue.
+
+First, we would really appreciate it if you could **make sure the bug was not
+already reported** (use the search bar on Github under Issues).
+
+Did not find it? :( So we can act quickly on it, please follow these steps:
+
+* Include your **OS type and version**, the versions of **Python** and **PyTorch**.
+* A short, self-contained, code snippet that allows us to reproduce the bug in
+  less than 30s;
+* Provide the *full* traceback if an exception is raised.
+
+### Do you want a new feature (that is not a model)?
+
+A world-class feature request addresses the following points:
+
+1. Motivation first:
+* Is it related to a problem/frustration with the library? If so, please explain
+  why. Providing a code snippet that demonstrates the problem is best.
+* Is it related to something you would need for a project? We'd love to hear
+  about it!
+* Is it something you worked on and think could benefit the community?
+  Awesome! Tell us what problem it solved for you.
+2. Write a *full paragraph* describing the feature;
+3. Provide a **code snippet** that demonstrates its future use;
+4. In case this is related to a paper, please attach a link;
+5. Attach any additional information (drawings, screenshots, etc.) you think may help.
+
+If your issue is well written we're already 80% of the way there by the time you
+post it.
+
+## Start contributing! (Pull Requests)
+
+Before writing code, we strongly advise you to search through the existing PRs or
+issues to make sure that nobody is already working on the same thing. If you are
+unsure, it is always a good idea to open an issue to get some feedback.
+
+You will need basic `git` proficiency to be able to contribute to
+`accelerate`. `git` is not the easiest tool to use but it has the greatest
```suggestion
   $ pip install -e ".[quality]"
```
There is no dev install for now ;-)
@@ -0,0 +1,220 @@
+<!---
+Copyright 2020 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.
+-->
+
+# How to contribute to Accelerate?
+
+Everyone is welcome to contribute, and we value everybody's contribution. Code
+is thus not the only way to help the community. Answering questions, helping
+others, reaching out and improving the documentations are immensely valuable to
+the community.
+
+It also helps us if you spread the word: reference the library from blog posts
+on the awesome projects it made possible, shout out on Twitter every time it has
+helped you, or simply star the repo to say "thank you".
+
+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md).
+
+## You can contribute in so many ways!
+
+There are 4 ways you can contribute to Accelerate:
+* Fixing outstanding issues with the existing code;
+* Contributing to the examples or to the documentation;
+* Submitting issues related to bugs or desired new features.
+
+Please ping @sgugger on your contribution.
+
+## Submitting a new issue or feature request
+
+Do your best to follow these guidelines when submitting an issue or a feature
+request. It will make it easier for us to come back to you quickly and with good
+feedback.
+
+### Did you find a bug?
+
+The 🤗 Accelerate library is robust and reliable thanks to the users who notify us of
+the problems they encounter. So thank you for reporting an issue.
+
+First, we would really appreciate it if you could **make sure the bug was not
+already reported** (use the search bar on Github under Issues).
+
+Did not find it? :( So we can act quickly on it, please follow these steps:
+
+* Include your **OS type and version**, the versions of **Python** and **PyTorch**.
+* A short, self-contained, code snippet that allows us to reproduce the bug in
+  less than 30s;
+* Provide the *full* traceback if an exception is raised.
+
+### Do you want a new feature (that is not a model)?
+
+A world-class feature request addresses the following points:
+
+1. Motivation first:
+* Is it related to a problem/frustration with the library? If so, please explain
+  why. Providing a code snippet that demonstrates the problem is best.
+* Is it related to something you would need for a project? We'd love to hear
+  about it!
+* Is it something you worked on and think could benefit the community?
+  Awesome! Tell us what problem it solved for you.
+2. Write a *full paragraph* describing the feature;
+3. Provide a **code snippet** that demonstrates its future use;
+4. In case this is related to a paper, please attach a link;
+5. Attach any additional information (drawings, screenshots, etc.) you think may help.
+
+If your issue is well written we're already 80% of the way there by the time you
+post it.
+
+## Start contributing! (Pull Requests)
+
+Before writing code, we strongly advise you to search through the existing PRs or
+issues to make sure that nobody is already working on the same thing. If you are
+unsure, it is always a good idea to open an issue to get some feedback.
+
+You will need basic `git` proficiency to be able to contribute to
+`accelerate`. `git` is not the easiest tool to use but it has the greatest
+manual. Type `git --help` in a shell and enjoy. If you prefer books, [Pro
+Git](https://git-scm.com/book/en/v2) is a very good reference.
+
+Follow these steps to start contributing:
+
+1. Fork the [repository](https://github.com/huggingface/accelerate) by
+   clicking on the 'Fork' button on the repository's page. This creates a copy of the code
+   under your GitHub user account.
+
+2. Clone your fork to your local disk, and add the base repository as a remote:
+
+   ```bash
+   $ git clone git@github.com:<your Github handle>/accelerate.git
+   $ cd accelerate
+   $ git remote add upstream https://github.com/huggingface/accelerate.git
+   ```
+
+3. Create a new branch to hold your development changes:
+
+   ```bash
+   $ git checkout -b a-descriptive-name-for-my-changes
+   ```
+
+   **Do not** work on the `master` branch.
+
+4. Set up a development environment by running the following command in a virtual environment:
+
+   ```bash
+   $ pip install -e ".[dev]"
shouldn't this go under `LICENSE` file instead?
@@ -0,0 +1,220 @@
+<!---
+Copyright 2020 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.
+-->
```suggestion
Some of the ways you can contribute to Accelerate:
```
@@ -0,0 +1,220 @@
+<!---
+Copyright 2020 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.
+-->
+
+# How to contribute to Accelerate?
+
+Everyone is welcome to contribute, and we value everybody's contribution. Code
+is thus not the only way to help the community. Answering questions, helping
+others, reaching out and improving the documentations are immensely valuable to
+the community.
+
+It also helps us if you spread the word: reference the library from blog posts
+on the awesome projects it made possible, shout out on Twitter every time it has
+helped you, or simply star the repo to say "thank you".
+
+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md).
+
+## You can contribute in so many ways!
+
+There are 4 ways you can contribute to Accelerate:
Some or most of this information could also go to an issue template. The templates also support _form_ format ([example](https://github.com/scikit-learn/scikit-learn/issues/new?assignees=&labels=Bug%2CNeeds+Triage&template=bug_report.yml)) and people reporting a bug are probably not reading this document, but they would see that form or that template.
@@ -0,0 +1,220 @@
+<!---
+Copyright 2020 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.
+-->
+
+# How to contribute to Accelerate?
+
+Everyone is welcome to contribute, and we value everybody's contribution. Code
+is thus not the only way to help the community. Answering questions, helping
+others, reaching out and improving the documentations are immensely valuable to
+the community.
+
+It also helps us if you spread the word: reference the library from blog posts
+on the awesome projects it made possible, shout out on Twitter every time it has
+helped you, or simply star the repo to say "thank you".
+
+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md).
+
+## You can contribute in so many ways!
+
+There are 4 ways you can contribute to Accelerate:
+* Fixing outstanding issues with the existing code;
+* Contributing to the examples or to the documentation;
+* Submitting issues related to bugs or desired new features.
+
+Please ping @sgugger on your contribution.
+
+## Submitting a new issue or feature request
+
+Do your best to follow these guidelines when submitting an issue or a feature
+request. It will make it easier for us to come back to you quickly and with good
+feedback.
+
+### Did you find a bug?
we found that having a `get_versions()` which returns all relevant versions is helpful. Users would just paste the output of that in the issue.
@@ -0,0 +1,220 @@
+<!---
+Copyright 2020 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.
+-->
+
+# How to contribute to Accelerate?
+
+Everyone is welcome to contribute, and we value everybody's contribution. Code
+is thus not the only way to help the community. Answering questions, helping
+others, reaching out and improving the documentations are immensely valuable to
+the community.
+
+It also helps us if you spread the word: reference the library from blog posts
+on the awesome projects it made possible, shout out on Twitter every time it has
+helped you, or simply star the repo to say "thank you".
+
+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md).
+
+## You can contribute in so many ways!
+
+There are 4 ways you can contribute to Accelerate:
+* Fixing outstanding issues with the existing code;
+* Contributing to the examples or to the documentation;
+* Submitting issues related to bugs or desired new features.
+
+Please ping @sgugger on your contribution.
+
+## Submitting a new issue or feature request
+
+Do your best to follow these guidelines when submitting an issue or a feature
+request. It will make it easier for us to come back to you quickly and with good
+feedback.
+
+### Did you find a bug?
+
+The 🤗 Accelerate library is robust and reliable thanks to the users who notify us of
+the problems they encounter. So thank you for reporting an issue.
+
+First, we would really appreciate it if you could **make sure the bug was not
+already reported** (use the search bar on Github under Issues).
+
+Did not find it? :( So we can act quickly on it, please follow these steps:
+
+* Include your **OS type and version**, the versions of **Python** and **PyTorch**.
```suggestion
A good feature request addresses the following points:
```
@@ -0,0 +1,220 @@
+<!---
+Copyright 2020 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.
+-->
+
+# How to contribute to Accelerate?
+
+Everyone is welcome to contribute, and we value everybody's contribution. Code
+is thus not the only way to help the community. Answering questions, helping
+others, reaching out and improving the documentations are immensely valuable to
+the community.
+
+It also helps us if you spread the word: reference the library from blog posts
+on the awesome projects it made possible, shout out on Twitter every time it has
+helped you, or simply star the repo to say "thank you".
+
+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md).
+
+## You can contribute in so many ways!
+
+There are 4 ways you can contribute to Accelerate:
+* Fixing outstanding issues with the existing code;
+* Contributing to the examples or to the documentation;
+* Submitting issues related to bugs or desired new features.
+
+Please ping @sgugger on your contribution.
+
+## Submitting a new issue or feature request
+
+Do your best to follow these guidelines when submitting an issue or a feature
+request. It will make it easier for us to come back to you quickly and with good
+feedback.
+
+### Did you find a bug?
+
+The 🤗 Accelerate library is robust and reliable thanks to the users who notify us of
+the problems they encounter. So thank you for reporting an issue.
+
+First, we would really appreciate it if you could **make sure the bug was not
+already reported** (use the search bar on Github under Issues).
+
+Did not find it? :( So we can act quickly on it, please follow these steps:
+
+* Include your **OS type and version**, the versions of **Python** and **PyTorch**.
+* A short, self-contained, code snippet that allows us to reproduce the bug in
+  less than 30s;
+* Provide the *full* traceback if an exception is raised.
+
+### Do you want a new feature (that is not a model)?
+
+A world-class feature request addresses the following points:
```suggestion
## Submitting a pull request (PR)
```

"Start contributing" implies the other ones were not contributions.
@@ -0,0 +1,220 @@
+<!---
+Copyright 2020 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.
+-->
+
+# How to contribute to Accelerate?
+
+Everyone is welcome to contribute, and we value everybody's contribution. Code
+is thus not the only way to help the community. Answering questions, helping
+others, reaching out and improving the documentations are immensely valuable to
+the community.
+
+It also helps us if you spread the word: reference the library from blog posts
+on the awesome projects it made possible, shout out on Twitter every time it has
+helped you, or simply star the repo to say "thank you".
+
+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md).
+
+## You can contribute in so many ways!
+
+There are 4 ways you can contribute to Accelerate:
+* Fixing outstanding issues with the existing code;
+* Contributing to the examples or to the documentation;
+* Submitting issues related to bugs or desired new features.
+
+Please ping @sgugger on your contribution.
+
+## Submitting a new issue or feature request
+
+Do your best to follow these guidelines when submitting an issue or a feature
+request. It will make it easier for us to come back to you quickly and with good
+feedback.
+
+### Did you find a bug?
+
+The 🤗 Accelerate library is robust and reliable thanks to the users who notify us of
+the problems they encounter. So thank you for reporting an issue.
+
+First, we would really appreciate it if you could **make sure the bug was not
+already reported** (use the search bar on Github under Issues).
+
+Did not find it? :( So we can act quickly on it, please follow these steps:
+
+* Include your **OS type and version**, the versions of **Python** and **PyTorch**.
+* A short, self-contained, code snippet that allows us to reproduce the bug in
+  less than 30s;
+* Provide the *full* traceback if an exception is raised.
+
+### Do you want a new feature (that is not a model)?
+
+A world-class feature request addresses the following points:
+
+1. Motivation first:
+* Is it related to a problem/frustration with the library? If so, please explain
+  why. Providing a code snippet that demonstrates the problem is best.
+* Is it related to something you would need for a project? We'd love to hear
+  about it!
+* Is it something you worked on and think could benefit the community?
+  Awesome! Tell us what problem it solved for you.
+2. Write a *full paragraph* describing the feature;
+3. Provide a **code snippet** that demonstrates its future use;
+4. In case this is related to a paper, please attach a link;
+5. Attach any additional information (drawings, screenshots, etc.) you think may help.
+
+If your issue is well written we're already 80% of the way there by the time you
+post it.
+
+## Start contributing! (Pull Requests)
this assumes the user has their public ssh key uploaded to github, it deserves a note.
@@ -0,0 +1,220 @@
+<!---
+Copyright 2020 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.
+-->
+
+# How to contribute to Accelerate?
+
+Everyone is welcome to contribute, and we value everybody's contribution. Code
+is thus not the only way to help the community. Answering questions, helping
+others, reaching out and improving the documentations are immensely valuable to
+the community.
+
+It also helps us if you spread the word: reference the library from blog posts
+on the awesome projects it made possible, shout out on Twitter every time it has
+helped you, or simply star the repo to say "thank you".
+
+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md).
+
+## You can contribute in so many ways!
+
+There are 4 ways you can contribute to Accelerate:
+* Fixing outstanding issues with the existing code;
+* Contributing to the examples or to the documentation;
+* Submitting issues related to bugs or desired new features.
+
+Please ping @sgugger on your contribution.
+
+## Submitting a new issue or feature request
+
+Do your best to follow these guidelines when submitting an issue or a feature
+request. It will make it easier for us to come back to you quickly and with good
+feedback.
+
+### Did you find a bug?
+
+The 🤗 Accelerate library is robust and reliable thanks to the users who notify us of
+the problems they encounter. So thank you for reporting an issue.
+
+First, we would really appreciate it if you could **make sure the bug was not
+already reported** (use the search bar on Github under Issues).
+
+Did not find it? :( So we can act quickly on it, please follow these steps:
+
+* Include your **OS type and version**, the versions of **Python** and **PyTorch**.
+* A short, self-contained, code snippet that allows us to reproduce the bug in
+  less than 30s;
+* Provide the *full* traceback if an exception is raised.
+
+### Do you want a new feature (that is not a model)?
+
+A world-class feature request addresses the following points:
+
+1. Motivation first:
+* Is it related to a problem/frustration with the library? If so, please explain
+  why. Providing a code snippet that demonstrates the problem is best.
+* Is it related to something you would need for a project? We'd love to hear
+  about it!
+* Is it something you worked on and think could benefit the community?
+  Awesome! Tell us what problem it solved for you.
+2. Write a *full paragraph* describing the feature;
+3. Provide a **code snippet** that demonstrates its future use;
+4. In case this is related to a paper, please attach a link;
+5. Attach any additional information (drawings, screenshots, etc.) you think may help.
+
+If your issue is well written we're already 80% of the way there by the time you
+post it.
+
+## Start contributing! (Pull Requests)
+
+Before writing code, we strongly advise you to search through the existing PRs or
+issues to make sure that nobody is already working on the same thing. If you are
+unsure, it is always a good idea to open an issue to get some feedback.
+
+You will need basic `git` proficiency to be able to contribute to
+`accelerate`. `git` is not the easiest tool to use but it has the greatest
+manual. Type `git --help` in a shell and enjoy. If you prefer books, [Pro
+Git](https://git-scm.com/book/en/v2) is a very good reference.
+
+Follow these steps to start contributing:
+
+1. Fork the [repository](https://github.com/huggingface/accelerate) by
+   clicking on the 'Fork' button on the repository's page. This creates a copy of the code
+   under your GitHub user account.
+
+2. Clone your fork to your local disk, and add the base repository as a remote:
+
+   ```bash
+   $ git clone git@github.com:<your Github handle>/accelerate.git
```suggestion
3. Create a new branch to hold your development changes, and do this for every new PR you work on:
```
@@ -0,0 +1,220 @@
+<!---
+Copyright 2020 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.
+-->
+
+# How to contribute to Accelerate?
+
+Everyone is welcome to contribute, and we value everybody's contribution. Code
+is thus not the only way to help the community. Answering questions, helping
+others, reaching out and improving the documentations are immensely valuable to
+the community.
+
+It also helps us if you spread the word: reference the library from blog posts
+on the awesome projects it made possible, shout out on Twitter every time it has
+helped you, or simply star the repo to say "thank you".
+
+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md).
+
+## You can contribute in so many ways!
+
+There are 4 ways you can contribute to Accelerate:
+* Fixing outstanding issues with the existing code;
+* Contributing to the examples or to the documentation;
+* Submitting issues related to bugs or desired new features.
+
+Please ping @sgugger on your contribution.
+
+## Submitting a new issue or feature request
+
+Do your best to follow these guidelines when submitting an issue or a feature
+request. It will make it easier for us to come back to you quickly and with good
+feedback.
+
+### Did you find a bug?
+
+The 🤗 Accelerate library is robust and reliable thanks to the users who notify us of
+the problems they encounter. So thank you for reporting an issue.
+
+First, we would really appreciate it if you could **make sure the bug was not
+already reported** (use the search bar on Github under Issues).
+
+Did not find it? :( So we can act quickly on it, please follow these steps:
+
+* Include your **OS type and version**, the versions of **Python** and **PyTorch**.
+* A short, self-contained, code snippet that allows us to reproduce the bug in
+  less than 30s;
+* Provide the *full* traceback if an exception is raised.
+
+### Do you want a new feature (that is not a model)?
+
+A world-class feature request addresses the following points:
+
+1. Motivation first:
+* Is it related to a problem/frustration with the library? If so, please explain
+  why. Providing a code snippet that demonstrates the problem is best.
+* Is it related to something you would need for a project? We'd love to hear
+  about it!
+* Is it something you worked on and think could benefit the community?
+  Awesome! Tell us what problem it solved for you.
+2. Write a *full paragraph* describing the feature;
+3. Provide a **code snippet** that demonstrates its future use;
+4. In case this is related to a paper, please attach a link;
+5. Attach any additional information (drawings, screenshots, etc.) you think may help.
+
+If your issue is well written we're already 80% of the way there by the time you
+post it.
+
+## Start contributing! (Pull Requests)
+
+Before writing code, we strongly advise you to search through the existing PRs or
+issues to make sure that nobody is already working on the same thing. If you are
+unsure, it is always a good idea to open an issue to get some feedback.
+
+You will need basic `git` proficiency to be able to contribute to
+`accelerate`. `git` is not the easiest tool to use but it has the greatest
+manual. Type `git --help` in a shell and enjoy. If you prefer books, [Pro
+Git](https://git-scm.com/book/en/v2) is a very good reference.
+
+Follow these steps to start contributing:
+
+1. Fork the [repository](https://github.com/huggingface/accelerate) by
+   clicking on the 'Fork' button on the repository's page. This creates a copy of the code
+   under your GitHub user account.
+
+2. Clone your fork to your local disk, and add the base repository as a remote:
+
+   ```bash
+   $ git clone git@github.com:<your Github handle>/accelerate.git
+   $ cd accelerate
+   $ git remote add upstream https://github.com/huggingface/accelerate.git
+   ```
+
+3. Create a new branch to hold your development changes:
either checkout from `upstream/master`, or have them sync those two first, e.g.:

- Synchronize your ``main`` branch with the ``upstream/main`` branch,
   more details on `GitHub Docs <https://docs.github.com/en/github/collaborating-with-issues-and-pull-requests/syncing-a-fork>`_:

   .. prompt:: bash $

        git checkout main
        git fetch upstream
        git merge upstream/main

@@ -0,0 +1,220 @@
+<!---
+Copyright 2020 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.
+-->
+
+# How to contribute to Accelerate?
+
+Everyone is welcome to contribute, and we value everybody's contribution. Code
+is thus not the only way to help the community. Answering questions, helping
+others, reaching out and improving the documentations are immensely valuable to
+the community.
+
+It also helps us if you spread the word: reference the library from blog posts
+on the awesome projects it made possible, shout out on Twitter every time it has
+helped you, or simply star the repo to say "thank you".
+
+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md).
+
+## You can contribute in so many ways!
+
+There are 4 ways you can contribute to Accelerate:
+* Fixing outstanding issues with the existing code;
+* Contributing to the examples or to the documentation;
+* Submitting issues related to bugs or desired new features.
+
+Please ping @sgugger on your contribution.
+
+## Submitting a new issue or feature request
+
+Do your best to follow these guidelines when submitting an issue or a feature
+request. It will make it easier for us to come back to you quickly and with good
+feedback.
+
+### Did you find a bug?
+
+The 🤗 Accelerate library is robust and reliable thanks to the users who notify us of
+the problems they encounter. So thank you for reporting an issue.
+
+First, we would really appreciate it if you could **make sure the bug was not
+already reported** (use the search bar on Github under Issues).
+
+Did not find it? :( So we can act quickly on it, please follow these steps:
+
+* Include your **OS type and version**, the versions of **Python** and **PyTorch**.
+* A short, self-contained, code snippet that allows us to reproduce the bug in
+  less than 30s;
+* Provide the *full* traceback if an exception is raised.
+
+### Do you want a new feature (that is not a model)?
+
+A world-class feature request addresses the following points:
+
+1. Motivation first:
+* Is it related to a problem/frustration with the library? If so, please explain
+  why. Providing a code snippet that demonstrates the problem is best.
+* Is it related to something you would need for a project? We'd love to hear
+  about it!
+* Is it something you worked on and think could benefit the community?
+  Awesome! Tell us what problem it solved for you.
+2. Write a *full paragraph* describing the feature;
+3. Provide a **code snippet** that demonstrates its future use;
+4. In case this is related to a paper, please attach a link;
+5. Attach any additional information (drawings, screenshots, etc.) you think may help.
+
+If your issue is well written we're already 80% of the way there by the time you
+post it.
+
+## Start contributing! (Pull Requests)
+
+Before writing code, we strongly advise you to search through the existing PRs or
+issues to make sure that nobody is already working on the same thing. If you are
+unsure, it is always a good idea to open an issue to get some feedback.
+
+You will need basic `git` proficiency to be able to contribute to
+`accelerate`. `git` is not the easiest tool to use but it has the greatest
+manual. Type `git --help` in a shell and enjoy. If you prefer books, [Pro
+Git](https://git-scm.com/book/en/v2) is a very good reference.
+
+Follow these steps to start contributing:
+
+1. Fork the [repository](https://github.com/huggingface/accelerate) by
+   clicking on the 'Fork' button on the repository's page. This creates a copy of the code
+   under your GitHub user account.
+
+2. Clone your fork to your local disk, and add the base repository as a remote:
+
+   ```bash
+   $ git clone git@github.com:<your Github handle>/accelerate.git
+   $ cd accelerate
+   $ git remote add upstream https://github.com/huggingface/accelerate.git
+   ```
+
+3. Create a new branch to hold your development changes:
+
+   ```bash
+   $ git checkout -b a-descriptive-name-for-my-changes
```suggestion
4. Set up a development environment by running the following command in a conda or a virtual environment you've created for working on this library:
```
@@ -0,0 +1,220 @@
+<!---
+Copyright 2020 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.
+-->
+
+# How to contribute to Accelerate?
+
+Everyone is welcome to contribute, and we value everybody's contribution. Code
+is thus not the only way to help the community. Answering questions, helping
+others, reaching out and improving the documentations are immensely valuable to
+the community.
+
+It also helps us if you spread the word: reference the library from blog posts
+on the awesome projects it made possible, shout out on Twitter every time it has
+helped you, or simply star the repo to say "thank you".
+
+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md).
+
+## You can contribute in so many ways!
+
+There are 4 ways you can contribute to Accelerate:
+* Fixing outstanding issues with the existing code;
+* Contributing to the examples or to the documentation;
+* Submitting issues related to bugs or desired new features.
+
+Please ping @sgugger on your contribution.
+
+## Submitting a new issue or feature request
+
+Do your best to follow these guidelines when submitting an issue or a feature
+request. It will make it easier for us to come back to you quickly and with good
+feedback.
+
+### Did you find a bug?
+
+The 🤗 Accelerate library is robust and reliable thanks to the users who notify us of
+the problems they encounter. So thank you for reporting an issue.
+
+First, we would really appreciate it if you could **make sure the bug was not
+already reported** (use the search bar on Github under Issues).
+
+Did not find it? :( So we can act quickly on it, please follow these steps:
+
+* Include your **OS type and version**, the versions of **Python** and **PyTorch**.
+* A short, self-contained, code snippet that allows us to reproduce the bug in
+  less than 30s;
+* Provide the *full* traceback if an exception is raised.
+
+### Do you want a new feature (that is not a model)?
+
+A world-class feature request addresses the following points:
+
+1. Motivation first:
+* Is it related to a problem/frustration with the library? If so, please explain
+  why. Providing a code snippet that demonstrates the problem is best.
+* Is it related to something you would need for a project? We'd love to hear
+  about it!
+* Is it something you worked on and think could benefit the community?
+  Awesome! Tell us what problem it solved for you.
+2. Write a *full paragraph* describing the feature;
+3. Provide a **code snippet** that demonstrates its future use;
+4. In case this is related to a paper, please attach a link;
+5. Attach any additional information (drawings, screenshots, etc.) you think may help.
+
+If your issue is well written we're already 80% of the way there by the time you
+post it.
+
+## Start contributing! (Pull Requests)
+
+Before writing code, we strongly advise you to search through the existing PRs or
+issues to make sure that nobody is already working on the same thing. If you are
+unsure, it is always a good idea to open an issue to get some feedback.
+
+You will need basic `git` proficiency to be able to contribute to
+`accelerate`. `git` is not the easiest tool to use but it has the greatest
+manual. Type `git --help` in a shell and enjoy. If you prefer books, [Pro
+Git](https://git-scm.com/book/en/v2) is a very good reference.
+
+Follow these steps to start contributing:
+
+1. Fork the [repository](https://github.com/huggingface/accelerate) by
+   clicking on the 'Fork' button on the repository's page. This creates a copy of the code
+   under your GitHub user account.
+
+2. Clone your fork to your local disk, and add the base repository as a remote:
+
+   ```bash
+   $ git clone git@github.com:<your Github handle>/accelerate.git
+   $ cd accelerate
+   $ git remote add upstream https://github.com/huggingface/accelerate.git
+   ```
+
+3. Create a new branch to hold your development changes:
+
+   ```bash
+   $ git checkout -b a-descriptive-name-for-my-changes
+   ```
+
+   **Do not** work on the `master` branch.
+
+4. Set up a development environment by running the following command in a virtual environment:
on a windows machine they'd need a git bash or similar to run this I think.
@@ -0,0 +1,220 @@
+<!---
+Copyright 2020 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.
+-->
+
+# How to contribute to Accelerate?
+
+Everyone is welcome to contribute, and we value everybody's contribution. Code
+is thus not the only way to help the community. Answering questions, helping
+others, reaching out and improving the documentations are immensely valuable to
+the community.
+
+It also helps us if you spread the word: reference the library from blog posts
+on the awesome projects it made possible, shout out on Twitter every time it has
+helped you, or simply star the repo to say "thank you".
+
+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md).
+
+## You can contribute in so many ways!
+
+There are 4 ways you can contribute to Accelerate:
+* Fixing outstanding issues with the existing code;
+* Contributing to the examples or to the documentation;
+* Submitting issues related to bugs or desired new features.
+
+Please ping @sgugger on your contribution.
+
+## Submitting a new issue or feature request
+
+Do your best to follow these guidelines when submitting an issue or a feature
+request. It will make it easier for us to come back to you quickly and with good
+feedback.
+
+### Did you find a bug?
+
+The 🤗 Accelerate library is robust and reliable thanks to the users who notify us of
+the problems they encounter. So thank you for reporting an issue.
+
+First, we would really appreciate it if you could **make sure the bug was not
+already reported** (use the search bar on Github under Issues).
+
+Did not find it? :( So we can act quickly on it, please follow these steps:
+
+* Include your **OS type and version**, the versions of **Python** and **PyTorch**.
+* A short, self-contained, code snippet that allows us to reproduce the bug in
+  less than 30s;
+* Provide the *full* traceback if an exception is raised.
+
+### Do you want a new feature (that is not a model)?
+
+A world-class feature request addresses the following points:
+
+1. Motivation first:
+* Is it related to a problem/frustration with the library? If so, please explain
+  why. Providing a code snippet that demonstrates the problem is best.
+* Is it related to something you would need for a project? We'd love to hear
+  about it!
+* Is it something you worked on and think could benefit the community?
+  Awesome! Tell us what problem it solved for you.
+2. Write a *full paragraph* describing the feature;
+3. Provide a **code snippet** that demonstrates its future use;
+4. In case this is related to a paper, please attach a link;
+5. Attach any additional information (drawings, screenshots, etc.) you think may help.
+
+If your issue is well written we're already 80% of the way there by the time you
+post it.
+
+## Start contributing! (Pull Requests)
+
+Before writing code, we strongly advise you to search through the existing PRs or
+issues to make sure that nobody is already working on the same thing. If you are
+unsure, it is always a good idea to open an issue to get some feedback.
+
+You will need basic `git` proficiency to be able to contribute to
+`accelerate`. `git` is not the easiest tool to use but it has the greatest
+manual. Type `git --help` in a shell and enjoy. If you prefer books, [Pro
+Git](https://git-scm.com/book/en/v2) is a very good reference.
+
+Follow these steps to start contributing:
+
+1. Fork the [repository](https://github.com/huggingface/accelerate) by
+   clicking on the 'Fork' button on the repository's page. This creates a copy of the code
+   under your GitHub user account.
+
+2. Clone your fork to your local disk, and add the base repository as a remote:
+
+   ```bash
+   $ git clone git@github.com:<your Github handle>/accelerate.git
+   $ cd accelerate
+   $ git remote add upstream https://github.com/huggingface/accelerate.git
+   ```
+
+3. Create a new branch to hold your development changes:
+
+   ```bash
+   $ git checkout -b a-descriptive-name-for-my-changes
+   ```
+
+   **Do not** work on the `master` branch.
+
+4. Set up a development environment by running the following command in a virtual environment:
+
+   ```bash
+   $ pip install -e ".[dev]"
+   ```
+
+   (If accelerate was already installed in the virtual environment, remove
+   it with `pip uninstall accelerate` before reinstalling it in editable
+   mode with the `-e` flag.)
+
+5. Develop the features on your branch.
+
+   As you work on the features, you should make sure that the test suite
+   passes. You should run the tests impacted by your changes like this (see 
+   below an explanation regarding the environment variable):
+
+   ```bash
+   $ pytest tests/<TEST_TO_RUN>.py
+   ```
+
+   You can also run the full suite with the following command.
+
+   ```bash
+   $ make test
+   ```
+
+   `accelerate` relies on `black` and `isort` to format its source code
+   consistently. After you make changes, apply automatic style corrections and code verifications
+   that can't be automated in one go with:
+
+   This target is also optimized to only work with files modified by the PR you're working on.
+
+   If you prefer to run the checks one after the other, the following command apply the
+   style corrections:
+
+   ```bash
+   $ make style
you could also add them as git hooks, which have helped us with users' contributions.
@@ -0,0 +1,220 @@
+<!---
+Copyright 2020 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.
+-->
+
+# How to contribute to Accelerate?
+
+Everyone is welcome to contribute, and we value everybody's contribution. Code
+is thus not the only way to help the community. Answering questions, helping
+others, reaching out and improving the documentations are immensely valuable to
+the community.
+
+It also helps us if you spread the word: reference the library from blog posts
+on the awesome projects it made possible, shout out on Twitter every time it has
+helped you, or simply star the repo to say "thank you".
+
+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md).
+
+## You can contribute in so many ways!
+
+There are 4 ways you can contribute to Accelerate:
+* Fixing outstanding issues with the existing code;
+* Contributing to the examples or to the documentation;
+* Submitting issues related to bugs or desired new features.
+
+Please ping @sgugger on your contribution.
+
+## Submitting a new issue or feature request
+
+Do your best to follow these guidelines when submitting an issue or a feature
+request. It will make it easier for us to come back to you quickly and with good
+feedback.
+
+### Did you find a bug?
+
+The 🤗 Accelerate library is robust and reliable thanks to the users who notify us of
+the problems they encounter. So thank you for reporting an issue.
+
+First, we would really appreciate it if you could **make sure the bug was not
+already reported** (use the search bar on Github under Issues).
+
+Did not find it? :( So we can act quickly on it, please follow these steps:
+
+* Include your **OS type and version**, the versions of **Python** and **PyTorch**.
+* A short, self-contained, code snippet that allows us to reproduce the bug in
+  less than 30s;
+* Provide the *full* traceback if an exception is raised.
+
+### Do you want a new feature (that is not a model)?
+
+A world-class feature request addresses the following points:
+
+1. Motivation first:
+* Is it related to a problem/frustration with the library? If so, please explain
+  why. Providing a code snippet that demonstrates the problem is best.
+* Is it related to something you would need for a project? We'd love to hear
+  about it!
+* Is it something you worked on and think could benefit the community?
+  Awesome! Tell us what problem it solved for you.
+2. Write a *full paragraph* describing the feature;
+3. Provide a **code snippet** that demonstrates its future use;
+4. In case this is related to a paper, please attach a link;
+5. Attach any additional information (drawings, screenshots, etc.) you think may help.
+
+If your issue is well written we're already 80% of the way there by the time you
+post it.
+
+## Start contributing! (Pull Requests)
+
+Before writing code, we strongly advise you to search through the existing PRs or
+issues to make sure that nobody is already working on the same thing. If you are
+unsure, it is always a good idea to open an issue to get some feedback.
+
+You will need basic `git` proficiency to be able to contribute to
+`accelerate`. `git` is not the easiest tool to use but it has the greatest
+manual. Type `git --help` in a shell and enjoy. If you prefer books, [Pro
+Git](https://git-scm.com/book/en/v2) is a very good reference.
+
+Follow these steps to start contributing:
+
+1. Fork the [repository](https://github.com/huggingface/accelerate) by
+   clicking on the 'Fork' button on the repository's page. This creates a copy of the code
+   under your GitHub user account.
+
+2. Clone your fork to your local disk, and add the base repository as a remote:
+
+   ```bash
+   $ git clone git@github.com:<your Github handle>/accelerate.git
+   $ cd accelerate
+   $ git remote add upstream https://github.com/huggingface/accelerate.git
+   ```
+
+3. Create a new branch to hold your development changes:
+
+   ```bash
+   $ git checkout -b a-descriptive-name-for-my-changes
+   ```
+
+   **Do not** work on the `master` branch.
+
+4. Set up a development environment by running the following command in a virtual environment:
+
+   ```bash
+   $ pip install -e ".[dev]"
+   ```
+
+   (If accelerate was already installed in the virtual environment, remove
+   it with `pip uninstall accelerate` before reinstalling it in editable
+   mode with the `-e` flag.)
+
+5. Develop the features on your branch.
+
+   As you work on the features, you should make sure that the test suite
+   passes. You should run the tests impacted by your changes like this (see 
+   below an explanation regarding the environment variable):
+
+   ```bash
+   $ pytest tests/<TEST_TO_RUN>.py
+   ```
+
+   You can also run the full suite with the following command.
+
+   ```bash
+   $ make test
+   ```
+
+   `accelerate` relies on `black` and `isort` to format its source code
+   consistently. After you make changes, apply automatic style corrections and code verifications
+   that can't be automated in one go with:
+
+   This target is also optimized to only work with files modified by the PR you're working on.
+
+   If you prefer to run the checks one after the other, the following command apply the
+   style corrections:
+
+   ```bash
+   $ make style
+   ```
+
+   `accelerate` also uses `flake8` and a few custom scripts to check for coding mistakes. Quality
+   control runs in CI, however you can also run the same checks with:
in longer PRs (time-wise), rebase can be very tricky for beginner users, merge usually makes it much easier for them, and if you're squash and merging, the merge commit wouldn't show up in your history.
@@ -0,0 +1,220 @@
+<!---
+Copyright 2020 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.
+-->
+
+# How to contribute to Accelerate?
+
+Everyone is welcome to contribute, and we value everybody's contribution. Code
+is thus not the only way to help the community. Answering questions, helping
+others, reaching out and improving the documentations are immensely valuable to
+the community.
+
+It also helps us if you spread the word: reference the library from blog posts
+on the awesome projects it made possible, shout out on Twitter every time it has
+helped you, or simply star the repo to say "thank you".
+
+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md).
+
+## You can contribute in so many ways!
+
+There are 4 ways you can contribute to Accelerate:
+* Fixing outstanding issues with the existing code;
+* Contributing to the examples or to the documentation;
+* Submitting issues related to bugs or desired new features.
+
+Please ping @sgugger on your contribution.
+
+## Submitting a new issue or feature request
+
+Do your best to follow these guidelines when submitting an issue or a feature
+request. It will make it easier for us to come back to you quickly and with good
+feedback.
+
+### Did you find a bug?
+
+The 🤗 Accelerate library is robust and reliable thanks to the users who notify us of
+the problems they encounter. So thank you for reporting an issue.
+
+First, we would really appreciate it if you could **make sure the bug was not
+already reported** (use the search bar on Github under Issues).
+
+Did not find it? :( So we can act quickly on it, please follow these steps:
+
+* Include your **OS type and version**, the versions of **Python** and **PyTorch**.
+* A short, self-contained, code snippet that allows us to reproduce the bug in
+  less than 30s;
+* Provide the *full* traceback if an exception is raised.
+
+### Do you want a new feature (that is not a model)?
+
+A world-class feature request addresses the following points:
+
+1. Motivation first:
+* Is it related to a problem/frustration with the library? If so, please explain
+  why. Providing a code snippet that demonstrates the problem is best.
+* Is it related to something you would need for a project? We'd love to hear
+  about it!
+* Is it something you worked on and think could benefit the community?
+  Awesome! Tell us what problem it solved for you.
+2. Write a *full paragraph* describing the feature;
+3. Provide a **code snippet** that demonstrates its future use;
+4. In case this is related to a paper, please attach a link;
+5. Attach any additional information (drawings, screenshots, etc.) you think may help.
+
+If your issue is well written we're already 80% of the way there by the time you
+post it.
+
+## Start contributing! (Pull Requests)
+
+Before writing code, we strongly advise you to search through the existing PRs or
+issues to make sure that nobody is already working on the same thing. If you are
+unsure, it is always a good idea to open an issue to get some feedback.
+
+You will need basic `git` proficiency to be able to contribute to
+`accelerate`. `git` is not the easiest tool to use but it has the greatest
+manual. Type `git --help` in a shell and enjoy. If you prefer books, [Pro
+Git](https://git-scm.com/book/en/v2) is a very good reference.
+
+Follow these steps to start contributing:
+
+1. Fork the [repository](https://github.com/huggingface/accelerate) by
+   clicking on the 'Fork' button on the repository's page. This creates a copy of the code
+   under your GitHub user account.
+
+2. Clone your fork to your local disk, and add the base repository as a remote:
+
+   ```bash
+   $ git clone git@github.com:<your Github handle>/accelerate.git
+   $ cd accelerate
+   $ git remote add upstream https://github.com/huggingface/accelerate.git
+   ```
+
+3. Create a new branch to hold your development changes:
+
+   ```bash
+   $ git checkout -b a-descriptive-name-for-my-changes
+   ```
+
+   **Do not** work on the `master` branch.
+
+4. Set up a development environment by running the following command in a virtual environment:
+
+   ```bash
+   $ pip install -e ".[dev]"
+   ```
+
+   (If accelerate was already installed in the virtual environment, remove
+   it with `pip uninstall accelerate` before reinstalling it in editable
+   mode with the `-e` flag.)
+
+5. Develop the features on your branch.
+
+   As you work on the features, you should make sure that the test suite
+   passes. You should run the tests impacted by your changes like this (see 
+   below an explanation regarding the environment variable):
+
+   ```bash
+   $ pytest tests/<TEST_TO_RUN>.py
+   ```
+
+   You can also run the full suite with the following command.
+
+   ```bash
+   $ make test
+   ```
+
+   `accelerate` relies on `black` and `isort` to format its source code
+   consistently. After you make changes, apply automatic style corrections and code verifications
+   that can't be automated in one go with:
+
+   This target is also optimized to only work with files modified by the PR you're working on.
+
+   If you prefer to run the checks one after the other, the following command apply the
+   style corrections:
+
+   ```bash
+   $ make style
+   ```
+
+   `accelerate` also uses `flake8` and a few custom scripts to check for coding mistakes. Quality
+   control runs in CI, however you can also run the same checks with:
+
+   ```bash
+   $ make quality
+   ```
+
+   Once you're happy with your changes, add changed files using `git add` and
+   make a commit with `git commit` to record your changes locally:
+
+   ```bash
+   $ git add modified_file.py
+   $ git commit
+   ```
+
+   Please write [good commit messages](https://chris.beams.io/posts/git-commit/).
+
+   It is a good idea to sync your copy of the code with the original
+   repository regularly. This way you can quickly account for changes:
+
+   ```bash
+   $ git fetch upstream
+   $ git rebase upstream/master
or create a _draft_ PR?
@@ -0,0 +1,220 @@
+<!---
+Copyright 2020 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.
+-->
+
+# How to contribute to Accelerate?
+
+Everyone is welcome to contribute, and we value everybody's contribution. Code
+is thus not the only way to help the community. Answering questions, helping
+others, reaching out and improving the documentations are immensely valuable to
+the community.
+
+It also helps us if you spread the word: reference the library from blog posts
+on the awesome projects it made possible, shout out on Twitter every time it has
+helped you, or simply star the repo to say "thank you".
+
+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md).
+
+## You can contribute in so many ways!
+
+There are 4 ways you can contribute to Accelerate:
+* Fixing outstanding issues with the existing code;
+* Contributing to the examples or to the documentation;
+* Submitting issues related to bugs or desired new features.
+
+Please ping @sgugger on your contribution.
+
+## Submitting a new issue or feature request
+
+Do your best to follow these guidelines when submitting an issue or a feature
+request. It will make it easier for us to come back to you quickly and with good
+feedback.
+
+### Did you find a bug?
+
+The 🤗 Accelerate library is robust and reliable thanks to the users who notify us of
+the problems they encounter. So thank you for reporting an issue.
+
+First, we would really appreciate it if you could **make sure the bug was not
+already reported** (use the search bar on Github under Issues).
+
+Did not find it? :( So we can act quickly on it, please follow these steps:
+
+* Include your **OS type and version**, the versions of **Python** and **PyTorch**.
+* A short, self-contained, code snippet that allows us to reproduce the bug in
+  less than 30s;
+* Provide the *full* traceback if an exception is raised.
+
+### Do you want a new feature (that is not a model)?
+
+A world-class feature request addresses the following points:
+
+1. Motivation first:
+* Is it related to a problem/frustration with the library? If so, please explain
+  why. Providing a code snippet that demonstrates the problem is best.
+* Is it related to something you would need for a project? We'd love to hear
+  about it!
+* Is it something you worked on and think could benefit the community?
+  Awesome! Tell us what problem it solved for you.
+2. Write a *full paragraph* describing the feature;
+3. Provide a **code snippet** that demonstrates its future use;
+4. In case this is related to a paper, please attach a link;
+5. Attach any additional information (drawings, screenshots, etc.) you think may help.
+
+If your issue is well written we're already 80% of the way there by the time you
+post it.
+
+## Start contributing! (Pull Requests)
+
+Before writing code, we strongly advise you to search through the existing PRs or
+issues to make sure that nobody is already working on the same thing. If you are
+unsure, it is always a good idea to open an issue to get some feedback.
+
+You will need basic `git` proficiency to be able to contribute to
+`accelerate`. `git` is not the easiest tool to use but it has the greatest
+manual. Type `git --help` in a shell and enjoy. If you prefer books, [Pro
+Git](https://git-scm.com/book/en/v2) is a very good reference.
+
+Follow these steps to start contributing:
+
+1. Fork the [repository](https://github.com/huggingface/accelerate) by
+   clicking on the 'Fork' button on the repository's page. This creates a copy of the code
+   under your GitHub user account.
+
+2. Clone your fork to your local disk, and add the base repository as a remote:
+
+   ```bash
+   $ git clone git@github.com:<your Github handle>/accelerate.git
+   $ cd accelerate
+   $ git remote add upstream https://github.com/huggingface/accelerate.git
+   ```
+
+3. Create a new branch to hold your development changes:
+
+   ```bash
+   $ git checkout -b a-descriptive-name-for-my-changes
+   ```
+
+   **Do not** work on the `master` branch.
+
+4. Set up a development environment by running the following command in a virtual environment:
+
+   ```bash
+   $ pip install -e ".[dev]"
+   ```
+
+   (If accelerate was already installed in the virtual environment, remove
+   it with `pip uninstall accelerate` before reinstalling it in editable
+   mode with the `-e` flag.)
+
+5. Develop the features on your branch.
+
+   As you work on the features, you should make sure that the test suite
+   passes. You should run the tests impacted by your changes like this (see 
+   below an explanation regarding the environment variable):
+
+   ```bash
+   $ pytest tests/<TEST_TO_RUN>.py
+   ```
+
+   You can also run the full suite with the following command.
+
+   ```bash
+   $ make test
+   ```
+
+   `accelerate` relies on `black` and `isort` to format its source code
+   consistently. After you make changes, apply automatic style corrections and code verifications
+   that can't be automated in one go with:
+
+   This target is also optimized to only work with files modified by the PR you're working on.
+
+   If you prefer to run the checks one after the other, the following command apply the
+   style corrections:
+
+   ```bash
+   $ make style
+   ```
+
+   `accelerate` also uses `flake8` and a few custom scripts to check for coding mistakes. Quality
+   control runs in CI, however you can also run the same checks with:
+
+   ```bash
+   $ make quality
+   ```
+
+   Once you're happy with your changes, add changed files using `git add` and
+   make a commit with `git commit` to record your changes locally:
+
+   ```bash
+   $ git add modified_file.py
+   $ git commit
+   ```
+
+   Please write [good commit messages](https://chris.beams.io/posts/git-commit/).
+
+   It is a good idea to sync your copy of the code with the original
+   repository regularly. This way you can quickly account for changes:
+
+   ```bash
+   $ git fetch upstream
+   $ git rebase upstream/master
+   ```
+
+   Push the changes to your account using:
+
+   ```bash
+   $ git push -u origin a-descriptive-name-for-my-changes
+   ```
+
+6. Once you are satisfied (**and the checklist below is happy too**), go to the
+   webpage of your fork on GitHub. Click on 'Pull request' to send your changes
+   to the project maintainers for review.
+
+7. It's ok if maintainers ask you for changes. It happens to core contributors
+   too! So everyone can see the changes in the Pull request, work in your local
+   branch and push the changes to your fork. They will automatically appear in
+   the pull request.
+
+
+### Checklist
+
+1. The title of your pull request should be a summary of its contribution;
+2. If your pull request addresses an issue, please mention the issue number in
+   the pull request description to make sure they are linked (and people
+   consulting the issue know you are working on it);
+3. To indicate a work in progress please prefix the title with `[WIP]`. These
+   are useful to avoid duplicated work, and to differentiate it from PRs ready
+   to be merged;
adding a link to a _good_ PR may help them by seeing a good example.
@@ -0,0 +1,220 @@
+<!---
+Copyright 2020 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.
+-->
+
+# How to contribute to Accelerate?
+
+Everyone is welcome to contribute, and we value everybody's contribution. Code
+is thus not the only way to help the community. Answering questions, helping
+others, reaching out and improving the documentations are immensely valuable to
+the community.
+
+It also helps us if you spread the word: reference the library from blog posts
+on the awesome projects it made possible, shout out on Twitter every time it has
+helped you, or simply star the repo to say "thank you".
+
+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md).
+
+## You can contribute in so many ways!
+
+There are 4 ways you can contribute to Accelerate:
+* Fixing outstanding issues with the existing code;
+* Contributing to the examples or to the documentation;
+* Submitting issues related to bugs or desired new features.
+
+Please ping @sgugger on your contribution.
+
+## Submitting a new issue or feature request
+
+Do your best to follow these guidelines when submitting an issue or a feature
+request. It will make it easier for us to come back to you quickly and with good
+feedback.
+
+### Did you find a bug?
+
+The 🤗 Accelerate library is robust and reliable thanks to the users who notify us of
+the problems they encounter. So thank you for reporting an issue.
+
+First, we would really appreciate it if you could **make sure the bug was not
+already reported** (use the search bar on Github under Issues).
+
+Did not find it? :( So we can act quickly on it, please follow these steps:
+
+* Include your **OS type and version**, the versions of **Python** and **PyTorch**.
+* A short, self-contained, code snippet that allows us to reproduce the bug in
+  less than 30s;
+* Provide the *full* traceback if an exception is raised.
+
+### Do you want a new feature (that is not a model)?
+
+A world-class feature request addresses the following points:
+
+1. Motivation first:
+* Is it related to a problem/frustration with the library? If so, please explain
+  why. Providing a code snippet that demonstrates the problem is best.
+* Is it related to something you would need for a project? We'd love to hear
+  about it!
+* Is it something you worked on and think could benefit the community?
+  Awesome! Tell us what problem it solved for you.
+2. Write a *full paragraph* describing the feature;
+3. Provide a **code snippet** that demonstrates its future use;
+4. In case this is related to a paper, please attach a link;
+5. Attach any additional information (drawings, screenshots, etc.) you think may help.
+
+If your issue is well written we're already 80% of the way there by the time you
+post it.
+
+## Start contributing! (Pull Requests)
+
+Before writing code, we strongly advise you to search through the existing PRs or
+issues to make sure that nobody is already working on the same thing. If you are
+unsure, it is always a good idea to open an issue to get some feedback.
+
+You will need basic `git` proficiency to be able to contribute to
+`accelerate`. `git` is not the easiest tool to use but it has the greatest
+manual. Type `git --help` in a shell and enjoy. If you prefer books, [Pro
+Git](https://git-scm.com/book/en/v2) is a very good reference.
+
+Follow these steps to start contributing:
+
+1. Fork the [repository](https://github.com/huggingface/accelerate) by
+   clicking on the 'Fork' button on the repository's page. This creates a copy of the code
+   under your GitHub user account.
+
+2. Clone your fork to your local disk, and add the base repository as a remote:
+
+   ```bash
+   $ git clone git@github.com:<your Github handle>/accelerate.git
+   $ cd accelerate
+   $ git remote add upstream https://github.com/huggingface/accelerate.git
+   ```
+
+3. Create a new branch to hold your development changes:
+
+   ```bash
+   $ git checkout -b a-descriptive-name-for-my-changes
+   ```
+
+   **Do not** work on the `master` branch.
+
+4. Set up a development environment by running the following command in a virtual environment:
+
+   ```bash
+   $ pip install -e ".[dev]"
+   ```
+
+   (If accelerate was already installed in the virtual environment, remove
+   it with `pip uninstall accelerate` before reinstalling it in editable
+   mode with the `-e` flag.)
+
+5. Develop the features on your branch.
+
+   As you work on the features, you should make sure that the test suite
+   passes. You should run the tests impacted by your changes like this (see 
+   below an explanation regarding the environment variable):
+
+   ```bash
+   $ pytest tests/<TEST_TO_RUN>.py
+   ```
+
+   You can also run the full suite with the following command.
+
+   ```bash
+   $ make test
+   ```
+
+   `accelerate` relies on `black` and `isort` to format its source code
+   consistently. After you make changes, apply automatic style corrections and code verifications
+   that can't be automated in one go with:
+
+   This target is also optimized to only work with files modified by the PR you're working on.
+
+   If you prefer to run the checks one after the other, the following command apply the
+   style corrections:
+
+   ```bash
+   $ make style
+   ```
+
+   `accelerate` also uses `flake8` and a few custom scripts to check for coding mistakes. Quality
+   control runs in CI, however you can also run the same checks with:
+
+   ```bash
+   $ make quality
+   ```
+
+   Once you're happy with your changes, add changed files using `git add` and
+   make a commit with `git commit` to record your changes locally:
+
+   ```bash
+   $ git add modified_file.py
+   $ git commit
+   ```
+
+   Please write [good commit messages](https://chris.beams.io/posts/git-commit/).
+
+   It is a good idea to sync your copy of the code with the original
+   repository regularly. This way you can quickly account for changes:
+
+   ```bash
+   $ git fetch upstream
+   $ git rebase upstream/master
+   ```
+
+   Push the changes to your account using:
+
+   ```bash
+   $ git push -u origin a-descriptive-name-for-my-changes
+   ```
+
+6. Once you are satisfied (**and the checklist below is happy too**), go to the
+   webpage of your fork on GitHub. Click on 'Pull request' to send your changes
+   to the project maintainers for review.
+
+7. It's ok if maintainers ask you for changes. It happens to core contributors
+   too! So everyone can see the changes in the Pull request, work in your local
+   branch and push the changes to your fork. They will automatically appear in
+   the pull request.
+
+
+### Checklist
There is also a `LICENSE` at the root of the repo. We have an additional license at the top of every file, with the appropriate copyrights for the code in that specific file. If some code was taken from somewhere else or from another org, we add a mention to that org in the copyright here.
@@ -0,0 +1,220 @@
+<!---
+Copyright 2020 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.
+-->
Sounds good. Ideally this would be done in a follow-up PR.
@@ -0,0 +1,220 @@
+<!---
+Copyright 2020 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.
+-->
+
+# How to contribute to Accelerate?
+
+Everyone is welcome to contribute, and we value everybody's contribution. Code
+is thus not the only way to help the community. Answering questions, helping
+others, reaching out and improving the documentations are immensely valuable to
+the community.
+
+It also helps us if you spread the word: reference the library from blog posts
+on the awesome projects it made possible, shout out on Twitter every time it has
+helped you, or simply star the repo to say "thank you".
+
+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md).
+
+## You can contribute in so many ways!
+
+There are 4 ways you can contribute to Accelerate:
+* Fixing outstanding issues with the existing code;
+* Contributing to the examples or to the documentation;
+* Submitting issues related to bugs or desired new features.
+
+Please ping @sgugger on your contribution.
+
+## Submitting a new issue or feature request
+
+Do your best to follow these guidelines when submitting an issue or a feature
+request. It will make it easier for us to come back to you quickly and with good
+feedback.
+
+### Did you find a bug?
Nice! That's what we use in `transformers`, with `transformers-cli env`. I agree it would be nice to port it to `accelerate`.
@@ -0,0 +1,220 @@
+<!---
+Copyright 2020 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.
+-->
+
+# How to contribute to Accelerate?
+
+Everyone is welcome to contribute, and we value everybody's contribution. Code
+is thus not the only way to help the community. Answering questions, helping
+others, reaching out and improving the documentations are immensely valuable to
+the community.
+
+It also helps us if you spread the word: reference the library from blog posts
+on the awesome projects it made possible, shout out on Twitter every time it has
+helped you, or simply star the repo to say "thank you".
+
+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md).
+
+## You can contribute in so many ways!
+
+There are 4 ways you can contribute to Accelerate:
+* Fixing outstanding issues with the existing code;
+* Contributing to the examples or to the documentation;
+* Submitting issues related to bugs or desired new features.
+
+Please ping @sgugger on your contribution.
+
+## Submitting a new issue or feature request
+
+Do your best to follow these guidelines when submitting an issue or a feature
+request. It will make it easier for us to come back to you quickly and with good
+feedback.
+
+### Did you find a bug?
+
+The 🤗 Accelerate library is robust and reliable thanks to the users who notify us of
+the problems they encounter. So thank you for reporting an issue.
+
+First, we would really appreciate it if you could **make sure the bug was not
+already reported** (use the search bar on Github under Issues).
+
+Did not find it? :( So we can act quickly on it, please follow these steps:
+
+* Include your **OS type and version**, the versions of **Python** and **PyTorch**.
I won't have the bandwidth right now to spin a tutorial about adding git hooks for our checks, but I agree that it would definitely have its place here. 

Feel free to push a commit directly to the `contributing-guide` branch with it, otherwise happy to do it in a follow-up PR when I have the bandwidth.
@@ -0,0 +1,220 @@
+<!---
+Copyright 2020 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.
+-->
+
+# How to contribute to Accelerate?
+
+Everyone is welcome to contribute, and we value everybody's contribution. Code
+is thus not the only way to help the community. Answering questions, helping
+others, reaching out and improving the documentations are immensely valuable to
+the community.
+
+It also helps us if you spread the word: reference the library from blog posts
+on the awesome projects it made possible, shout out on Twitter every time it has
+helped you, or simply star the repo to say "thank you".
+
+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md).
+
+## You can contribute in so many ways!
+
+There are 4 ways you can contribute to Accelerate:
+* Fixing outstanding issues with the existing code;
+* Contributing to the examples or to the documentation;
+* Submitting issues related to bugs or desired new features.
+
+Please ping @sgugger on your contribution.
+
+## Submitting a new issue or feature request
+
+Do your best to follow these guidelines when submitting an issue or a feature
+request. It will make it easier for us to come back to you quickly and with good
+feedback.
+
+### Did you find a bug?
+
+The 🤗 Accelerate library is robust and reliable thanks to the users who notify us of
+the problems they encounter. So thank you for reporting an issue.
+
+First, we would really appreciate it if you could **make sure the bug was not
+already reported** (use the search bar on Github under Issues).
+
+Did not find it? :( So we can act quickly on it, please follow these steps:
+
+* Include your **OS type and version**, the versions of **Python** and **PyTorch**.
+* A short, self-contained, code snippet that allows us to reproduce the bug in
+  less than 30s;
+* Provide the *full* traceback if an exception is raised.
+
+### Do you want a new feature (that is not a model)?
+
+A world-class feature request addresses the following points:
+
+1. Motivation first:
+* Is it related to a problem/frustration with the library? If so, please explain
+  why. Providing a code snippet that demonstrates the problem is best.
+* Is it related to something you would need for a project? We'd love to hear
+  about it!
+* Is it something you worked on and think could benefit the community?
+  Awesome! Tell us what problem it solved for you.
+2. Write a *full paragraph* describing the feature;
+3. Provide a **code snippet** that demonstrates its future use;
+4. In case this is related to a paper, please attach a link;
+5. Attach any additional information (drawings, screenshots, etc.) you think may help.
+
+If your issue is well written we're already 80% of the way there by the time you
+post it.
+
+## Start contributing! (Pull Requests)
+
+Before writing code, we strongly advise you to search through the existing PRs or
+issues to make sure that nobody is already working on the same thing. If you are
+unsure, it is always a good idea to open an issue to get some feedback.
+
+You will need basic `git` proficiency to be able to contribute to
+`accelerate`. `git` is not the easiest tool to use but it has the greatest
+manual. Type `git --help` in a shell and enjoy. If you prefer books, [Pro
+Git](https://git-scm.com/book/en/v2) is a very good reference.
+
+Follow these steps to start contributing:
+
+1. Fork the [repository](https://github.com/huggingface/accelerate) by
+   clicking on the 'Fork' button on the repository's page. This creates a copy of the code
+   under your GitHub user account.
+
+2. Clone your fork to your local disk, and add the base repository as a remote:
+
+   ```bash
+   $ git clone git@github.com:<your Github handle>/accelerate.git
+   $ cd accelerate
+   $ git remote add upstream https://github.com/huggingface/accelerate.git
+   ```
+
+3. Create a new branch to hold your development changes:
+
+   ```bash
+   $ git checkout -b a-descriptive-name-for-my-changes
+   ```
+
+   **Do not** work on the `master` branch.
+
+4. Set up a development environment by running the following command in a virtual environment:
+
+   ```bash
+   $ pip install -e ".[dev]"
+   ```
+
+   (If accelerate was already installed in the virtual environment, remove
+   it with `pip uninstall accelerate` before reinstalling it in editable
+   mode with the `-e` flag.)
+
+5. Develop the features on your branch.
+
+   As you work on the features, you should make sure that the test suite
+   passes. You should run the tests impacted by your changes like this (see 
+   below an explanation regarding the environment variable):
+
+   ```bash
+   $ pytest tests/<TEST_TO_RUN>.py
+   ```
+
+   You can also run the full suite with the following command.
+
+   ```bash
+   $ make test
+   ```
+
+   `accelerate` relies on `black` and `isort` to format its source code
+   consistently. After you make changes, apply automatic style corrections and code verifications
+   that can't be automated in one go with:
+
+   This target is also optimized to only work with files modified by the PR you're working on.
+
+   If you prefer to run the checks one after the other, the following command apply the
+   style corrections:
+
+   ```bash
+   $ make style
+   ```
+
+   `accelerate` also uses `flake8` and a few custom scripts to check for coding mistakes. Quality
+   control runs in CI, however you can also run the same checks with:
What we've seen is that merging can cause a lot of issues with github not understanding what was done with the branch, and users ending up with all the commits on their branch (which of course came with all file changes accompanying these commits).

We then ask the users to close the PR and to open a new one with the same branch which generally solves the issue, but it adds more friction than necessary, so up to now we've mostly asked users to rebase their branches.

If you feel strongly, happy to add an additional mention for merge.
@@ -0,0 +1,220 @@
+<!---
+Copyright 2020 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.
+-->
+
+# How to contribute to Accelerate?
+
+Everyone is welcome to contribute, and we value everybody's contribution. Code
+is thus not the only way to help the community. Answering questions, helping
+others, reaching out and improving the documentations are immensely valuable to
+the community.
+
+It also helps us if you spread the word: reference the library from blog posts
+on the awesome projects it made possible, shout out on Twitter every time it has
+helped you, or simply star the repo to say "thank you".
+
+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md).
+
+## You can contribute in so many ways!
+
+There are 4 ways you can contribute to Accelerate:
+* Fixing outstanding issues with the existing code;
+* Contributing to the examples or to the documentation;
+* Submitting issues related to bugs or desired new features.
+
+Please ping @sgugger on your contribution.
+
+## Submitting a new issue or feature request
+
+Do your best to follow these guidelines when submitting an issue or a feature
+request. It will make it easier for us to come back to you quickly and with good
+feedback.
+
+### Did you find a bug?
+
+The 🤗 Accelerate library is robust and reliable thanks to the users who notify us of
+the problems they encounter. So thank you for reporting an issue.
+
+First, we would really appreciate it if you could **make sure the bug was not
+already reported** (use the search bar on Github under Issues).
+
+Did not find it? :( So we can act quickly on it, please follow these steps:
+
+* Include your **OS type and version**, the versions of **Python** and **PyTorch**.
+* A short, self-contained, code snippet that allows us to reproduce the bug in
+  less than 30s;
+* Provide the *full* traceback if an exception is raised.
+
+### Do you want a new feature (that is not a model)?
+
+A world-class feature request addresses the following points:
+
+1. Motivation first:
+* Is it related to a problem/frustration with the library? If so, please explain
+  why. Providing a code snippet that demonstrates the problem is best.
+* Is it related to something you would need for a project? We'd love to hear
+  about it!
+* Is it something you worked on and think could benefit the community?
+  Awesome! Tell us what problem it solved for you.
+2. Write a *full paragraph* describing the feature;
+3. Provide a **code snippet** that demonstrates its future use;
+4. In case this is related to a paper, please attach a link;
+5. Attach any additional information (drawings, screenshots, etc.) you think may help.
+
+If your issue is well written we're already 80% of the way there by the time you
+post it.
+
+## Start contributing! (Pull Requests)
+
+Before writing code, we strongly advise you to search through the existing PRs or
+issues to make sure that nobody is already working on the same thing. If you are
+unsure, it is always a good idea to open an issue to get some feedback.
+
+You will need basic `git` proficiency to be able to contribute to
+`accelerate`. `git` is not the easiest tool to use but it has the greatest
+manual. Type `git --help` in a shell and enjoy. If you prefer books, [Pro
+Git](https://git-scm.com/book/en/v2) is a very good reference.
+
+Follow these steps to start contributing:
+
+1. Fork the [repository](https://github.com/huggingface/accelerate) by
+   clicking on the 'Fork' button on the repository's page. This creates a copy of the code
+   under your GitHub user account.
+
+2. Clone your fork to your local disk, and add the base repository as a remote:
+
+   ```bash
+   $ git clone git@github.com:<your Github handle>/accelerate.git
+   $ cd accelerate
+   $ git remote add upstream https://github.com/huggingface/accelerate.git
+   ```
+
+3. Create a new branch to hold your development changes:
+
+   ```bash
+   $ git checkout -b a-descriptive-name-for-my-changes
+   ```
+
+   **Do not** work on the `master` branch.
+
+4. Set up a development environment by running the following command in a virtual environment:
+
+   ```bash
+   $ pip install -e ".[dev]"
+   ```
+
+   (If accelerate was already installed in the virtual environment, remove
+   it with `pip uninstall accelerate` before reinstalling it in editable
+   mode with the `-e` flag.)
+
+5. Develop the features on your branch.
+
+   As you work on the features, you should make sure that the test suite
+   passes. You should run the tests impacted by your changes like this (see 
+   below an explanation regarding the environment variable):
+
+   ```bash
+   $ pytest tests/<TEST_TO_RUN>.py
+   ```
+
+   You can also run the full suite with the following command.
+
+   ```bash
+   $ make test
+   ```
+
+   `accelerate` relies on `black` and `isort` to format its source code
+   consistently. After you make changes, apply automatic style corrections and code verifications
+   that can't be automated in one go with:
+
+   This target is also optimized to only work with files modified by the PR you're working on.
+
+   If you prefer to run the checks one after the other, the following command apply the
+   style corrections:
+
+   ```bash
+   $ make style
+   ```
+
+   `accelerate` also uses `flake8` and a few custom scripts to check for coding mistakes. Quality
+   control runs in CI, however you can also run the same checks with:
+
+   ```bash
+   $ make quality
+   ```
+
+   Once you're happy with your changes, add changed files using `git add` and
+   make a commit with `git commit` to record your changes locally:
+
+   ```bash
+   $ git add modified_file.py
+   $ git commit
+   ```
+
+   Please write [good commit messages](https://chris.beams.io/posts/git-commit/).
+
+   It is a good idea to sync your copy of the code with the original
+   repository regularly. This way you can quickly account for changes:
+
+   ```bash
+   $ git fetch upstream
+   $ git rebase upstream/master
```suggestion
Copyright 2022 The HuggingFace Team. All rights reserved.
```
@@ -0,0 +1,235 @@
+<!---
+Copyright 2020 The HuggingFace Team. All rights reserved.
If you mean including the git hooks in the repository itself, I think a lot of us in the team are averse to having them in our own setups so it would be polarising :)
@@ -0,0 +1,220 @@
+<!---
+Copyright 2020 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.
+-->
+
+# How to contribute to Accelerate?
+
+Everyone is welcome to contribute, and we value everybody's contribution. Code
+is thus not the only way to help the community. Answering questions, helping
+others, reaching out and improving the documentations are immensely valuable to
+the community.
+
+It also helps us if you spread the word: reference the library from blog posts
+on the awesome projects it made possible, shout out on Twitter every time it has
+helped you, or simply star the repo to say "thank you".
+
+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md).
+
+## You can contribute in so many ways!
+
+There are 4 ways you can contribute to Accelerate:
+* Fixing outstanding issues with the existing code;
+* Contributing to the examples or to the documentation;
+* Submitting issues related to bugs or desired new features.
+
+Please ping @sgugger on your contribution.
+
+## Submitting a new issue or feature request
+
+Do your best to follow these guidelines when submitting an issue or a feature
+request. It will make it easier for us to come back to you quickly and with good
+feedback.
+
+### Did you find a bug?
+
+The 🤗 Accelerate library is robust and reliable thanks to the users who notify us of
+the problems they encounter. So thank you for reporting an issue.
+
+First, we would really appreciate it if you could **make sure the bug was not
+already reported** (use the search bar on Github under Issues).
+
+Did not find it? :( So we can act quickly on it, please follow these steps:
+
+* Include your **OS type and version**, the versions of **Python** and **PyTorch**.
+* A short, self-contained, code snippet that allows us to reproduce the bug in
+  less than 30s;
+* Provide the *full* traceback if an exception is raised.
+
+### Do you want a new feature (that is not a model)?
+
+A world-class feature request addresses the following points:
+
+1. Motivation first:
+* Is it related to a problem/frustration with the library? If so, please explain
+  why. Providing a code snippet that demonstrates the problem is best.
+* Is it related to something you would need for a project? We'd love to hear
+  about it!
+* Is it something you worked on and think could benefit the community?
+  Awesome! Tell us what problem it solved for you.
+2. Write a *full paragraph* describing the feature;
+3. Provide a **code snippet** that demonstrates its future use;
+4. In case this is related to a paper, please attach a link;
+5. Attach any additional information (drawings, screenshots, etc.) you think may help.
+
+If your issue is well written we're already 80% of the way there by the time you
+post it.
+
+## Start contributing! (Pull Requests)
+
+Before writing code, we strongly advise you to search through the existing PRs or
+issues to make sure that nobody is already working on the same thing. If you are
+unsure, it is always a good idea to open an issue to get some feedback.
+
+You will need basic `git` proficiency to be able to contribute to
+`accelerate`. `git` is not the easiest tool to use but it has the greatest
+manual. Type `git --help` in a shell and enjoy. If you prefer books, [Pro
+Git](https://git-scm.com/book/en/v2) is a very good reference.
+
+Follow these steps to start contributing:
+
+1. Fork the [repository](https://github.com/huggingface/accelerate) by
+   clicking on the 'Fork' button on the repository's page. This creates a copy of the code
+   under your GitHub user account.
+
+2. Clone your fork to your local disk, and add the base repository as a remote:
+
+   ```bash
+   $ git clone git@github.com:<your Github handle>/accelerate.git
+   $ cd accelerate
+   $ git remote add upstream https://github.com/huggingface/accelerate.git
+   ```
+
+3. Create a new branch to hold your development changes:
+
+   ```bash
+   $ git checkout -b a-descriptive-name-for-my-changes
+   ```
+
+   **Do not** work on the `master` branch.
+
+4. Set up a development environment by running the following command in a virtual environment:
+
+   ```bash
+   $ pip install -e ".[dev]"
+   ```
+
+   (If accelerate was already installed in the virtual environment, remove
+   it with `pip uninstall accelerate` before reinstalling it in editable
+   mode with the `-e` flag.)
+
+5. Develop the features on your branch.
+
+   As you work on the features, you should make sure that the test suite
+   passes. You should run the tests impacted by your changes like this (see 
+   below an explanation regarding the environment variable):
+
+   ```bash
+   $ pytest tests/<TEST_TO_RUN>.py
+   ```
+
+   You can also run the full suite with the following command.
+
+   ```bash
+   $ make test
+   ```
+
+   `accelerate` relies on `black` and `isort` to format its source code
+   consistently. After you make changes, apply automatic style corrections and code verifications
+   that can't be automated in one go with:
+
+   This target is also optimized to only work with files modified by the PR you're working on.
+
+   If you prefer to run the checks one after the other, the following command apply the
+   style corrections:
+
+   ```bash
+   $ make style
+   ```
+
+   `accelerate` also uses `flake8` and a few custom scripts to check for coding mistakes. Quality
+   control runs in CI, however you can also run the same checks with:
Yeah I saw. We can follow up on that later.
@@ -0,0 +1,220 @@
+<!---
+Copyright 2020 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.
+-->
+
+# How to contribute to Accelerate?
+
+Everyone is welcome to contribute, and we value everybody's contribution. Code
+is thus not the only way to help the community. Answering questions, helping
+others, reaching out and improving the documentations are immensely valuable to
+the community.
+
+It also helps us if you spread the word: reference the library from blog posts
+on the awesome projects it made possible, shout out on Twitter every time it has
+helped you, or simply star the repo to say "thank you".
+
+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md).
+
+## You can contribute in so many ways!
+
+There are 4 ways you can contribute to Accelerate:
+* Fixing outstanding issues with the existing code;
+* Contributing to the examples or to the documentation;
+* Submitting issues related to bugs or desired new features.
+
+Please ping @sgugger on your contribution.
+
+## Submitting a new issue or feature request
+
+Do your best to follow these guidelines when submitting an issue or a feature
+request. It will make it easier for us to come back to you quickly and with good
+feedback.
+
+### Did you find a bug?
+
+The 🤗 Accelerate library is robust and reliable thanks to the users who notify us of
+the problems they encounter. So thank you for reporting an issue.
+
+First, we would really appreciate it if you could **make sure the bug was not
+already reported** (use the search bar on Github under Issues).
+
+Did not find it? :( So we can act quickly on it, please follow these steps:
+
+* Include your **OS type and version**, the versions of **Python** and **PyTorch**.
+* A short, self-contained, code snippet that allows us to reproduce the bug in
+  less than 30s;
+* Provide the *full* traceback if an exception is raised.
+
+### Do you want a new feature (that is not a model)?
+
+A world-class feature request addresses the following points:
+
+1. Motivation first:
+* Is it related to a problem/frustration with the library? If so, please explain
+  why. Providing a code snippet that demonstrates the problem is best.
+* Is it related to something you would need for a project? We'd love to hear
+  about it!
+* Is it something you worked on and think could benefit the community?
+  Awesome! Tell us what problem it solved for you.
+2. Write a *full paragraph* describing the feature;
+3. Provide a **code snippet** that demonstrates its future use;
+4. In case this is related to a paper, please attach a link;
+5. Attach any additional information (drawings, screenshots, etc.) you think may help.
+
+If your issue is well written we're already 80% of the way there by the time you
+post it.
+
+## Start contributing! (Pull Requests)
+
+Before writing code, we strongly advise you to search through the existing PRs or
+issues to make sure that nobody is already working on the same thing. If you are
+unsure, it is always a good idea to open an issue to get some feedback.
+
+You will need basic `git` proficiency to be able to contribute to
+`accelerate`. `git` is not the easiest tool to use but it has the greatest
+manual. Type `git --help` in a shell and enjoy. If you prefer books, [Pro
+Git](https://git-scm.com/book/en/v2) is a very good reference.
+
+Follow these steps to start contributing:
+
+1. Fork the [repository](https://github.com/huggingface/accelerate) by
+   clicking on the 'Fork' button on the repository's page. This creates a copy of the code
+   under your GitHub user account.
+
+2. Clone your fork to your local disk, and add the base repository as a remote:
+
+   ```bash
+   $ git clone git@github.com:<your Github handle>/accelerate.git
+   $ cd accelerate
+   $ git remote add upstream https://github.com/huggingface/accelerate.git
+   ```
+
+3. Create a new branch to hold your development changes:
+
+   ```bash
+   $ git checkout -b a-descriptive-name-for-my-changes
+   ```
+
+   **Do not** work on the `master` branch.
+
+4. Set up a development environment by running the following command in a virtual environment:
+
+   ```bash
+   $ pip install -e ".[dev]"
+   ```
+
+   (If accelerate was already installed in the virtual environment, remove
+   it with `pip uninstall accelerate` before reinstalling it in editable
+   mode with the `-e` flag.)
+
+5. Develop the features on your branch.
+
+   As you work on the features, you should make sure that the test suite
+   passes. You should run the tests impacted by your changes like this (see 
+   below an explanation regarding the environment variable):
+
+   ```bash
+   $ pytest tests/<TEST_TO_RUN>.py
+   ```
+
+   You can also run the full suite with the following command.
+
+   ```bash
+   $ make test
+   ```
+
+   `accelerate` relies on `black` and `isort` to format its source code
+   consistently. After you make changes, apply automatic style corrections and code verifications
+   that can't be automated in one go with:
+
+   This target is also optimized to only work with files modified by the PR you're working on.
+
+   If you prefer to run the checks one after the other, the following command apply the
+   style corrections:
+
+   ```bash
+   $ make style
+   ```
+
+   `accelerate` also uses `flake8` and a few custom scripts to check for coding mistakes. Quality
+   control runs in CI, however you can also run the same checks with:
Just a little thing, it's `git rebase upstream/main`
@@ -0,0 +1,235 @@
+<!---
+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.
+-->
+
+# How to contribute to 🤗 Accelerate?
+
+Everyone is welcome to contribute, and we value everybody's contribution. Code
+is thus not the only way to help the community. Answering questions, helping
+others, reaching out and improving the documentations are immensely valuable to
+the community.
+
+It also helps us if you spread the word: reference the library from blog posts
+on the awesome projects it made possible, shout out on Twitter every time it has
+helped you, or simply star the repo to say "thank you".
+
+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md).
+
+## You can contribute in so many ways!
+
+Some of the ways you can contribute to Accelerate:
+* Fixing outstanding issues with the existing code;
+* Contributing to the examples or to the documentation;
+* Submitting issues related to bugs or desired new features.
+
+## Submitting a new issue or feature request
+
+Do your best to follow these guidelines when submitting an issue or a feature
+request. It will make it easier for us to come back to you quickly and with good
+feedback.
+
+### Did you find a bug?
+
+The 🤗 Accelerate library is robust and reliable thanks to the users who notify us of
+the problems they encounter. So thank you for reporting an issue.
+
+First, we would really appreciate it if you could **make sure the bug was not
+already reported** (use the search bar on Github under Issues).
+
+Did not find it? :( So we can act quickly on it, please follow these steps:
+
+* Include your **OS type and version**, the versions of **Python** and **PyTorch**.
+* A short, self-contained, code snippet that allows us to reproduce the bug in
+  less than 30s;
+* Provide the with your Accelerate configuration (located by default in `~/.cache/huggingface/accelerate/default_congig.yml`)
+
+### Do you want a new feature?
+
+A good feature request addresses the following points:
+
+1. Motivation first:
+* Is it related to a problem/frustration with the library? If so, please explain
+  why. Providing a code snippet that demonstrates the problem is best.
+* Is it related to something you would need for a project? We'd love to hear
+  about it!
+* Is it something you worked on and think could benefit the community?
+  Awesome! Tell us what problem it solved for you.
+2. Write a *full paragraph* describing the feature;
+3. Provide a **code snippet** that demonstrates its future use;
+4. In case this is related to a paper, please attach a link;
+5. Attach any additional information (drawings, screenshots, etc.) you think may help.
+
+If your issue is well written we're already 80% of the way there by the time you
+post it.
+
+## Submitting a pull request (PR)
+
+Before writing code, we strongly advise you to search through the existing PRs or
+issues to make sure that nobody is already working on the same thing. If you are
+unsure, it is always a good idea to open an issue to get some feedback.
+
+You will need basic `git` proficiency to be able to contribute to
+🤗 Accelerate. `git` is not the easiest tool to use but it has the greatest
+manual. Type `git --help` in a shell and enjoy. If you prefer books, [Pro
+Git](https://git-scm.com/book/en/v2) is a very good reference.
+
+Follow these steps to start contributing:
+
+1. Fork the [repository](https://github.com/huggingface/accelerate) by
+   clicking on the 'Fork' button on the repository's page. This creates a copy of the code
+   under your GitHub user account.
+
+2. Clone your fork to your local disk, and add the base repository as a remote. The following command
+   assumes you have your public SSH key uploaded to GitHub. See the following guide for more
+   [information](https://docs.github.com/en/repositories/creating-and-managing-repositories/cloning-a-repository).
+
+   ```bash
+   $ git clone git@github.com:<your Github handle>/accelerate.git
+   $ cd accelerate
+   $ git remote add upstream https://github.com/huggingface/accelerate.git
+   ```
+
+3. Create a new branch to hold your development changes, and do this for every new PR you work on.
+
+   Start by synchronizing your `main` branch with the `upstream/main` branch (ore details in the [GitHub Docs](https://docs.github.com/en/github/collaborating-with-issues-and-pull-requests/syncing-a-fork)):
+
+   ```bash
+   $ git checkout main
+   $ git fetch upstream
+   $ git merge upstream/main
+   ```
+
+   Once your `main` branch is synchronized, create a new branch from it:
+
+   ```bash
+   $ git checkout -b a-descriptive-name-for-my-changes
+   ```
+
+   **Do not** work on the `master` branch.
+
+4. Set up a development environment by running the following command in a conda or a virtual environment you've created for working on this library:
+
+   ```bash
+   $ pip install -e ".[quality]"
+   ```
+
+   (If accelerate was already installed in the virtual environment, remove
+   it with `pip uninstall accelerate` before reinstalling it in editable
+   mode with the `-e` flag.)
+
+5. Develop the features on your branch.
+
+   As you work on the features, you should make sure that the test suite
+   passes. You should run the tests impacted by your changes like this (see 
+   below an explanation regarding the environment variable):
+
+   ```bash
+   $ pytest tests/<TEST_TO_RUN>.py
+   ```
+   
+   > For the following commands leveraging the `make` utility, we recommend using the WSL system when running on
+   > Windows. More information [here](https://docs.microsoft.com/en-us/windows/wsl/about).
+
+   You can also run the full suite with the following command.
+
+   ```bash
+   $ make test
+   ```
+
+   `accelerate` relies on `black` and `isort` to format its source code
+   consistently. After you make changes, apply automatic style corrections and code verifications
+   that can't be automated in one go with:
+
+   This target is also optimized to only work with files modified by the PR you're working on.
+
+   If you prefer to run the checks one after the other, the following command apply the
+   style corrections:
+
+   ```bash
+   $ make style
+   ```
+
+   `accelerate` also uses `flake8` and a few custom scripts to check for coding mistakes. Quality
+   control runs in CI, however you can also run the same checks with:
+
+   ```bash
+   $ make quality
+   ```
+
+   Once you're happy with your changes, add changed files using `git add` and
+   make a commit with `git commit` to record your changes locally:
+
+   ```bash
+   $ git add modified_file.py
+   $ git commit
+   ```
+
+   Please write [good commit messages](https://chris.beams.io/posts/git-commit/).
+
+   It is a good idea to sync your copy of the code with the original
+   repository regularly. This way you can quickly account for changes:
+
+   ```bash
+   $ git fetch upstream
+   $ git rebase upstream/master
Good catch! Do you want to make a PR to fix this?
@@ -0,0 +1,235 @@
+<!---
+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.
+-->
+
+# How to contribute to 🤗 Accelerate?
+
+Everyone is welcome to contribute, and we value everybody's contribution. Code
+is thus not the only way to help the community. Answering questions, helping
+others, reaching out and improving the documentations are immensely valuable to
+the community.
+
+It also helps us if you spread the word: reference the library from blog posts
+on the awesome projects it made possible, shout out on Twitter every time it has
+helped you, or simply star the repo to say "thank you".
+
+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md).
+
+## You can contribute in so many ways!
+
+Some of the ways you can contribute to Accelerate:
+* Fixing outstanding issues with the existing code;
+* Contributing to the examples or to the documentation;
+* Submitting issues related to bugs or desired new features.
+
+## Submitting a new issue or feature request
+
+Do your best to follow these guidelines when submitting an issue or a feature
+request. It will make it easier for us to come back to you quickly and with good
+feedback.
+
+### Did you find a bug?
+
+The 🤗 Accelerate library is robust and reliable thanks to the users who notify us of
+the problems they encounter. So thank you for reporting an issue.
+
+First, we would really appreciate it if you could **make sure the bug was not
+already reported** (use the search bar on Github under Issues).
+
+Did not find it? :( So we can act quickly on it, please follow these steps:
+
+* Include your **OS type and version**, the versions of **Python** and **PyTorch**.
+* A short, self-contained, code snippet that allows us to reproduce the bug in
+  less than 30s;
+* Provide the with your Accelerate configuration (located by default in `~/.cache/huggingface/accelerate/default_congig.yml`)
+
+### Do you want a new feature?
+
+A good feature request addresses the following points:
+
+1. Motivation first:
+* Is it related to a problem/frustration with the library? If so, please explain
+  why. Providing a code snippet that demonstrates the problem is best.
+* Is it related to something you would need for a project? We'd love to hear
+  about it!
+* Is it something you worked on and think could benefit the community?
+  Awesome! Tell us what problem it solved for you.
+2. Write a *full paragraph* describing the feature;
+3. Provide a **code snippet** that demonstrates its future use;
+4. In case this is related to a paper, please attach a link;
+5. Attach any additional information (drawings, screenshots, etc.) you think may help.
+
+If your issue is well written we're already 80% of the way there by the time you
+post it.
+
+## Submitting a pull request (PR)
+
+Before writing code, we strongly advise you to search through the existing PRs or
+issues to make sure that nobody is already working on the same thing. If you are
+unsure, it is always a good idea to open an issue to get some feedback.
+
+You will need basic `git` proficiency to be able to contribute to
+🤗 Accelerate. `git` is not the easiest tool to use but it has the greatest
+manual. Type `git --help` in a shell and enjoy. If you prefer books, [Pro
+Git](https://git-scm.com/book/en/v2) is a very good reference.
+
+Follow these steps to start contributing:
+
+1. Fork the [repository](https://github.com/huggingface/accelerate) by
+   clicking on the 'Fork' button on the repository's page. This creates a copy of the code
+   under your GitHub user account.
+
+2. Clone your fork to your local disk, and add the base repository as a remote. The following command
+   assumes you have your public SSH key uploaded to GitHub. See the following guide for more
+   [information](https://docs.github.com/en/repositories/creating-and-managing-repositories/cloning-a-repository).
+
+   ```bash
+   $ git clone git@github.com:<your Github handle>/accelerate.git
+   $ cd accelerate
+   $ git remote add upstream https://github.com/huggingface/accelerate.git
+   ```
+
+3. Create a new branch to hold your development changes, and do this for every new PR you work on.
+
+   Start by synchronizing your `main` branch with the `upstream/main` branch (ore details in the [GitHub Docs](https://docs.github.com/en/github/collaborating-with-issues-and-pull-requests/syncing-a-fork)):
+
+   ```bash
+   $ git checkout main
+   $ git fetch upstream
+   $ git merge upstream/main
+   ```
+
+   Once your `main` branch is synchronized, create a new branch from it:
+
+   ```bash
+   $ git checkout -b a-descriptive-name-for-my-changes
+   ```
+
+   **Do not** work on the `master` branch.
+
+4. Set up a development environment by running the following command in a conda or a virtual environment you've created for working on this library:
+
+   ```bash
+   $ pip install -e ".[quality]"
+   ```
+
+   (If accelerate was already installed in the virtual environment, remove
+   it with `pip uninstall accelerate` before reinstalling it in editable
+   mode with the `-e` flag.)
+
+5. Develop the features on your branch.
+
+   As you work on the features, you should make sure that the test suite
+   passes. You should run the tests impacted by your changes like this (see 
+   below an explanation regarding the environment variable):
+
+   ```bash
+   $ pytest tests/<TEST_TO_RUN>.py
+   ```
+   
+   > For the following commands leveraging the `make` utility, we recommend using the WSL system when running on
+   > Windows. More information [here](https://docs.microsoft.com/en-us/windows/wsl/about).
+
+   You can also run the full suite with the following command.
+
+   ```bash
+   $ make test
+   ```
+
+   `accelerate` relies on `black` and `isort` to format its source code
+   consistently. After you make changes, apply automatic style corrections and code verifications
+   that can't be automated in one go with:
+
+   This target is also optimized to only work with files modified by the PR you're working on.
+
+   If you prefer to run the checks one after the other, the following command apply the
+   style corrections:
+
+   ```bash
+   $ make style
+   ```
+
+   `accelerate` also uses `flake8` and a few custom scripts to check for coding mistakes. Quality
+   control runs in CI, however you can also run the same checks with:
+
+   ```bash
+   $ make quality
+   ```
+
+   Once you're happy with your changes, add changed files using `git add` and
+   make a commit with `git commit` to record your changes locally:
+
+   ```bash
+   $ git add modified_file.py
+   $ git commit
+   ```
+
+   Please write [good commit messages](https://chris.beams.io/posts/git-commit/).
+
+   It is a good idea to sync your copy of the code with the original
+   repository regularly. This way you can quickly account for changes:
+
+   ```bash
+   $ git fetch upstream
+   $ git rebase upstream/master
Will merge it like that, with the following action items:
- Work on adding an issue template (for both `huggingface_hub` and `accelerate`)
- Have a similar method to `transformers-cli env` for `accelerate` (can be a `get_versions()` like Adrin mentions)

@sgugger, ok for you if I open issues for that on `accelerate`?
Fine by me! It can be a good first issue!
https://api.github.com/repos/huggingface/accelerate/pulls/254
diff --git a/setup.py b/setup.py
index f06ab3c99..bec3b9805 100644
--- a/setup.py
+++ b/setup.py
@@ -16,7 +16,7 @@
 from setuptools import find_packages
 
 extras = {}
-extras["quality"] = ["black == 21.4b0", "isort >= 5.5.4", "flake8 >= 3.8.3"]
+extras["quality"] = ["black ~= 22.0", "isort >= 5.5.4", "flake8 >= 3.8.3"]
 extras["docs"] = [
     "docutils==0.16.0",
     "recommonmark",

https://api.github.com/repos/huggingface/accelerate/pulls/250
diff --git a/README.md b/README.md
index cbc23ed10..c27131808 100644
--- a/README.md
+++ b/README.md
@@ -210,8 +210,11 @@ You shouldn't use 🤗 Accelerate if you don't want to write a training loop you
 
 If you like the simplicity of 🤗 Accelerate but would prefer a higher-level abstraction around your training loop, some frameworks that are built on top of 🤗 Accelerate are listed below:
 
-* [pytorch-accelerated](https://github.com/Chris-hughes10/pytorch-accelerated) is a lightweight training library, with a streamlined feature set centred around a general-purpose [Trainer](https://pytorch-accelerated.readthedocs.io/en/latest/trainer.html), that places a huge emphasis on simplicity and transparency; enabling users to understand exactly what is going on under the hood, but without having to write and maintain the boilerplate themselves!
+* [Animus](https://github.com/Scitator/animus) is a minimalistic framework to run machine learning experiments. Animus highlights common "breakpoints" in ML experiments and provides a unified interface for them within [IExperiment](https://github.com/Scitator/animus/blob/main/animus/core.py#L76).
+* [Catalyst](https://github.com/catalyst-team/catalyst#getting-started) is a PyTorch framework for Deep Learning Research and Development. It focuses on reproducibility, rapid experimentation, and codebase reuse so you can create something new rather than write yet another train loop. Catalyst provides a [Runner](https://catalyst-team.github.io/catalyst/api/core.html#runner) to connect all parts of the experiment: hardware backend, data transformations, model train, and inference logic.
 * [Kornia](https://kornia.readthedocs.io/en/latest/get-started/introduction.html) is a differentiable library that allows classical computer vision to be integrated into deep learning models. Kornia provides a [Trainer](https://kornia.readthedocs.io/en/latest/x.html#kornia.x.Trainer) with the specific purpose to train and fine-tune the supported deep learning algorithms within the library.
+* [pytorch-accelerated](https://github.com/Chris-hughes10/pytorch-accelerated) is a lightweight training library, with a streamlined feature set centred around a general-purpose [Trainer](https://pytorch-accelerated.readthedocs.io/en/latest/trainer.html), that places a huge emphasis on simplicity and transparency; enabling users to understand exactly what is going on under the hood, but without having to write and maintain the boilerplate themselves!
+
 
 ## Installation
 

btw, it would be great if you could comment on my "Accelerator mixin" - [Engine](https://github.com/Scitator/animus/blob/main/animus/torch/engine.py) in terms of "a bit more user-friendly Accelerator API". I found it useful for me to run pipelines within notebooks or python scripts without CLI.
Ah thanks! I am wrapping up some things in Transformers and will go back to Accelerate pretty soon (haven't forgotten #215 !)

You mixins look great! Having one per kind of distributed setup is a bit against the philosophy of Accelerate (where the same exact code runs on any setup) but I like the abstraction.

To run things in notebooks, we also have `notebook_launcher`, I don't know if you tried it or not.
yeah, I found `notebook_launcher` an interesting functional alternative... nevertheless, from my previous experience I knew a few limitations of pure functional way, so... a small class-based API was created ;)
https://api.github.com/repos/huggingface/accelerate/pulls/249
diff --git a/.gitignore b/.gitignore
index 3b89c957d..f41377320 100644
--- a/.gitignore
+++ b/.gitignore
@@ -130,3 +130,6 @@ dmypy.json
 
 # VSCode
 .vscode
+
+# IntelliJ
+.idea
\ No newline at end of file
diff --git a/examples/cv_example.py b/examples/cv_example.py
index f4164613f..bb7c7f5b1 100644
--- a/examples/cv_example.py
+++ b/examples/cv_example.py
@@ -73,7 +73,7 @@ def __getitem__(self, idx):
 
 def training_function(config, args):
     # Initialize accelerator
-    accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu)
+    accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mix_precision)
 
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
     lr = config["lr"]
@@ -187,6 +187,15 @@ def main():
     parser = argparse.ArgumentParser(description="Simple example of training script.")
     parser.add_argument("--data_dir", required=True, help="The data folder on disk.")
     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()
     config = {"lr": 3e-2, "num_epochs": 3, "seed": 42, "batch_size": 64, "image_size": 224}
diff --git a/examples/nlp_example.py b/examples/nlp_example.py
index a1f310271..2f869bd0f 100644
--- a/examples/nlp_example.py
+++ b/examples/nlp_example.py
@@ -51,7 +51,7 @@
 
 def training_function(config, args):
     # Initialize accelerator
-    accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu)
+    accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision)
 
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
     lr = config["lr"]
@@ -163,6 +163,15 @@ def collate_fn(examples):
 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()
     config = {"lr": 2e-5, "num_epochs": 3, "correct_bias": True, "seed": 42, "batch_size": 16}
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 76c3067a3..ed3b2acb7 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -14,6 +14,7 @@
 
 import gc
 import os
+import warnings
 from contextlib import contextmanager
 from typing import List, Optional, Union
 
@@ -60,10 +61,11 @@ class Accelerator:
             :obj:`True` the actual batch size used will be the same on any kind of distributed processes, but it must
             be a round multiple of the :obj:`num_processes` you are using. If :obj:`False`, actual batch size used will
             be the one set in your script multiplied by the number of processes.
-        fp16 (:obj:`bool`, `optional`):
-            Whether or not to use mixed precision training. Will default to the value in the environment variable
-            :obj:`USE_FP16`, which will use the default value in the accelerate config of the current system or the
-            flag passed with the :obj:`accelerate.launch` command.
+        mixed_precision (:obj:`str`, `optional`):
+            Whether or not to use mixed precision training (fp16 or bfloat16). Choose from 'no','fp16','bf16'. Will
+            default to the value in the environment variable :obj:`MIXED_PRECISION`, which will use the default value
+            in the accelerate config of the current system or the flag passed with the :obj:`accelerate.launch`
+            command. 'fp16' requires pytorch 1.6 or higher. 'bf16' requires pytorch 1.10 or higher.
         cpu (:obj:`bool`, `optional`):
             Whether or not to force the script to execute on CPU. Will ignore GPU available if set to :obj:`True` and
             force the execution on one process only.
@@ -101,12 +103,25 @@ def __init__(
         device_placement: bool = True,
         split_batches: bool = False,
         fp16: bool = None,
+        mixed_precision: str = None,
         cpu: bool = False,
         deepspeed_plugin: DeepSpeedPlugin = None,
         rng_types: Optional[List[Union[str, RNGType]]] = None,
         dispatch_batches: Optional[bool] = None,
         kwargs_handlers: Optional[List[KwargsHandler]] = None,
     ):
+
+        if mixed_precision is not None:
+            mixed_precision = mixed_precision.lower()
+            if mixed_precision not in ["no", "fp16", "bf16"]:
+                raise ValueError(
+                    f"Unknown mixed_precision mode: {mixed_precision}. Choose between 'no', 'fp16' and 'bf16'."
+                )
+
+        if fp16:
+            warnings.warn('fp16=True is deprecated. Use mixed_precision="fp16" instead.', DeprecationWarning)
+            mixed_precision = "fp16"
+
         if deepspeed_plugin is None:  # init from env variables
             deepspeed_plugin = DeepSpeedPlugin() if os.environ.get("USE_DEEPSPEED", "false") == "true" else None
         else:
@@ -139,7 +154,11 @@ def __init__(
 
         kwargs = self.init_handler.to_kwargs() if self.init_handler is not None else {}
         self.state = AcceleratorState(
-            fp16=fp16, cpu=cpu, deepspeed_plugin=deepspeed_plugin, _from_accelerator=True, **kwargs
+            mixed_precision=mixed_precision,
+            cpu=cpu,
+            deepspeed_plugin=deepspeed_plugin,
+            _from_accelerator=True,
+            **kwargs,
         )
 
         self.device_placement = device_placement
@@ -149,8 +168,18 @@ def __init__(
         # Mixed precision attributes
         self.scaler = None
         self.native_amp = False
-        if self.state.use_fp16:
+        if self.state.mixed_precision == "fp16":
             self.native_amp = version.parse(torch.__version__) >= version.parse("1.6")
+            if version.parse(torch.__version__) < version.parse("1.6"):
+                raise ValueError("fp16 mixed precision requires PyTorch >= 1.6")
+
+            kwargs = self.scaler_handler.to_kwargs() if self.scaler_handler is not None else {}
+            self.scaler = torch.cuda.amp.GradScaler(**kwargs)
+        elif self.state.mixed_precision == "bf16":
+            self.native_amp = version.parse(torch.__version__) >= version.parse("1.10")
+            if mixed_precision == "bf16" and version.parse(torch.__version__) < version.parse("1.10"):
+                raise ValueError("bf16 mixed precision requires PyTorch >= 1.10")
+
             kwargs = self.scaler_handler.to_kwargs() if self.scaler_handler is not None else {}
             self.scaler = torch.cuda.amp.GradScaler(**kwargs)
 
@@ -195,11 +224,20 @@ def is_local_main_process(self):
 
     @property
     def use_fp16(self):
+        return self.mixed_precision != "no"
+
+    @property
+    def mixed_precision(self):
         if self.distributed_type == DistributedType.DEEPSPEED:
-            use_fp16 = self.state.deepspeed_plugin.deepspeed_config["fp16"]["enabled"]
+            if self.state.deepspeed_plugin.deepspeed_config["fp16"]["enabled"]:
+                mixed_precision = "fp16"
+            elif self.state.deepspeed_plugin.deepspeed_config["bf16"]["enabled"]:
+                mixed_precision = "bf16"
+            else:
+                mixed_precision = "no"
         else:
-            use_fp16 = self.state.use_fp16
-        return use_fp16
+            mixed_precision = self.state.mixed_precision
+        return mixed_precision
 
     @contextmanager
     def local_main_process_first(self):
@@ -311,7 +349,12 @@ def prepare_model(self, model):
             kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {}
             model = torch.nn.parallel.DistributedDataParallel(model, **kwargs)
         if self.native_amp:
-            model.forward = torch.cuda.amp.autocast()(model.forward)
+            if self.mixed_precision == "fp16" and version.parse(torch.__version__) >= version.parse("1.10"):
+                model.forward = torch.cuda.amp.autocast(dtype=torch.float16)(model.forward)
+            elif self.mixed_precision == "bf16":
+                model.forward = torch.cuda.amp.autocast(dtype=torch.bfloat16)(model.forward)
+            else:
+                model.forward = torch.cuda.amp.autocast()(model.forward)
             model.forward = convert_outputs_to_fp32(model.forward)
         return model
 
@@ -577,7 +620,13 @@ def autocast(self):
         different will happen otherwise.
         """
         if self.native_amp:
-            autocast_context = torch.cuda.amp.autocast()
+            if self.mixed_precision == "fp16" and version.parse(torch.__version__) >= version.parse("1.10"):
+                autocast_context = torch.cuda.amp.autocast(dtype=torch.float16)
+            elif self.mixed_precision == "bf16":
+                autocast_context = torch.cuda.amp.autocast(dtype=torch.bfloat16)
+            else:
+                autocast_context = torch.cuda.amp.autocast()
+
             autocast_context.__enter__()
             yield
             autocast_context.__exit__()
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index eaaad20f2..534ea1526 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -103,20 +103,19 @@ def get_cluster_input():
     )
 
     if distributed_type != DistributedType.TPU:
-        fp16 = _ask_field(
-            "Do you wish to use FP16 (mixed precision)? [yes/NO]: ",
-            _convert_yes_no_to_bool,
-            default=False,
-            error_message="Please enter yes or no.",
+        mixed_precision = _ask_field(
+            "Do you wish to use FP16 or BF16 (mixed precision)? [NO/fp16/bf16]: ",
+            lambda x: str(x).lower(),
+            default="no",
         )
     else:
-        fp16 = False
+        mixed_precision = "no"
 
     return ClusterConfig(
         compute_environment=ComputeEnvironment.LOCAL_MACHINE,
         distributed_type=distributed_type,
         num_processes=num_processes,
-        fp16=fp16,
+        mixed_precision=mixed_precision,
         machine_rank=machine_rank,
         num_machines=num_machines,
         main_process_ip=main_process_ip,
diff --git a/src/accelerate/commands/config/config_args.py b/src/accelerate/commands/config/config_args.py
index 39c00c280..891580ab5 100644
--- a/src/accelerate/commands/config/config_args.py
+++ b/src/accelerate/commands/config/config_args.py
@@ -66,7 +66,7 @@ def load_config_from_file(config_file):
 class BaseConfig:
     compute_environment: ComputeEnvironment
     distributed_type: Union[DistributedType, SageMakerDistributedType]
-    fp16: bool
+    mixed_precision: str
 
     def to_dict(self):
         result = self.__dict__
@@ -83,6 +83,10 @@ def from_json_file(cls, json_file=None):
             config_dict = json.load(f)
         if "compute_environment" not in config_dict:
             config_dict["compute_environment"] = ComputeEnvironment.LOCAL_MACHINE
+        if "mixed_precision" not in config_dict:
+            config_dict["mixed_precision"] = "fp16" if ("fp16" in config_dict and config_dict["fp16"]) else "no"
+        if "fp16" in config_dict:  # Convert the config to the new format.
+            del config_dict["fp16"]
         return cls(**config_dict)
 
     def to_json_file(self, json_file):
@@ -97,6 +101,12 @@ def from_yaml_file(cls, yaml_file=None):
             config_dict = yaml.safe_load(f)
         if "compute_environment" not in config_dict:
             config_dict["compute_environment"] = ComputeEnvironment.LOCAL_MACHINE
+
+        if "mixed_precision" not in config_dict:
+            config_dict["mixed_precision"] = "fp16" if ("fp16" in config_dict and config_dict["fp16"]) else "no"
+        if "fp16" in config_dict:  # Convert the config to the new format.
+            del config_dict["fp16"]
+
         return cls(**config_dict)
 
     def to_yaml_file(self, yaml_file):
diff --git a/src/accelerate/commands/config/sagemaker.py b/src/accelerate/commands/config/sagemaker.py
index feff0aecf..3b0fc78dd 100644
--- a/src/accelerate/commands/config/sagemaker.py
+++ b/src/accelerate/commands/config/sagemaker.py
@@ -20,7 +20,7 @@
 from accelerate.utils import is_boto3_available
 
 from .config_args import SageMakerConfig
-from .config_utils import _ask_field, _convert_sagemaker_distributed_mode, _convert_yes_no_to_bool
+from .config_utils import _ask_field, _convert_sagemaker_distributed_mode
 
 
 if is_boto3_available():
@@ -139,11 +139,11 @@ def get_sagemaker_input():
             lambda x: int(x),
             default=2,
         )
-    fp16 = _ask_field(
-        "Do you wish to use FP16 (mixed precision)? [yes/NO]: ",
-        _convert_yes_no_to_bool,
-        default=False,
-        error_message="Please enter yes or no.",
+
+    mixed_precision = _ask_field(
+        "Do you wish to use FP16 or BF16 (mixed precision)? [No/FP16/BF16]: ",
+        lambda x: str(x),
+        default="No",
     )
 
     return SageMakerConfig(
@@ -153,6 +153,6 @@ def get_sagemaker_input():
         profile=aws_profile,
         region=aws_region,
         iam_role_name=iam_role_name,
-        fp16=fp16,
+        mixed_precision=mixed_precision,
         num_machines=num_machines,
     )
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 4f392d6fa..77db4546b 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -19,6 +19,7 @@
 import os
 import subprocess
 import sys
+import warnings
 from ast import literal_eval
 from pathlib import Path
 from typing import Dict, List
@@ -53,6 +54,16 @@ def launch_command_parser(subparsers=None):
     parser.add_argument(
         "--tpu", default=False, action="store_true", help="Whether or not this should launch a TPU training."
     )
+    parser.add_argument(
+        "--mixed_precision",
+        default="no",
+        type=str,
+        choices=["no", "fp16", "bf16"],
+        help="Whether or not to use mixed precision training. "
+        "Choose between FP16 and BF16 (bfloat16) training. "
+        "BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.",
+    )
+
     parser.add_argument(
         "--fp16", default=False, action="store_true", help="Whether or not to use mixed precision training."
     )
@@ -134,7 +145,16 @@ def simple_launcher(args):
 
     current_env = os.environ.copy()
     current_env["USE_CPU"] = str(args.cpu)
-    current_env["USE_FP16"] = str(args.fp16)
+
+    mixed_precision = args.mixed_precision.lower()
+    if mixed_precision not in ["no", "fp16", "bf16"]:
+        raise ValueError(f"Unknown mixed_precision mode: {mixed_precision}. Choose between 'no', 'fp16' and 'bf16'.")
+
+    if args.fp16:
+        warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', DeprecationWarning)
+        mixed_precision = "fp16"
+
+    current_env["MIXED_PRECISION"] = str(mixed_precision)
 
     process = subprocess.Popen(cmd, env=current_env)
     process.wait()
@@ -168,7 +188,16 @@ def multi_gpu_launcher(args):
     cmd.extend(args.training_script_args)
 
     current_env = os.environ.copy()
-    current_env["USE_FP16"] = str(args.fp16)
+    mixed_precision = args.mixed_precision.lower()
+
+    if mixed_precision not in ["no", "fp16", "bf16"]:
+        raise ValueError(f"Unknown mixed_precision mode: {mixed_precision}. Choose between 'no', 'fp16' and 'bf16'.")
+
+    if args.fp16:
+        warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', DeprecationWarning)
+        mixed_precision = "fp16"
+
+    current_env["MIXED_PRECISION"] = str(mixed_precision)
 
     process = subprocess.Popen(cmd, env=current_env)
     process.wait()
@@ -201,7 +230,16 @@ def deepspeed_launcher(args):
     cmd.extend(args.training_script_args)
 
     current_env = os.environ.copy()
-    current_env["USE_FP16"] = str(args.fp16)
+    mixed_precision = args.mixed_precision.lower()
+
+    if mixed_precision not in ["no", "fp16", "bf16"]:
+        raise ValueError(f"Unknown mixed_precision mode: {mixed_precision}. Choose between 'no', 'fp16' and 'bf16'.")
+
+    if args.fp16:
+        warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', DeprecationWarning)
+        mixed_precision = "fp16"
+
+    current_env["MIXED_PRECISION"] = str(mixed_precision)
     current_env["USE_DEEPSPEED"] = "true"
     current_env["DEEPSPEED_ZERO_STAGE"] = str(args.zero_stage)
     current_env["GRADIENT_ACCUMULATION_STEPS"] = str(args.gradient_accumulation_steps)
@@ -309,8 +347,17 @@ def sagemaker_launcher(sagemaker_config: SageMakerConfig, args):
     print("Converting Arguments to Hyperparameters")
     hyperparameters = _convert_nargs_to_dict(args.training_script_args)
 
-    environment = {"USE_FP16": args.fp16}  # Environment variables to be set for use during training job
+    mixed_precision = args.mixed_precision.lower()
+
+    if mixed_precision not in ["no", "fp16", "bf16"]:
+        raise ValueError(f"Unknown mixed_precision mode: {mixed_precision}. Choose between 'no', 'fp16' and 'bf16'.")
 
+    if args.fp16:
+        warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', DeprecationWarning)
+        mixed_precision = "fp16"
+
+    # Environment variables to be set for use during training job
+    environment = {"MIXED_PRECISION": str(mixed_precision)}
     # configure distribution set up
     distribution = None  # TODO: not yet implemented
 
@@ -360,13 +407,16 @@ def launch_command(args):
 
                 # Those args are handled separately
                 if (
-                    name not in ["compute_environment", "fp16", "distributed_type"]
+                    name not in ["compute_environment", "fp16", "mixed_precision", "distributed_type"]
                     and getattr(args, name, None) is None
                 ):
                     setattr(args, name, attr)
 
-        if not args.fp16:
-            args.fp16 = defaults.fp16
+        if not args.mixed_precision:
+            if args.fp16:
+                args.mixed_precision = "fp16"
+            else:
+                args.mixed_precision = defaults.mixed_precision
     else:
         if args.num_processes is None:
             args.num_processes = 1
diff --git a/src/accelerate/notebook_launcher.py b/src/accelerate/notebook_launcher.py
index b848e13d4..90b4f826a 100644
--- a/src/accelerate/notebook_launcher.py
+++ b/src/accelerate/notebook_launcher.py
@@ -14,6 +14,7 @@
 
 import os
 import sys
+import warnings
 
 import torch
 
@@ -23,7 +24,7 @@
 from .utils import PrepareForLaunch
 
 
-def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, use_port="29500"):
+def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, mixed_precision="no", use_port="29500"):
     """
     Launches a training function, using several processes if it's possible in the current environment (TPU with
     multiple cores for instance).
@@ -37,8 +38,8 @@ def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, use
         num_processes (:obj:`int`, `optional`):
             The number of processes to use for training. Will default to 8 in Colab/Kaggle if a TPU is available, to
             the number of GPUs available otherwise.
-        use_fp16 (:obj:`bool`, `optional`, defaults to :obj:`False`):
-            If :obj:`True`, will use mixed precision training on multi-GPU.
+        mixed_precision (:obj:`str`, `optional`, defaults to :obj:`no`):
+            If :obj:`fp16` or :obj:`bf16`, will use mixed precision training on multi-GPU.
         use_port (:obj:`str`, `optional`, defaults to :obj:`"29500"`):
             The port to use to communicate between processes when launching a multi-GPU training.
     """
@@ -110,7 +111,18 @@ def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, use
             os.environ["WORLD_SIZE"] = str(num_processes)
             os.environ["MASTER_ADDR"] = "127.0.0.1"
             os.environ["MASTER_PORT"] = str(use_port)
-            os.environ["USE_FP16"] = str(use_fp16)
+
+            mixed_precision = mixed_precision.lower()
+            if mixed_precision not in ["no", "fp16", "bf16"]:
+                raise ValueError(
+                    f"Unknown mixed_precision: {mixed_precision}. Choose between 'no', 'fp16' and 'bf16'."
+                )
+
+            if use_fp16:
+                warnings.warn('use_fp16=True is deprecated. Use mixed_precision="fp16" instead.', DeprecationWarning)
+                mixed_precision = "fp16"
+
+            os.environ["MIXED_PRECISION"] = str(mixed_precision)
 
             launcher = PrepareForLaunch(function, distributed_type="MULTI_GPU")
             try:
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index e4d536e14..6693a34ca 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -66,6 +66,11 @@ def parse_flag_from_env(key, default=False):
     return strtobool(value) == 1  # As its name indicates `strtobool` actually returns an int...
 
 
+def parse_choice_from_env(key, default="no"):
+    value = os.environ.get(key, str(default))
+    return value
+
+
 class DistributedType(str, Enum):
     """
     Represents a type of distributed environment.
@@ -133,18 +138,25 @@ class AcceleratorState:
         - **num_processes** (:obj:`int`) -- The number of processes currently launched in parallel.
         - **process_index** (:obj:`int`) -- The index of the current process.
         - **local_process_index** (:obj:`int`) -- The index of the current process on the current server.
-        - **use_fp16** (:obj:`bool`) -- Whether or not the current script will use mixed precision.
+        - **mixed_precision** (:obj:`str`) -- Whether or not the current script will use mixed precision. If you are
+          using mixed precision, define if you want to use FP16 or BF16 (bfloat16) as the floating point.
     """
 
     _shared_state = {}
 
     def __init__(
-        self, fp16: bool = None, cpu: bool = False, deepspeed_plugin=None, _from_accelerator: bool = False, **kwargs
+        self,
+        mixed_precision: str = None,
+        cpu: bool = False,
+        deepspeed_plugin=None,
+        _from_accelerator: bool = False,
+        **kwargs,
     ):
         self.__dict__ = self._shared_state
         if not getattr(self, "initialized", False):
             self.backend = None
             self.deepspeed_plugin = None
+            mixed_precision = mixed_precision.lower() if mixed_precision else None
             if not _from_accelerator:
                 raise ValueError(
                     "Please make sure to properly initialize your accelerator via `accelerator = Accelerator()` "
@@ -156,7 +168,7 @@ def __init__(
                 self.process_index = xm.get_ordinal()
                 self.local_process_index = xm.get_local_ordinal()
                 self.device = xm.xla_device()
-                self.use_fp16 = False
+                self.mixed_precision = "no"
             elif os.environ.get("USE_DEEPSPEED", "false") == "true" and not cpu:
                 assert (
                     is_deepspeed_available()
@@ -170,9 +182,14 @@ def __init__(
                 self.local_process_index = int(os.environ.get("LOCAL_RANK", -1))
                 self.device = torch.device("cuda", self.local_process_index)
                 torch.cuda.set_device(self.device)
-                self.use_fp16 = False  # deepspeed handles fp16 using deepspeed_config
-                fp16 = parse_flag_from_env("USE_FP16", False) if fp16 is None else fp16
-                deepspeed_plugin.deepspeed_config.update({"fp16": {"enabled": fp16}})
+                self.mixed_precision = "no"  # deepspeed handles mixed_precision using deepspeed_config
+                mixed_precision = (
+                    parse_choice_from_env("MIXED_PRECISION", "no") if mixed_precision is None else mixed_precision
+                )
+                if mixed_precision == "fp16":
+                    deepspeed_plugin.deepspeed_config.update({"fp16": {"enabled": True}})
+                elif mixed_precision == "bf16":
+                    deepspeed_plugin.deepspeed_config.update({"bfloat16": {"enabled": True}})
                 self.deepspeed_plugin = deepspeed_plugin
             elif int(os.environ.get("LOCAL_RANK", -1)) != -1 and not cpu:
                 self.distributed_type = DistributedType.MULTI_GPU
@@ -184,7 +201,10 @@ def __init__(
                 self.local_process_index = int(os.environ.get("LOCAL_RANK", -1))
                 self.device = torch.device("cuda", self.local_process_index)
                 torch.cuda.set_device(self.device)
-                self.use_fp16 = parse_flag_from_env("USE_FP16", False) if fp16 is None else fp16
+                self.mixed_precision = (
+                    parse_choice_from_env("MIXED_PRECISION", "no") if mixed_precision is None else mixed_precision
+                )
+
             elif get_int_from_env(["PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE", "WORLD_SIZE"], 1) > 1:
                 self.distributed_type = DistributedType.MULTI_CPU
                 if is_ccl_available() and get_int_from_env(["CCL_WORKER_COUNT"], 0) > 0:
@@ -221,24 +241,32 @@ def __init__(
                 self.process_index = torch.distributed.get_rank()
                 self.local_process_index = local_rank
                 self.device = torch.device("cpu")
-                self.use_fp16 = False
+                self.mixed_precision = "no"
             else:
                 self.distributed_type = DistributedType.NO
                 self.num_processes = 1
                 self.process_index = self.local_process_index = 0
                 self.device = torch.device("cuda" if torch.cuda.is_available() and not cpu else "cpu")
-                self.use_fp16 = parse_flag_from_env("USE_FP16", False) if fp16 is None else fp16
+                self.mixed_precision = (
+                    parse_choice_from_env("MIXED_PRECISION", "no") if mixed_precision is None else mixed_precision
+                )
             self.initialized = True
 
     def __repr__(self):
-        use_fp16 = self.deepspeed_plugin.fp16 if self.distributed_type == DistributedType.DEEPSPEED else self.use_fp16
+        mixed_precision = self.mixed_precision
+        if self.distributed_type == DistributedType.DEEPSPEED:
+            if self.deepspeed_plugin.fp16:
+                mixed_precision = "fp16"
+            if self.deepspeed_plugin.bflaot16:
+                mixed_precision = "bf16"
+
         repr = (
             f"Distributed environment: {self.distributed_type}{('  Backend: ' + self.backend) if self.backend else ''}\n"
             f"Num processes: {self.num_processes}\n"
             f"Process index: {self.process_index}\n"
             f"Local process index: {self.local_process_index}\n"
             f"Device: {self.device}\n"
-            f"Use FP16 precision: {use_fp16}\n"
+            f"Mixed precision type: {mixed_precision}\n"
         )
         if self.distributed_type == DistributedType.DEEPSPEED:
             repr += f"ds_config: {self.deepspeed_plugin.ds_config}\n"
diff --git a/src/accelerate/test_utils/test_script.py b/src/accelerate/test_utils/test_script.py
index 41efcf61c..ecd8e4f02 100644
--- a/src/accelerate/test_utils/test_script.py
+++ b/src/accelerate/test_utils/test_script.py
@@ -245,6 +245,29 @@ def training_check():
     accelerator.print("Training yielded the same results on one CPU or distributes setup with batch split.")
 
     # Mostly a test that FP16 doesn't crash as the operation inside the model is not converted to FP16
+    print("FP16 training check.")
+    accelerator = Accelerator(mixed_precision="fp16")
+    train_dl = DataLoader(train_set, batch_size=batch_size, shuffle=True, generator=generator)
+    model = RegressionModel()
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
+
+    train_dl, model, optimizer = accelerator.prepare(train_dl, model, optimizer)
+    set_seed(42)
+    generator.manual_seed(42)
+    for _ in range(3):
+        for batch in train_dl:
+            model.zero_grad()
+            output = model(batch["x"])
+            loss = torch.nn.functional.mse_loss(output, batch["y"])
+            accelerator.backward(loss)
+            optimizer.step()
+
+    model = accelerator.unwrap_model(model).cpu()
+    assert torch.allclose(old_model.a, model.a), "Did not obtain the same model on CPU or distributed training."
+    assert torch.allclose(old_model.b, model.b), "Did not obtain the same model on CPU or distributed training."
+
+    # TEST that previous fp16 flag still works
+    print("Legacy FP16 training check.")
     accelerator = Accelerator(fp16=True)
     train_dl = DataLoader(train_set, batch_size=batch_size, shuffle=True, generator=generator)
     model = RegressionModel()
@@ -265,6 +288,28 @@ def training_check():
     assert torch.allclose(old_model.a, model.a), "Did not obtain the same model on CPU or distributed training."
     assert torch.allclose(old_model.b, model.b), "Did not obtain the same model on CPU or distributed training."
 
+    # Mostly a test that BF16 doesn't crash as the operation inside the model is not converted to BF16
+    print("BF16 training check.")
+    accelerator = Accelerator(mixed_precision="bf16")
+    train_dl = DataLoader(train_set, batch_size=batch_size, shuffle=True, generator=generator)
+    model = RegressionModel()
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
+
+    train_dl, model, optimizer = accelerator.prepare(train_dl, model, optimizer)
+    set_seed(42)
+    generator.manual_seed(42)
+    for _ in range(3):
+        for batch in train_dl:
+            model.zero_grad()
+            output = model(batch["x"])
+            loss = torch.nn.functional.mse_loss(output, batch["y"])
+            accelerator.backward(loss)
+            optimizer.step()
+
+    model = accelerator.unwrap_model(model).cpu()
+    assert torch.allclose(old_model.a, model.a), "Did not obtain the same model on CPU or distributed training."
+    assert torch.allclose(old_model.b, model.b), "Did not obtain the same model on CPU or distributed training."
+
 
 def main():
     accelerator = Accelerator()
diff --git a/src/accelerate/test_utils/training.py b/src/accelerate/test_utils/training.py
index 3cdea4517..2ea86fbb1 100644
--- a/src/accelerate/test_utils/training.py
+++ b/src/accelerate/test_utils/training.py
@@ -36,6 +36,10 @@ def __init__(self, a=0, b=0, double_output=False):
         super().__init__()
         self.a = torch.nn.Parameter(torch.tensor(a).float())
         self.b = torch.nn.Parameter(torch.tensor(b).float())
+        self.first_batch = True
 
     def forward(self, x=None):
+        if self.first_batch:
+            print(f"Model dtype: {self.a.dtype}, {self.b.dtype}. Input dtype: {x.dtype}")
+            self.first_batch = False
         return x * self.a + self.b
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 991997282..3af792f1f 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -23,6 +23,8 @@
 import numpy as np
 import torch
 
+from packaging import version
+
 from .state import AcceleratorState, DistributedType, is_deepspeed_available, is_tpu_available
 
 
@@ -240,23 +242,26 @@ def _initialize_tensor(tensor_info):
 
 def convert_to_fp32(tensor):
     """
-    Recursively converts the elements nested list/tuple/dictionary of tensors in FP16 precision to FP32.
+    Recursively converts the elements nested list/tuple/dictionary of tensors in FP16/BF16 precision to FP32.
 
     Args:
         tensor (nested list/tuple/dictionary of :obj:`torch.Tensor`):
-            The data to convert from FP16 to FP32.
+            The data to convert from FP16/BF16 to FP32.
 
     Returns:
-        The same data structure as :obj:`tensor` with all tensors that were in FP16 precision converted to FP32.
+        The same data structure as :obj:`tensor` with all tensors that were in FP16/BF16 precision converted to FP32.
     """
 
     def _convert_to_fp32(tensor):
         return tensor.float()
 
-    def _is_fp16_tensor(tensor):
-        return hasattr(tensor, "dtype") and tensor.dtype == torch.float16
+    def _is_fp16_bf16_tensor(tensor):
+        return hasattr(tensor, "dtype") and (
+            tensor.dtype == torch.float16
+            or (version.parse(torch.__version__) >= version.parse("1.10") and tensor.dtype == torch.bfloat16)
+        )
 
-    return recursively_apply(_convert_to_fp32, tensor, test_type=_is_fp16_tensor)
+    return recursively_apply(_convert_to_fp32, tensor, test_type=_is_fp16_bf16_tensor)
 
 
 def convert_outputs_to_fp32(model_forward):

Default here should be `None`, and then get the value of `MIXED_PRECISION` (as is done for `fp16`).
@@ -101,12 +106,24 @@ def __init__(
         device_placement: bool = True,
         split_batches: bool = False,
         fp16: bool = None,
+        mixed_precision: str = "no",
```suggestion
            warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', DeprecationWarning)
```
We're not inside a script, so I would adapt to use arguments: `fp16=True`is deprecated, use `mixed_precision="fp16"` instead.
@@ -101,12 +106,24 @@ def __init__(
         device_placement: bool = True,
         split_batches: bool = False,
         fp16: bool = None,
+        mixed_precision: str = "no",
         cpu: bool = False,
         deepspeed_plugin: DeepSpeedPlugin = None,
         rng_types: Optional[List[Union[str, RNGType]]] = None,
         dispatch_batches: Optional[bool] = None,
         kwargs_handlers: Optional[List[KwargsHandler]] = None,
     ):
+        mixed_precision = mixed_precision.lower()
+        assert mixed_precision in [
+            "no",
+            "fp16",
+            "bf16",
+        ], f"Unknown mixed_precision: {mixed_precision}. Choose between 'no', 'fp16' and 'bf16'."
+
+        if fp16:
+            warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', DeprecationWarning)
Rephrase as a test and raise a ValueError:
```
if mixed_precision not in [...]:
    raise ValueError(...)
```
@@ -101,12 +106,24 @@ def __init__(
         device_placement: bool = True,
         split_batches: bool = False,
         fp16: bool = None,
+        mixed_precision: str = "no",
         cpu: bool = False,
         deepspeed_plugin: DeepSpeedPlugin = None,
         rng_types: Optional[List[Union[str, RNGType]]] = None,
         dispatch_batches: Optional[bool] = None,
         kwargs_handlers: Optional[List[KwargsHandler]] = None,
     ):
+        mixed_precision = mixed_precision.lower()
+        assert mixed_precision in [
+            "no",
+            "fp16",
+            "bf16",
+        ], f"Unknown mixed_precision: {mixed_precision}. Choose between 'no', 'fp16' and 'bf16'."
```suggestion
        return self.mixed_precision != "no":
```
Seems easier this way.
@@ -195,11 +220,23 @@ def is_local_main_process(self):
 
     @property
     def use_fp16(self):
+        if self.mixed_precision != "no":
+            return True
+        else:
+            return False
We have to be a bit more careful here and add a version test, as the argument `dtype` was only introduce in PyTorch 1.10.  So this first test should include
```
and version.parse(torch.__version__) >= version.parse("1.10")
```
otherwise we will fall in the else branch at the end.
@@ -311,7 +348,12 @@ def prepare_model(self, model):
             kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {}
             model = torch.nn.parallel.DistributedDataParallel(model, **kwargs)
         if self.native_amp:
-            model.forward = torch.cuda.amp.autocast()(model.forward)
+            if self.mixed_precision == "fp16":
+                model.forward = torch.cuda.amp.autocast(dtype=torch.float16)(model.forward)
Same here.
@@ -577,7 +619,13 @@ def autocast(self):
         different will happen otherwise.
         """
         if self.native_amp:
-            autocast_context = torch.cuda.amp.autocast()
+            if self.mixed_precision == "fp16":
It should also be added in the `from_json_file` method above for backward compatibility, and I think we should remove the fp16 key as well, to fully convert the config to the new format.
@@ -97,6 +98,12 @@ def from_yaml_file(cls, yaml_file=None):
             config_dict = yaml.safe_load(f)
         if "compute_environment" not in config_dict:
             config_dict["compute_environment"] = ComputeEnvironment.LOCAL_MACHINE
+        if "mixed_precision" not in config_dict:
As before, use a test and raise a ValueError here.
@@ -134,7 +146,19 @@ def simple_launcher(args):
 
     current_env = os.environ.copy()
     current_env["USE_CPU"] = str(args.cpu)
-    current_env["USE_FP16"] = str(args.fp16)
+
+    mixed_precision = args.mixed_precision.lower()
+    assert mixed_precision in [
+        "no",
+        "fp16",
+        "bf16",
+    ], f"Unknown mixed_precision: {mixed_precision}. Choose between 'no', 'fp16' and 'bf16'."
Same here.
@@ -168,7 +192,18 @@ def multi_gpu_launcher(args):
     cmd.extend(args.training_script_args)
 
     current_env = os.environ.copy()
-    current_env["USE_FP16"] = str(args.fp16)
+    mixed_precision = args.mixed_precision.lower()
+    assert mixed_precision in [
+        "no",
+        "fp16",
+        "bf16",
+    ], f"Unknown mixed_precision: {mixed_precision}. Choose between 'no', 'fp16' and 'bf16'."
Same here.
@@ -201,7 +236,18 @@ def deepspeed_launcher(args):
     cmd.extend(args.training_script_args)
 
     current_env = os.environ.copy()
-    current_env["USE_FP16"] = str(args.fp16)
+    mixed_precision = args.mixed_precision.lower()
+    assert mixed_precision in [
+        "no",
+        "fp16",
+        "bf16",
+    ], f"Unknown mixed_precision: {mixed_precision}. Choose between 'no', 'fp16' and 'bf16'."
And same here.
@@ -309,8 +355,20 @@ def sagemaker_launcher(sagemaker_config: SageMakerConfig, args):
     print("Converting Arguments to Hyperparameters")
     hyperparameters = _convert_nargs_to_dict(args.training_script_args)
 
-    environment = {"USE_FP16": args.fp16}  # Environment variables to be set for use during training job
+    mixed_precision = args.mixed_precision.lower()
+    assert mixed_precision in [
+        "no",
+        "fp16",
+        "bf16",
+    ], f"Unknown mixed_precision: {mixed_precision}. Choose between 'no', 'fp16' and 'bf16'."
Nit, put the comment before and have the dict on one line:
```suggestion
    # Environment variables to be set for use during training job
    environment = {"MIXED_PRECISION": str(mixed_precision)}
```
@@ -309,8 +355,20 @@ def sagemaker_launcher(sagemaker_config: SageMakerConfig, args):
     print("Converting Arguments to Hyperparameters")
     hyperparameters = _convert_nargs_to_dict(args.training_script_args)
 
-    environment = {"USE_FP16": args.fp16}  # Environment variables to be set for use during training job
+    mixed_precision = args.mixed_precision.lower()
+    assert mixed_precision in [
+        "no",
+        "fp16",
+        "bf16",
+    ], f"Unknown mixed_precision: {mixed_precision}. Choose between 'no', 'fp16' and 'bf16'."
+
+    if args.fp16:
+        warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', DeprecationWarning)
+        mixed_precision = "fp16"
 
+    environment = {
+        "MIXED_PRECISION": str(mixed_precision)
+    }  # Environment variables to be set for use during training job
As before, raise a value error here.
@@ -110,7 +113,19 @@ def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, use
             os.environ["WORLD_SIZE"] = str(num_processes)
             os.environ["MASTER_ADDR"] = "127.0.0.1"
             os.environ["MASTER_PORT"] = str(use_port)
-            os.environ["USE_FP16"] = str(use_fp16)
+
+            mixed_precision = mixed_precision.lower()
+            assert mixed_precision in [
+                "no",
+                "fp16",
+                "bf16",
+            ], f"Unknown mixed_precision: {mixed_precision}. Choose between 'no', 'fp16' and 'bf16'."
Should be `None` here
@@ -133,18 +138,25 @@ class AcceleratorState:
         - **num_processes** (:obj:`int`) -- The number of processes currently launched in parallel.
         - **process_index** (:obj:`int`) -- The index of the current process.
         - **local_process_index** (:obj:`int`) -- The index of the current process on the current server.
-        - **use_fp16** (:obj:`bool`) -- Whether or not the current script will use mixed precision.
+        - **mixed_precision** (:obj:`str`) -- Whether or not the current script will use mixed precision. If
+          you are using mixed precision, define if you want to use FP16 or BF16 (bfloat16) as the floating point.
     """
 
     _shared_state = {}
 
     def __init__(
-        self, fp16: bool = None, cpu: bool = False, deepspeed_plugin=None, _from_accelerator: bool = False, **kwargs
+        self,
+        mixed_precision: str = "no",
Test for None before lowering there.
@@ -133,18 +138,25 @@ class AcceleratorState:
         - **num_processes** (:obj:`int`) -- The number of processes currently launched in parallel.
         - **process_index** (:obj:`int`) -- The index of the current process.
         - **local_process_index** (:obj:`int`) -- The index of the current process on the current server.
-        - **use_fp16** (:obj:`bool`) -- Whether or not the current script will use mixed precision.
+        - **mixed_precision** (:obj:`str`) -- Whether or not the current script will use mixed precision. If
+          you are using mixed precision, define if you want to use FP16 or BF16 (bfloat16) as the floating point.
     """
 
     _shared_state = {}
 
     def __init__(
-        self, fp16: bool = None, cpu: bool = False, deepspeed_plugin=None, _from_accelerator: bool = False, **kwargs
+        self,
+        mixed_precision: str = "no",
+        cpu: bool = False,
+        deepspeed_plugin=None,
+        _from_accelerator: bool = False,
+        **kwargs,
     ):
         self.__dict__ = self._shared_state
         if not getattr(self, "initialized", False):
             self.backend = None
             self.deepspeed_plugin = None
+            mixed_precision = mixed_precision.lower()
Get the env variable when `mixed_precision` is None
@@ -170,9 +182,14 @@ def __init__(
                 self.local_process_index = int(os.environ.get("LOCAL_RANK", -1))
                 self.device = torch.device("cuda", self.local_process_index)
                 torch.cuda.set_device(self.device)
-                self.use_fp16 = False  # deepspeed handles fp16 using deepspeed_config
-                fp16 = parse_flag_from_env("USE_FP16", False) if fp16 is None else fp16
-                deepspeed_plugin.deepspeed_config.update({"fp16": {"enabled": fp16}})
+                self.mixed_precision = "no"  # deepspeed handles mixed_precision using deepspeed_config
+                mixed_precision = (
+                    parse_choice_from_env("MIXED_PRECISION", "no") if mixed_precision == "no" else mixed_precision
This would also need a version check, as I'm not sure the `torch.bfloat16` was there in PyTorch 1.6.
@@ -240,23 +240,23 @@ def _initialize_tensor(tensor_info):
 
 def convert_to_fp32(tensor):
     """
-    Recursively converts the elements nested list/tuple/dictionary of tensors in FP16 precision to FP32.
+    Recursively converts the elements nested list/tuple/dictionary of tensors in FP16/BF16 precision to FP32.
 
     Args:
         tensor (nested list/tuple/dictionary of :obj:`torch.Tensor`):
-            The data to convert from FP16 to FP32.
+            The data to convert from FP16/BF16 to FP32.
 
     Returns:
-        The same data structure as :obj:`tensor` with all tensors that were in FP16 precision converted to FP32.
+        The same data structure as :obj:`tensor` with all tensors that were in FP16/BF16 precision converted to FP32.
     """
 
     def _convert_to_fp32(tensor):
         return tensor.float()
 
-    def _is_fp16_tensor(tensor):
-        return hasattr(tensor, "dtype") and tensor.dtype == torch.float16
+    def _is_fp16_bf16_tensor(tensor):
+        return hasattr(tensor, "dtype") and (tensor.dtype == torch.float16 or tensor.dtype == torch.bfloat16)
Version test added 
@@ -311,7 +348,12 @@ def prepare_model(self, model):
             kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {}
             model = torch.nn.parallel.DistributedDataParallel(model, **kwargs)
         if self.native_amp:
-            model.forward = torch.cuda.amp.autocast()(model.forward)
+            if self.mixed_precision == "fp16":
+                model.forward = torch.cuda.amp.autocast(dtype=torch.float16)(model.forward)
I forgot about the "from_json_file" method, i have updated it too and I added a line to delete the fp16 key from old config files
@@ -97,6 +98,12 @@ def from_yaml_file(cls, yaml_file=None):
             config_dict = yaml.safe_load(f)
         if "compute_environment" not in config_dict:
             config_dict["compute_environment"] = ComputeEnvironment.LOCAL_MACHINE
+        if "mixed_precision" not in config_dict:
```suggestion
        if mixed_precision is not None:
```
We usually prefer explicit tests instead of relying on Python's bool conversion.
@@ -101,12 +106,29 @@ def __init__(
         device_placement: bool = True,
         split_batches: bool = False,
         fp16: bool = None,
+        mixed_precision: str = None,
         cpu: bool = False,
         deepspeed_plugin: DeepSpeedPlugin = None,
         rng_types: Optional[List[Union[str, RNGType]]] = None,
         dispatch_batches: Optional[bool] = None,
         kwargs_handlers: Optional[List[KwargsHandler]] = None,
     ):
+
+        if mixed_precision:
```suggestion
            if mixed_precision not in ["no", "fp16", "bf16"]:
```
Should fit in one line now.
@@ -101,12 +106,29 @@ def __init__(
         device_placement: bool = True,
         split_batches: bool = False,
         fp16: bool = None,
+        mixed_precision: str = None,
         cpu: bool = False,
         deepspeed_plugin: DeepSpeedPlugin = None,
         rng_types: Optional[List[Union[str, RNGType]]] = None,
         dispatch_batches: Optional[bool] = None,
         kwargs_handlers: Optional[List[KwargsHandler]] = None,
     ):
+
+        if mixed_precision:
+            mixed_precision = mixed_precision.lower()
+            if mixed_precision not in [
+                "no",
+                "fp16",
+                "bf16",
+            ]:
I realize we are silently not training in mixed precision when:
- `mixed_precision="fp16"` and torch version is < 1.6
- `mixed_precision="bf16"` and torch version is < 1.10

I think we should raise an error message in both cases, or at least a warning.
@@ -149,10 +175,14 @@ def __init__(
         # Mixed precision attributes
         self.scaler = None
         self.native_amp = False
-        if self.state.use_fp16:
+        if self.state.mixed_precision == "fp16":
             self.native_amp = version.parse(torch.__version__) >= version.parse("1.6")
             kwargs = self.scaler_handler.to_kwargs() if self.scaler_handler is not None else {}
             self.scaler = torch.cuda.amp.GradScaler(**kwargs)
+        elif self.state.mixed_precision == "bf16":
+            self.native_amp = version.parse(torch.__version__) >= version.parse("1.10")
+            kwargs = self.scaler_handler.to_kwargs() if self.scaler_handler is not None else {}
+            self.scaler = torch.cuda.amp.GradScaler(**kwargs)
Nice!
@@ -97,6 +98,12 @@ def from_yaml_file(cls, yaml_file=None):
             config_dict = yaml.safe_load(f)
         if "compute_environment" not in config_dict:
             config_dict["compute_environment"] = ComputeEnvironment.LOCAL_MACHINE
+        if "mixed_precision" not in config_dict:
No need for this new line I think.
```suggestion
```
@@ -53,9 +54,20 @@ def launch_command_parser(subparsers=None):
     parser.add_argument(
         "--tpu", default=False, action="store_true", help="Whether or not this should launch a TPU training."
     )
+    parser.add_argument(
+        "--mixed_precision",
+        default="no",
+        type=str,
+        choices=["no", "fp16", "bf16"],
+        help="Whether or not to use mixed precision training. "
+        "Choose between FP16 and BF16 (bfloat16) training. "
+        "BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.",
+    )
+
     parser.add_argument(
         "--fp16", default=False, action="store_true", help="Whether or not to use mixed precision training."
     )
+
```suggestion
    if mixed_precision not in ["no", "fp16", "bf16"]:
```
@@ -134,7 +146,20 @@ def simple_launcher(args):
 
     current_env = os.environ.copy()
     current_env["USE_CPU"] = str(args.cpu)
-    current_env["USE_FP16"] = str(args.fp16)
+
+    mixed_precision = args.mixed_precision.lower()
+    if mixed_precision not in [
+        "no",
+        "fp16",
+        "bf16",
+    ]:
```suggestion
    if mixed_precision not in ["no", "fp16", "bf16"]:
```
@@ -168,7 +193,20 @@ def multi_gpu_launcher(args):
     cmd.extend(args.training_script_args)
 
     current_env = os.environ.copy()
-    current_env["USE_FP16"] = str(args.fp16)
+    mixed_precision = args.mixed_precision.lower()
+
+    if mixed_precision not in [
+        "no",
+        "fp16",
+        "bf16",
+    ]:
```suggestion
    if mixed_precision not in ["no", "fp16", "bf16"]:
```
@@ -201,7 +239,20 @@ def deepspeed_launcher(args):
     cmd.extend(args.training_script_args)
 
     current_env = os.environ.copy()
-    current_env["USE_FP16"] = str(args.fp16)
+    mixed_precision = args.mixed_precision.lower()
+
+    if mixed_precision not in [
+        "no",
+        "fp16",
+        "bf16",
+    ]:
```suggestion
    if mixed_precision not in ["no", "fp16", "bf16"]:
```
@@ -309,8 +360,21 @@ def sagemaker_launcher(sagemaker_config: SageMakerConfig, args):
     print("Converting Arguments to Hyperparameters")
     hyperparameters = _convert_nargs_to_dict(args.training_script_args)
 
-    environment = {"USE_FP16": args.fp16}  # Environment variables to be set for use during training job
+    mixed_precision = args.mixed_precision.lower()
+
+    if mixed_precision not in [
+        "no",
+        "fp16",
+        "bf16",
+    ]:
You can remove this from the documentation (we try to avoid documenting deprecated args so users don't start using them).
@@ -39,6 +40,8 @@ def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, use
             the number of GPUs available otherwise.
         use_fp16 (:obj:`bool`, `optional`, defaults to :obj:`False`):
             If :obj:`True`, will use mixed precision training on multi-GPU.
```suggestion
            if mixed_precision not in ["no", "fp16", "bf16"]:
```
@@ -110,7 +113,22 @@ def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, use
             os.environ["WORLD_SIZE"] = str(num_processes)
             os.environ["MASTER_ADDR"] = "127.0.0.1"
             os.environ["MASTER_PORT"] = str(use_port)
-            os.environ["USE_FP16"] = str(use_fp16)
+
+            mixed_precision = mixed_precision.lower()
+            if mixed_precision not in [
+                "no",
+                "fp16",
+                "bf16",
+            ]:
I have updated the PR with the fixes you mention 
Let me know if you need any help with the last standing comments. When you're finished, make sure to run `make style` on your branch so that the quality check passes :-)
I think they are solved now. Maybe we should add a warning in 'cluster.py' to warn if the user chooses the "fp16" or "bf16" option but doesn't have the correct version of Pytorch installed?
This [comment](https://github.com/huggingface/accelerate/pull/247/files#r800669214) and this [one](https://github.com/huggingface/accelerate/pull/247/files#r800671521) don't seem to have been addressed. Could you take care of them?
Quality check is also still failing.

We can also add warnings when the options are chosen, yes!
Sorry the previous commit to solve the problem of silently not training in mixed precision was incorrect. I have a question about implementation. If  the version is correct we set self.native_amp to True. 
```Python
self.native_amp = version.parse(torch.__version__) >= version.parse("1.6")
```
If the version is incorrect we disable mixed precision by setting that variable to False. Is it used for anything else? If we are going to raise an error, should we remove it? 
Should we raise a error always if the pytorch version is not correct? I'm not sure I understand this part of the code 100% and I don't know if I've implemented right the errors if the pytorch version is not compatible with the chosen mixed precision. 


We should at least raise a warning, to tell the user that we are not delivering what they asked for.
What I mean is that what I have done is adding the following code (same for bf16 but it checks pytorch 1.10).  This raises an error if you try to run "fp16" with Pytorch < 1.6. However, the previous line sets self.native_amp=True with the oposite condition. 
We will never run the training if native_amp = False because we raise an error. Is this the correct behavior? If it is, should we keep the native_amp variable?

```Python
if self.state.mixed_precision == "fp16":
    self.native_amp = version.parse(torch.__version__) >= version.parse("1.6")
    if version.parse(torch.__version__) < version.parse("1.6"):
        raise ValueError("fp16 mixed precision requires PyTorch >= 1.6")
```
The way you did it is great! Thanks again for your awesome contribution :-)
https://api.github.com/repos/huggingface/accelerate/pulls/247
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 76c3067a3..d3e4fa603 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -302,10 +302,7 @@ def prepare_model(self, model):
         if self.distributed_type == DistributedType.MULTI_GPU:
             kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {}
             model = torch.nn.parallel.DistributedDataParallel(
-                model,
-                device_ids=[self.local_process_index],
-                output_device=self.local_process_index,
-                **kwargs,
+                model, device_ids=[self.local_process_index], output_device=self.local_process_index, **kwargs
             )
         elif self.distributed_type == DistributedType.MULTI_CPU:
             kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {}
@@ -352,8 +349,12 @@ def _prepare_deepspeed(self, *args):
             is_adamw = isinstance(optimizer, torch.optim.AdamW)
             if (is_adam or is_adamw) and deepspeed_plugin.offload_optimizer_device == "cpu":
                 defaults = optimizer.defaults
+                params = []
+                for group in optimizer.param_groups:
+                    params.extend(group["params"])
+
                 optimizer = deepspeed.ops.adam.DeepSpeedCPUAdam(
-                    model.parameters(),
+                    params,
                     lr=defaults["lr"],
                     bias_correction=True,
                     betas=defaults["betas"],

Could you just run `make style` on your branch to fix the quality issue? Thanks!
is this okay? it looks like there are a few other lines in there that also got style fixed
Looks good, thanks again!
https://api.github.com/repos/huggingface/accelerate/pulls/246
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 443c55622..6eba4b295 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -118,11 +118,11 @@ def honor_type(obj, generator):
     """
     Cast a generator to the same type as obj (list, tuple or namedtuple)
     """
-    # There is no direct check whether an object if of type namedtuple sadly, this is a workaround.
-    if isinstance(obj, tuple) and hasattr(obj, "_fields"):
-        # Can instantiate a namedtuple from a generator directly, contrary to a tuple/list.
+    try:
+        return type(obj)(generator)
+    except TypeError:
+        # Some objects may not be able to instantiate from a generator directly
         return type(obj)(*list(generator))
-    return type(obj)(generator)
 
 
 def is_torch_tensor(tensor):

Sorry I was on my Chinese new year leave and just get back.
Yeah I do agree this should raise exact error with enough information for debugging.
Could you just run `male style` on your branch so we can merge this?
hi, I just noticed the existence of this pr......
I've fixed the linting issue related to this pr, but there seems to be more on other code
Seems like the checks are all passing now, thanks again!
https://api.github.com/repos/huggingface/accelerate/pulls/241
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 3af792f1f..36eaea592 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -18,7 +18,7 @@
 from collections.abc import Mapping
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import List, Optional, Union
+from typing import Any, List, Optional, Union
 
 import numpy as np
 import torch
@@ -350,6 +350,39 @@ def gather(tensor):
         return tensor
 
 
+def _gpu_gather_object(object: Any):
+    def _gpu_gather_object_one(object: Any):
+        output_objects = [None for _ in range(AcceleratorState().num_processes)]
+        torch.distributed.all_gather_object(output_objects, object)
+        return output_objects
+
+    return recursively_apply(_gpu_gather_object_one, object)
+
+
+_cpu_gather_object = _gpu_gather_object
+
+
+def gather_object(object: Any):
+    """
+    Recursively gather object in a nested list/tuple/dictionary of objects from all devices.
+
+    Args:
+        object (nested list/tuple/dictionary of picklable object):
+            The data to gather.
+
+    Returns:
+        The same data structure as :obj:`object` with all the objects sent to every device.
+    """
+    if AcceleratorState().distributed_type == DistributedType.TPU:
+        raise NotImplementedError("gather objects in TPU is not supported")
+    elif AcceleratorState().distributed_type == DistributedType.MULTI_GPU:
+        return _gpu_gather_object(object)
+    elif AcceleratorState().distributed_type == DistributedType.MULTI_CPU:
+        return _cpu_gather_object(object)
+    else:
+        return object
+
+
 def _gpu_broadcast(data, src=0):
     def _gpu_broadcast_one(tensor, src=0):
         torch.distributed.broadcast(tensor, src=src)

```suggestion
        object (nested list/tuple/dictionary of any object):
```
@@ -345,6 +345,40 @@ def gather(tensor):
         return tensor
 
 
+
+def _gpu_gather_object(object: Any):
+    def _gpu_gather_object_one(object: Any):
+        output_objects = [None for _ in range(torch.distributed.get_world_size())]
+        torch.distributed.all_gather_object(output_objects, object)
+        return [i for j in output_objects for i in j]
+
+    return recursively_apply(_gpu_gather_one, tensor, error_on_other_type=True)
+
+
+_cpu_gather_object = _gpu_gather_object
+
+
+def gather_object(object: Any):
+    """
+    Recursively gather object in a nested list/tuple/dictionary of objects from all devices.
+
+    Args:
+        object (nested list/tuple/dictionary of :obj:`Any`):
Here we should use the `state.num_processes` instead.
@@ -345,6 +345,40 @@ def gather(tensor):
         return tensor
 
 
+
+def _gpu_gather_object(object: Any):
+    def _gpu_gather_object_one(object: Any):
+        output_objects = [None for _ in range(torch.distributed.get_world_size())]
I don't understand this line. This presupposes that every element in `output_objects` is iterable, but `object` is supposed to be a primitive object at this point (tensor or other if we remove `error_on_type=True`).
@@ -345,6 +345,40 @@ def gather(tensor):
         return tensor
 
 
+
+def _gpu_gather_object(object: Any):
+    def _gpu_gather_object_one(object: Any):
+        output_objects = [None for _ in range(torch.distributed.get_world_size())]
+        torch.distributed.all_gather_object(output_objects, object)
+        return [i for j in output_objects for i in j]
`torch.all_gather_object` works for all picklable objects, so we can skip the `error_on_other_type=True` I believe.
@@ -345,6 +345,40 @@ def gather(tensor):
         return tensor
 
 
+
+def _gpu_gather_object(object: Any):
+    def _gpu_gather_object_one(object: Any):
+        output_objects = [None for _ in range(torch.distributed.get_world_size())]
+        torch.distributed.all_gather_object(output_objects, object)
+        return [i for j in output_objects for i in j]
+
+    return recursively_apply(_gpu_gather_one, tensor, error_on_other_type=True)
```suggestion
        The same data structure as :obj:`object` with all the objects sent to every device.
```
I think we need to explain where the list arrives (since each tensor is replaced by a list of tensors normally).
@@ -345,6 +345,40 @@ def gather(tensor):
         return tensor
 
 
+
+def _gpu_gather_object(object: Any):
+    def _gpu_gather_object_one(object: Any):
+        output_objects = [None for _ in range(torch.distributed.get_world_size())]
+        torch.distributed.all_gather_object(output_objects, object)
+        return [i for j in output_objects for i in j]
+
+    return recursively_apply(_gpu_gather_one, tensor, error_on_other_type=True)
+
+
+_cpu_gather_object = _gpu_gather_object
+
+
+def gather_object(object: Any):
+    """
+    Recursively gather object in a nested list/tuple/dictionary of objects from all devices.
+
+    Args:
+        object (nested list/tuple/dictionary of :obj:`Any`):
+            The data to gather.
+
+    Returns:
+        The same data structure as :obj:`Any` with all objects sent to the proper device.
```suggestion
        output_objects = [None for _ in range(AcceleratorState().num_processes)]
```
`state` is not defined.
@@ -350,6 +350,39 @@ def gather(tensor):
         return tensor
 
 
+def _gpu_gather_object(object: Any):
+    def _gpu_gather_object_one(object: Any):
+        output_objects = [None for _ in range(state.num_processes)]
```suggestion
    return recursively_apply(_gpu_gather_object_one, object)
```
@@ -350,6 +350,39 @@ def gather(tensor):
         return tensor
 
 
+def _gpu_gather_object(object: Any):
+    def _gpu_gather_object_one(object: Any):
+        output_objects = [None for _ in range(state.num_processes)]
+        torch.distributed.all_gather_object(output_objects, object)
+        return output_objects
+
+    return recursively_apply(_gpu_gather_one, tensor)
> I don't understand this line. This presupposes that every element in output_objects is iterable, but object is supposed to be a primitive object at this point (tensor or other if we remove error_on_type=True).

Should we separate `gather_object` to two functions, `gather_object_list` perform `[i for j in output_objects for i in j]` so it works much like original `gather`, while `gather_object` doesn not perform concatenation so it has better support for every objects

Or maybe we could add an argument to control it's behavior like

```python
def _gpu_gather_object(object: Any, concat: bool = True):
    def _gpu_gather_object_one(object: Any):
        output_objects = [None for _ in range(torch.distributed.get_world_size())]
        torch.distributed.all_gather_object(output_objects, object)
        if concat:
            return [i for j in output_objects for i in j]
        else:
            return output_objects 
```
It can't really concatenate if the tensors are not of the same size, so `gather_object` should not concatenate IMO.
> It can't really concatenate if the tensors are not of the same size, so `gather_object` should not concatenate IMO.

Yes we can leave it to the user to concat manually. 
I have resolved all conversations, could you please review it again?
> Thanks for addressing all comments, this looks good to merge! Could you just run `make style` on your branch to fix the quality failure?

Sorry for this late reply.... I believe I've resolved the quality issue.
> Here is how to solve the last quality issues (variables not defined).

Thank you for your kind suggestions! I have resolved these
(Sorry I'm recently experiencing some computer issues and `make style` is not giving me full outputs. 
Hmmmm, is it possible for us to determine the type of input data and auto call corresponding gather function? I just noticed [torchvison](https://github.com/pytorch/vision/blob/main/references/detection/utils.py#L70-L83) uses `gather_object` at all time
That's too magical for Accelerate's philosophy, users should pick which function they want to use.

All green now so we can merge your PR :-) Thanks again!
> That's too magical for Accelerate's philosophy, users should pick which function they want to use.

Haha, thats the reason why i love this repo so much. Users should be trained programmer and should know what they're doing/supposed to do. 

> All green now so we can merge your PR :-) Thanks again!

Thank you for all the review and comments~
https://api.github.com/repos/huggingface/accelerate/pulls/238
diff --git a/docs/source/installation.md b/docs/source/installation.md
index 6dac579e9..ca2f88286 100644
--- a/docs/source/installation.md
+++ b/docs/source/installation.md
@@ -73,7 +73,7 @@ If you want to constantly use the bleeding edge `master` version of the source c
 
 ``` bash
 git clone https://github.com/huggingface/accelerate.git
-cd transformers
+cd accelerate
 pip install -e .
 ```
 

Sorry I forgot to merge for some reason...
https://api.github.com/repos/huggingface/accelerate/pulls/234
diff --git a/examples/nlp_example.py b/examples/nlp_example.py
index aad2b8b65..a1f310271 100644
--- a/examples/nlp_example.py
+++ b/examples/nlp_example.py
@@ -125,7 +125,7 @@ def collate_fn(examples):
     lr_scheduler = get_linear_schedule_with_warmup(
         optimizer=optimizer,
         num_warmup_steps=100,
-        num_training_steps=len(train_dataloader) * num_epochs,
+        num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,
     )
 
     # Now we train the model
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index f94b2cd99..76c3067a3 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -22,7 +22,7 @@
 from packaging import version
 
 from .data_loader import prepare_data_loader
-from .kwargs_handlers import DistributedDataParallelKwargs, GradScalerKwargs, KwargsHandler
+from .kwargs_handlers import DistributedDataParallelKwargs, GradScalerKwargs, InitProcessGroupKwargs, KwargsHandler
 from .optimizer import AcceleratedOptimizer
 from .state import AcceleratorState, DistributedType, is_deepspeed_available
 from .utils import (
@@ -114,15 +114,10 @@ def __init__(
                 deepspeed_plugin, DeepSpeedPlugin
             ), "`deepspeed_plugin` must be a DeepSpeedPlugin object."
 
-        self.state = AcceleratorState(fp16=fp16, cpu=cpu, deepspeed_plugin=deepspeed_plugin, _from_accelerator=True)
-
-        self.device_placement = device_placement
-        self.split_batches = split_batches
-        self.dispatch_batches = dispatch_batches
-
         # Kwargs handlers
         self.ddp_handler = None
         self.scaler_handler = None
+        self.init_handler = None
         if kwargs_handlers is not None:
             for handler in kwargs_handlers:
                 assert isinstance(handler, KwargsHandler), f"Unsupported kwargs handler passed: {handler}."
@@ -136,6 +131,20 @@ def __init__(
                         raise ValueError("You can only pass one `GradScalerKwargs` in `kwargs_handler`.")
                     else:
                         self.scaler_handler = handler
+                elif isinstance(handler, InitProcessGroupKwargs):
+                    if self.init_handler is not None:
+                        raise ValueError("You can only pass one `InitProcessGroupKwargs` in `kwargs_handler`.")
+                    else:
+                        self.init_handler = handler
+
+        kwargs = self.init_handler.to_kwargs() if self.init_handler is not None else {}
+        self.state = AcceleratorState(
+            fp16=fp16, cpu=cpu, deepspeed_plugin=deepspeed_plugin, _from_accelerator=True, **kwargs
+        )
+
+        self.device_placement = device_placement
+        self.split_batches = split_batches
+        self.dispatch_batches = dispatch_batches
 
         # Mixed precision attributes
         self.scaler = None
diff --git a/src/accelerate/kwargs_handlers.py b/src/accelerate/kwargs_handlers.py
index 7f52d6639..51760efdc 100644
--- a/src/accelerate/kwargs_handlers.py
+++ b/src/accelerate/kwargs_handlers.py
@@ -14,6 +14,8 @@
 
 import copy
 from dataclasses import dataclass
+from datetime import timedelta
+from typing import Optional
 
 
 class KwargsHandler:
@@ -71,3 +73,16 @@ class GradScalerKwargs(KwargsHandler):
     backoff_factor: float = 0.5
     growth_interval: int = 2000
     enabled: bool = True
+
+
+@dataclass
+class InitProcessGroupKwargs(KwargsHandler):
+    """
+    Use this object in your :class:`~accelerate.Accelerator` to customize the initialization of the distributed
+    processes. Please refer to the documentation of this `method
+    <https://pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group>`__ for more information on
+    each argument.
+    """
+
+    init_method: Optional[str] = None
+    timeout: timedelta = timedelta(seconds=1800)
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index 5b349fc38..e4d536e14 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -138,7 +138,9 @@ class AcceleratorState:
 
     _shared_state = {}
 
-    def __init__(self, fp16: bool = None, cpu: bool = False, deepspeed_plugin=None, _from_accelerator: bool = False):
+    def __init__(
+        self, fp16: bool = None, cpu: bool = False, deepspeed_plugin=None, _from_accelerator: bool = False, **kwargs
+    ):
         self.__dict__ = self._shared_state
         if not getattr(self, "initialized", False):
             self.backend = None
@@ -161,7 +163,7 @@ def __init__(self, fp16: bool = None, cpu: bool = False, deepspeed_plugin=None,
                 ), "DeepSpeed is not available => install it using `pip3 install deepspeed` or build it from source"
                 self.distributed_type = DistributedType.DEEPSPEED
                 if not torch.distributed.is_initialized():
-                    torch.distributed.init_process_group(backend="nccl")
+                    torch.distributed.init_process_group(backend="nccl", **kwargs)
                     self.backend = "nccl"
                 self.num_processes = torch.distributed.get_world_size()
                 self.process_index = torch.distributed.get_rank()
@@ -175,7 +177,7 @@ def __init__(self, fp16: bool = None, cpu: bool = False, deepspeed_plugin=None,
             elif int(os.environ.get("LOCAL_RANK", -1)) != -1 and not cpu:
                 self.distributed_type = DistributedType.MULTI_GPU
                 if not torch.distributed.is_initialized():
-                    torch.distributed.init_process_group(backend="nccl")
+                    torch.distributed.init_process_group(backend="nccl", **kwargs)
                     self.backend = "nccl"
                 self.num_processes = torch.distributed.get_world_size()
                 self.process_index = torch.distributed.get_rank()
@@ -213,7 +215,7 @@ def __init__(self, fp16: bool = None, cpu: bool = False, deepspeed_plugin=None,
                             "please try exporting rank 0's hostname as MASTER_ADDR"
                         )
                 if not torch.distributed.is_initialized():
-                    torch.distributed.init_process_group(backend, rank=rank, world_size=size)
+                    torch.distributed.init_process_group(backend, rank=rank, world_size=size, **kwargs)
                     self.backend = backend
                 self.num_processes = torch.distributed.get_world_size()
                 self.process_index = torch.distributed.get_rank()

https://api.github.com/repos/huggingface/accelerate/pulls/228
diff --git a/examples/nlp_example.py b/examples/nlp_example.py
index aad2b8b65..a1f310271 100644
--- a/examples/nlp_example.py
+++ b/examples/nlp_example.py
@@ -125,7 +125,7 @@ def collate_fn(examples):
     lr_scheduler = get_linear_schedule_with_warmup(
         optimizer=optimizer,
         num_warmup_steps=100,
-        num_training_steps=len(train_dataloader) * num_epochs,
+        num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,
     )
 
     # Now we train the model

:+1: 
https://api.github.com/repos/huggingface/accelerate/pulls/227
diff --git a/src/accelerate/optimizer.py b/src/accelerate/optimizer.py
index 4269242b4..52217ba6a 100644
--- a/src/accelerate/optimizer.py
+++ b/src/accelerate/optimizer.py
@@ -53,19 +53,27 @@ class AcceleratedOptimizer(torch.optim.Optimizer):
     def __init__(self, optimizer, device_placement=True, scaler=None):
         self.optimizer = optimizer
         self.scaler = scaler
-        self.state = AcceleratorState()
+        self.accelerator_state = AcceleratorState()
         self.device_placement = device_placement
         self._is_overflow = False
 
         # Handle device placement
         if device_placement:
             state_dict = self.optimizer.state_dict()
-            if self.state.distributed_type == DistributedType.TPU:
-                xm.send_cpu_data_to_device(state_dict, self.state.device)
+            if self.accelerator_state.distributed_type == DistributedType.TPU:
+                xm.send_cpu_data_to_device(state_dict, self.accelerator_state.device)
             else:
-                state_dict = move_to_device(state_dict, self.state.device)
+                state_dict = move_to_device(state_dict, self.accelerator_state.device)
             self.optimizer.load_state_dict(state_dict)
 
+    @property
+    def state(self):
+        return self.optimizer.state
+
+    @state.setter
+    def state(self, state):
+        self.optimizer.state = state
+
     @property
     def param_groups(self):
         return self.optimizer.param_groups
@@ -86,8 +94,8 @@ def add_param_group(self, param_group):
         self.optimizer.add_param_group(param_group)
 
     def load_state_dict(self, state_dict):
-        if self.state.distributed_type == DistributedType.TPU and self.device_placement:
-            xm.send_cpu_data_to_device(state_dict, self.state.device)
+        if self.accelerator_state.distributed_type == DistributedType.TPU and self.device_placement:
+            xm.send_cpu_data_to_device(state_dict, self.accelerator_state.device)
         self.optimizer.load_state_dict(state_dict)
 
     def state_dict(self):
@@ -113,7 +121,7 @@ def zero_grad(self, set_to_none=None):
                 self.optimizer.zero_grad()
 
     def step(self, closure=None):
-        if self.state.distributed_type == DistributedType.TPU:
+        if self.accelerator_state.distributed_type == DistributedType.TPU:
             optimizer_args = {"closure": closure} if closure is not None else {}
             xm.optimizer_step(self.optimizer, optimizer_args=optimizer_args)
         elif self.scaler is not None:

https://api.github.com/repos/huggingface/accelerate/pulls/224
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index 73f8ecc0b..b28f24fc6 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -335,6 +335,7 @@ def __iter__(self):
             # We only iterate through the DataLoader on process 0.
             main_iterator = super().__iter__()
         stop_iteration = False
+        first_batch = None
         while not stop_iteration:
             # On process 0, we gather the batch to dispatch.
             if state.process_index == 0:
@@ -382,7 +383,18 @@ def __iter__(self):
             # Broadcast the batch before splitting it.
             batch = broadcast(batch, from_process=0)
 
-            batch_size = find_batch_size(batch) // state.num_processes
+            if not self.drop_last and first_batch is None:
+                # We keep at least num processes elements of the first batch to be able to complete the last batch
+                first_batch = slice_tensors(batch, slice(0, state.num_processes))
+
+            observed_batch_size = find_batch_size(batch)
+            batch_size = observed_batch_size // state.num_processes
+
+            if not self.drop_last and stop_iteration and observed_batch_size % state.num_processes != 0:
+                # If the last batch is not complete, let's add the first batch to it.
+                batch = concatenate([batch, first_batch], dim=0)
+                batch_size += 1
+
             data_slice = slice(state.process_index * batch_size, (state.process_index + 1) * batch_size)
 
             if state.distributed_type == DistributedType.TPU:
@@ -534,6 +546,7 @@ def prepare_data_loader(
 
     # Need to provide batch_size as batch_sampler is None for Iterable dataset
     if new_batch_sampler is None:
+        kwargs["drop_last"] = dataloader.drop_last
         kwargs["batch_size"] = dataloader.batch_size // num_processes if split_batches else dataloader.batch_size
 
     if dispatch_batches:

Will this run twice on the first batch in case `drop_last` is `False`?
@@ -382,7 +383,18 @@ def __iter__(self):
             # Broadcast the batch before splitting it.
             batch = broadcast(batch, from_process=0)
 
-            batch_size = find_batch_size(batch) // state.num_processes
+            if not self.drop_last and first_batch is None:
+                # We keep at least num processes elements of the first batch to be able to complete the last batch
+                first_batch = slice_tensors(batch, slice(0, state.num_processes))
+
+            observed_batch_size = find_batch_size(batch)
+            batch_size = observed_batch_size // state.num_processes
+
+            if not self.drop_last and stop_iteration and observed_batch_size % state.num_processes != 0:
+                # If the last batch is not complete, let's add the first batch to it.
+                batch = concatenate([batch, first_batch], dim=0)
Yes, like all the dataloading utils, the batches need to be the same length on all processes, which means we have to add some data from the beginning to completet the last batch when `drop_last=False`.
@@ -382,7 +383,18 @@ def __iter__(self):
             # Broadcast the batch before splitting it.
             batch = broadcast(batch, from_process=0)
 
-            batch_size = find_batch_size(batch) // state.num_processes
+            if not self.drop_last and first_batch is None:
+                # We keep at least num processes elements of the first batch to be able to complete the last batch
+                first_batch = slice_tensors(batch, slice(0, state.num_processes))
+
+            observed_batch_size = find_batch_size(batch)
+            batch_size = observed_batch_size // state.num_processes
+
+            if not self.drop_last and stop_iteration and observed_batch_size % state.num_processes != 0:
+                # If the last batch is not complete, let's add the first batch to it.
+                batch = concatenate([batch, first_batch], dim=0)
The `DispatchDataLoader` in general has no test, so we would need to add those before adding a test for the fix. I can do that, but probably not before next week so think we should merge now to at least fix the issue.
Sounds good!
https://api.github.com/repos/huggingface/accelerate/pulls/212
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index ac79d3e0a..f94b2cd99 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -150,7 +150,8 @@ def __init__(
         self._models = []
 
         # RNG Types
-        if rng_types is None:
+        self.rng_types = rng_types
+        if self.rng_types is None:
             self.rng_types = ["torch"] if version.parse(torch.__version__) <= version.parse("1.5.1") else ["generator"]
 
     @property

https://api.github.com/repos/huggingface/accelerate/pulls/206
diff --git a/README.md b/README.md
index f944eea2d..cbc23ed10 100644
--- a/README.md
+++ b/README.md
@@ -206,6 +206,13 @@ You should use 🤗 Accelerate when you want to easily run your training scripts
 
 You shouldn't use 🤗 Accelerate if you don't want to write a training loop yourself. There are plenty of high-level libraries above PyTorch that will offer you that, 🤗 Accelerate is not one of them.
 
+## Frameworks using 🤗 Accelerate
+
+If you like the simplicity of 🤗 Accelerate but would prefer a higher-level abstraction around your training loop, some frameworks that are built on top of 🤗 Accelerate are listed below:
+
+* [pytorch-accelerated](https://github.com/Chris-hughes10/pytorch-accelerated) is a lightweight training library, with a streamlined feature set centred around a general-purpose [Trainer](https://pytorch-accelerated.readthedocs.io/en/latest/trainer.html), that places a huge emphasis on simplicity and transparency; enabling users to understand exactly what is going on under the hood, but without having to write and maintain the boilerplate themselves!
+* [Kornia](https://kornia.readthedocs.io/en/latest/get-started/introduction.html) is a differentiable library that allows classical computer vision to be integrated into deep learning models. Kornia provides a [Trainer](https://kornia.readthedocs.io/en/latest/x.html#kornia.x.Trainer) with the specific purpose to train and fine-tune the supported deep learning algorithms within the library.
+
 ## Installation
 
 This repository is tested on Python 3.6+ and PyTorch 1.4.0+
diff --git a/src/accelerate/notebook_launcher.py b/src/accelerate/notebook_launcher.py
index fef2f0c08..b848e13d4 100644
--- a/src/accelerate/notebook_launcher.py
+++ b/src/accelerate/notebook_launcher.py
@@ -114,7 +114,7 @@ def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, use
 
             launcher = PrepareForLaunch(function, distributed_type="MULTI_GPU")
             try:
-                print(f"Launching a training on {num_processes} GPUs.")
+                print(f"Launching training on {num_processes} GPUs.")
                 start_processes(launcher, args=args, nprocs=num_processes, start_method="fork")
             finally:
                 # Clean up the environment variables set.

No need for those two new lines here :-)
@@ -220,6 +227,8 @@ Then, you will need to install PyTorch: refer to the [official installation page
 pip install accelerate
 ```
 
+
+
Removed.
@@ -220,6 +227,8 @@ Then, you will need to install PyTorch: refer to the [official installation page
 pip install accelerate
 ```
 
+
+
> Thanks for your PR! I was more thinking of having a separate section named "Frameworks using Accelerate" where you could put the same content. There is also kornia using Accelerate internally in their trainer I would add there :-)

Sure thing, I can create an additional section, migrate the content, and will also look into and add Kornia. 
@sgugger I have updated the README, hopefully that is closer to what you had in mind. Would you also like to me add a corresponding section to the docs?
https://api.github.com/repos/huggingface/accelerate/pulls/204
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 3b1281c84..384db658a 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -69,7 +69,7 @@ Changing it to work with accelerate is really easy and only adds a few lines of
           loss = my_loss_function(outputs, targets)
           # Just a small change for the backward instruction
     -     loss.backward()
-    +     accelerate.backward(loss)
+    +     accelerator.backward(loss)
           my_optimizer.step()
 
 and with this, your script can now run in a distributed environment (multi-GPU, TPU).
@@ -97,7 +97,7 @@ especially for TPU training):
           loss = my_loss_function(outputs, targets)
           # Just a small change for the backward instruction
     -     loss.backward()
-    +     accelerate.backward(loss)
+    +     accelerator.backward(loss)
           my_optimizer.step()
 
 

https://api.github.com/repos/huggingface/accelerate/pulls/199
diff --git a/src/accelerate/optimizer.py b/src/accelerate/optimizer.py
index 37ca78833..4269242b4 100644
--- a/src/accelerate/optimizer.py
+++ b/src/accelerate/optimizer.py
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import inspect
+
 import torch
 
 from packaging import version
@@ -100,9 +102,15 @@ def zero_grad(self, set_to_none=None):
                 )
             self.optimizer.zero_grad()
         else:
-            if set_to_none is None:
-                set_to_none = False
-            self.optimizer.zero_grad(set_to_none=set_to_none)
+            accept_arg = "set_to_none" in inspect.signature(self.optimizer.zero_grad).parameters
+            if accept_arg:
+                if set_to_none is None:
+                    set_to_none = False
+                self.optimizer.zero_grad(set_to_none=set_to_none)
+            else:
+                if set_to_none is not None:
+                    raise ValueError("`set_to_none` for Optimizer.zero_grad` is not supported by this optimizer.")
+                self.optimizer.zero_grad()
 
     def step(self, closure=None):
         if self.state.distributed_type == DistributedType.TPU:

https://api.github.com/repos/huggingface/accelerate/pulls/189
diff --git a/README.md b/README.md
index 281325bfe..f944eea2d 100644
--- a/README.md
+++ b/README.md
@@ -155,6 +155,8 @@ For instance, here is how you would run the GLUE example on the MRPC task (from
 accelerate launch examples/nlp_example.py
 ```
 
+This CLI tool is **optional**, and you can still use `python my_script.py` or `python -m torch.distributed.launch my_script.py` at your convenance.
+
 ## Launching multi-CPU run using MPI
 
 🤗 Here is another way to launch multi-CPU run using MPI. You can learn how to install Open MPI on [this page](https://www.open-mpi.org/faq/?category=building#easy-build). You can use Intel MPI or MVAPICH as well.

https://api.github.com/repos/huggingface/accelerate/pulls/187
diff --git a/examples/nlp_example.py b/examples/nlp_example.py
index 220c6ddcc..aad2b8b65 100644
--- a/examples/nlp_example.py
+++ b/examples/nlp_example.py
@@ -162,8 +162,8 @@ def collate_fn(examples):
 
 def main():
     parser = argparse.ArgumentParser(description="Simple example of training script.")
-    parser.add_argument("--fp16", type=bool, default=False, help="If passed, will use FP16 training.")
-    parser.add_argument("--cpu", type=bool, default=False, help="If passed, will train on the CPU.")
+    parser.add_argument("--fp16", action="store_true", help="If passed, will use FP16 training.")
+    parser.add_argument("--cpu", action="store_true", help="If passed, will train on the CPU.")
     args = parser.parse_args()
     config = {"lr": 2e-5, "num_epochs": 3, "correct_bias": True, "seed": 42, "batch_size": 16}
     training_function(config, args)

https://api.github.com/repos/huggingface/accelerate/pulls/183
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index b8428733b..991997282 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -15,7 +15,7 @@
 import importlib
 import os
 import random
-from collections import UserDict
+from collections.abc import Mapping
 from dataclasses import dataclass, field
 from enum import Enum
 from typing import List, Optional, Union
@@ -163,7 +163,7 @@ def recursively_apply(func, data, *args, test_type=is_torch_tensor, error_on_oth
                 for o in data
             ),
         )
-    elif isinstance(data, UserDict):
+    elif isinstance(data, Mapping):
         return type(data)(
             {
                 k: recursively_apply(
@@ -172,15 +172,6 @@ def recursively_apply(func, data, *args, test_type=is_torch_tensor, error_on_oth
                 for k, v in data.items()
             }
         )
-    elif isinstance(data, dict):
-        return type(data)(
-            **{
-                k: recursively_apply(
-                    func, v, *args, test_type=test_type, error_on_other_type=error_on_other_type, **kwargs
-                )
-                for k, v in data.items()
-            }
-        )
     elif test_type(data):
         return func(data, *args, **kwargs)
     elif error_on_other_type:
@@ -310,7 +301,7 @@ def extract_model_from_parallel(model):
 def _tpu_gather(tensor, name="gather tensor"):
     if isinstance(tensor, (list, tuple)):
         return honor_type(tensor, (_tpu_gather(t, name=f"{name}_{i}") for i, t in enumerate(tensor)))
-    elif isinstance(tensor, (dict, UserDict)):
+    elif isinstance(tensor, Mapping):
         return type(tensor)({k: _tpu_gather(v, name=f"{name}_{k}") for k, v in tensor.items()})
     elif not isinstance(tensor, torch.Tensor):
         raise TypeError(f"Can't gather the values of type {type(tensor)}, only of nested list/tuple/dicts of tensors.")
@@ -365,7 +356,7 @@ def _gpu_broadcast_one(tensor, src=0):
 def _tpu_broadcast(tensor, src=0, name="broadcast tensor"):
     if isinstance(tensor, (list, tuple)):
         return honor_type(tensor, (_tpu_broadcast(t, name=f"{name}_{i}") for i, t in enumerate(tensor)))
-    elif isinstance(tensor, (dict, UserDict)):
+    elif isinstance(tensor, Mapping):
         return type(tensor)({k: _tpu_broadcast(v, name=f"{name}_{k}") for k, v in tensor.items()})
     return xm.mesh_reduce(name, tensor, lambda x: x[src])
 
@@ -448,7 +439,7 @@ def find_batch_size(data):
     """
     if isinstance(data, (tuple, list)):
         return find_batch_size(data[0])
-    elif isinstance(data, (dict, UserDict)):
+    elif isinstance(data, Mapping):
         for k in data.keys():
             return find_batch_size(data[k])
     elif not isinstance(data, torch.Tensor):
@@ -471,10 +462,8 @@ def concatenate(data, dim=0):
     """
     if isinstance(data[0], (tuple, list)):
         return honor_type(data[0], (concatenate([d[i] for d in data], dim=dim) for i in range(len(data[0]))))
-    elif isinstance(data[0], UserDict):
+    elif isinstance(data[0], Mapping):
         return type(data[0])({k: concatenate([d[k] for d in data], dim=dim) for k in data[0].keys()})
-    elif isinstance(data[0], dict):
-        return type(data[0])(**{k: concatenate([d[k] for d in data], dim=dim) for k in data[0].keys()})
     elif not isinstance(data[0], torch.Tensor):
         raise TypeError(f"Can only concatenate tensors but got {type(data[0])}")
     return torch.cat(data, dim=dim)

You would need to remove the `**` for this to work with UserDict I believe.
```suggestion
            {
```
@@ -163,16 +163,7 @@ def recursively_apply(func, data, *args, test_type=is_torch_tensor, error_on_oth
                 for o in data
             ),
         )
-    elif isinstance(data, UserDict):
-        return type(data)(
-            {
-                k: recursively_apply(
-                    func, v, *args, test_type=test_type, error_on_other_type=error_on_other_type, **kwargs
-                )
-                for k, v in data.items()
-            }
-        )
-    elif isinstance(data, dict):
+    elif isinstance(data, Mapping):
         return type(data)(
             **{
```suggestion
        return type(data[0])({k: concatenate([d[k] for d in data], dim=dim) for k in data[0].keys()})
```
Same here
@@ -471,9 +462,7 @@ def concatenate(data, dim=0):
     """
     if isinstance(data[0], (tuple, list)):
         return honor_type(data[0], (concatenate([d[i] for d in data], dim=dim) for i in range(len(data[0]))))
-    elif isinstance(data[0], UserDict):
-        return type(data[0])({k: concatenate([d[k] for d in data], dim=dim) for k in data[0].keys()})
-    elif isinstance(data[0], dict):
+    elif isinstance(data[0], Mapping):
         return type(data[0])(**{k: concatenate([d[k] for d in data], dim=dim) for k in data[0].keys()})
`UserDict` also works with the `**` syntax. You can test this with the following code or by looking at the source [here](https://github.com/python/cpython/blob/2c56c97f015a7ea81719615ddcf3c745fba5b4f3/Lib/collections/__init__.py#L964-L984) (warning you that the code is a bit messy in Python 3.6-3.8 compared to Python 3.9 😄):
```python
from collections import UserDict
assert UserDict(**{"a": 1, "b": 2}) == UserDict({"a": 1, "b": 2})
```
@@ -163,16 +163,7 @@ def recursively_apply(func, data, *args, test_type=is_torch_tensor, error_on_oth
                 for o in data
             ),
         )
-    elif isinstance(data, UserDict):
-        return type(data)(
-            {
-                k: recursively_apply(
-                    func, v, *args, test_type=test_type, error_on_other_type=error_on_other_type, **kwargs
-                )
-                for k, v in data.items()
-            }
-        )
-    elif isinstance(data, dict):
+    elif isinstance(data, Mapping):
         return type(data)(
             **{
But not the `BatchEncoding` subclass we are doing all of this for :-)

You can try and see that:
```
from transformers import BatchEncoding
assert BatchEncoding(**{"a": 1, "b": 2}) == BatchEncoding({"a": 1, "b": 2})
```
fails
@@ -163,16 +163,7 @@ def recursively_apply(func, data, *args, test_type=is_torch_tensor, error_on_oth
                 for o in data
             ),
         )
-    elif isinstance(data, UserDict):
-        return type(data)(
-            {
-                k: recursively_apply(
-                    func, v, *args, test_type=test_type, error_on_other_type=error_on_other_type, **kwargs
-                )
-                for k, v in data.items()
-            }
-        )
-    elif isinstance(data, dict):
+    elif isinstance(data, Mapping):
         return type(data)(
             **{
Oh, I see. Thanks.
@@ -163,16 +163,7 @@ def recursively_apply(func, data, *args, test_type=is_torch_tensor, error_on_oth
                 for o in data
             ),
         )
-    elif isinstance(data, UserDict):
-        return type(data)(
-            {
-                k: recursively_apply(
-                    func, v, *args, test_type=test_type, error_on_other_type=error_on_other_type, **kwargs
-                )
-                for k, v in data.items()
-            }
-        )
-    elif isinstance(data, dict):
+    elif isinstance(data, Mapping):
         return type(data)(
             **{
This one still needs to be changed.
@@ -471,9 +462,7 @@ def concatenate(data, dim=0):
     """
     if isinstance(data[0], (tuple, list)):
         return honor_type(data[0], (concatenate([d[i] for d in data], dim=dim) for i in range(len(data[0]))))
-    elif isinstance(data[0], UserDict):
-        return type(data[0])({k: concatenate([d[k] for d in data], dim=dim) for k in data[0].keys()})
-    elif isinstance(data[0], dict):
+    elif isinstance(data[0], Mapping):
         return type(data[0])(**{k: concatenate([d[k] for d in data], dim=dim) for k in data[0].keys()})
But besides replacing `(dict, UserDict)` with `Mapping` in the `isinstance` checks, the code right now is identical to the code you pushed in the initial PR. However, I see an additional simplification we can use if we abandon the `**` syntax (maybe that is what you are aiming for).
@@ -471,9 +462,7 @@ def concatenate(data, dim=0):
     """
     if isinstance(data[0], (tuple, list)):
         return honor_type(data[0], (concatenate([d[i] for d in data], dim=dim) for i in range(len(data[0]))))
-    elif isinstance(data[0], UserDict):
-        return type(data[0])({k: concatenate([d[k] for d in data], dim=dim) for k in data[0].keys()})
-    elif isinstance(data[0], dict):
+    elif isinstance(data[0], Mapping):
         return type(data[0])(**{k: concatenate([d[k] for d in data], dim=dim) for k in data[0].keys()})
cc @sgugger 
@sgugger I've reverted the changes in these two functions, so now the code should work with `BatchEncoding`. However, even though this would add some complexity, maybe a cleaner approach would be to use a specific check for `BatchEncoding`. But for now, I think the current solution is good enough.
Thanks again for your contribution!
https://api.github.com/repos/huggingface/accelerate/pulls/180
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 28a38539a..b8428733b 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -15,6 +15,7 @@
 import importlib
 import os
 import random
+from collections import UserDict
 from dataclasses import dataclass, field
 from enum import Enum
 from typing import List, Optional, Union
@@ -162,6 +163,15 @@ def recursively_apply(func, data, *args, test_type=is_torch_tensor, error_on_oth
                 for o in data
             ),
         )
+    elif isinstance(data, UserDict):
+        return type(data)(
+            {
+                k: recursively_apply(
+                    func, v, *args, test_type=test_type, error_on_other_type=error_on_other_type, **kwargs
+                )
+                for k, v in data.items()
+            }
+        )
     elif isinstance(data, dict):
         return type(data)(
             **{
@@ -300,7 +310,7 @@ def extract_model_from_parallel(model):
 def _tpu_gather(tensor, name="gather tensor"):
     if isinstance(tensor, (list, tuple)):
         return honor_type(tensor, (_tpu_gather(t, name=f"{name}_{i}") for i, t in enumerate(tensor)))
-    elif isinstance(tensor, dict):
+    elif isinstance(tensor, (dict, UserDict)):
         return type(tensor)({k: _tpu_gather(v, name=f"{name}_{k}") for k, v in tensor.items()})
     elif not isinstance(tensor, torch.Tensor):
         raise TypeError(f"Can't gather the values of type {type(tensor)}, only of nested list/tuple/dicts of tensors.")
@@ -355,7 +365,7 @@ def _gpu_broadcast_one(tensor, src=0):
 def _tpu_broadcast(tensor, src=0, name="broadcast tensor"):
     if isinstance(tensor, (list, tuple)):
         return honor_type(tensor, (_tpu_broadcast(t, name=f"{name}_{i}") for i, t in enumerate(tensor)))
-    elif isinstance(tensor, dict):
+    elif isinstance(tensor, (dict, UserDict)):
         return type(tensor)({k: _tpu_broadcast(v, name=f"{name}_{k}") for k, v in tensor.items()})
     return xm.mesh_reduce(name, tensor, lambda x: x[src])
 
@@ -438,7 +448,7 @@ def find_batch_size(data):
     """
     if isinstance(data, (tuple, list)):
         return find_batch_size(data[0])
-    elif isinstance(data, dict):
+    elif isinstance(data, (dict, UserDict)):
         for k in data.keys():
             return find_batch_size(data[k])
     elif not isinstance(data, torch.Tensor):
@@ -461,6 +471,8 @@ def concatenate(data, dim=0):
     """
     if isinstance(data[0], (tuple, list)):
         return honor_type(data[0], (concatenate([d[i] for d in data], dim=dim) for i in range(len(data[0]))))
+    elif isinstance(data[0], UserDict):
+        return type(data[0])({k: concatenate([d[k] for d in data], dim=dim) for k in data[0].keys()})
     elif isinstance(data[0], dict):
         return type(data[0])(**{k: concatenate([d[k] for d in data], dim=dim) for k in data[0].keys()})
     elif not isinstance(data[0], torch.Tensor):
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 9b16aba7b..e515df8fe 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import unittest
-from collections import namedtuple
+from collections import UserDict, namedtuple
 
 import torch
 
@@ -54,3 +54,11 @@ def test_send_to_device(self):
         self.assertTrue(torch.equal(result3.b[0].cpu(), tensor))
         self.assertTrue(torch.equal(result3.b[1].cpu(), tensor))
         self.assertEqual(result3.c, 1)
+
+        result4 = send_to_device(UserDict({"a": tensor, "b": [tensor, tensor], "c": 1}), device)
+        self.assertIsInstance(result4, UserDict)
+        self.assertTrue(torch.equal(result4["a"].cpu(), tensor))
+        self.assertIsInstance(result4["b"], list)
+        self.assertTrue(torch.equal(result4["b"][0].cpu(), tensor))
+        self.assertTrue(torch.equal(result4["b"][1].cpu(), tensor))
+        self.assertEqual(result4["c"], 1)

@sgugger @LysandreJik Instead of adding `UserDict` to the `isinstance` class tuple, you can use `collections.abc.Mapping` to cover both the `dict` and the `UserDict` types (e.g. PyTorch DataLoader uses such check in its default `collate_fn` to stack `dict`-like objects). 

Oh thanks for the tip! I will clean the code to use that when I have a bit of time (unless you want to suggest a PR!)
I'm busy currently, so probably won't find time to fix this myself.
I did find time to fix this, so I've opened a PR.
https://api.github.com/repos/huggingface/accelerate/pulls/179
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index e25b80c49..28a38539a 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -201,7 +201,7 @@ def _send_to_device(t, device):
     def _has_to_method(t):
         return hasattr(t, "to")
 
-    return recursively_apply(_send_to_device, tensor, device, test_type=_has_to_method, error_on_other_type=True)
+    return recursively_apply(_send_to_device, tensor, device, test_type=_has_to_method)
 
 
 def get_data_structure(data):
diff --git a/tests/test_utils.py b/tests/test_utils.py
index ca617634d..9b16aba7b 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -20,7 +20,7 @@
 from accelerate.utils import send_to_device
 
 
-TestNamedTuple = namedtuple("TestNamedTuple", "a b")
+TestNamedTuple = namedtuple("TestNamedTuple", "a b c")
 
 
 class UtilsTester(unittest.TestCase):
@@ -31,23 +31,26 @@ def test_send_to_device(self):
         result1 = send_to_device(tensor, device)
         self.assertTrue(torch.equal(result1.cpu(), tensor))
 
-        result2 = send_to_device((tensor, [tensor, tensor]), device)
+        result2 = send_to_device((tensor, [tensor, tensor], 1), device)
         self.assertIsInstance(result2, tuple)
         self.assertTrue(torch.equal(result2[0].cpu(), tensor))
         self.assertIsInstance(result2[1], list)
         self.assertTrue(torch.equal(result2[1][0].cpu(), tensor))
         self.assertTrue(torch.equal(result2[1][1].cpu(), tensor))
+        self.assertEqual(result2[2], 1)
 
-        result2 = send_to_device({"a": tensor, "b": [tensor, tensor]}, device)
+        result2 = send_to_device({"a": tensor, "b": [tensor, tensor], "c": 1}, device)
         self.assertIsInstance(result2, dict)
         self.assertTrue(torch.equal(result2["a"].cpu(), tensor))
         self.assertIsInstance(result2["b"], list)
         self.assertTrue(torch.equal(result2["b"][0].cpu(), tensor))
         self.assertTrue(torch.equal(result2["b"][1].cpu(), tensor))
+        self.assertEqual(result2["c"], 1)
 
-        result3 = send_to_device(TestNamedTuple(a=tensor, b=[tensor, tensor]), device)
+        result3 = send_to_device(TestNamedTuple(a=tensor, b=[tensor, tensor], c=1), device)
         self.assertIsInstance(result3, TestNamedTuple)
         self.assertTrue(torch.equal(result3.a.cpu(), tensor))
         self.assertIsInstance(result3.b, list)
         self.assertTrue(torch.equal(result3.b[0].cpu(), tensor))
         self.assertTrue(torch.equal(result3.b[1].cpu(), tensor))
+        self.assertEqual(result3.c, 1)

https://api.github.com/repos/huggingface/accelerate/pulls/177
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index 582066414..73f8ecc0b 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import math
 from typing import List, Optional, Union
 
 import torch
@@ -388,6 +389,14 @@ def __iter__(self):
                 xm.mark_step()
             yield slice_tensors(batch, data_slice)
 
+    def __len__(self):
+        state = AcceleratorState()
+        whole_length = super().__len__()
+        if self.drop_last:
+            return whole_length // state.num_processes
+        else:
+            return math.ceil(whole_length / state.num_processes)
+
 
 def prepare_data_loader(
     dataloader: DataLoader,

https://api.github.com/repos/huggingface/accelerate/pulls/175
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 8df05578a..e25b80c49 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -255,7 +255,7 @@ def _convert_to_fp32(tensor):
     def _is_fp16_tensor(tensor):
         return hasattr(tensor, "dtype") and tensor.dtype == torch.float16
 
-    return recursively_apply(_is_fp16_tensor, tensor, test_type=_is_fp16_tensor)
+    return recursively_apply(_convert_to_fp32, tensor, test_type=_is_fp16_tensor)
 
 
 def convert_outputs_to_fp32(model_forward):

Thanks for fixing this!
https://api.github.com/repos/huggingface/accelerate/pulls/173
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index 03e38baf8..582066414 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import warnings
 from typing import List, Optional, Union
 
 import torch
@@ -228,8 +227,8 @@ def __init__(
         process_index: int = 0,
         split_batches: bool = False,
     ):
-        if split_batches and batch_size % num_processes != 0:
-            warnings.warn(
+        if split_batches and batch_size > 1 and batch_size % num_processes != 0:
+            raise ValueError(
                 f"To use `IterableDatasetShard` in `split_batches` mode, the batch size ({batch_size}) "
                 f"needs to be a round multiple of the number of processes ({num_processes})."
             )
@@ -466,9 +465,9 @@ def prepare_data_loader(
         process_index = state.process_index
 
     # Sanity check
-    if split_batches and dataloader.batch_size % num_processes != 0:
-        warnings.warn(
-            f"To use `IterableDatasetShard` in `split_batches` mode, the batch size ({dataloader.batch_size}) "
+    if split_batches and dataloader.batch_size > 1 and dataloader.batch_size % num_processes != 0:
+        raise ValueError(
+            f"To use a `DataLoader` in `split_batches` mode, the batch size ({dataloader.batch_size}) "
             f"needs to be a round multiple of the number of processes ({num_processes})."
         )
 

https://api.github.com/repos/huggingface/accelerate/pulls/170
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index a70ebdaba..ac79d3e0a 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -82,9 +82,10 @@ class Accelerator:
 
             Will default to :obj:`["torch"]` for PyTorch versions <=1.5.1 and :obj:`["generator"]` for PyTorch versions
             >= 1.6.
-        dispatch_batches (:obj:`bool`, `optional`, defaults to :obj:`False`):
+        dispatch_batches (:obj:`bool`, `optional`):
             If set to :obj:`True`, the dataloader prepared by the Accelerator is only iterated through on the main
-            process and then the batches are split and broadcast to each process.
+            process and then the batches are split and broadcast to each process. Will default to :obj:`True` for
+            :obj:`DataLoader` whose underlying dataset is an :obj:`IterableDataset`, :obj:`False` otherwise.
         kwargs_handlers (list of kwargs handlers, `optional`)
             A list of :obj:`KwargHandler` to customize how the objects related to distributed training or mixed
             precision are created. See :doc:`kwargs` for more information.
@@ -103,7 +104,7 @@ def __init__(
         cpu: bool = False,
         deepspeed_plugin: DeepSpeedPlugin = None,
         rng_types: Optional[List[Union[str, RNGType]]] = None,
-        dispatch_batches: bool = False,
+        dispatch_batches: Optional[bool] = None,
         kwargs_handlers: Optional[List[KwargsHandler]] = None,
     ):
         if deepspeed_plugin is None:  # init from env variables
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index 1d464c55f..03e38baf8 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -398,7 +398,7 @@ def prepare_data_loader(
     split_batches: bool = False,
     put_on_device: bool = False,
     rng_types: Optional[List[Union[str, RNGType]]] = None,
-    dispatch_batches: bool = False,
+    dispatch_batches: Optional[bool] = None,
 ) -> DataLoader:
     """
     Wraps a PyTorch :obj:`DataLoader` to generate batches for one of the processes only.
@@ -441,9 +441,10 @@ def prepare_data_loader(
             - :obj:`"generator"`: the :obj:`torch.Generator` of the sampler (or batch sampler if there is no sampler in
               your dataloader) or of the iterable dataset (if it exists) if the underlying dataset is of that type.
 
-        dispatch_batches (:obj:`bool`, `optional`, defaults to :obj:`False`):
+        dispatch_batches (:obj:`bool`, `optional`):
             If set to :obj:`True`, the datalaoder prepared is only iterated through on the main process and then the
-            batches are split and broadcast to each process.
+            batches are split and broadcast to each process. Will default to :obj:`True` when the underlying dataset is
+            an :obj:`IterableDataset`, :obj:`False` otherwise.
 
     Returns:
         :obj:`torch.utils.data.dataloader.DataLoader`: A new data loader that will yield the portion of the batches
@@ -452,6 +453,9 @@ def prepare_data_loader(
 
         This does not support :obj:`BatchSampler` with varying batch size yet.
     """
+    if dispatch_batches is None:
+        dispatch_batches = False if not put_on_device else isinstance(dataloader.dataset, IterableDataset)
+
     if dispatch_batches and not put_on_device:
         raise ValueError("Using `dispatch_batches=True` requires `put_on_device=True`.")
     # Grab defaults from AcceleratorState

Clean!
@@ -452,6 +453,9 @@ def prepare_data_loader(
 
         This does not support :obj:`BatchSampler` with varying batch size yet.
     """
+    if dispatch_batches is None:
+        dispatch_batches = False if not put_on_device else isinstance(new_dataset, IterableDataset)
https://api.github.com/repos/huggingface/accelerate/pulls/168
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index a35f7e268..a70ebdaba 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -82,6 +82,9 @@ class Accelerator:
 
             Will default to :obj:`["torch"]` for PyTorch versions <=1.5.1 and :obj:`["generator"]` for PyTorch versions
             >= 1.6.
+        dispatch_batches (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            If set to :obj:`True`, the dataloader prepared by the Accelerator is only iterated through on the main
+            process and then the batches are split and broadcast to each process.
         kwargs_handlers (list of kwargs handlers, `optional`)
             A list of :obj:`KwargHandler` to customize how the objects related to distributed training or mixed
             precision are created. See :doc:`kwargs` for more information.
@@ -100,6 +103,7 @@ def __init__(
         cpu: bool = False,
         deepspeed_plugin: DeepSpeedPlugin = None,
         rng_types: Optional[List[Union[str, RNGType]]] = None,
+        dispatch_batches: bool = False,
         kwargs_handlers: Optional[List[KwargsHandler]] = None,
     ):
         if deepspeed_plugin is None:  # init from env variables
@@ -113,6 +117,7 @@ def __init__(
 
         self.device_placement = device_placement
         self.split_batches = split_batches
+        self.dispatch_batches = dispatch_batches
 
         # Kwargs handlers
         self.ddp_handler = None
@@ -384,6 +389,7 @@ def prepare_data_loader(self, data_loader):
             split_batches=self.split_batches,
             put_on_device=self.device_placement,
             rng_types=self.rng_types.copy(),
+            dispatch_batches=self.dispatch_batches,
         )
 
     def prepare_optimizer(self, optimizer):
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index 96648ca65..fd10140b6 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -20,7 +20,18 @@
 from packaging import version
 
 from .state import AcceleratorState, DistributedType, is_tpu_available
-from .utils import RNGType, send_to_device, synchronize_rng_states
+from .utils import (
+    RNGType,
+    broadcast,
+    broadcast_object_list,
+    concatenate,
+    find_batch_size,
+    get_data_structure,
+    initialize_tensors,
+    send_to_device,
+    slice_tensors,
+    synchronize_rng_states,
+)
 
 
 if is_tpu_available():
@@ -294,6 +305,90 @@ def __iter__(self):
             yield batch if self.device is None else send_to_device(batch, self.device)
 
 
+class DataLoaderDispatcher(DataLoader):
+    """
+    Subclass of a PyTorch :obj:`DataLoader` that will iterate and preprocess on process 0 only, then dispatch on each
+    process their part of the batch.
+
+    Args:
+        split_batches (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            Whether the resulting :obj:`DataLoader` should split the batches of the original data loader across devices
+            or yield full batches (in which case it will yield batches starting at the :obj:`process_index`-th and
+            advancing of :obj:`num_processes` batches at each iteration).
+
+            Another way to see this is that the observed batch size will be the same as the initial :obj:`dataloader`
+            if this option is set to :obj:`True`, the batch size of the initial :obj:`dataloader` multiplied by
+            :obj:`num_processes` otherwise.
+
+            Setting this option to :obj:`True` requires that the batch size of the :obj:`dataloader` is a round
+            multiple of :obj:`batch_size`.
+    """
+
+    def __init__(self, dataset, split_batches: bool = False, **kwargs):
+        super().__init__(dataset, **kwargs)
+        self.split_batches = split_batches
+
+    def __iter__(self):
+        state = AcceleratorState()
+        if state.process_index == 0:
+            # We only iterate through the DataLoader on process 0.
+            main_iterator = super().__iter__()
+        stop_iteration = False
+        while not stop_iteration:
+            # On process 0, we gather the batch to dispatch.
+            if state.process_index == 0:
+                try:
+                    if self.split_batches:
+                        # One batch of the main iterator is dispatched and split.
+                        batch = next(main_iterator)
+                    else:
+                        # num_processes batches of the main iterator are concatenated then dispatched and split.
+                        # We add the batches one by one so we have the remainder available when drop_last=False.
+                        batches = []
+                        for _ in range(state.num_processes):
+                            batches.append(next(main_iterator))
+                        batch = concatenate(batches, dim=0)
+                    # In both cases, we need to get the structure of the batch that we will broadcast on other
+                    # processes to initialize the tensors with the right shape.
+                    # data_structure, stop_iteration
+                    batch_info = [get_data_structure(batch), False]
+                except StopIteration:
+                    batch_info = [None, True]
+            else:
+                batch_info = [None, stop_iteration]
+
+            # This is inplace, so after this instruction, every process has the same `batch_info` as process 0.
+            broadcast_object_list(batch_info)
+            stop_iteration = batch_info[1]
+            if stop_iteration:
+                # If drop_last is False and split_batches is False, we may have a remainder to take care of.
+                if not self.split_batches and not self.drop_last:
+                    if state.process_index == 0 and len(batches) > 0:
+                        batch = concatenate(batches, dim=0)
+                        batch_info = [get_data_structure(batch), False]
+                    else:
+                        batch_info = [None, True]
+                    broadcast_object_list(batch_info)
+                    if batch_info[1]:
+                        continue
+                else:
+                    continue
+
+            if state.process_index != 0:
+                # Initialize tensors on other processes than process 0.
+                batch = initialize_tensors(batch_info[0])
+            batch = send_to_device(batch, state.device)
+            # Broadcast the batch before splitting it.
+            batch = broadcast(batch, from_process=0)
+
+            batch_size = find_batch_size(batch) // state.num_processes
+            data_slice = slice(state.process_index * batch_size, (state.process_index + 1) * batch_size)
+
+            if state.distributed_type == DistributedType.TPU:
+                xm.mark_step()
+            yield slice_tensors(batch, data_slice)
+
+
 def prepare_data_loader(
     dataloader: DataLoader,
     device: Optional[torch.device] = None,
@@ -302,6 +397,7 @@ def prepare_data_loader(
     split_batches: bool = False,
     put_on_device: bool = False,
     rng_types: Optional[List[Union[str, RNGType]]] = None,
+    dispatch_batches: bool = False,
 ) -> DataLoader:
     """
     Wraps a PyTorch :obj:`DataLoader` to generate batches for one of the processes only.
@@ -344,6 +440,10 @@ def prepare_data_loader(
             - :obj:`"generator"`: the :obj:`torch.Generator` of the sampler (or batch sampler if there is no sampler in
               your dataloader) or of the iterable dataset (if it exists) if the underlying dataset is of that type.
 
+        dispatch_batches (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            If set to :obj:`True`, the datalaoder prepared is only iterated through on the main process and then the
+            batches are split and broadcast to each process.
+
     Returns:
         :obj:`torch.utils.data.dataloader.DataLoader`: A new data loader that will yield the portion of the batches
 
@@ -351,6 +451,8 @@ def prepare_data_loader(
 
         This does not support :obj:`BatchSampler` with varying batch size yet.
     """
+    if dispatch_batches and not put_on_device:
+        raise ValueError("Using `dispatch_batches=True` requires `put_on_device=True`.")
     # Grab defaults from AcceleratorState
     state = AcceleratorState()
     if num_processes is None:
@@ -370,7 +472,7 @@ def prepare_data_loader(
     new_batch_sampler = dataloader.batch_sampler if not isinstance(new_dataset, IterableDataset) else None
     generator = getattr(dataloader, "generator", None)
     # No change if no multiprocess
-    if num_processes != 1:
+    if num_processes != 1 and not dispatch_batches:
         if isinstance(new_dataset, IterableDataset):
             if getattr(dataloader.dataset, "generator", None) is not None:
                 generator = dataloader.dataset.generator
@@ -421,6 +523,11 @@ def prepare_data_loader(
     if new_batch_sampler is None:
         kwargs["batch_size"] = dataloader.batch_size // num_processes if split_batches else dataloader.batch_size
 
+    if dispatch_batches:
+        return DataLoaderDispatcher(
+            new_dataset, split_batches=split_batches, batch_sampler=new_batch_sampler, **kwargs
+        )
+
     return DataLoaderShard(
         new_dataset,
         device=device if put_on_device else None,
diff --git a/src/accelerate/test_utils/test_script.py b/src/accelerate/test_utils/test_script.py
index 2e5095533..41efcf61c 100644
--- a/src/accelerate/test_utils/test_script.py
+++ b/src/accelerate/test_utils/test_script.py
@@ -36,14 +36,14 @@ def init_state_check():
 def rng_sync_check():
     state = AcceleratorState()
     synchronize_rng_states(["torch"])
-    assert are_the_same_tensors(torch.get_rng_state())
+    assert are_the_same_tensors(torch.get_rng_state()), "RNG states improperly synchronized on CPU."
     if state.distributed_type == DistributedType.MULTI_GPU:
         synchronize_rng_states(["cuda"])
-        assert are_the_same_tensors(torch.cuda.get_rng_state())
+        assert are_the_same_tensors(torch.cuda.get_rng_state()), "RNG states improperly synchronized on GPU."
     if version.parse(torch.__version__) >= version.parse("1.6.0"):
         generator = torch.Generator()
         synchronize_rng_states(["generator"], generator=generator)
-        assert are_the_same_tensors(generator.get_state())
+        assert are_the_same_tensors(generator.get_state()), "RNG states improperly synchronized in generator."
 
     if state.local_process_index == 0:
         print("All rng are properly synched.")
@@ -59,7 +59,7 @@ def dl_preparation_check():
     for batch in dl:
         result.append(gather(batch))
     result = torch.cat(result)
-    assert torch.equal(result.cpu(), torch.arange(0, length).long())
+    assert torch.equal(result.cpu(), torch.arange(0, length).long()), "Wrong non-shuffled dataloader result."
 
     dl = DataLoader(range(length), batch_size=8)
     dl = prepare_data_loader(
@@ -74,7 +74,7 @@ def dl_preparation_check():
     for batch in dl:
         result.append(gather(batch))
     result = torch.cat(result)
-    assert torch.equal(result.cpu(), torch.arange(0, length).long())
+    assert torch.equal(result.cpu(), torch.arange(0, length).long()), "Wrong non-shuffled dataloader result."
 
     if state.process_index == 0:
         print("Non-shuffled dataloader passing.")
@@ -86,7 +86,7 @@ def dl_preparation_check():
         result.append(gather(batch))
     result = torch.cat(result).tolist()
     result.sort()
-    assert result == list(range(length))
+    assert result == list(range(length)), "Wrong shuffled dataloader result."
 
     dl = DataLoader(range(length), batch_size=8, shuffle=True)
     dl = prepare_data_loader(
@@ -102,12 +102,77 @@ def dl_preparation_check():
         result.append(gather(batch))
     result = torch.cat(result).tolist()
     result.sort()
-    assert result == list(range(length))
+    assert result == list(range(length)), "Wrong shuffled dataloader result."
 
     if state.local_process_index == 0:
         print("Shuffled dataloader passing.")
 
 
+def central_dl_preparation_check():
+    state = AcceleratorState()
+    length = 32 * state.num_processes
+
+    dl = DataLoader(range(length), batch_size=8)
+    dl = prepare_data_loader(
+        dl, state.device, state.num_processes, state.process_index, put_on_device=True, dispatch_batches=True
+    )
+    result = []
+    for batch in dl:
+        result.append(gather(batch))
+    result = torch.cat(result)
+    assert torch.equal(result.cpu(), torch.arange(0, length).long()), "Wrong non-shuffled dataloader result."
+
+    dl = DataLoader(range(length), batch_size=8)
+    dl = prepare_data_loader(
+        dl,
+        state.device,
+        state.num_processes,
+        state.process_index,
+        put_on_device=True,
+        split_batches=True,
+        dispatch_batches=True,
+    )
+    result = []
+    for batch in dl:
+        result.append(gather(batch))
+    result = torch.cat(result)
+    assert torch.equal(result.cpu(), torch.arange(0, length).long()), "Wrong non-shuffled dataloader result."
+
+    if state.process_index == 0:
+        print("Non-shuffled central dataloader passing.")
+
+    dl = DataLoader(range(length), batch_size=8, shuffle=True)
+    dl = prepare_data_loader(
+        dl, state.device, state.num_processes, state.process_index, put_on_device=True, dispatch_batches=True
+    )
+    result = []
+    for batch in dl:
+        result.append(gather(batch))
+    result = torch.cat(result).tolist()
+    result.sort()
+    assert result == list(range(length)), "Wrong shuffled dataloader result."
+
+    dl = DataLoader(range(length), batch_size=8, shuffle=True)
+    dl = prepare_data_loader(
+        dl,
+        state.device,
+        state.num_processes,
+        state.process_index,
+        put_on_device=True,
+        split_batches=True,
+        dispatch_batches=True,
+    )
+    result = []
+    for batch in dl:
+        result.append(gather(batch))
+    result = torch.cat(result).tolist()
+    result.sort()
+    assert result == list(range(length)), "Wrong shuffled dataloader result."
+
+    if state.local_process_index == 0:
+        print("Shuffled central dataloader passing.")
+
+
 def mock_training(length, batch_size, generator):
     set_seed(42)
     generator.manual_seed(42)
@@ -132,8 +197,8 @@ def training_check():
     length = batch_size * 4 * state.num_processes
 
     train_set, old_model = mock_training(length, batch_size * state.num_processes, generator)
-    assert are_the_same_tensors(old_model.a)
-    assert are_the_same_tensors(old_model.b)
+    assert are_the_same_tensors(old_model.a), "Did not obtain the same model on both processes."
+    assert are_the_same_tensors(old_model.b), "Did not obtain the same model on both processes."
 
     accelerator = Accelerator()
     train_dl = DataLoader(train_set, batch_size=batch_size, shuffle=True, generator=generator)
@@ -152,8 +217,8 @@ def training_check():
             optimizer.step()
 
     model = accelerator.unwrap_model(model).cpu()
-    assert torch.allclose(old_model.a, model.a)
-    assert torch.allclose(old_model.b, model.b)
+    assert torch.allclose(old_model.a, model.a), "Did not obtain the same model on CPU or distributed training."
+    assert torch.allclose(old_model.b, model.b), "Did not obtain the same model on CPU or distributed training."
 
     accelerator.print("Training yielded the same results on one CPU or distributed setup with no batch split.")
 
@@ -174,8 +239,8 @@ def training_check():
             optimizer.step()
 
     model = accelerator.unwrap_model(model).cpu()
-    assert torch.allclose(old_model.a, model.a)
-    assert torch.allclose(old_model.b, model.b)
+    assert torch.allclose(old_model.a, model.a), "Did not obtain the same model on CPU or distributed training."
+    assert torch.allclose(old_model.b, model.b), "Did not obtain the same model on CPU or distributed training."
 
     accelerator.print("Training yielded the same results on one CPU or distributes setup with batch split.")
 
@@ -197,8 +262,8 @@ def training_check():
             optimizer.step()
 
     model = accelerator.unwrap_model(model).cpu()
-    assert torch.allclose(old_model.a, model.a)
-    assert torch.allclose(old_model.b, model.b)
+    assert torch.allclose(old_model.a, model.a), "Did not obtain the same model on CPU or distributed training."
+    assert torch.allclose(old_model.b, model.b), "Did not obtain the same model on CPU or distributed training."
 
 
 def main():
@@ -215,6 +280,7 @@ def main():
     if state.local_process_index == 0:
         print("\n**DataLoader integration test**")
     dl_preparation_check()
+    central_dl_preparation_check()
 
     if state.local_process_index == 0:
         print("\n**Training integration test**")
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 026ab9501..8df05578a 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -48,6 +48,12 @@ class RNGType(Enum):
     GENERATOR = "generator"
 
 
+@dataclass
+class TensorInformation:
+    shape: torch.Size
+    dtype: torch.dtype
+
+
 def set_seed(seed: int):
     """
     Helper function for reproducible behavior to set the seed in ``random``, ``numpy``, ``torch``.
@@ -115,6 +121,66 @@ def honor_type(obj, generator):
     return type(obj)(generator)
 
 
+def is_torch_tensor(tensor):
+    return isinstance(tensor, torch.Tensor)
+
+
+def is_tensor_information(tensor_info):
+    return isinstance(tensor_info, TensorInformation)
+
+
+def recursively_apply(func, data, *args, test_type=is_torch_tensor, error_on_other_type=False, **kwargs):
+    """
+    Recursively apply a function on a data structure that is a nested list/tuple/dictionary of a given base type.
+
+    Args:
+        func (:obj:`callable`):
+            The function to recursively apply.
+        data (nested list/tuple/dictionary of :obj:`main_type`):
+            The data on which to apply :obj:`func`
+        *args:
+            Positional arguments that will be passed to :obj:`func` when applied on the unpacked data.
+        main_type (:obj:`type`, `optional`, defaults to :obj:`torch.Tensor`):
+            The base type of the objects to which apply :obj:`func`.
+        error_on_other_type (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            Whether to return an error or not if after unpacking :obj:`data`, we get on an object that is not of type
+            :obj:`main_type`. If :obj:`False`, the function will leave objects of types different than :obj:`main_type`
+            unchanged.
+        **kwargs:
+            Keyword arguments that will be passed to :obj:`func` when applied on the unpacked data.
+
+    Returns:
+        The same data structure as :obj:`data` with :obj:`func` applied to every object of type :obj:`main_type`.
+    """
+    if isinstance(data, (tuple, list)):
+        return honor_type(
+            data,
+            (
+                recursively_apply(
+                    func, o, *args, test_type=test_type, error_on_other_type=error_on_other_type, **kwargs
+                )
+                for o in data
+            ),
+        )
+    elif isinstance(data, dict):
+        return type(data)(
+            **{
+                k: recursively_apply(
+                    func, v, *args, test_type=test_type, error_on_other_type=error_on_other_type, **kwargs
+                )
+                for k, v in data.items()
+            }
+        )
+    elif test_type(data):
+        return func(data, *args, **kwargs)
+    elif error_on_other_type:
+        raise TypeError(
+            f"Can't apply {func.__name__} on object of type {type(data)}, only of nested list/tuple/dicts of objects "
+            f"that satisfy {test_type.__name__}."
+        )
+    return data
+
+
 def send_to_device(tensor, device):
     """
     Recursively sends the elements in a nested list/tuple/dictionary of tensors to a given device.
@@ -128,18 +194,52 @@ def send_to_device(tensor, device):
     Returns:
         The same data structure as :obj:`tensor` with all tensors sent to the proper device.
     """
-    if isinstance(tensor, (list, tuple)):
-        return honor_type(tensor, (send_to_device(t, device) for t in tensor))
-    elif isinstance(tensor, dict):
-        return type(tensor)({k: send_to_device(v, device) for k, v in tensor.items()})
-    elif not hasattr(tensor, "to"):
-        return tensor
-    return tensor.to(device)
+
+    def _send_to_device(t, device):
+        return t.to(device)
+
+    def _has_to_method(t):
+        return hasattr(t, "to")
+
+    return recursively_apply(_send_to_device, tensor, device, test_type=_has_to_method, error_on_other_type=True)
+
+
+def get_data_structure(data):
+    """
+    Recursively gathers the information needed to rebuild a nested list/tuple/dictionary of tensors.
+
+    Args:
+        data (nested list/tuple/dictionary of :obj:`torch.Tensor`):
+            The data to send to analyze.
+
+    Returns:
+        The same data structure as :obj:`data` with :class:`~accelerate.utils.TensorInformation` instead of tensors.
+    """
+
+    def _get_data_structure(tensor):
+        return TensorInformation(shape=tensor.shape, dtype=tensor.dtype)
+
+    return recursively_apply(_get_data_structure, data)
+
+
+def initialize_tensors(data_structure):
+    """
+    Recursively initializes tensors from a nested list/tuple/dictionary of
+    :class:`~accelerate.utils.TensorInformation`.
+
+    Returns:
+        The same data structure as :obj:`data` with tensors instead of :class:`~accelerate.utils.TensorInformation`.
+    """
+
+    def _initialize_tensor(tensor_info):
+        return torch.empty(*tensor_info.shape, dtype=tensor_info.dtype)
+
+    return recursively_apply(_initialize_tensor, data_structure, test_type=is_tensor_information)
 
 
 def convert_to_fp32(tensor):
     """
-    Recursively converts the lements nested list/tuple/dictionary of tensors in FP16 precision to FP32.
+    Recursively converts the elements nested list/tuple/dictionary of tensors in FP16 precision to FP32.
 
     Args:
         tensor (nested list/tuple/dictionary of :obj:`torch.Tensor`):
@@ -148,13 +248,14 @@ def convert_to_fp32(tensor):
     Returns:
         The same data structure as :obj:`tensor` with all tensors that were in FP16 precision converted to FP32.
     """
-    if isinstance(tensor, (list, tuple)):
-        return honor_type(tensor, (convert_to_fp32(t) for t in tensor))
-    elif isinstance(tensor, dict):
-        return type(tensor)(**{k: convert_to_fp32(v) for k, v in tensor.items()})
-    elif not hasattr(tensor, "dtype") or tensor.dtype != torch.float16:
-        return tensor
-    return tensor.float()
+
+    def _convert_to_fp32(tensor):
+        return tensor.float()
+
+    def _is_fp16_tensor(tensor):
+        return hasattr(tensor, "dtype") and tensor.dtype == torch.float16
+
+    return recursively_apply(_is_fp16_tensor, tensor, test_type=_is_fp16_tensor)
 
 
 def convert_outputs_to_fp32(model_forward):
@@ -196,7 +297,7 @@ def extract_model_from_parallel(model):
     return model
 
 
-def _tpu_gather(tensor, name="tensor"):
+def _tpu_gather(tensor, name="gather tensor"):
     if isinstance(tensor, (list, tuple)):
         return honor_type(tensor, (_tpu_gather(t, name=f"{name}_{i}") for i, t in enumerate(tensor)))
     elif isinstance(tensor, dict):
@@ -209,17 +310,14 @@ def _tpu_gather(tensor, name="tensor"):
 
 
 def _gpu_gather(tensor):
-    if isinstance(tensor, (list, tuple)):
-        return honor_type(tensor, (_gpu_gather(t) for t in tensor))
-    elif isinstance(tensor, dict):
-        return type(tensor)({k: _gpu_gather(v) for k, v in tensor.items()})
-    elif not isinstance(tensor, torch.Tensor):
-        raise TypeError(f"Can't gather the values of type {type(tensor)}, only of nested list/tuple/dicts of tensors.")
-    if tensor.ndim == 0:
-        tensor = tensor.clone()[None]
-    output_tensors = [tensor.clone() for _ in range(torch.distributed.get_world_size())]
-    torch.distributed.all_gather(output_tensors, tensor)
-    return torch.cat(output_tensors, dim=0)
+    def _gpu_gather_one(tensor):
+        if tensor.ndim == 0:
+            tensor = tensor.clone()[None]
+        output_tensors = [tensor.clone() for _ in range(torch.distributed.get_world_size())]
+        torch.distributed.all_gather(output_tensors, tensor)
+        return torch.cat(output_tensors, dim=0)
+
+    return recursively_apply(_gpu_gather_one, tensor, error_on_other_type=True)
 
 
 _cpu_gather = _gpu_gather
@@ -246,6 +344,130 @@ def gather(tensor):
         return tensor
 
 
+def _gpu_broadcast(data, src=0):
+    def _gpu_broadcast_one(tensor, src=0):
+        torch.distributed.broadcast(tensor, src=src)
+        return tensor
+
+    return recursively_apply(_gpu_broadcast_one, data, error_on_other_type=True, src=src)
+
+
+def _tpu_broadcast(tensor, src=0, name="broadcast tensor"):
+    if isinstance(tensor, (list, tuple)):
+        return honor_type(tensor, (_tpu_broadcast(t, name=f"{name}_{i}") for i, t in enumerate(tensor)))
+    elif isinstance(tensor, dict):
+        return type(tensor)({k: _tpu_broadcast(v, name=f"{name}_{k}") for k, v in tensor.items()})
+    return xm.mesh_reduce(name, tensor, lambda x: x[src])
+
+
+def broadcast(tensor, from_process: int = 0):
+    """
+    Recursively broadcast tensor in a nested list/tuple/dictionary of tensors to all devices.
+
+    Args:
+        tensor (nested list/tuple/dictionary of :obj:`torch.Tensor`):
+            The data to gather.
+        from_process (:obj:`int`, `optional`, defaults to 0):
+            The process from which to send the data
+
+    Returns:
+        The same data structure as :obj:`tensor` with all tensors broadcasted to the proper device.
+    """
+    if AcceleratorState().distributed_type == DistributedType.TPU:
+        return _tpu_broadcast(tensor, src=from_process, name="accelerate.utils.broadcast")
+    elif AcceleratorState().distributed_type == DistributedType.MULTI_GPU:
+        return _gpu_broadcast(tensor, src=from_process)
+    elif AcceleratorState().distributed_type == DistributedType.MULTI_CPU:
+        return _gpu_broadcast(tensor, src=from_process)
+    else:
+        return tensor
+
+
+def broadcast_object_list(object_list, from_process: int = 0):
+    """
+    Broadcast a list of picklable objects form one process to the others.
+
+    Args:
+        object_list (list of picklable objects):
+            The list of objects to broadcast. This list will be modified inplace.
+        from_process (:obj:`int`, `optional`, defaults to 0):
+            The process from which to send the data.
+
+    Returns:
+        The same list containing the objects from process 0.
+    """
+    if AcceleratorState().distributed_type == DistributedType.TPU:
+        for i, obj in enumerate(object_list):
+            object_list[i] = xm.mesh_reduce("accelerate.utils.broadcast_object_list", obj, lambda x: x[from_process])
+    elif AcceleratorState().distributed_type == DistributedType.MULTI_GPU:
+        torch.distributed.broadcast_object_list(object_list, src=from_process)
+    elif AcceleratorState().distributed_type == DistributedType.MULTI_CPU:
+        torch.distributed.broadcast_object_list(object_list, src=from_process)
+    return object_list
+
+
+def slice_tensors(data, tensor_slice):
+    """
+    Recursively takes a slice in a nested list/tuple/dictionary of tensors.
+
+    Args:
+        data (nested list/tuple/dictionary of :obj:`torch.Tensor`):
+            The data to slice.
+        tensor_slice (:obj:`slice`):
+            The slice to take.
+
+    Returns:
+        The same data structure as :obj:`data` with all the tensors slices.
+    """
+
+    def _slice_tensor(tensor, tensor_slice):
+        return tensor[tensor_slice]
+
+    return recursively_apply(_slice_tensor, data, tensor_slice)
+
+
+def find_batch_size(data):
+    """
+    Recursively finds the batch size in a nested list/tuple/dictionary of lists of tensors.
+
+    Args:
+        data (nested list/tuple/dictionary of :obj:`torch.Tensor`): The data from which to find the batch size.
+
+    Returns:
+        :obj:`int`: The batch size.
+    """
+    if isinstance(data, (tuple, list)):
+        return find_batch_size(data[0])
+    elif isinstance(data, dict):
+        for k in data.keys():
+            return find_batch_size(data[k])
+    elif not isinstance(data, torch.Tensor):
+        raise TypeError(f"Can only find the batch size of tensors but got {type(data)}.")
+    return data.shape[0]
+
+
+def concatenate(data, dim=0):
+    """
+    Recursively concatenate the tensors in a nested list/tuple/dictionary of lists of tensors with the same shape.
+
+    Args:
+        data (nested list/tuple/dictionary of lists of tensors :obj:`torch.Tensor`):
+            The data to concatenate.
+        dim (:obj:`int`, `optional`, defaults to 0):
+            The dimension on which to concatenate.
+
+    Returns:
+        The same data structure as :obj:`data` with all the tensors concatenated.
+    """
+    if isinstance(data[0], (tuple, list)):
+        return honor_type(data[0], (concatenate([d[i] for d in data], dim=dim) for i in range(len(data[0]))))
+    elif isinstance(data[0], dict):
+        return type(data[0])(**{k: concatenate([d[k] for d in data], dim=dim) for k in data[0].keys()})
+    elif not isinstance(data[0], torch.Tensor):
+        raise TypeError(f"Can only concatenate tensors but got {type(data[0])}")
+    return torch.cat(data, dim=dim)
+
+
 def pad_across_processes(tensor, dim=0, pad_index=0, pad_first=False):
     """
     Recursively pad the tensors in a nested list/tuple/dictionary of tensors from all devices to the same size so they
@@ -261,36 +483,35 @@ def pad_across_processes(tensor, dim=0, pad_index=0, pad_first=False):
         pad_first (:obj:`bool`, `optional`, defaults to :obj:`False`):
             Whether to pad at the beginning or the end.
     """
-    if isinstance(tensor, (list, tuple)):
-        return honor_type(tensor, (pad_across_processes(t, dim=dim, pad_index=pad_index) for t in tensor))
-    elif isinstance(tensor, dict):
-        return type(tensor)({k: pad_across_processes(v, dim=dim, pad_index=pad_index) for k, v in tensor.items()})
-    elif not isinstance(tensor, torch.Tensor):
-        raise TypeError(f"Can't pad the values of type {type(tensor)}, only of nested list/tuple/dicts of tensors.")
 
-    if dim >= len(tensor.shape):
-        return tensor
-
-    # Gather all sizes
-    size = torch.tensor(tensor.shape, device=tensor.device)[None]
-    sizes = gather(size).cpu()
-    # Then pad to the maximum size
-    max_size = max(s[dim] for s in sizes)
-    if max_size == tensor.shape[dim]:
-        return tensor
-
-    old_size = tensor.shape
-    new_size = list(old_size)
-    new_size[dim] = max_size
-    new_tensor = tensor.new_zeros(tuple(new_size)) + pad_index
-    if pad_first:
-        indices = tuple(
-            slice(max_size - old_size[dim], max_size) if i == dim else slice(None) for i in range(len(new_size))
-        )
-    else:
-        indices = tuple(slice(0, old_size[dim]) if i == dim else slice(None) for i in range(len(new_size)))
-    new_tensor[indices] = tensor
-    return new_tensor
+    def _pad_across_processes(tensor, dim=0, pad_index=0, pad_first=False):
+        if dim >= len(tensor.shape):
+            return tensor
+
+        # Gather all sizes
+        size = torch.tensor(tensor.shape, device=tensor.device)[None]
+        sizes = gather(size).cpu()
+        # Then pad to the maximum size
+        max_size = max(s[dim] for s in sizes)
+        if max_size == tensor.shape[dim]:
+            return tensor
+
+        old_size = tensor.shape
+        new_size = list(old_size)
+        new_size[dim] = max_size
+        new_tensor = tensor.new_zeros(tuple(new_size)) + pad_index
+        if pad_first:
+            indices = tuple(
+                slice(max_size - old_size[dim], max_size) if i == dim else slice(None) for i in range(len(new_size))
+            )
+        else:
+            indices = tuple(slice(0, old_size[dim]) if i == dim else slice(None) for i in range(len(new_size)))
+        new_tensor[indices] = tensor
+        return new_tensor
+
+    return recursively_apply(
+        _pad_across_processes, tensor, error_on_other_type=True, dim=dim, pad_index=pad_index, pad_first=pad_first
+    )
 
 
 def wait_for_everyone():

```suggestion
            If set to :obj:`True`, the dataloader prepared by the Accelerator is only iterated through on the main
```
@@ -82,6 +82,9 @@ class Accelerator:
 
             Will default to :obj:`["torch"]` for PyTorch versions <=1.5.1 and :obj:`["generator"]` for PyTorch versions
             >= 1.6.
+        central_dataloader (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            If set to :obj:`True`, the datalaoder prepared by the Accelerator is only iterated through on the main
I don't like `central` as I don't really understand what it means from a quick read, but I spent 10 minutes trying to find a better name without success, so fine to leave it like it is :)
@@ -82,6 +82,9 @@ class Accelerator:
 
             Will default to :obj:`["torch"]` for PyTorch versions <=1.5.1 and :obj:`["generator"]` for PyTorch versions
             >= 1.6.
+        central_dataloader (:obj:`bool`, `optional`, defaults to :obj:`False`):
What is this `my_batch_size` used for?
@@ -421,6 +524,13 @@ def prepare_data_loader(
     if new_batch_sampler is None:
         kwargs["batch_size"] = dataloader.batch_size // num_processes if split_batches else dataloader.batch_size
 
+    if central_dataloader:
+        data_loader = DataLoaderDispatcher(
+            new_dataset, split_batches=split_batches, batch_sampler=new_batch_sampler, **kwargs
+        )
+        data_loader.my_batch_size = dataloader.batch_size
Please add an error message if it fails :)
@@ -108,6 +108,71 @@ def dl_preparation_check():
         print("Shuffled dataloader passing.")
 
 
+def central_dl_preparation_check():
+    state = AcceleratorState()
+    length = 32 * state.num_processes
+
+    dl = DataLoader(range(length), batch_size=8)
+    dl = prepare_data_loader(
+        dl, state.device, state.num_processes, state.process_index, put_on_device=True, central_dataloader=True
+    )
+    result = []
+    for batch in dl:
+        result.append(gather(batch))
+    result = torch.cat(result)
+    assert torch.equal(result.cpu(), torch.arange(0, length).long())
Oh actually this applies to all `assert` in this file!
@@ -108,6 +108,71 @@ def dl_preparation_check():
         print("Shuffled dataloader passing.")
 
 
+def central_dl_preparation_check():
+    state = AcceleratorState()
+    length = 32 * state.num_processes
+
+    dl = DataLoader(range(length), batch_size=8)
+    dl = prepare_data_loader(
+        dl, state.device, state.num_processes, state.process_index, put_on_device=True, central_dataloader=True
+    )
+    result = []
+    for batch in dl:
+        result.append(gather(batch))
+    result = torch.cat(result)
+    assert torch.equal(result.cpu(), torch.arange(0, length).long())
Woopsie, this is some leftover to clean up! Thanks for flagging!
@@ -421,6 +524,13 @@ def prepare_data_loader(
     if new_batch_sampler is None:
         kwargs["batch_size"] = dataloader.batch_size // num_processes if split_batches else dataloader.batch_size
 
+    if central_dataloader:
+        data_loader = DataLoaderDispatcher(
+            new_dataset, split_batches=split_batches, batch_sampler=new_batch_sampler, **kwargs
+        )
+        data_loader.my_batch_size = dataloader.batch_size
This is just a test script ;-)
@@ -108,6 +108,71 @@ def dl_preparation_check():
         print("Shuffled dataloader passing.")
 
 
+def central_dl_preparation_check():
+    state = AcceleratorState()
+    length = 32 * state.num_processes
+
+    dl = DataLoader(range(length), batch_size=8)
+    dl = prepare_data_loader(
+        dl, state.device, state.num_processes, state.process_index, put_on_device=True, central_dataloader=True
+    )
+    result = []
+    for batch in dl:
+        result.append(gather(batch))
+    result = torch.cat(result)
+    assert torch.equal(result.cpu(), torch.arange(0, length).long())
Let's go for `dispath_batches` to look like `split_batches`
@@ -82,6 +82,9 @@ class Accelerator:
 
             Will default to :obj:`["torch"]` for PyTorch versions <=1.5.1 and :obj:`["generator"]` for PyTorch versions
             >= 1.6.
+        central_dataloader (:obj:`bool`, `optional`, defaults to :obj:`False`):
https://api.github.com/repos/huggingface/accelerate/pulls/164
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index d81e4ab61..ca49adf45 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -203,6 +203,8 @@ def _tpu_gather(tensor, name="tensor"):
         return type(tensor)({k: _tpu_gather(v, name=f"{name}_{k}") for k, v in tensor.items()})
     elif not isinstance(tensor, torch.Tensor):
         raise TypeError(f"Can't gather the values of type {type(tensor)}, only of nested list/tuple/dicts of tensors.")
+    if tensor.ndim == 0:
+        tensor = tensor.clone()[None]
     return xm.mesh_reduce(name, tensor, torch.cat)
 
 
@@ -213,6 +215,8 @@ def _gpu_gather(tensor):
         return type(tensor)({k: _gpu_gather(v) for k, v in tensor.items()})
     elif not isinstance(tensor, torch.Tensor):
         raise TypeError(f"Can't gather the values of type {type(tensor)}, only of nested list/tuple/dicts of tensors.")
+    if tensor.ndim == 0:
+        tensor = tensor.clone()[None]
     output_tensors = [tensor.clone() for _ in range(torch.distributed.get_world_size())]
     torch.distributed.all_gather(output_tensors, tensor)
     return torch.cat(output_tensors, dim=0)

https://api.github.com/repos/huggingface/accelerate/pulls/152
diff --git a/docs/source/quicktour.rst b/docs/source/quicktour.rst
index 3e9f5f470..585d2aa88 100644
--- a/docs/source/quicktour.rst
+++ b/docs/source/quicktour.rst
@@ -152,7 +152,7 @@ just run
 
     accelerate config
 
-on your machine and reply to the questions asked. This will save a `default_config.json` file in your cache folder for
+on your machine and reply to the questions asked. This will save a `default_config.yaml` file in your cache folder for
 🤗 Accelerate. That cache folder is (with decreasing order of priority):
 
     - The content of your environment variable ``HF_HOME`` suffixed with `accelerate`.
@@ -176,7 +176,7 @@ Note that if you specified a location for the config file in the previous step,
 
 .. code-block:: bash
 
-    accelerate test --config_file path_to_config.json
+    accelerate test --config_file path_to_config.yaml
 
 
 Now that this is done, you can run your script with the following command:
@@ -190,7 +190,7 @@ If you stored the config file in a non-default location, you can indicate it to
 
 .. code-block:: bash
 
-    accelerate launch --config_file path_to_config.json path_to_script.py --args_for_the_script
+    accelerate launch --config_file path_to_config.yaml path_to_script.py --args_for_the_script
 
 You can also override any of the arguments determined by your config file, see TODO: insert ref here.
 
diff --git a/src/accelerate/commands/config/__init__.py b/src/accelerate/commands/config/__init__.py
index 48ea96ab1..5bafb3be2 100644
--- a/src/accelerate/commands/config/__init__.py
+++ b/src/accelerate/commands/config/__init__.py
@@ -48,7 +48,7 @@ def config_command_parser(subparsers=None):
         "--config_file",
         default=None,
         help=(
-            "The path to use to store the config file. Will default to a file named default_config.json in the cache "
+            "The path to use to store the config file. Will default to a file named default_config.yaml in the cache "
             "location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have "
             "such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed "
             "with 'huggingface'."
diff --git a/src/accelerate/commands/config/config_args.py b/src/accelerate/commands/config/config_args.py
index 55b019fa4..39c00c280 100644
--- a/src/accelerate/commands/config/config_args.py
+++ b/src/accelerate/commands/config/config_args.py
@@ -28,7 +28,7 @@
     os.getenv("HF_HOME", os.path.join(os.getenv("XDG_CACHE_HOME", "~/.cache"), "huggingface"))
 )
 cache_dir = os.path.join(hf_cache_home, "accelerate")
-default_json_config_file = os.path.join(cache_dir, "default_config.json")
+default_json_config_file = os.path.join(cache_dir, "default_config.yaml")
 default_yaml_config_file = os.path.join(cache_dir, "default_config.yaml")
 
 # For backward compatibility: the default config is the json one if it's the only existing file.
diff --git a/src/accelerate/commands/test.py b/src/accelerate/commands/test.py
index ec6d2524b..379233679 100644
--- a/src/accelerate/commands/test.py
+++ b/src/accelerate/commands/test.py
@@ -30,7 +30,7 @@ def test_command_parser(subparsers=None):
         "--config_file",
         default=None,
         help=(
-            "The path to use to store the config file. Will default to a file named default_config.json in the cache "
+            "The path to use to store the config file. Will default to a file named default_config.yaml in the cache "
             "location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have "
             "such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed "
             "with 'huggingface'."

https://api.github.com/repos/huggingface/accelerate/pulls/151
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index a6b7a7d74..c5238369c 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -388,4 +388,5 @@ def __post_init__(self):
                 },
             },
             "steps_per_print": float("inf"),  # this will stop deepspeed from logging @ stdout
+            "zero_allow_untested_optimizer": True,
         }

FYI, I'm not involved with accelerate's deepspeed integration. I suggested that it will be done in same way as transformers when it was first contributed, but it was still committed in a different way. Thus now the user has to use a completely different setup when switching from transformers and developers have to maintain something totally different. So I'm not participating in this particular endeavour.
https://api.github.com/repos/huggingface/accelerate/pulls/150
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index a6b7a7d74..d81e4ab61 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -151,7 +151,7 @@ def convert_to_fp32(tensor):
     if isinstance(tensor, (list, tuple)):
         return honor_type(tensor, (convert_to_fp32(t) for t in tensor))
     elif isinstance(tensor, dict):
-        return type(tensor)({k: convert_to_fp32(v) for k, v in tensor.items()})
+        return type(tensor)(**{k: convert_to_fp32(v) for k, v in tensor.items()})
     elif not hasattr(tensor, "dtype") or tensor.dtype != torch.float16:
         return tensor
     return tensor.float()

https://api.github.com/repos/huggingface/accelerate/pulls/149
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 8707ed373..3b1281c84 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -136,7 +136,7 @@ Supported integrations
 - multi-GPU on several nodes (machines)
 - TPU
 - FP16 with native AMP (apex on the roadmap)
-- DeepSpeed (experimental)
+- DeepSpeed (experimental support)
 
 .. toctree::
     :maxdepth: 2
diff --git a/docs/source/quicktour.rst b/docs/source/quicktour.rst
index 68a23fff5..3e9f5f470 100644
--- a/docs/source/quicktour.rst
+++ b/docs/source/quicktour.rst
@@ -381,6 +381,22 @@ this:
         lr_scheduler.step()
 
 
+DeepSpeed
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+DeepSpeed support is experimental, so the underlying API will evolve in the near future and may have some slight
+breaking changes. In particular, 🤗 Accelerate does not support DeepSpeed config you have written yourself yet, this
+will be added in a next version.
+
+One main caveat for the DeepSpeed integration is that the DeepSpeed launcher always passes a ``local_rank`` variable to
+the training script, so your training script should accept it (whether you launch training with the DeepSpeed launcher
+or ``accelerate launch``).
+
+.. Warning::
+
+    The :func:`~accelerate.notebook_launcher` does not support the DeepSpeed integration yet.
+
+
 Internal mechanism
 -----------------------------------------------------------------------------------------------------------------------
 

```suggestion
or ``accelerate launch``).
```
@@ -381,6 +381,22 @@ this:
         lr_scheduler.step()
 
 
+DeepSpeed
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+DeepSpeed support is experimental, so the underlying API will evolve in the near future and may have some slight
+breaking changes. In particular, 🤗 Accelerate does not support DeepSpeed config you have written yourself yet, this
+will be added in a next version.
+
+One main caveat for the DeepSpeed integration is that the DeepSpeed launcher always passes a ``local_rank`` variable to
+the training script, so your training script should accept it (whether you launch training with the DeepSpeed launcher
+or ``accelerate launch`` Accelerate.
https://api.github.com/repos/huggingface/accelerate/pulls/140
diff --git a/docs/source/quicktour.rst b/docs/source/quicktour.rst
index 7ccdf5daf..07035a427 100644
--- a/docs/source/quicktour.rst
+++ b/docs/source/quicktour.rst
@@ -358,6 +358,20 @@ softmax). However you might want to put your loss computation inside the `accele
     with accelerator.autocast():
         loss = complex_loss_function(outputs, target):
 
+Another caveat with Mixed Precision training is that the gradient will skip a few updates at the beginning and
+sometimes during training: because of the dynamic loss scaling strategy, there are points during training where the
+gradients have overflown, and the loss scaling factor is reduced to avoid this happening again at the next step.
+
+This means that you may update your learning rate scheduler when there was no update, which is fine in general, but may
+have an impact when you have very little training data, or if the first learning rate values of your scheduler are very
+important. In this case, you can skip the learning rate scheduler updates when the optimizer step was not done like
+this:
+
+.. codeblock::
+
+    if not accelerator.optimizer_step_was_skipped:
+        lr_scheduler.step()
+
 
 Internal mechanism
 -----------------------------------------------------------------------------------------------------------------------
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index a11b7a7a5..a35f7e268 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -566,3 +566,14 @@ def autocast(self):
             autocast_context.__exit__()
         else:
             yield
+
+    @property
+    def optimizer_step_was_skipped(self):
+        """
+        Whether or not the optimizer update was skipped (because of gradient overflow in mixed precision), in which
+        case the learning rate should not be changed.
+        """
+        for optimizer in self._optimizers:
+            if optimizer.is_overflow:
+                return True
+        return False
diff --git a/src/accelerate/deepspeed_utils.py b/src/accelerate/deepspeed_utils.py
index e52c0c030..07450a87e 100644
--- a/src/accelerate/deepspeed_utils.py
+++ b/src/accelerate/deepspeed_utils.py
@@ -89,7 +89,7 @@ def step(self):
 
     @property
     def is_overflow(self):
-        """This must be called before lr_scheduler.step() when using deepspeed with fp16"""
+        """Whether or not the optimizer step was done, or skipped because of gradient overflow."""
         overflow = False
         if hasattr(self.optimizer, "overflow"):
             overflow = self.optimizer.overflow
diff --git a/src/accelerate/optimizer.py b/src/accelerate/optimizer.py
index b2909dbb5..37ca78833 100644
--- a/src/accelerate/optimizer.py
+++ b/src/accelerate/optimizer.py
@@ -53,6 +53,7 @@ def __init__(self, optimizer, device_placement=True, scaler=None):
         self.scaler = scaler
         self.state = AcceleratorState()
         self.device_placement = device_placement
+        self._is_overflow = False
 
         # Handle device placement
         if device_placement:
@@ -108,8 +109,12 @@ def step(self, closure=None):
             optimizer_args = {"closure": closure} if closure is not None else {}
             xm.optimizer_step(self.optimizer, optimizer_args=optimizer_args)
         elif self.scaler is not None:
+            scale_before = self.scaler.get_scale()
             self.scaler.step(self.optimizer, closure)
             self.scaler.update()
+            scale_after = self.scaler.get_scale()
+            # If we reduced the loss scale, it means the optimizer step was skipped because of gradient overflow.
+            self._is_overflow = scale_after < scale_before
         else:
             self.optimizer.step(closure)
 
@@ -119,5 +124,5 @@ def _switch_parameters(self, parameters_map):
 
     @property
     def is_overflow(self):
-        """This needs to be implemented at the end"""
-        return False  # TODO: implement it
+        """Whether or not the optimizer step was done, or skipped because of gradient overflow."""
+        return self._is_overflow

```suggestion
sometimes during training: because of the dynamic loss scaling strategy, there are points during training where the
```
@@ -358,6 +358,20 @@ softmax). However you might want to put your loss computation inside the `accele
     with accelerator.autocast():
         loss = complex_loss_function(outputs, target):
 
+Another caveat with Mixed Precision training is that the gradient will skip a few updates at the beginning and
+sometimes during training: because of the dynamic loss scaling strategy, there are points at training where the
```suggestion
have an impact when you have very little training data, or if the first learning rate values of your scheduler are very
```
@@ -358,6 +358,20 @@ softmax). However you might want to put your loss computation inside the `accele
     with accelerator.autocast():
         loss = complex_loss_function(outputs, target):
 
+Another caveat with Mixed Precision training is that the gradient will skip a few updates at the beginning and
+sometimes during training: because of the dynamic loss scaling strategy, there are points at training where the
+gradients have overflown, and the loss scaling factor is reduced to avoid this happening again at the next step.
+
+This means that you may update your learning rate scheduler when there was no update, which is fine in general, but may
+have an impace when you have very little training data, or if the first learning rate values of your scheduler are very
https://api.github.com/repos/huggingface/accelerate/pulls/139
diff --git a/docs/source/quicktour.rst b/docs/source/quicktour.rst
index 7ccdf5daf..2efdf6a82 100644
--- a/docs/source/quicktour.rst
+++ b/docs/source/quicktour.rst
@@ -251,6 +251,13 @@ To introduce special behavior in your script for TPUs you can check the :obj:`di
 The `NLP example <https://github.com/huggingface/accelerate/blob/main/examples/nlp_example.py>`__ shows an example in
 situation with dynamic padding.
 
+One last thing to pay close attnetion to: if your model has tied weights (such as language models which tie the weights
+of the embedding matrix with the weights of the decoder), moving this model to the TPU (either yourself or after you
+passed your model to :meth:`~accelerate.Accelerator.prepare`) will break the tying. You will need to retie the weights
+after. You can find an example of this in the `run_clm_no_trainer
+<https://github.com/huggingface/transformers/blob/master/examples/pytorch/language-modeling/run_clm.py>`__ script in
+the Transformers repository.
+
 
 Other caveats
 -----------------------------------------------------------------------------------------------------------------------
@@ -345,6 +352,7 @@ If you are using gradient clipping in your script, you should replace the calls
 :obj:`torch.nn.utils.clip_grad_norm_` or :obj:`torch.nn.utils.clip_grad_value_` with :obj:`accelerator.clip_grad_norm_`
 and :obj:`accelerator.clip_grad_value_` respectively.
 
+
 Mixed Precision training
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

https://api.github.com/repos/huggingface/accelerate/pulls/138
diff --git a/docs/source/quicktour.rst b/docs/source/quicktour.rst
index 1e28fed85..7ccdf5daf 100644
--- a/docs/source/quicktour.rst
+++ b/docs/source/quicktour.rst
@@ -345,6 +345,19 @@ If you are using gradient clipping in your script, you should replace the calls
 :obj:`torch.nn.utils.clip_grad_norm_` or :obj:`torch.nn.utils.clip_grad_value_` with :obj:`accelerator.clip_grad_norm_`
 and :obj:`accelerator.clip_grad_value_` respectively.
 
+Mixed Precision training
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+If you are running your training in Mixed Precision with Accelerate, you will get the best result with your loss being
+computed inside your model (like in Transformer models for instance). Every computation outside of the model will be
+executed in full precision (which is generally what you want for loss computation, expecially if it involves a
+softmax). However you might want to put your loss computation inside the `accelerator.autocast` context manager:
+
+.. codeblock::
+
+    with accelerator.autocast():
+        loss = complex_loss_function(outputs, target):
+
 
 Internal mechanism
 -----------------------------------------------------------------------------------------------------------------------
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 281f4e9a7..a11b7a7a5 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -28,6 +28,7 @@
 from .utils import (
     DeepSpeedPlugin,
     RNGType,
+    convert_outputs_to_fp32,
     extract_model_from_parallel,
     gather,
     pad_across_processes,
@@ -295,6 +296,7 @@ def prepare_model(self, model):
             model = torch.nn.parallel.DistributedDataParallel(model, **kwargs)
         if self.native_amp:
             model.forward = torch.cuda.amp.autocast()(model.forward)
+            model.forward = convert_outputs_to_fp32(model.forward)
         return model
 
     def _prepare_deepspeed(self, *args):
@@ -550,3 +552,17 @@ def get_state_dict(self, model):
                 state_dict[k] = state_dict[k].float()
 
         return state_dict
+
+    @contextmanager
+    def autocast(self):
+        """
+        Will apply automatic mixed-precision inside the block inside this context manager, if it is enabled. Nothing
+        different will happen otherwise.
+        """
+        if self.native_amp:
+            autocast_context = torch.cuda.amp.autocast()
+            autocast_context.__enter__()
+            yield
+            autocast_context.__exit__()
+        else:
+            yield
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 6edd1f1c0..a6b7a7d74 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -137,6 +137,46 @@ def send_to_device(tensor, device):
     return tensor.to(device)
 
 
+def convert_to_fp32(tensor):
+    """
+    Recursively converts the lements nested list/tuple/dictionary of tensors in FP16 precision to FP32.
+
+    Args:
+        tensor (nested list/tuple/dictionary of :obj:`torch.Tensor`):
+            The data to convert from FP16 to FP32.
+
+    Returns:
+        The same data structure as :obj:`tensor` with all tensors that were in FP16 precision converted to FP32.
+    """
+    if isinstance(tensor, (list, tuple)):
+        return honor_type(tensor, (convert_to_fp32(t) for t in tensor))
+    elif isinstance(tensor, dict):
+        return type(tensor)({k: convert_to_fp32(v) for k, v in tensor.items()})
+    elif not hasattr(tensor, "dtype") or tensor.dtype != torch.float16:
+        return tensor
+    return tensor.float()
+
+
+def convert_outputs_to_fp32(model_forward):
+    """
+    Decorator to apply to a function outputing tensors (like a model forward pass) that ensures the outputs in FP16
+    precision will be convert back to FP32.
+
+    Args:
+        model_forward (:obj:`Callable`):
+            The function which outputs we want to treat.
+
+    Returns:
+        The same function as :obj:`model_forward` but with converted outputs.
+    """
+
+    def convert_outputs(*args, **kwargs):
+        outputs = model_forward(*args, **kwargs)
+        return convert_to_fp32(outputs)
+
+    return convert_outputs
+
+
 def extract_model_from_parallel(model):
     """
     Extract a model from its distributed containers.

```suggestion
        Will apply automatic mixed-precision inside the block inside this context manager, if it is enabled. Nothing
```
@@ -550,3 +552,17 @@ def get_state_dict(self, model):
                 state_dict[k] = state_dict[k].float()
 
         return state_dict
+
+    @contextmanager
+    def autocast(self):
+        """
+        Will apply automatic mixed precision inside the block inside thics context manager, if it is enabled. Nothing
https://api.github.com/repos/huggingface/accelerate/pulls/134
diff --git a/src/accelerate/optimizer.py b/src/accelerate/optimizer.py
index 30fd8314e..57ea686f9 100644
--- a/src/accelerate/optimizer.py
+++ b/src/accelerate/optimizer.py
@@ -52,6 +52,7 @@ def __init__(self, optimizer, device_placement=True, scaler=None):
         self.optimizer = optimizer
         self.scaler = scaler
         self.state = AcceleratorState()
+        self.device_placement = device_placement
 
         # Handle device placement
         if device_placement:

https://api.github.com/repos/huggingface/accelerate/pulls/127
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index 44909ab3c..96648ca65 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -104,6 +104,8 @@ def __init__(
         self.drop_last = batch_sampler.drop_last
 
     def __len__(self):
+        if self.split_batches:
+            return len(self.batch_sampler)
         if len(self.batch_sampler) % self.num_processes == 0:
             return len(self.batch_sampler) // self.num_processes
         length = len(self.batch_sampler) // self.num_processes

https://api.github.com/repos/huggingface/accelerate/pulls/121
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 20274b671..281f4e9a7 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -413,8 +413,10 @@ def unscale_gradients(self, optimizer=None):
                 optimizer = self._optimizers
             elif not isinstance(optimizer, (tuple, list)):
                 optimizer = [optimizer]
-            for optimizer in optimizer:
-                self.scaler.unscale_(optimizer)
+            for opt in optimizer:
+                while isinstance(opt, AcceleratedOptimizer):
+                    opt = opt.optimizer
+                self.scaler.unscale_(opt)
 
     def clip_grad_norm_(self, parameters, max_norm, norm_type=2):
         """

The tests are not properly setup :see_no_evil: 
https://api.github.com/repos/huggingface/accelerate/pulls/115
diff --git a/src/accelerate/notebook_launcher.py b/src/accelerate/notebook_launcher.py
index 4c785b12f..fef2f0c08 100644
--- a/src/accelerate/notebook_launcher.py
+++ b/src/accelerate/notebook_launcher.py
@@ -16,7 +16,8 @@
 import sys
 
 import torch
-from torch.multiprocessing import start_processes
+
+from packaging import version
 
 from .state import AcceleratorState
 from .utils import PrepareForLaunch
@@ -82,6 +83,14 @@ def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, use
 
         if num_processes > 1:
             # Multi-GPU launch
+            if version.parse(torch.__version__) < version.parse("1.5.0"):
+                raise ImportError(
+                    "Using `notebook_launcher` for distributed training on GPUs require torch >= 1.5.0, got "
+                    f"{torch.__version__}."
+                )
+
+            from torch.multiprocessing import start_processes
+
             if len(AcceleratorState._shared_state) > 0:
                 raise ValueError(
                     "To launch a multi-GPU training from your notebook, the `Accelerator` should only be initialized "

https://api.github.com/repos/huggingface/accelerate/pulls/108
diff --git a/README.md b/README.md
index ffbf36bbd..78a2ceb2a 100644
--- a/README.md
+++ b/README.md
@@ -44,6 +44,10 @@ limitations under the License.
 <p>Run your *raw* PyTorch training script on any kind of device
 </h3>
 
+<h3 align="center">
+    <a href="https://hf.co/course"><img src="https://raw.githubusercontent.com/huggingface/transformers/master/docs/source/imgs/course_banner.png"></a>
+</h3>
+
 ## Easy to integrate
 
 🤗 Accelerate was created for PyTorch users who like to write the training loop of PyTorch models but are reluctant to write and maintain the boilerplate code needed to use multi-GPUs/TPU/fp16.
diff --git a/docs/source/imgs/course_banner.png b/docs/source/imgs/course_banner.png
new file mode 100644
index 000000000..45773d164
Binary files /dev/null and b/docs/source/imgs/course_banner.png differ

https://api.github.com/repos/huggingface/accelerate/pulls/107
diff --git a/src/accelerate/optimizer.py b/src/accelerate/optimizer.py
index 0bbba3644..30fd8314e 100644
--- a/src/accelerate/optimizer.py
+++ b/src/accelerate/optimizer.py
@@ -102,14 +102,15 @@ def zero_grad(self, set_to_none=None):
                 set_to_none = False
             self.optimizer.zero_grad(set_to_none=set_to_none)
 
-    def step(self):
+    def step(self, closure=None):
         if self.state.distributed_type == DistributedType.TPU:
-            xm.optimizer_step(self.optimizer)
+            optimizer_args = {"closure": closure} if closure is not None else {}
+            xm.optimizer_step(self.optimizer, optimizer_args=optimizer_args)
         elif self.scaler is not None:
-            self.scaler.step(self.optimizer)
+            self.scaler.step(self.optimizer, closure)
             self.scaler.update()
         else:
-            self.optimizer.step()
+            self.optimizer.step(closure)
 
     def _switch_parameters(self, parameters_map):
         for param_group in self.optimizer.param_groups:

```suggestion
            xm.optimizer_step(self.optimizer, optimizer_args={"closure": closure})
```
The [PyTorch XLA doc](https://pytorch.org/xla/release/1.8/index.html#torch_xla.core.xla_model.optimizer_step) seems to suggest we have to use `optimizer_args` to pass this along and it's not the second argument.
@@ -102,14 +102,14 @@ def zero_grad(self, set_to_none=None):
                 set_to_none = False
             self.optimizer.zero_grad(set_to_none=set_to_none)
 
-    def step(self):
+    def step(self, closure=None):
         if self.state.distributed_type == DistributedType.TPU:
-            xm.optimizer_step(self.optimizer)
+            xm.optimizer_step(self.optimizer, dict(closure=closure))
good catch!
@@ -102,14 +102,14 @@ def zero_grad(self, set_to_none=None):
                 set_to_none = False
             self.optimizer.zero_grad(set_to_none=set_to_none)
 
-    def step(self):
+    def step(self, closure=None):
         if self.state.distributed_type == DistributedType.TPU:
-            xm.optimizer_step(self.optimizer)
+            xm.optimizer_step(self.optimizer, dict(closure=closure))
Actually thinking a little bit more about it, let's not pass anything if there is no closure:
```suggestion
            optimizer_args = {"closure": closure} if closure is not None else {}
            xm.optimizer_step(self.optimizer, optimizer_args=optimizer_args)
```
@@ -102,14 +102,14 @@ def zero_grad(self, set_to_none=None):
                 set_to_none = False
             self.optimizer.zero_grad(set_to_none=set_to_none)
 
-    def step(self):
+    def step(self, closure=None):
         if self.state.distributed_type == DistributedType.TPU:
-            xm.optimizer_step(self.optimizer)
+            xm.optimizer_step(self.optimizer, optimizer_args={"closure": closure})
I concur, this is a little more explicit, but it shouldn't matter (famous last words...) because the base `torch.optim.Optimizer` has a `closure` argument. For the other options, we pass that along anyway. 
@@ -102,14 +102,14 @@ def zero_grad(self, set_to_none=None):
                 set_to_none = False
             self.optimizer.zero_grad(set_to_none=set_to_none)
 
-    def step(self):
+    def step(self, closure=None):
         if self.state.distributed_type == DistributedType.TPU:
-            xm.optimizer_step(self.optimizer)
+            xm.optimizer_step(self.optimizer, optimizer_args={"closure": closure})
If you don't mind committing the suggestion then, we're all good to merge :-)
@@ -102,14 +102,14 @@ def zero_grad(self, set_to_none=None):
                 set_to_none = False
             self.optimizer.zero_grad(set_to_none=set_to_none)
 
-    def step(self):
+    def step(self, closure=None):
         if self.state.distributed_type == DistributedType.TPU:
-            xm.optimizer_step(self.optimizer)
+            xm.optimizer_step(self.optimizer, optimizer_args={"closure": closure})
https://api.github.com/repos/huggingface/accelerate/pulls/105
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 2bc1459fd..a700cad62 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -387,16 +387,16 @@ def prepare_data_loader(self, data_loader):
     def prepare_optimizer(self, optimizer):
         return AcceleratedOptimizer(optimizer, device_placement=self.device_placement, scaler=self.scaler)
 
-    def backward(self, loss):
+    def backward(self, loss, **kwargs):
         """
         Use :obj:`accelerator.backward(loss)` in lieu of :obj:`loss.backward()`.
         """
         if self.distributed_type == DistributedType.DEEPSPEED:
-            self.deepspeed_engine.backward(loss)
+            self.deepspeed_engine.backward(loss, **kwargs)
         elif self.scaler is not None:
-            self.scaler.scale(loss).backward()
+            self.scaler.scale(loss).backward(**kwargs)
         else:
-            loss.backward()
+            loss.backward(**kwargs)
 
     def unscale_gradients(self, optimizer=None):
         """

https://api.github.com/repos/huggingface/accelerate/pulls/104
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index d119013d2..2bc1459fd 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -14,6 +14,7 @@
 
 import gc
 import os
+from contextlib import contextmanager
 from typing import List, Optional, Union
 
 import torch
@@ -183,6 +184,33 @@ def use_fp16(self):
             use_fp16 = self.state.use_fp16
         return use_fp16
 
+    @contextmanager
+    def local_main_process_first(self):
+        """
+        Lets the local main process go inside a with block.
+
+        The other processes will enter the with block after the main process exits.
+        """
+        yield from self._goes_first(self.is_local_main_process)
+
+    @contextmanager
+    def main_process_first(self):
+        """
+        Lets the main process go first inside a with block.
+
+        The other processes will enter the with block after the main process exits.
+        """
+        yield from self._goes_first(self.is_main_process)
+
+    def _goes_first(self, is_main):
+        if not is_main:
+            self.wait_for_everyone()
+
+        yield
+
+        if is_main:
+            self.wait_for_everyone()
+
     def print(self, *args, **kwargs):
         """
         Use in replacement of :obj:`print()` to only print once per server.

Small nit: can we use `is_main` for consistency here (as well as at line 211 below)?
@@ -183,6 +184,33 @@ def use_fp16(self):
             use_fp16 = self.state.use_fp16
         return use_fp16
 
+    @contextmanager
+    def local_main_process_first(self):
+        """
+        Lets the local main process go inside a with block.
+
+        The other processes will enter the with block after the main process exits.
+        """
+        yield from self._goes_first(self.is_local_main_process)
+
+    @contextmanager
+    def main_process_first(self):
+        """
+        Lets the main process go first inside a with block.
+
+        The other processes will enter the with block after the main process exits.
+        """
+        yield from self._goes_first(self.is_main_process)
+
+    def _goes_first(self, is_master):
+        if not is_master:
Yes sure!
@@ -183,6 +184,33 @@ def use_fp16(self):
             use_fp16 = self.state.use_fp16
         return use_fp16
 
+    @contextmanager
+    def local_main_process_first(self):
+        """
+        Lets the local main process go inside a with block.
+
+        The other processes will enter the with block after the main process exits.
+        """
+        yield from self._goes_first(self.is_local_main_process)
+
+    @contextmanager
+    def main_process_first(self):
+        """
+        Lets the main process go first inside a with block.
+
+        The other processes will enter the with block after the main process exits.
+        """
+        yield from self._goes_first(self.is_main_process)
+
+    def _goes_first(self, is_master):
+        if not is_master:
Addressed the request.
Thanks!
https://api.github.com/repos/huggingface/accelerate/pulls/98
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index c7706d390..7c12b18e6 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import gc
 from typing import List, Optional, Union
 
 import torch
@@ -338,6 +339,15 @@ def save(self, obj, f):
         """
         save(obj, f)
 
+    def free_memory(self):
+        """
+        Will release all references to the internal objects stored and call the garbage collector. You should call this
+        method between two trainings with different models/optimizers.
+        """
+        self._optimizers = []
+        gc.collect()
+        torch.cuda.empty_cache()
+
     def _get_named_parameters(self, *args):
         named_parameters = {}
         for obj in args:

https://api.github.com/repos/huggingface/accelerate/pulls/89
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index c7706d390..5a4bbe217 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -253,24 +253,36 @@ def backward(self, loss):
         else:
             loss.backward()
 
-    def clip_grad_norm_(self, parameters, max_norm, norm_type=2):
+    def unscale_gradients(self, optimizer=None):
         """
-        Should be used in place of :func:`torch.nn.utils.clip_grad_norm_`.
+        Unscale the gradients in mixed precision training with AMP. This is a noop in all other settings.
+
+        Args:
+            optimizer (:obj:`torch.optim.Optimizer` or :obj:`List[torch.optim.Optimizer]`, `optional`):
+                The optimizer(s) for which to unscale gradients. If not set, will unscale gradients on all optimizers
+                that were passed to :meth:`~accelerate.Accelerator.prepare`.
         """
-        # TODO: this unscales all optimizers where we should only unscale the one where parameters are.
         if self.state.use_fp16 and self.native_amp:
-            for optimizer in self._optimizers:
+            if optimizer is None:
+                # TODO: this unscales all optimizers where we should only unscale the one where parameters are.
+                optimizer = self._optimizers
+            elif not isinstance(optimizer, (tuple, list)):
+                optimizer = [optimizer]
+            for optimizer in optimizer:
                 self.scaler.unscale_(optimizer)
+
+    def clip_grad_norm_(self, parameters, max_norm, norm_type=2):
+        """
+        Should be used in place of :func:`torch.nn.utils.clip_grad_norm_`.
+        """
+        self.unscale_gradients()
         torch.nn.utils.clip_grad_norm_(parameters, max_norm, norm_type=norm_type)
 
     def clip_grad_value_(self, parameters, clip_value):
         """
         Should be used in place of :func:`torch.nn.utils.clip_grad_value_`.
         """
-        # TODO: this unscales all optimizers where we should only unscale the one where parameters are.
-        if self.state.use_fp16 and self.native_amp:
-            for optimizer in self._optimizers:
-                self.scaler.unscale_(optimizer)
+        self.unscale_gradients()
         torch.nn.utils.clip_grad_value_(parameters, clip_value)
 
     def gather(self, tensor):

https://api.github.com/repos/huggingface/accelerate/pulls/88
diff --git a/examples/README.md b/examples/README.md
index 4d941b5d3..19995fc0a 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -23,7 +23,7 @@ The [nlp_example.py](./nlp_example.py) script is a simple example to train a Ber
 Prior to running it you should install 🤗 Dataset and 🤗 Transformers:
 
 ```bash
-pip install datasets, transformers
+pip install datasets transformers
 ```
 
 The same script can be run in any of the following configurations:

https://api.github.com/repos/huggingface/accelerate/pulls/87
diff --git a/README.md b/README.md
index 4f6cf5f40..ffbf36bbd 100644
--- a/README.md
+++ b/README.md
@@ -160,6 +160,26 @@ Once you have MPI setup on your cluster, just run:
 mpirun -np 2 python examples/nlp_example.py
 ```
 
+## Launching training using DeepSpeed
+
+🤗 Accelerate supports training on single/multiple GPUs using DeepSpeed. to use it, you don't need to change anything in your training code; you can set everything using just `accelerate config`. However, if you desire to tweak your DeepSpeed related args from your python script, we provide you the `DeepSpeedPlugin`.
+
+```python
+from accelerator import Accelerator, DeepSpeedPlugin
+
+# deepspeed needs to know your gradient accumulation steps before hand, so don't forget to pass it
+# Remember you still need to do gradient accumulation by yourself, just like you would have done without deepspeed
+deepspeed_plugin = DeepSpeedPlugin(zero_stage=2, gradient_accumulation_steps=2)
+accelerator = Accelerator(fp16=True, deepspeed_plugin=deepspeed_plugin)
+
+# How to save your 🤗 Transformer?
+accelerator.wait_for_everyone()
+unwrapped_model = accelerator.unwrap_model(model)
+unwrapped_model.save_pretrained(save_dir, save_function=accelerator.save, state_dict=accelerator.get_state_dict(model))
+```
+
+Note: DeepSpeed support is experimental for now. In case you get into some problem, please open an issue.
+
 ## Launching your training from a notebook
 
 🤗 Accelerate also provides a `notebook_launcher` function you can use in a notebook to launch a distributed training. This is especially useful for Colab or Kaggle notebooks with a TPU backend. Just define your training loop in a `training_function` then in your last cell, add:
@@ -204,3 +224,4 @@ pip install accelerate
 - multi-GPU on several nodes (machines)
 - TPU
 - FP16 with native AMP (apex on the roadmap)
+- DeepSpeed support (experimental)
diff --git a/src/accelerate/__init__.py b/src/accelerate/__init__.py
index f793b38c1..f6206df02 100644
--- a/src/accelerate/__init__.py
+++ b/src/accelerate/__init__.py
@@ -8,4 +8,4 @@
 from .kwargs_handlers import DistributedDataParallelKwargs, GradScalerKwargs
 from .notebook_launcher import notebook_launcher
 from .state import DistributedType
-from .utils import synchronize_rng_states
+from .utils import DeepSpeedPlugin, synchronize_rng_states
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 7b09375df..d119013d2 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import gc
+import os
 from typing import List, Optional, Union
 
 import torch
@@ -22,8 +23,26 @@
 from .data_loader import prepare_data_loader
 from .kwargs_handlers import DistributedDataParallelKwargs, GradScalerKwargs, KwargsHandler
 from .optimizer import AcceleratedOptimizer
-from .state import AcceleratorState, DistributedType
-from .utils import RNGType, extract_model_from_parallel, gather, pad_across_processes, save, wait_for_everyone
+from .state import AcceleratorState, DistributedType, is_deepspeed_available
+from .utils import (
+    DeepSpeedPlugin,
+    RNGType,
+    extract_model_from_parallel,
+    gather,
+    pad_across_processes,
+    save,
+    wait_for_everyone,
+)
+
+
+if is_deepspeed_available():
+    import deepspeed
+    from .deepspeed_utils import DeepSpeedEngineWrapper, DeepSpeedOptimizerWrapper
+
+import logging
+
+
+logger = logging.getLogger(__name__)
 
 
 class Accelerator:
@@ -46,6 +65,9 @@ class Accelerator:
         cpu (:obj:`bool`, `optional`):
             Whether or not to force the script to execute on CPU. Will ignore GPU available if set to :obj:`True` and
             force the execution on one process only.
+        deepspeed_plugin (:obj:`DeepSpeedPlugin`, `optional`):
+            Tweak your DeepSpeed related args using this argument. This argument is optional and can be configured
+            directly using `accelerate config`
         rng_types (list of :obj:`str` or :class:`~accelerate.utils.RNGType`):
             The list of random number generators to synchronize at the beginning of each iteration in your prepared
             dataloaders. Should be one or several of:
@@ -74,10 +96,18 @@ def __init__(
         split_batches: bool = False,
         fp16: bool = None,
         cpu: bool = False,
+        deepspeed_plugin: DeepSpeedPlugin = None,
         rng_types: Optional[List[Union[str, RNGType]]] = None,
         kwargs_handlers: Optional[List[KwargsHandler]] = None,
     ):
-        self.state = AcceleratorState(fp16=fp16, cpu=cpu, _from_accelerator=True)
+        if deepspeed_plugin is None:  # init from env variables
+            deepspeed_plugin = DeepSpeedPlugin() if os.environ.get("USE_DEEPSPEED", "false") == "true" else None
+        else:
+            assert isinstance(
+                deepspeed_plugin, DeepSpeedPlugin
+            ), "`deepspeed_plugin` must be a DeepSpeedPlugin object."
+
+        self.state = AcceleratorState(fp16=fp16, cpu=cpu, deepspeed_plugin=deepspeed_plugin, _from_accelerator=True)
 
         self.device_placement = device_placement
         self.split_batches = split_batches
@@ -109,6 +139,7 @@ def __init__(
 
         # Internal references to the training objects
         self._optimizers = []
+        self._models = []
 
         # RNG Types
         if rng_types is None:
@@ -146,7 +177,11 @@ def is_local_main_process(self):
 
     @property
     def use_fp16(self):
-        return self.state.use_fp16
+        if self.distributed_type == DistributedType.DEEPSPEED:
+            use_fp16 = self.state.deepspeed_plugin.deepspeed_config["fp16"]["enabled"]
+        else:
+            use_fp16 = self.state.use_fp16
+        return use_fp16
 
     def print(self, *args, **kwargs):
         """
@@ -159,6 +194,7 @@ def _prepare_one(self, obj):
         if isinstance(obj, torch.utils.data.DataLoader):
             return self.prepare_data_loader(obj)
         elif isinstance(obj, torch.nn.Module):
+            self._models.append(obj)
             return self.prepare_model(obj)
         elif isinstance(obj, torch.optim.Optimizer):
             optimizer = self.prepare_optimizer(obj)
@@ -177,7 +213,6 @@ def prepare(self, *args):
             - :obj:`torch.utils.data.DataLoader`: PyTorch Dataloader
             - :obj:`torch.nn.Module`: PyTorch Module
             - :obj:`torch.optim.Optimizer`: PyTorch Optimizer
-
         """
         # On TPUs, putting the model on the XLA device will create new parameters, so the corresponding optimizer will
         # have parameters disconnected from the model (so no training :-( ).
@@ -199,7 +234,10 @@ def prepare(self, *args):
             # 1. grabbing old model parameters
             old_named_params = self._get_named_parameters(*args)
 
-        result = tuple(self._prepare_one(obj) for obj in args)
+        if self.distributed_type == DistributedType.DEEPSPEED:
+            result = self._prepare_deepspeed(*args)
+        else:
+            result = tuple(self._prepare_one(obj) for obj in args)
 
         if tpu_should_fix_optimizer:
             # 2. grabbing new model parameters
@@ -231,6 +269,82 @@ def prepare_model(self, model):
             model.forward = torch.cuda.amp.autocast()(model.forward)
         return model
 
+    def _prepare_deepspeed(self, *args):
+
+        deepspeed_plugin = self.state.deepspeed_plugin
+        self.deepspeed_config = deepspeed_plugin.deepspeed_config
+
+        batch_sizes = [obj.batch_size for obj in args if hasattr(obj, "batch_size")]
+        if len(batch_sizes) == 0:
+            raise ValueError(
+                "You must specify a training or evaluation dataloader in `accelerate.prepare()` when using DeepSpeed."
+            )
+
+        batch_size_per_device = min(batch_sizes) if deepspeed_plugin.is_train_batch_min else max(batch_sizes)
+        if len(batch_sizes) > 1:
+            logger.info(
+                f"Since you passed both train and evaluation dataloader, `is_train_batch_min` (here \
+                {deepspeed_plugin.is_train_batch_min} will decide the `train_batch_size` ({batch_size_per_device})."
+            )
+
+        self.deepspeed_config["train_batch_size"] = (
+            batch_size_per_device * deepspeed_plugin.gradient_accumulation_steps * self.num_processes
+        )
+
+        result = [self._prepare_one(obj) if isinstance(obj, torch.utils.data.DataLoader) else obj for obj in args]
+
+        model = None
+        optimizer = None
+        for obj in result:
+            if isinstance(obj, torch.nn.Module):
+                model = obj
+            elif isinstance(obj, (torch.optim.Optimizer, dict)):
+                optimizer = obj
+
+        if deepspeed_plugin.auto_opt_mapping:
+            is_adam = isinstance(optimizer, torch.optim.Adam)
+            is_adamw = isinstance(optimizer, torch.optim.AdamW)
+            if (is_adam or is_adamw) and deepspeed_plugin.offload_optimizer_device == "cpu":
+                defaults = optimizer.defaults
+                optimizer = deepspeed.ops.adam.DeepSpeedCPUAdam(
+                    model.parameters(),
+                    lr=defaults["lr"],
+                    bias_correction=True,
+                    betas=defaults["betas"],
+                    eps=defaults["eps"],
+                    weight_decay=defaults["weight_decay"],
+                    amsgrad=defaults["amsgrad"],
+                    adamw_mode=is_adamw,
+                )
+
+        # useful when only eval_dataloader is given into `accelerator.prepare()`
+        if model is not None:
+            engine = DeepSpeedEngineWrapper(
+                args=None,
+                model=model,
+                optimizer=optimizer,
+                config_params=self.deepspeed_config,
+                dist_init_required=False,
+            )
+            for i in range(len(result)):
+                if isinstance(result[i], torch.nn.Module):
+                    result[i] = engine
+                elif isinstance(result[i], torch.optim.Optimizer):
+                    result[i] = DeepSpeedOptimizerWrapper(engine.optimizer, engine)
+            self.deepspeed_engine = engine  # pointing for deepspeed_engine.backward()
+            self._models.append(engine)
+            self._optimizers.append(engine.optimizer)
+            assert (
+                len(self._models) == 1
+            ), "You can't use same `Accelerator()` instance with 2 models when using DeepSpeed"
+
+        if self.distributed_type == DistributedType.DEEPSPEED:
+            assert hasattr(
+                self, "deepspeed_engine"
+            ), "You need to pass the model along the optimizer when using Deepspeed."
+
+        return tuple(result)
+
     def prepare_data_loader(self, data_loader):
         return prepare_data_loader(
             data_loader,
@@ -249,7 +363,9 @@ def backward(self, loss):
         """
         Use :obj:`accelerator.backward(loss)` in lieu of :obj:`loss.backward()`.
         """
-        if self.scaler is not None:
+        if self.distributed_type == DistributedType.DEEPSPEED:
+            self.deepspeed_engine.backward(loss)
+        elif self.scaler is not None:
             self.scaler.scale(loss).backward()
         else:
             loss.backward()
@@ -357,6 +473,8 @@ def free_memory(self):
         method between two trainings with different models/optimizers.
         """
         self._optimizers = []
+        self._models = []
+        self.deepspeed_engine = None
         gc.collect()
         torch.cuda.empty_cache()
 
@@ -384,3 +502,21 @@ def _get_devices(self, *args):
                         optimizer_device = param_group["params"][0].device
                         break
         return (model_device, optimizer_device)
+
+    def get_state_dict(self, model):
+        is_zero_3 = False
+        if is_deepspeed_available():
+            if isinstance(model, DeepSpeedEngineWrapper) and self.distributed_type == DistributedType.DEEPSPEED:
+                is_zero_3 = self.state.deepspeed_plugin.zero_stage == 3
+
+        if is_zero_3:
+            state_dict = model._zero3_consolidated_fp16_state_dict()
+        else:
+            model = self.unwrap_model(model)
+            state_dict = model.state_dict()
+
+        for k in state_dict:
+            if state_dict[k].dtype == torch.float16:
+                state_dict[k] = state_dict[k].float()
+
+        return state_dict
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index ebc176089..eaaad20f2 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -16,6 +16,7 @@
 
 from accelerate.state import ComputeEnvironment, DistributedType
 
+from ...utils import is_deepspeed_available
 from .config_args import ClusterConfig
 from .config_utils import _ask_field, _convert_distributed_mode, _convert_yes_no_to_bool
 
@@ -31,7 +32,7 @@ def get_cluster_input():
     num_machines = 1
     main_process_ip = None
     main_process_port = None
-    if distributed_type == DistributedType.MULTI_GPU or distributed_type == DistributedType.MULTI_CPU:
+    if distributed_type in [DistributedType.MULTI_GPU, DistributedType.MULTI_CPU]:
         num_machines = _ask_field(
             "How many different machines will you use (use more than 1 for multi-node training)? [1]: ",
             lambda x: int(x),
@@ -50,6 +51,42 @@ def get_cluster_input():
                 "What is the port you will use to communicate with the main process? ",
                 lambda x: int(x),
             )
+
+    deepspeed_config = None
+    if distributed_type in [DistributedType.MULTI_GPU, DistributedType.NO]:
+        use_deepspeed = _ask_field(
+            "Do you want to use DeepSpeed? [yes/NO]: ",
+            _convert_yes_no_to_bool,
+            default=False,
+            error_message="Please enter yes or no.",
+        )
+        if use_deepspeed:
+            distributed_type = DistributedType.DEEPSPEED
+            assert (
+                is_deepspeed_available()
+            ), "DeepSpeed is not installed => run `pip3 install deepspeed` or build it from source"
+
+        deepspeed_config = {}
+        if distributed_type == DistributedType.DEEPSPEED:
+            deepspeed_config["zero_stage"] = _ask_field(
+                "What should be your DeepSpeed's ZeRO optimization stage (0, 1, 2, 3)? [2]: ",
+                lambda x: int(x),
+                default=2,
+            )
+
+            if deepspeed_config["zero_stage"] >= 2:
+                deepspeed_config["offload_optimizer_device"] = _ask_field(
+                    "Where to offload optimizer states? [NONE/cpu/nvme]: ",
+                    lambda x: str(x),
+                    default="none",
+                )
+
+            deepspeed_config["gradient_accumulation_steps"] = _ask_field(
+                "How many gradient accumulation steps you're passing in your script? [1]: ",
+                lambda x: int(x),
+                default=1,
+            )
+
     if distributed_type == DistributedType.TPU:
         main_training_function = _ask_field(
             "What is the name of the function in your script that should be launched in all parallel scripts? [main]: ",
@@ -85,4 +122,5 @@ def get_cluster_input():
         main_process_ip=main_process_ip,
         main_process_port=main_process_port,
         main_training_function=main_training_function,
+        deepspeed_config=deepspeed_config,
     )
diff --git a/src/accelerate/commands/config/config_args.py b/src/accelerate/commands/config/config_args.py
index 8cc993b59..55b019fa4 100644
--- a/src/accelerate/commands/config/config_args.py
+++ b/src/accelerate/commands/config/config_args.py
@@ -119,6 +119,9 @@ class ClusterConfig(BaseConfig):
     main_process_port: Optional[int] = None
     main_training_function: str = "main"
 
+    # args for deepspeed_plugin
+    deepspeed_config: dict = None
+
 
 @dataclass
 class SageMakerConfig(BaseConfig):
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index e27c6ed5c..4f392d6fa 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -44,6 +44,12 @@ def launch_command_parser(subparsers=None):
         action="store_true",
         help="Whether or not this should launch a distributed GPU training.",
     )
+    parser.add_argument(
+        "--use_deepspeed",
+        default=False,
+        action="store_true",
+        help="Whether to use deepspeed.",
+    )
     parser.add_argument(
         "--tpu", default=False, action="store_true", help="Whether or not this should launch a TPU training."
     )
@@ -95,6 +101,25 @@ def launch_command_parser(subparsers=None):
             "script."
         ),
     )
+    parser.add_argument(
+        "--zero_stage",
+        default=None,
+        type=int,
+        help="DeepSpeed's ZeRO optimization stage (useful only when `use_deepspeed` flag is passed).",
+    )
+    parser.add_argument(
+        "--offload_optimizer_device",
+        default=None,
+        type=str,
+        help="Decides where (none|cpu|nvme) to offload optimizer states (useful only when `use_deepspeed` flag is passed).",
+    )
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        default=None,
+        type=int,
+        help="No of gradient_accumulation_steps used in your training script (useful only when `use_deepspeed` flag is passed).",
+    )
+
     # Other arguments of the training scripts
     parser.add_argument("training_script_args", nargs=argparse.REMAINDER, help="Arguments of the training script.")
 
@@ -151,6 +176,43 @@ def multi_gpu_launcher(args):
         raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)
 
 
+def deepspeed_launcher(args):
+
+    cmd = ["deepspeed"]
+    if args.num_machines > 1:
+        cmd.extend(
+            [
+                "--num_gpus",
+                str(args.num_processes // args.num_machines),
+                "--num_nodes",
+                str(args.num_machines),
+                "--node_rank",
+                str(args.machine_rank),
+                "--master_addr",
+                args.main_process_ip,
+                "--master_port",
+                str(args.main_process_port),
+            ]
+        )
+    else:
+        cmd.extend(["--num_gpus", str(args.num_processes)])
+
+    cmd.append(args.training_script)
+    cmd.extend(args.training_script_args)
+
+    current_env = os.environ.copy()
+    current_env["USE_FP16"] = str(args.fp16)
+    current_env["USE_DEEPSPEED"] = "true"
+    current_env["DEEPSPEED_ZERO_STAGE"] = str(args.zero_stage)
+    current_env["GRADIENT_ACCUMULATION_STEPS"] = str(args.gradient_accumulation_steps)
+    current_env["DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE"] = str(args.offload_optimizer_device)
+
+    process = subprocess.Popen(cmd, env=current_env)
+    process.wait()
+    if process.returncode != 0:
+        raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)
+
+
 def tpu_launcher(args):
     import torch_xla.distributed.xla_multiprocessing as xmp
 
@@ -276,19 +338,26 @@ def sagemaker_launcher(sagemaker_config: SageMakerConfig, args):
 
 def launch_command(args):
     # Sanity checks
-    if args.multi_gpu and args.tpu:
-        raise ValueError("You can only pick one between `--multi_gpu` and `--tpu`.")
+    if sum([args.multi_gpu, args.tpu, args.use_deepspeed]) > 1:
+        raise ValueError("You can only pick one between `--multi_gpu`, `--use_deepspeed`, `--tpu`.")
 
     defaults = None
     # Get the default from the config file.
     if args.config_file is not None or os.path.isfile(default_config_file) and not args.cpu:
         defaults = load_config_from_file(args.config_file)
-        if not args.multi_gpu and not args.tpu:
+        if not args.multi_gpu and not args.tpu and not args.use_deepspeed:
+            args.use_deepspeed = defaults.distributed_type == DistributedType.DEEPSPEED
             args.multi_gpu = defaults.distributed_type == DistributedType.MULTI_GPU
             args.tpu = defaults.distributed_type == DistributedType.TPU
         if defaults.compute_environment == ComputeEnvironment.LOCAL_MACHINE:
             # Update args with the defaults
             for name, attr in defaults.__dict__.items():
+                if isinstance(attr, dict):
+                    for k in defaults.deepspeed_config:
+                        if getattr(args, k) is None:
+                            setattr(args, k, defaults.deepspeed_config[k])
+                    continue
+
                 # Those args are handled separately
                 if (
                     name not in ["compute_environment", "fp16", "distributed_type"]
@@ -303,7 +372,9 @@ def launch_command(args):
             args.num_processes = 1
 
     # Use the proper launcher
-    if args.multi_gpu and not args.cpu:
+    if args.use_deepspeed and not args.cpu:
+        deepspeed_launcher(args)
+    elif args.multi_gpu and not args.cpu:
         multi_gpu_launcher(args)
     elif args.tpu and not args.cpu:
         tpu_launcher(args)
diff --git a/src/accelerate/deepspeed_utils.py b/src/accelerate/deepspeed_utils.py
new file mode 100644
index 000000000..11897ddc7
--- /dev/null
+++ b/src/accelerate/deepspeed_utils.py
@@ -0,0 +1,96 @@
+# 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.
+
+from .optimizer import AcceleratedOptimizer
+from .state import is_apex_available, is_deepspeed_available
+
+
+if is_deepspeed_available():
+    from deepspeed import DeepSpeedEngine
+
+if is_apex_available():
+    import amp
+
+
+class DeepSpeedEngineWrapper(DeepSpeedEngine):
+    """
+    Wrapper over deepspeed.DeepSpeedEngine object
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # overwriting micro_steps for user's gradient_accumulation
+        self.micro_steps = -1
+
+    def step(self, lr_kwargs=None):
+        """DeepSpeedEngine.step() without `micro_steps` update & no profiling"""
+        if self.is_gradient_accumulation_boundary():  # it shouldn't matter whether we keep this line or not
+            if self.progressive_layer_drop:
+                self.progressive_layer_drop.update_state(self.global_steps)
+
+            self._take_model_step(lr_kwargs)
+
+    def backward(self, loss):
+        """DeepSpeedEngine.backward() with with no loss scaling; no profiling but with `micro_steps` update"""
+
+        if self.zero_optimization():
+            self.optimizer.is_gradient_accumulation_boundary = self.is_gradient_accumulation_boundary()
+            self.optimizer.backward(loss)
+        elif self.amp_enabled():
+            # AMP requires delaying unscale when inside gradient accumulation boundaries
+            # https://nvidia.github.io/apex/advanced.html#gradient-accumulation-across-iterations
+            delay_unscale = not self.is_gradient_accumulation_boundary()
+            with amp.scale_loss(loss, self.optimizer, delay_unscale=delay_unscale) as scaled_loss:
+                scaled_loss.backward()
+        elif self.fp16_enabled():
+            self.optimizer.backward(loss)
+        else:
+            loss.backward()
+
+        if self.enable_backward_allreduce:
+            self.allreduce_gradients()
+
+        # this will ensure deepspeed gradient_accumulation matches user's accumulation
+        self.micro_steps += 1
+
+
+class DeepSpeedOptimizerWrapper(AcceleratedOptimizer):
+    """
+    Internal wrapper around a deepspeed optimizer.
+
+    Args:
+        optimizer (:obj:`torch.optim.optimizer.Optimizer`):
+            The optimizer to wrap.
+    """
+
+    def __init__(self, optimizer, model: DeepSpeedEngineWrapper):
+        super().__init__(optimizer, device_placement=False, scaler=None)
+
+        self.model = model
+
+    def zero_grad(self, set_to_none=None):
+        pass  # `model.step()` is doing that automatically. Therefore, it's implementation is not needed
+
+    def step(self):
+        """This will handle optimizer.step() & optimizer.zero_grad() with gradient_accumulation"""
+        self.model.step()
+
+    @property
+    def is_overflow(self):
+        """This must be called before lr_scheduler.step() when using deepspeed with fp16"""
+        overflow = False
+        if hasattr(self.optimizer, "overflow"):
+            overflow = self.optimizer.overflow
+        return overflow
diff --git a/src/accelerate/optimizer.py b/src/accelerate/optimizer.py
index 1dd9f506a..0bbba3644 100644
--- a/src/accelerate/optimizer.py
+++ b/src/accelerate/optimizer.py
@@ -114,3 +114,8 @@ def step(self):
     def _switch_parameters(self, parameters_map):
         for param_group in self.optimizer.param_groups:
             param_group["params"] = [parameters_map.get(p, p) for p in param_group["params"]]
+
+    @property
+    def is_overflow(self):
+        """This needs to be implemented at the end"""
+        return False  # TODO: implement it
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index 76376038b..5b349fc38 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import importlib
 import os
 from distutils.util import strtobool
 from enum import Enum
@@ -48,10 +49,18 @@ def is_ccl_available():
     return _ccl_available
 
 
+def is_apex_available():
+    return importlib.util.find_spec("apex") is not None
+
+
 def is_tpu_available():
     return _tpu_available
 
 
+def is_deepspeed_available():
+    return importlib.util.find_spec("deepspeed") is not None
+
+
 def parse_flag_from_env(key, default=False):
     value = os.environ.get(key, str(default))
     return strtobool(value) == 1  # As its name indicates `strtobool` actually returns an int...
@@ -66,6 +75,7 @@ class DistributedType(str, Enum):
         - **NO** -- Not a distributed environment, just a single process.
         - **MULTI_CPU** -- Distributed on multiple CPU nodes.
         - **MULTI_GPU** -- Distributed on multiple GPUs.
+        - **DEEPSPEED** -- Using DeepSpeed.
         - **TPU** -- Distributed on TPUs.
     """
 
@@ -73,6 +83,7 @@ class DistributedType(str, Enum):
     NO = "NO"
     MULTI_CPU = "MULTI_CPU"
     MULTI_GPU = "MULTI_GPU"
+    DEEPSPEED = "DEEPSPEED"
     TPU = "TPU"
 
 
@@ -127,10 +138,11 @@ class AcceleratorState:
 
     _shared_state = {}
 
-    def __init__(self, fp16: bool = None, cpu: bool = False, _from_accelerator: bool = False):
+    def __init__(self, fp16: bool = None, cpu: bool = False, deepspeed_plugin=None, _from_accelerator: bool = False):
         self.__dict__ = self._shared_state
         if not getattr(self, "initialized", False):
             self.backend = None
+            self.deepspeed_plugin = None
             if not _from_accelerator:
                 raise ValueError(
                     "Please make sure to properly initialize your accelerator via `accelerator = Accelerator()` "
@@ -143,6 +155,23 @@ def __init__(self, fp16: bool = None, cpu: bool = False, _from_accelerator: bool
                 self.local_process_index = xm.get_local_ordinal()
                 self.device = xm.xla_device()
                 self.use_fp16 = False
+            elif os.environ.get("USE_DEEPSPEED", "false") == "true" and not cpu:
+                assert (
+                    is_deepspeed_available()
+                ), "DeepSpeed is not available => install it using `pip3 install deepspeed` or build it from source"
+                self.distributed_type = DistributedType.DEEPSPEED
+                if not torch.distributed.is_initialized():
+                    torch.distributed.init_process_group(backend="nccl")
+                    self.backend = "nccl"
+                self.num_processes = torch.distributed.get_world_size()
+                self.process_index = torch.distributed.get_rank()
+                self.local_process_index = int(os.environ.get("LOCAL_RANK", -1))
+                self.device = torch.device("cuda", self.local_process_index)
+                torch.cuda.set_device(self.device)
+                self.use_fp16 = False  # deepspeed handles fp16 using deepspeed_config
+                fp16 = parse_flag_from_env("USE_FP16", False) if fp16 is None else fp16
+                deepspeed_plugin.deepspeed_config.update({"fp16": {"enabled": fp16}})
+                self.deepspeed_plugin = deepspeed_plugin
             elif int(os.environ.get("LOCAL_RANK", -1)) != -1 and not cpu:
                 self.distributed_type = DistributedType.MULTI_GPU
                 if not torch.distributed.is_initialized():
@@ -200,11 +229,15 @@ def __init__(self, fp16: bool = None, cpu: bool = False, _from_accelerator: bool
             self.initialized = True
 
     def __repr__(self):
-        return (
+        use_fp16 = self.deepspeed_plugin.fp16 if self.distributed_type == DistributedType.DEEPSPEED else self.use_fp16
+        repr = (
             f"Distributed environment: {self.distributed_type}{('  Backend: ' + self.backend) if self.backend else ''}\n"
             f"Num processes: {self.num_processes}\n"
             f"Process index: {self.process_index}\n"
             f"Local process index: {self.local_process_index}\n"
             f"Device: {self.device}\n"
-            f"Use FP16 precision: {self.use_fp16}\n"
+            f"Use FP16 precision: {use_fp16}\n"
         )
+        if self.distributed_type == DistributedType.DEEPSPEED:
+            repr += f"ds_config: {self.deepspeed_plugin.ds_config}\n"
+        return repr
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 7302ead15..6edd1f1c0 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -15,13 +15,14 @@
 import importlib
 import os
 import random
+from dataclasses import dataclass, field
 from enum import Enum
 from typing import List, Optional, Union
 
 import numpy as np
 import torch
 
-from .state import AcceleratorState, DistributedType, is_tpu_available
+from .state import AcceleratorState, DistributedType, is_deepspeed_available, is_tpu_available
 
 
 if is_tpu_available():
@@ -36,6 +37,10 @@ def is_sagemaker_available():
     return importlib.util.find_spec("sagemaker") is not None
 
 
+if is_deepspeed_available():
+    from deepspeed import DeepSpeedEngine
+
+
 class RNGType(Enum):
     TORCH = "torch"
     CUDA = "cuda"
@@ -142,7 +147,11 @@ def extract_model_from_parallel(model):
     Returns:
         :obj:`torch.nn.Module`: The extracted model.
     """
-    while isinstance(model, (torch.nn.parallel.DistributedDataParallel, torch.nn.DataParallel)):
+    options = (torch.nn.parallel.DistributedDataParallel, torch.nn.DataParallel)
+    if is_deepspeed_available():
+        options += (DeepSpeedEngine,)
+
+    while isinstance(model, options):
         model = model.module
     return model
 
@@ -251,6 +260,7 @@ def wait_for_everyone():
     if (
         AcceleratorState().distributed_type == DistributedType.MULTI_GPU
         or AcceleratorState().distributed_type == DistributedType.MULTI_CPU
+        or AcceleratorState().distributed_type == DistributedType.DEEPSPEED
     ):
         torch.distributed.barrier()
     elif AcceleratorState().distributed_type == DistributedType.TPU:
@@ -293,3 +303,49 @@ def __call__(self, index, *args):
             os.environ["RANK"] = str(index)
 
         self.launcher(*args)
+
+
+@dataclass
+class DeepSpeedPlugin:
+
+    gradient_accumulation_steps: int = field(
+        default=None, metadata={"help": "Number of steps to accumulate gradients before updating optimizer states"}
+    )
+    zero_stage: int = field(
+        default=None,
+        metadata={"help": "Possible options are 0,1,2,3; Default will be taken from environment variable"},
+    )
+    is_train_batch_min: str = field(
+        default=True,
+        metadata={"help": "If both train & eval dataloaders are specified, this will decide the train_batch_size"},
+    )
+
+    auto_opt_mapping: bool = field(
+        default=True,
+        metadata={"help": "whether to map torch.adam to deepspeed optimizer version of adam based on config"},
+    )
+
+    offload_optimizer_device: bool = field(default=None, metadata={"help": "Possible options are none|cpu|nvme"})
+
+    def __post_init__(self):
+
+        if self.gradient_accumulation_steps is None:
+            self.gradient_accumulation_steps = int(os.environ.get("GRADIENT_ACCUMULATION_STEPS", 1))
+
+        if self.zero_stage is None:
+            self.zero_stage = int(os.environ.get("DEEPSPEED_ZERO_STAGE", 2))
+
+        if self.offload_optimizer_device is None:
+            self.offload_optimizer_device = os.environ.get("DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE", "none")
+
+        self.deepspeed_config = {
+            "train_batch_size": None,
+            "gradient_accumulation_steps": self.gradient_accumulation_steps,
+            "zero_optimization": {
+                "stage": self.zero_stage,
+                "offload_optimizer": {
+                    "device": self.offload_optimizer_device,
+                },
+            },
+            "steps_per_print": float("inf"),  # this will stop deepspeed from logging @ stdout
+        }

```suggestion
            assert isinstance(deepspeed_plugin, DeepSpeedPlugin), "`deepspeed_plugin` must be a DeepSpeedPlugin object."
```
@@ -86,18 +85,10 @@ def __init__(
         rng_types: Optional[List[Union[str, RNGType]]] = None,
         kwargs_handlers: Optional[List[KwargsHandler]] = None,
     ):
+        if deepspeed_plugin is not None:
+            assert isinstance(deepspeed_plugin, DeepSpeedPlugin), "deepspeed_plugin must be instance of DeepSpeedPlugin"
```suggestion
        assert len(batch_size) > 0, "You must specify a training or evaluation dataloader in `accelerate.prepare()` when using DeepSpeed."
```
@@ -230,6 +245,51 @@ def prepare_model(self, model):
             model.forward = torch.cuda.amp.autocast()(model.forward)
         return model
 
+    def _prepare_deepspeed(self, *args):
+
+        ds_plugin = self.state.deepspeed_plugin
+        self.ds_config = ds_plugin.ds_config
+
+        batch_size = [obj.batch_size for obj in args if hasattr(obj, "batch_size")]
+        assert len(batch_size) > 0, "You must specify training_dataloader in `accelerate.prepare()` when using DeepSpeed"
Rename to batch_sizes (since there are several of them)?
@@ -230,6 +245,51 @@ def prepare_model(self, model):
             model.forward = torch.cuda.amp.autocast()(model.forward)
         return model
 
+    def _prepare_deepspeed(self, *args):
+
+        ds_plugin = self.state.deepspeed_plugin
+        self.ds_config = ds_plugin.ds_config
+
+        batch_size = [obj.batch_size for obj in args if hasattr(obj, "batch_size")]
```suggestion
        logger.info("Since you passed several dataloaders, `is_train_batch_min` will determine the `train_batch_size`".)
```

+ I think we should check the list has more than one element before issuing this info!
@@ -230,6 +245,51 @@ def prepare_model(self, model):
             model.forward = torch.cuda.amp.autocast()(model.forward)
         return model
 
+    def _prepare_deepspeed(self, *args):
+
+        ds_plugin = self.state.deepspeed_plugin
+        self.ds_config = ds_plugin.ds_config
+
+        batch_size = [obj.batch_size for obj in args if hasattr(obj, "batch_size")]
+        assert len(batch_size) > 0, "You must specify training_dataloader in `accelerate.prepare()` when using DeepSpeed"
+        logger.info("Since you passed both train & eval dataloader, `is_train_batch_min` will decide the `train_batch_size`")
Nothing needs to be done for the `eval_batch_size`?
@@ -230,6 +245,51 @@ def prepare_model(self, model):
             model.forward = torch.cuda.amp.autocast()(model.forward)
         return model
 
+    def _prepare_deepspeed(self, *args):
+
+        ds_plugin = self.state.deepspeed_plugin
+        self.ds_config = ds_plugin.ds_config
+
+        batch_size = [obj.batch_size for obj in args if hasattr(obj, "batch_size")]
+        assert len(batch_size) > 0, "You must specify training_dataloader in `accelerate.prepare()` when using DeepSpeed"
+        logger.info("Since you passed both train & eval dataloader, `is_train_batch_min` will decide the `train_batch_size`")
+        batch_size = min(batch_size) if ds_plugin.is_train_batch_min else max(batch_size)
+
+        self.ds_config["train_batch_size"] = batch_size * ds_plugin.gradient_accumulation_steps * self.num_processes
+
I don't understand this part, could you explain?
@@ -230,6 +245,51 @@ def prepare_model(self, model):
             model.forward = torch.cuda.amp.autocast()(model.forward)
         return model
 
+    def _prepare_deepspeed(self, *args):
+
+        ds_plugin = self.state.deepspeed_plugin
+        self.ds_config = ds_plugin.ds_config
+
+        batch_size = [obj.batch_size for obj in args if hasattr(obj, "batch_size")]
+        assert len(batch_size) > 0, "You must specify training_dataloader in `accelerate.prepare()` when using DeepSpeed"
+        logger.info("Since you passed both train & eval dataloader, `is_train_batch_min` will decide the `train_batch_size`")
+        batch_size = min(batch_size) if ds_plugin.is_train_batch_min else max(batch_size)
+
+        self.ds_config["train_batch_size"] = batch_size * ds_plugin.gradient_accumulation_steps * self.num_processes
+
+        result = [self._prepare_one(obj) if isinstance(obj, torch.utils.data.DataLoader) else obj for obj in args]
+
+        model = None
+        optimizer = None
+        model_parameters = None
+        for obj in result:
+            if isinstance(obj, torch.nn.Module):
+                model = obj
+                model_parameters = filter(lambda p: p.requires_grad, model.parameters())
+            elif isinstance(obj, (torch.optim.Optimizer, dict)):
+                optimizer = obj
+
+        # TODO: how dict based optimizer should be passed (in optimizer or ds_plugin) ??
+        # if isinstance(optimizer, dict):
+        #     self.ds_config.update({"optimizer": optimizer})
+        #     optimizer = None
The error message should be more along the lines of "You need to pass the model along the optimizer when using Deepspeed."
@@ -230,6 +245,51 @@ def prepare_model(self, model):
             model.forward = torch.cuda.amp.autocast()(model.forward)
         return model
 
+    def _prepare_deepspeed(self, *args):
+
+        ds_plugin = self.state.deepspeed_plugin
+        self.ds_config = ds_plugin.ds_config
+
+        batch_size = [obj.batch_size for obj in args if hasattr(obj, "batch_size")]
+        assert len(batch_size) > 0, "You must specify training_dataloader in `accelerate.prepare()` when using DeepSpeed"
+        logger.info("Since you passed both train & eval dataloader, `is_train_batch_min` will decide the `train_batch_size`")
+        batch_size = min(batch_size) if ds_plugin.is_train_batch_min else max(batch_size)
+
+        self.ds_config["train_batch_size"] = batch_size * ds_plugin.gradient_accumulation_steps * self.num_processes
+
+        result = [self._prepare_one(obj) if isinstance(obj, torch.utils.data.DataLoader) else obj for obj in args]
+
+        model = None
+        optimizer = None
+        model_parameters = None
+        for obj in result:
+            if isinstance(obj, torch.nn.Module):
+                model = obj
+                model_parameters = filter(lambda p: p.requires_grad, model.parameters())
+            elif isinstance(obj, (torch.optim.Optimizer, dict)):
+                optimizer = obj
+
+        # TODO: how dict based optimizer should be passed (in optimizer or ds_plugin) ??
+        # if isinstance(optimizer, dict):
+        #     self.ds_config.update({"optimizer": optimizer})
+        #     optimizer = None
+
+        # useful when only eval_dataloader is given into `accelerator.prepare()`
+        if model is not None:
+            model, optimizer, _, _ = deepspeed.initialize(model=model, optimizer=optimizer, config_params=self.ds_config, dist_init_required=False, model_parameters=model_parameters)
+            assert isinstance(model, deepspeed.DeepSpeedEngine), "PipelineEngine not supported currently"
+            for i in range(len(result)):
+                if isinstance(result[i], torch.nn.Module):
+                    result[i] = DeepSpeedEngineWrapper(model)
+                elif isinstance(result[i], torch.optim.Optimizer):
+                    result[i] = DeepSpeedOptimizerWrapper(optimizer, model)
+            self.model = model # pointing for model.backward()
+        
+        if self.distributed_type == DistributedType.DEEPSPEED:
+            assert hasattr(self, "model"), "Accelerator instance must have model as its attribute"
I'm not sure I understand this error message.
@@ -230,6 +245,51 @@ def prepare_model(self, model):
             model.forward = torch.cuda.amp.autocast()(model.forward)
         return model
 
+    def _prepare_deepspeed(self, *args):
+
+        ds_plugin = self.state.deepspeed_plugin
+        self.ds_config = ds_plugin.ds_config
+
+        batch_size = [obj.batch_size for obj in args if hasattr(obj, "batch_size")]
+        assert len(batch_size) > 0, "You must specify training_dataloader in `accelerate.prepare()` when using DeepSpeed"
+        logger.info("Since you passed both train & eval dataloader, `is_train_batch_min` will decide the `train_batch_size`")
+        batch_size = min(batch_size) if ds_plugin.is_train_batch_min else max(batch_size)
+
+        self.ds_config["train_batch_size"] = batch_size * ds_plugin.gradient_accumulation_steps * self.num_processes
+
+        result = [self._prepare_one(obj) if isinstance(obj, torch.utils.data.DataLoader) else obj for obj in args]
+
+        model = None
+        optimizer = None
+        model_parameters = None
+        for obj in result:
+            if isinstance(obj, torch.nn.Module):
+                model = obj
+                model_parameters = filter(lambda p: p.requires_grad, model.parameters())
+            elif isinstance(obj, (torch.optim.Optimizer, dict)):
+                optimizer = obj
+
+        # TODO: how dict based optimizer should be passed (in optimizer or ds_plugin) ??
+        # if isinstance(optimizer, dict):
+        #     self.ds_config.update({"optimizer": optimizer})
+        #     optimizer = None
+
+        # useful when only eval_dataloader is given into `accelerator.prepare()`
+        if model is not None:
+            model, optimizer, _, _ = deepspeed.initialize(model=model, optimizer=optimizer, config_params=self.ds_config, dist_init_required=False, model_parameters=model_parameters)
+            assert isinstance(model, deepspeed.DeepSpeedEngine), "PipelineEngine not supported currently"
Great to have this handled here :-) 
@@ -248,7 +308,9 @@ def backward(self, loss):
         """
         Use :obj:`accelerator.backward(loss)` in lieu of :obj:`loss.backward()`.
         """
-        if self.scaler is not None:
+        if self.distributed_type == DistributedType.DEEPSPEED:
+            self.model.backward(loss)
I think deepspeed is a subcase of multi-CPU/multi-GPU no? I would rather ask the question after.
@@ -16,22 +16,22 @@
 
 from accelerate.state import ComputeEnvironment, DistributedType
 
-from .config_args import ClusterConfig
+from .config_args import ClusterConfig, DeepSpeedConfig
 from .config_utils import _ask_field, _convert_distributed_mode, _convert_yes_no_to_bool
 
 
 def get_cluster_input():
     distributed_type = _ask_field(
-        "Which type of machine are you using? ([0] No distributed training, [1] multi-CPU, [2] multi-GPU, [3] TPU): ",
+        "Which type of machine are you using? ([0] No distributed training, [1] multi-CPU, [2] multi-GPU, [3] DeepSpeed [4] TPU): ",
         _convert_distributed_mode,
-        error_message="Please enter 0, 1, 2 or 3.",
+        error_message="Please enter 0, 1, 2, 3 or 4.",
I think this one should just be left at True: there are going to be tons of parameters for deepspeed so let's keep the essential ones in the question (and use the default for the others). Users can edit the config after!
@@ -50,6 +50,25 @@ def get_cluster_input():
                 "What is the port you will use to communicate with the main process? ",
                 lambda x: int(x),
             )
+
+        # TODO: add other deepspeed args
+        if distributed_type == DistributedType.DEEPSPEED:
+            zero_stage = _ask_field(
+                "What should be your DeepSpeed's ZeRO optimization stage (0, 1, 2)? [2]: ",
+                lambda x: int(x),
+                default=2,
+            )
+            gradient_accumulation_steps = _ask_field(
+                "What should be your number of gradient accumulation steps? [1]: ",
+                lambda x: int(x),
+                default=1,
+            )
+            is_train_batch_min = _ask_field(
+                "Should accelerate choose minimum batch size incase both train & eval dataloaders are passed? [True]: ",
Why add this default?
@@ -111,13 +111,25 @@ def __post_init__(self):
 
 @dataclass
 class ClusterConfig(BaseConfig):
-    num_processes: int
+    num_processes: int = None
     machine_rank: int = 0
     num_machines: int = 1
     main_process_ip: Optional[str] = None
     main_process_port: Optional[int] = None
     main_training_function: str = "main"
 
+    def __post_init__(self):
+        super().__post_init__()
+        assert self.num_processes is not None, "num_processes can't be None"
Are those all expected by deepspeed?
@@ -151,6 +151,44 @@ def multi_gpu_launcher(args):
         raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)
 
 
+def deepspeed_launcher(args):
+    # TODO: need to complete this
+
+    cmd = ["deepspeed"]
+    if args.num_machines > 1:
+        cmd.extend(
+            [
+                "--num_gpus",
+                str(args.num_processes // args.num_machines),
+                "--num_nodes",
+                str(args.num_machines),
+                "--node_rank",
+                str(args.machine_rank),
+                "--master_addr",
+                args.main_process_ip,
+                "--master_port",
+                str(args.main_process_port),
+            ]
+        )
+    else:
+        cmd.extend(["--num_gpus", str(args.num_processes)])
+
+    cmd.append(args.training_script)
+    cmd.extend(args.training_script_args)
+
+    current_env = os.environ.copy()
+    current_env["USE_FP16"] = str(args.fp16)
+    current_env["ZERO_STAGE"] = int(args.zero_stage)
+    current_env["GRADIENT_ACCUMULATION_STEPS"] = int(args.gradient_accumulation_steps)
+    current_env["IS_TRAIN_BATCH_MIN"] = bool(args.is_train_batch_min)
There is something done in Transformers for that, if you need help.
@@ -276,14 +314,15 @@ def sagemaker_launcher(sagemaker_config: SageMakerConfig, args):
 
 def launch_command(args):
     # Sanity checks
-    if args.multi_gpu and args.tpu:
+    if args.multi_gpu and args.tpu: # TODO: fix this for deepspeed
Make sure to use the `make style` command, this is not super readable here :-)
@@ -0,0 +1,177 @@
+# 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 torch
+import torch.nn as nn
+from dataclasses import dataclass, field
+
+from .state import is_apex_available
+
+@dataclass
+class DeepSpeedPlugin:
+
+    gradient_accumulation_steps: int = field(default=1, metadata={"help": "Number of steps to accumulate gradients before updating optimizer states"})
+    zero_stage: int = field(default=0, metadata={"help": "Possible options are 0,1,2"})
+    is_train_batch_min: str = field(default=True, metadata={"help": "If both train & eval dataloaders are specified, this will decide the train_batch_size"})
This is copied from `deepspeed`? Could we subclass their class then and only change the method that requires it?
@@ -0,0 +1,177 @@
+# 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 torch
+import torch.nn as nn
+from dataclasses import dataclass, field
+
+from .state import is_apex_available
+
+@dataclass
+class DeepSpeedPlugin:
+
+    gradient_accumulation_steps: int = field(default=1, metadata={"help": "Number of steps to accumulate gradients before updating optimizer states"})
+    zero_stage: int = field(default=0, metadata={"help": "Possible options are 0,1,2"})
+    is_train_batch_min: str = field(default=True, metadata={"help": "If both train & eval dataloaders are specified, this will decide the train_batch_size"})
+
+    fp16: bool = field(repr=False, default=False, metadata={"help": "You need not define this here"})
+
+    def __post_init__(self):
+        self.ds_config = {
+            "train_batch_size": None,
+            "fp16": {"enabled": False if self.fp16 is None else self.fp16},
+            "gradient_accumulation_steps": self.gradient_accumulation_steps,
+            "zero_optimization": {"stage": self.zero_stage},
+            "steps_per_print": float("inf"), # this will stop deepspeed from logging @ stdout
+        }
+
+
+class DeepSpeedEngineWrapper(nn.Module):
Should this subclass the `AcceleratedOptimizer` for most properties and only overrides the necessary ones?
@@ -0,0 +1,177 @@
+# 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 torch
+import torch.nn as nn
+from dataclasses import dataclass, field
+
+from .state import is_apex_available
+
+@dataclass
+class DeepSpeedPlugin:
+
+    gradient_accumulation_steps: int = field(default=1, metadata={"help": "Number of steps to accumulate gradients before updating optimizer states"})
+    zero_stage: int = field(default=0, metadata={"help": "Possible options are 0,1,2"})
+    is_train_batch_min: str = field(default=True, metadata={"help": "If both train & eval dataloaders are specified, this will decide the train_batch_size"})
+
+    fp16: bool = field(repr=False, default=False, metadata={"help": "You need not define this here"})
+
+    def __post_init__(self):
+        self.ds_config = {
+            "train_batch_size": None,
+            "fp16": {"enabled": False if self.fp16 is None else self.fp16},
+            "gradient_accumulation_steps": self.gradient_accumulation_steps,
+            "zero_optimization": {"stage": self.zero_stage},
+            "steps_per_print": float("inf"), # this will stop deepspeed from logging @ stdout
+        }
+
+
+class DeepSpeedEngineWrapper(nn.Module):
+    """
+    Internal wrapper around DeepSpeedEngine instance
+
+    Args:
+        model (:obj: `DeepSpeedEngine`):
+            DeepSpeedEngine instance created using deepspeed.initalize()
+    """
+
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+        self.optimizer = model.optimizer
+
+        self.progressive_layer_drop = model.progressive_layer_drop
+        self.global_steps = model.global_steps
+
+        # overwriting micro_steps for user's gradient_accumulation
+        self.model.micro_steps = -1
+
+    def forward(self, *args, **kwargs):
+        return self.model(*args, **kwargs)
+
+    def train(self, *args, **kwargs):
+        return self.model.train(*args, **kwargs)
+
+    def eval(self, *args, **kwargs):
+        return self.model.eval(*args, **kwargs)
+
+    def is_gradient_accumulation_boundary(self):
+        return self.model.is_gradient_accumulation_boundary()
+
+    def zero_optimization(self):
+        return self.model.zero_optimization()
+
+    def amp_enabled(self):
+        return self.model.amp_enabled()
+
+    def fp16_enabled(self):
+        return self.model.fp16_enabled()
+
+    def _take_model_step(self, *args, **kwargs):
+        return self.model._take_model_step(*args, **kwargs)
+
+    def step(self, lr_kwargs=None):
+        """ DeepSpeedEngine.step() without `micro_steps` update & no profiling """
+        if self.is_gradient_accumulation_boundary(): # it shouldn't matter whether we keep this line or not
+            if self.progressive_layer_drop:
+                self.progressive_layer_drop.update_state(self.global_steps)
+
+            self._take_model_step(lr_kwargs)
+
+    def backward(self, loss):
+        """ DeepSpeedEngine.backward() with no loss scaling & no profiling """
+        if self.zero_optimization():
+            self.optimizer.is_gradient_accumulation_boundary = self.is_gradient_accumulation_boundary(
+            )
+            self.optimizer.backward(loss)
+        elif self.amp_enabled():
+            assert is_apex_available, "You have enabled apex in deepspeed_plugin, but apex is unavailable in your machine"
+            # AMP requires delaying unscale when inside gradient accumulation boundaries
+            # https://nvidia.github.io/apex/advanced.html#gradient-accumulation-across-iterations
+            delay_unscale = not self.is_gradient_accumulation_boundary()
+            with amp.scale_loss(loss,
+                                self.optimizer,
+                                delay_unscale=delay_unscale) as scaled_loss:
+                scaled_loss.backward()
+        elif self.fp16_enabled():
+            self.optimizer.backward(loss)
+        else:
+            loss.backward()
+
+        # this will ensure deepspeed gradient_accumulation matches user's accumulation
+        self.model.micro_steps += 1
+
+    def save_checkpoint(self, *args, **kwargs):
+        return self.model.save_checkpoint(*args, **kwargs)
+
+    def load_checkpoint(self, *args, **kwargs):
+        return self.model.load_checkpoint(*args, **kwargs)
+
+
+class DeepSpeedOptimizerWrapper(torch.optim.Optimizer):
I can do this next week!
@@ -113,3 +113,8 @@ def step(self):
     def _switch_parameters(self, parameters_map):
         for param_group in self.optimizer.param_groups:
             param_group["params"] = [parameters_map.get(p, p) for p in param_group["params"]]
+
+    @property
+    def is_overflow(self):
+        """ This needs to be implemented at the end """
+        return False # TODO: implement it
Apex is not supported yet. Where do you need it?
@@ -34,6 +35,20 @@
 except ImportError:
     _tpu_available = False
 
+try:
+    import deepspeed
+
+    _deepspeed_available = True
+except:
+    _deepspeed_available = False
+
+try:
+    from apex import amp
+
+    _apex_available = True
+except ImportError:
+    _apex_available = False
+
I would prefer not importing directly (see [here](https://github.com/huggingface/accelerate/blob/a96fbaaf19e650da101bfadb886ecc1d1b1cb205/src/accelerate/utils.py#L31) for an example on how to this).
@@ -34,6 +35,20 @@
 except ImportError:
     _tpu_available = False
 
+try:
+    import deepspeed
+
+    _deepspeed_available = True
+except:
+    _deepspeed_available = False
Yes, we can do that I think. Then  we won't need `deepspeed.initialize()` since deepspeed.initialize() is doing just sending everything to `DeepSpeedEngine` (see [this](https://github.com/microsoft/DeepSpeed/blob/46f4573b1a8a9cda2b45f4de4e90b631cee3f80b/deepspeed/__init__.py#L119)).
@@ -0,0 +1,177 @@
+# 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 torch
+import torch.nn as nn
+from dataclasses import dataclass, field
+
+from .state import is_apex_available
+
+@dataclass
+class DeepSpeedPlugin:
+
+    gradient_accumulation_steps: int = field(default=1, metadata={"help": "Number of steps to accumulate gradients before updating optimizer states"})
+    zero_stage: int = field(default=0, metadata={"help": "Possible options are 0,1,2"})
+    is_train_batch_min: str = field(default=True, metadata={"help": "If both train & eval dataloaders are specified, this will decide the train_batch_size"})
+
+    fp16: bool = field(repr=False, default=False, metadata={"help": "You need not define this here"})
+
+    def __post_init__(self):
+        self.ds_config = {
+            "train_batch_size": None,
+            "fp16": {"enabled": False if self.fp16 is None else self.fp16},
+            "gradient_accumulation_steps": self.gradient_accumulation_steps,
+            "zero_optimization": {"stage": self.zero_stage},
+            "steps_per_print": float("inf"), # this will stop deepspeed from logging @ stdout
+        }
+
+
+class DeepSpeedEngineWrapper(nn.Module):
yes in `deepspeed.initialize()` as `config_params`
@@ -151,6 +151,44 @@ def multi_gpu_launcher(args):
         raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)
 
 
+def deepspeed_launcher(args):
+    # TODO: need to complete this
+
+    cmd = ["deepspeed"]
+    if args.num_machines > 1:
+        cmd.extend(
+            [
+                "--num_gpus",
+                str(args.num_processes // args.num_machines),
+                "--num_nodes",
+                str(args.num_machines),
+                "--node_rank",
+                str(args.machine_rank),
+                "--master_addr",
+                args.main_process_ip,
+                "--master_port",
+                str(args.main_process_port),
+            ]
+        )
+    else:
+        cmd.extend(["--num_gpus", str(args.num_processes)])
+
+    cmd.append(args.training_script)
+    cmd.extend(args.training_script_args)
+
+    current_env = os.environ.copy()
+    current_env["USE_FP16"] = str(args.fp16)
+    current_env["ZERO_STAGE"] = int(args.zero_stage)
+    current_env["GRADIENT_ACCUMULATION_STEPS"] = int(args.gradient_accumulation_steps)
+    current_env["IS_TRAIN_BATCH_MIN"] = bool(args.is_train_batch_min)
`deepspeed.initialize()` passes arguments to either of 2 class -  `DeepSpeedEngine` or  `PipelineEngine`. & I had to make changes to DeepSpeedEngine for preventing automatic gradient_accumulation inside deepspeed (in `DeepSpeedEngineWrapper`) & hence only `DeepSpeedEngine` is supported right now. But this line can be removed if we directly inherit `DeepSpeedEngineWrapper` from `DeepSpeedEngine` directly (as pointed out by you in a later comment).
@@ -230,6 +245,51 @@ def prepare_model(self, model):
             model.forward = torch.cuda.amp.autocast()(model.forward)
         return model
 
+    def _prepare_deepspeed(self, *args):
+
+        ds_plugin = self.state.deepspeed_plugin
+        self.ds_config = ds_plugin.ds_config
+
+        batch_size = [obj.batch_size for obj in args if hasattr(obj, "batch_size")]
+        assert len(batch_size) > 0, "You must specify training_dataloader in `accelerate.prepare()` when using DeepSpeed"
+        logger.info("Since you passed both train & eval dataloader, `is_train_batch_min` will decide the `train_batch_size`")
+        batch_size = min(batch_size) if ds_plugin.is_train_batch_min else max(batch_size)
+
+        self.ds_config["train_batch_size"] = batch_size * ds_plugin.gradient_accumulation_steps * self.num_processes
+
+        result = [self._prepare_one(obj) if isinstance(obj, torch.utils.data.DataLoader) else obj for obj in args]
+
+        model = None
+        optimizer = None
+        model_parameters = None
+        for obj in result:
+            if isinstance(obj, torch.nn.Module):
+                model = obj
+                model_parameters = filter(lambda p: p.requires_grad, model.parameters())
+            elif isinstance(obj, (torch.optim.Optimizer, dict)):
+                optimizer = obj
+
+        # TODO: how dict based optimizer should be passed (in optimizer or ds_plugin) ??
+        # if isinstance(optimizer, dict):
+        #     self.ds_config.update({"optimizer": optimizer})
+        #     optimizer = None
+
+        # useful when only eval_dataloader is given into `accelerator.prepare()`
+        if model is not None:
+            model, optimizer, _, _ = deepspeed.initialize(model=model, optimizer=optimizer, config_params=self.ds_config, dist_init_required=False, model_parameters=model_parameters)
+            assert isinstance(model, deepspeed.DeepSpeedEngine), "PipelineEngine not supported currently"
DeepSpeed can build optimizer automatically for us if optimizer key is included in `ds_config`, instead of giving `torch.optim.Optimizer` in `deepspeed.initialize()`. Passing optimizer dict is the recommended way by deepspeed since they have some fused implementation & it's more easy to switch their optimizers. So we need to include that somehow.

So one way is to let user pass optimizer as dict in accelerate.prepare() itself & later we can handle that; while other can be to include it in deepspeed_plugin. What way you would suggest??
@@ -230,6 +245,51 @@ def prepare_model(self, model):
             model.forward = torch.cuda.amp.autocast()(model.forward)
         return model
 
+    def _prepare_deepspeed(self, *args):
+
+        ds_plugin = self.state.deepspeed_plugin
+        self.ds_config = ds_plugin.ds_config
+
+        batch_size = [obj.batch_size for obj in args if hasattr(obj, "batch_size")]
+        assert len(batch_size) > 0, "You must specify training_dataloader in `accelerate.prepare()` when using DeepSpeed"
+        logger.info("Since you passed both train & eval dataloader, `is_train_batch_min` will decide the `train_batch_size`")
+        batch_size = min(batch_size) if ds_plugin.is_train_batch_min else max(batch_size)
+
+        self.ds_config["train_batch_size"] = batch_size * ds_plugin.gradient_accumulation_steps * self.num_processes
+
+        result = [self._prepare_one(obj) if isinstance(obj, torch.utils.data.DataLoader) else obj for obj in args]
+
+        model = None
+        optimizer = None
+        model_parameters = None
+        for obj in result:
+            if isinstance(obj, torch.nn.Module):
+                model = obj
+                model_parameters = filter(lambda p: p.requires_grad, model.parameters())
+            elif isinstance(obj, (torch.optim.Optimizer, dict)):
+                optimizer = obj
+
+        # TODO: how dict based optimizer should be passed (in optimizer or ds_plugin) ??
+        # if isinstance(optimizer, dict):
+        #     self.ds_config.update({"optimizer": optimizer})
+        #     optimizer = None
its needed in DeepSpeedEngineWrapper.
@@ -34,6 +35,20 @@
 except ImportError:
     _tpu_available = False
 
+try:
+    import deepspeed
+
+    _deepspeed_available = True
+except:
+    _deepspeed_available = False
+
+try:
+    from apex import amp
+
+    _apex_available = True
+except ImportError:
+    _apex_available = False
+
I need to fix it a bit. Will work on it.
@@ -111,13 +111,25 @@ def __post_init__(self):
 
 @dataclass
 class ClusterConfig(BaseConfig):
-    num_processes: int
+    num_processes: int = None
     machine_rank: int = 0
     num_machines: int = 1
     main_process_ip: Optional[str] = None
     main_process_port: Optional[int] = None
     main_training_function: str = "main"
 
+    def __post_init__(self):
+        super().__post_init__()
+        assert self.num_processes is not None, "num_processes can't be None"
no
@@ -230,6 +245,51 @@ def prepare_model(self, model):
             model.forward = torch.cuda.amp.autocast()(model.forward)
         return model
 
+    def _prepare_deepspeed(self, *args):
+
+        ds_plugin = self.state.deepspeed_plugin
+        self.ds_config = ds_plugin.ds_config
+
+        batch_size = [obj.batch_size for obj in args if hasattr(obj, "batch_size")]
+        assert len(batch_size) > 0, "You must specify training_dataloader in `accelerate.prepare()` when using DeepSpeed"
+        logger.info("Since you passed both train & eval dataloader, `is_train_batch_min` will decide the `train_batch_size`")
+        batch_size = min(batch_size) if ds_plugin.is_train_batch_min else max(batch_size)
+
+        self.ds_config["train_batch_size"] = batch_size * ds_plugin.gradient_accumulation_steps * self.num_processes
+
Let's focus on this on another PR. We could have accelerate support this in general (since the goal is to have the same API for everything) but it would require more work.
@@ -230,6 +245,51 @@ def prepare_model(self, model):
             model.forward = torch.cuda.amp.autocast()(model.forward)
         return model
 
+    def _prepare_deepspeed(self, *args):
+
+        ds_plugin = self.state.deepspeed_plugin
+        self.ds_config = ds_plugin.ds_config
+
+        batch_size = [obj.batch_size for obj in args if hasattr(obj, "batch_size")]
+        assert len(batch_size) > 0, "You must specify training_dataloader in `accelerate.prepare()` when using DeepSpeed"
+        logger.info("Since you passed both train & eval dataloader, `is_train_batch_min` will decide the `train_batch_size`")
+        batch_size = min(batch_size) if ds_plugin.is_train_batch_min else max(batch_size)
+
+        self.ds_config["train_batch_size"] = batch_size * ds_plugin.gradient_accumulation_steps * self.num_processes
+
+        result = [self._prepare_one(obj) if isinstance(obj, torch.utils.data.DataLoader) else obj for obj in args]
+
+        model = None
+        optimizer = None
+        model_parameters = None
+        for obj in result:
+            if isinstance(obj, torch.nn.Module):
+                model = obj
+                model_parameters = filter(lambda p: p.requires_grad, model.parameters())
+            elif isinstance(obj, (torch.optim.Optimizer, dict)):
+                optimizer = obj
+
+        # TODO: how dict based optimizer should be passed (in optimizer or ds_plugin) ??
+        # if isinstance(optimizer, dict):
+        #     self.ds_config.update({"optimizer": optimizer})
+        #     optimizer = None
Ok, then please use the same as below and not a direct import.
@@ -34,6 +35,20 @@
 except ImportError:
     _tpu_available = False
 
+try:
+    import deepspeed
+
+    _deepspeed_available = True
+except:
+    _deepspeed_available = False
+
+try:
+    from apex import amp
+
+    _apex_available = True
+except ImportError:
+    _apex_available = False
+
okay. another PR after this then :)
@@ -230,6 +245,51 @@ def prepare_model(self, model):
             model.forward = torch.cuda.amp.autocast()(model.forward)
         return model
 
+    def _prepare_deepspeed(self, *args):
+
+        ds_plugin = self.state.deepspeed_plugin
+        self.ds_config = ds_plugin.ds_config
+
+        batch_size = [obj.batch_size for obj in args if hasattr(obj, "batch_size")]
+        assert len(batch_size) > 0, "You must specify training_dataloader in `accelerate.prepare()` when using DeepSpeed"
+        logger.info("Since you passed both train & eval dataloader, `is_train_batch_min` will decide the `train_batch_size`")
+        batch_size = min(batch_size) if ds_plugin.is_train_batch_min else max(batch_size)
+
+        self.ds_config["train_batch_size"] = batch_size * ds_plugin.gradient_accumulation_steps * self.num_processes
+
+        result = [self._prepare_one(obj) if isinstance(obj, torch.utils.data.DataLoader) else obj for obj in args]
+
+        model = None
+        optimizer = None
+        model_parameters = None
+        for obj in result:
+            if isinstance(obj, torch.nn.Module):
+                model = obj
+                model_parameters = filter(lambda p: p.requires_grad, model.parameters())
+            elif isinstance(obj, (torch.optim.Optimizer, dict)):
+                optimizer = obj
+
+        # TODO: how dict based optimizer should be passed (in optimizer or ds_plugin) ??
+        # if isinstance(optimizer, dict):
+        #     self.ds_config.update({"optimizer": optimizer})
+        #     optimizer = None
Can we use `batch_sizes` here?
@@ -231,6 +255,70 @@ def prepare_model(self, model):
             model.forward = torch.cuda.amp.autocast()(model.forward)
         return model
 
+    def _prepare_deepspeed(self, *args):
+
+        ds_plugin = self.state.deepspeed_plugin
+        self.ds_config = ds_plugin.ds_config
+
+        bz = [obj.batch_size for obj in args if hasattr(obj, "batch_size")]
It think this one would read better as (including the suggested rename form before):
```suggestion
        if len(batch_sizes) == 0:
            raise ValueError("You must specify a training or evaluation dataloader in `accelerate.prepare()` when using DeepSpeed.")
```
@@ -231,6 +255,70 @@ def prepare_model(self, model):
             model.forward = torch.cuda.amp.autocast()(model.forward)
         return model
 
+    def _prepare_deepspeed(self, *args):
+
+        ds_plugin = self.state.deepspeed_plugin
+        self.ds_config = ds_plugin.ds_config
+
+        bz = [obj.batch_size for obj in args if hasattr(obj, "batch_size")]
+        assert (
+            len(bz) > 0
+        ), "You must specify a training or evaluation dataloader in `accelerate.prepare()` when using DeepSpeed."
Since we are computing it just after, can we log it? the log should be something like:
```
f"Since you passed both train and evaluation dataloader, `is_train_batch_min` (here {ds_plugin.is_train_batch_min} will decide the `train_batch_size` ({batch_size_per_device})."
```
This might need to be split in two lines as well to respect the 119 char limit.
@@ -231,6 +255,70 @@ def prepare_model(self, model):
             model.forward = torch.cuda.amp.autocast()(model.forward)
         return model
 
+    def _prepare_deepspeed(self, *args):
+
+        ds_plugin = self.state.deepspeed_plugin
+        self.ds_config = ds_plugin.ds_config
+
+        bz = [obj.batch_size for obj in args if hasattr(obj, "batch_size")]
+        assert (
+            len(bz) > 0
+        ), "You must specify a training or evaluation dataloader in `accelerate.prepare()` when using DeepSpeed."
+        if len(bz) > 1:
+            logger.info(
+                "Since you passed both train & eval dataloader, `is_train_batch_min` will decide the `train_batch_size`"
+            )
+        batch_size_per_device = min(bz) if ds_plugin.is_train_batch_min else max(bz)
There can be several models or several optimizers passed, so we should make a list of models and optimizers for that case here? If deepspeed only supports one, we should raise an error here.
@@ -231,6 +255,70 @@ def prepare_model(self, model):
             model.forward = torch.cuda.amp.autocast()(model.forward)
         return model
 
+    def _prepare_deepspeed(self, *args):
+
+        ds_plugin = self.state.deepspeed_plugin
+        self.ds_config = ds_plugin.ds_config
+
+        bz = [obj.batch_size for obj in args if hasattr(obj, "batch_size")]
+        assert (
+            len(bz) > 0
+        ), "You must specify a training or evaluation dataloader in `accelerate.prepare()` when using DeepSpeed."
+        if len(bz) > 1:
+            logger.info(
+                "Since you passed both train & eval dataloader, `is_train_batch_min` will decide the `train_batch_size`"
+            )
+        batch_size_per_device = min(bz) if ds_plugin.is_train_batch_min else max(bz)
+
+        self.ds_config["train_batch_size"] = (
+            batch_size_per_device * ds_plugin.gradient_accumulation_steps * self.num_processes
+        )
+
+        result = [self._prepare_one(obj) if isinstance(obj, torch.utils.data.DataLoader) else obj for obj in args]
+
+        model = None
+        optimizer = None
+        for obj in result:
+            if isinstance(obj, torch.nn.Module):
+                model = obj
+            elif isinstance(obj, (torch.optim.Optimizer, dict)):
+                optimizer = obj
Is this the only reference we create? It seems so but I want to be certain.
@@ -357,6 +447,7 @@ def free_memory(self):
         method between two trainings with different models/optimizers.
         """
         self._optimizers = []
+        self.deepspeed_engine = None
```suggestion
                "What should be your DeepSpeed's ZeRO optimization stage (0, 1, 2)? [0]: ",
```
Should include 3?
@@ -50,6 +50,39 @@ def get_cluster_input():
                 "What is the port you will use to communicate with the main process? ",
                 lambda x: int(x),
             )
+
+    ds_config = None
+    if distributed_type in [DistributedType.MULTI_GPU, DistributedType.NO]:
+        use_deepspeed = _ask_field(
+            "Do you want to use DeepSpeed (yes/no)? [NO]",
+            _convert_yes_no_to_bool,
+            default=False,
+            error_message="Please enter yes or no.",
+        )
+        if use_deepspeed:
+            distributed_type = DistributedType.DEEPSPEED
+
+        ds_config = {}
+        if distributed_type == DistributedType.DEEPSPEED:
+            ds_config["zero_stage"] = _ask_field(
+                "What should be your DeepSpeed's ZeRO optimization stage (0, 1, 2)? [0]: ",
Should this be a more general check of whether `attr` is a dict?
@@ -276,19 +338,26 @@ def sagemaker_launcher(sagemaker_config: SageMakerConfig, args):
 
 def launch_command(args):
     # Sanity checks
-    if args.multi_gpu and args.tpu:
-        raise ValueError("You can only pick one between `--multi_gpu` and `--tpu`.")
+    if sum([args.multi_gpu, args.tpu, args.use_deepspeed]) > 1:
+        raise ValueError("You can only pick one between `--multi_gpu`, `--use_deepspeed`, `--tpu`.")
 
     defaults = None
     # Get the default from the config file.
     if args.config_file is not None or os.path.isfile(default_config_file) and not args.cpu:
         defaults = load_config_from_file(args.config_file)
-        if not args.multi_gpu and not args.tpu:
+        if not args.multi_gpu and not args.tpu and not args.use_deepspeed:
+            args.use_deepspeed = defaults.distributed_type == DistributedType.DEEPSPEED
             args.multi_gpu = defaults.distributed_type == DistributedType.MULTI_GPU
             args.tpu = defaults.distributed_type == DistributedType.TPU
         if defaults.compute_environment == ComputeEnvironment.LOCAL_MACHINE:
             # Update args with the defaults
             for name, attr in defaults.__dict__.items():
+                if name == "ds_config":
If I am faster than you and implement it for more general usecases, please rebase. If I'm not I will do this after merging :-)
@@ -114,3 +114,8 @@ def step(self):
     def _switch_parameters(self, parameters_map):
         for param_group in self.optimizer.param_groups:
             param_group["params"] = [parameters_map.get(p, p) for p in param_group["params"]]
+
+    @property
+    def is_overflow(self):
+        """This needs to be implemented at the end"""
+        return False  # TODO: implement it
This file should be removed afterward :-) 
@@ -0,0 +1,461 @@
+# coding=utf-8
Can we say `self.deepspeed_config` to go with `self.state.deepspeed_plugin`?

Also would love if this function was using `deepspeed_plugin` instead of `ds_pluging`.
@@ -231,6 +255,70 @@ def prepare_model(self, model):
             model.forward = torch.cuda.amp.autocast()(model.forward)
         return model
 
+    def _prepare_deepspeed(self, *args):
+
+        ds_plugin = self.state.deepspeed_plugin
+        self.ds_config = ds_plugin.ds_config
value of `ds_config` will always be a dict since I am defining it to be `dict` when taking deepspeed arguments from `accelerate config`.
@@ -276,19 +338,26 @@ def sagemaker_launcher(sagemaker_config: SageMakerConfig, args):
 
 def launch_command(args):
     # Sanity checks
-    if args.multi_gpu and args.tpu:
-        raise ValueError("You can only pick one between `--multi_gpu` and `--tpu`.")
+    if sum([args.multi_gpu, args.tpu, args.use_deepspeed]) > 1:
+        raise ValueError("You can only pick one between `--multi_gpu`, `--use_deepspeed`, `--tpu`.")
 
     defaults = None
     # Get the default from the config file.
     if args.config_file is not None or os.path.isfile(default_config_file) and not args.cpu:
         defaults = load_config_from_file(args.config_file)
-        if not args.multi_gpu and not args.tpu:
+        if not args.multi_gpu and not args.tpu and not args.use_deepspeed:
+            args.use_deepspeed = defaults.distributed_type == DistributedType.DEEPSPEED
             args.multi_gpu = defaults.distributed_type == DistributedType.MULTI_GPU
             args.tpu = defaults.distributed_type == DistributedType.TPU
         if defaults.compute_environment == ComputeEnvironment.LOCAL_MACHINE:
             # Update args with the defaults
             for name, attr in defaults.__dict__.items():
+                if name == "ds_config":
will remove it in the very end (before merge)
@@ -0,0 +1,461 @@
+# coding=utf-8
I didn't understand the question. Could you please clarify??

I will rename `ds_plugin` to `deepspeed_plugin`.
@@ -231,6 +255,70 @@ def prepare_model(self, model):
             model.forward = torch.cuda.amp.autocast()(model.forward)
         return model
 
+    def _prepare_deepspeed(self, *args):
+
+        ds_plugin = self.state.deepspeed_plugin
+        self.ds_config = ds_plugin.ds_config
I would prefer we use `deepspeed_plugin` and `deepspeed_config` instead of `ds_plugin` and `ds_config`
@@ -231,6 +255,70 @@ def prepare_model(self, model):
             model.forward = torch.cuda.amp.autocast()(model.forward)
         return model
 
+    def _prepare_deepspeed(self, *args):
+
+        ds_plugin = self.state.deepspeed_plugin
+        self.ds_config = ds_plugin.ds_config
Yes, that's why I'm suggesting we check `if isinstance(attr, dict)` instead of the name.
@@ -276,19 +338,26 @@ def sagemaker_launcher(sagemaker_config: SageMakerConfig, args):
 
 def launch_command(args):
     # Sanity checks
-    if args.multi_gpu and args.tpu:
-        raise ValueError("You can only pick one between `--multi_gpu` and `--tpu`.")
+    if sum([args.multi_gpu, args.tpu, args.use_deepspeed]) > 1:
+        raise ValueError("You can only pick one between `--multi_gpu`, `--use_deepspeed`, `--tpu`.")
 
     defaults = None
     # Get the default from the config file.
     if args.config_file is not None or os.path.isfile(default_config_file) and not args.cpu:
         defaults = load_config_from_file(args.config_file)
-        if not args.multi_gpu and not args.tpu:
+        if not args.multi_gpu and not args.tpu and not args.use_deepspeed:
+            args.use_deepspeed = defaults.distributed_type == DistributedType.DEEPSPEED
             args.multi_gpu = defaults.distributed_type == DistributedType.MULTI_GPU
             args.tpu = defaults.distributed_type == DistributedType.TPU
         if defaults.compute_environment == ComputeEnvironment.LOCAL_MACHINE:
             # Update args with the defaults
             for name, attr in defaults.__dict__.items():
+                if name == "ds_config":
yes
@@ -357,6 +447,7 @@ def free_memory(self):
         method between two trainings with different models/optimizers.
         """
         self._optimizers = []
+        self.deepspeed_engine = None
This sentence is a bit ambiguous: do I have to change something in my code? I'd either say "You don't need to change anything. If you desire to tweak your deepspeed related args, we provide ...."
@@ -160,6 +160,26 @@ Once you have MPI setup on your cluster, just run:
 mpirun -np 2 python examples/nlp_example.py
 ```
 
+## Launching training using DeepSpeed
+
+🤗 Accelerate supports training on single/multiple GPUs using DeepSpeed. For using deepspeed, you don't really need to change anything in your code once you integrate accelerate in your code. You can set everything using just `accelerate config`.
```suggestion
1. DeepSpeed support is experimental for now. In case you get into some issue, please raise an issue.
```
Users should always feel free to open issues :)
@@ -160,6 +160,26 @@ Once you have MPI setup on your cluster, just run:
 mpirun -np 2 python examples/nlp_example.py
 ```
 
+## Launching training using DeepSpeed
+
+🤗 Accelerate supports training on single/multiple GPUs using DeepSpeed. For using deepspeed, you don't really need to change anything in your code once you integrate accelerate in your code. You can set everything using just `accelerate config`.
+
+In case you would like to change your deepspeed related args from python script, we provide you the `DeepSpeedPlugin`.
+
+```python
+from accelerator import Accelerator, DeepSpeedPlugin
+
+# deepspeed needs to know your gradient accumulation steps before hand, so don't forget to pass it
+# Remember you need to do gradient accumulation by yourself, just like you would have done without deepspeed
+deepspeed_plugin = DeepSpeedPlugin(zero_stage=2, gradient_accumulation_steps=2)
+accelerator = Accelerator(fp16=True, deepspeed_plugin=deepspeed_plugin)
+```
+
+Note:
+
+1. DeepSpeed support is experimental for now. In case you get into some issue, feel free to raise an issue.
The file `deepspeed_utils` has a hard requirement on `deepspeed` being available
@@ -20,12 +20,22 @@
 from packaging import version
 
 from .data_loader import prepare_data_loader
+from .deepspeed_utils import DeepSpeedEngineWrapper, DeepSpeedOptimizerWrapper, DeepSpeedPlugin
This should be conditionally imported if objects are to be imported from this file in others where deepspeed isn't necessarily installed
@@ -0,0 +1,150 @@
+# 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
+from dataclasses import dataclass, field
+
+from deepspeed import DeepSpeedEngine
I don't understand this argument
@@ -0,0 +1,150 @@
+# 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
+from dataclasses import dataclass, field
+
+from deepspeed import DeepSpeedEngine
+
+from .optimizer import AcceleratedOptimizer
+from .state import is_apex_available
+
+
+if is_apex_available():
+    import amp
+
+
+@dataclass
+class DeepSpeedPlugin:
+
+    gradient_accumulation_steps: int = field(
+        default=None, metadata={"help": "Number of steps to accumulate gradients before updating optimizer states"}
+    )
+    zero_stage: int = field(
+        default=None,
+        metadata={"help": "Possible options are 0,1,2,3; Default will be taken from environment variable"},
+    )
+    is_train_batch_min: str = field(
+        default=True,
+        metadata={"help": "If both train & eval dataloaders are specified, this will decide the train_batch_size"},
+    )
Is it necessary to have this here if it shouldn't be used?
@@ -0,0 +1,150 @@
+# 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
+from dataclasses import dataclass, field
+
+from deepspeed import DeepSpeedEngine
+
+from .optimizer import AcceleratedOptimizer
+from .state import is_apex_available
+
+
+if is_apex_available():
+    import amp
+
+
+@dataclass
+class DeepSpeedPlugin:
+
+    gradient_accumulation_steps: int = field(
+        default=None, metadata={"help": "Number of steps to accumulate gradients before updating optimizer states"}
+    )
+    zero_stage: int = field(
+        default=None,
+        metadata={"help": "Possible options are 0,1,2,3; Default will be taken from environment variable"},
+    )
+    is_train_batch_min: str = field(
+        default=True,
+        metadata={"help": "If both train & eval dataloaders are specified, this will decide the train_batch_size"},
+    )
+
+    auto_opt_mapping: bool = field(
+        default=True,
+        metadata={"help": "whether to map torch.adam to deepspeed optimizer version of adam based on config"},
+    )
+
+    offload_optimizer_device: bool = field(default=None, metadata={"help": "Possible options are none|cpu|nvme"})
+
+    fp16: bool = field(repr=False, default=None, metadata={"help": "You need not define this here"})
This would place the comment as a docstring, but is that really the goal? Shouldn't it rather be:
```suggestion
        pass  # `model.step()` is doing that automatically. Therefore, it's implementation is not needed
```
@@ -0,0 +1,150 @@
+# 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
+from dataclasses import dataclass, field
+
+from deepspeed import DeepSpeedEngine
+
+from .optimizer import AcceleratedOptimizer
+from .state import is_apex_available
+
+
+if is_apex_available():
+    import amp
+
+
+@dataclass
+class DeepSpeedPlugin:
+
+    gradient_accumulation_steps: int = field(
+        default=None, metadata={"help": "Number of steps to accumulate gradients before updating optimizer states"}
+    )
+    zero_stage: int = field(
+        default=None,
+        metadata={"help": "Possible options are 0,1,2,3; Default will be taken from environment variable"},
+    )
+    is_train_batch_min: str = field(
+        default=True,
+        metadata={"help": "If both train & eval dataloaders are specified, this will decide the train_batch_size"},
+    )
+
+    auto_opt_mapping: bool = field(
+        default=True,
+        metadata={"help": "whether to map torch.adam to deepspeed optimizer version of adam based on config"},
+    )
+
+    offload_optimizer_device: bool = field(default=None, metadata={"help": "Possible options are none|cpu|nvme"})
+
+    fp16: bool = field(repr=False, default=None, metadata={"help": "You need not define this here"})
+
+    def __post_init__(self):
+
+        if self.gradient_accumulation_steps is None:
+            self.gradient_accumulation_steps = int(os.environ.get("GRADIENT_ACCUMULATION_STEPS", 1))
+
+        if self.zero_stage is None:
+            self.zero_stage = int(os.environ.get("DEEPSPEED_ZERO_STAGE", 2))
+
+        if self.fp16 is None:
+            self.fp16 = bool(os.environ.get("USE_FP16", "False"))
+
+        if self.offload_optimizer_device is None:
+            self.offload_optimizer_device = os.environ.get("DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE", "none")
+
+        self.deepspeed_config = {
+            "train_batch_size": None,
+            "fp16": {"enabled": self.fp16},
+            "gradient_accumulation_steps": self.gradient_accumulation_steps,
+            "zero_optimization": {
+                "stage": self.zero_stage,
+                "offload_optimizer": {
+                    "device": self.offload_optimizer_device,
+                },
+            },
+            "steps_per_print": float("inf"),  # this will stop deepspeed from logging @ stdout
+        }
+
+
+class DeepSpeedEngineWrapper(DeepSpeedEngine):
+    """
+    Wrapper over deepspeed.DeepSpeedEngine object
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # overwriting micro_steps for user's gradient_accumulation
+        self.micro_steps = -1
+
+    def step(self, lr_kwargs=None):
+        """DeepSpeedEngine.step() without `micro_steps` update & no profiling"""
+        if self.is_gradient_accumulation_boundary():  # it shouldn't matter whether we keep this line or not
+            if self.progressive_layer_drop:
+                self.progressive_layer_drop.update_state(self.global_steps)
+
+            self._take_model_step(lr_kwargs)
+
+    def backward(self, loss):
+        """DeepSpeedEngine.backward() with with no loss scaling; no profiling but with `micro_steps` update"""
+
+        if self.zero_optimization():
+            self.optimizer.is_gradient_accumulation_boundary = self.is_gradient_accumulation_boundary()
+            self.optimizer.backward(loss)
+        elif self.amp_enabled():
+            # AMP requires delaying unscale when inside gradient accumulation boundaries
+            # https://nvidia.github.io/apex/advanced.html#gradient-accumulation-across-iterations
+            delay_unscale = not self.is_gradient_accumulation_boundary()
+            with amp.scale_loss(loss, self.optimizer, delay_unscale=delay_unscale) as scaled_loss:
+                scaled_loss.backward()
+        elif self.fp16_enabled():
+            self.optimizer.backward(loss)
+        else:
+            loss.backward()
+
+        if self.enable_backward_allreduce:
+            self.allreduce_gradients()
+
+        # this will ensure deepspeed gradient_accumulation matches user's accumulation
+        self.micro_steps += 1
+
+
+class DeepSpeedOptimizerWrapper(AcceleratedOptimizer):
+    """
+    Internal wrapper around a deepspeed optimizer.
+
+    Args:
+        optimizer (:obj:`torch.optim.optimizer.Optimizer`):
+            The optimizer to wrap.
+    """
+
+    def __init__(self, optimizer, model: DeepSpeedEngineWrapper):
+        super().__init__(optimizer, device_placement=False, scaler=None)
+
+        self.model = model
+
+    def zero_grad(self, set_to_none=None):
+        """`model.step()` is doing that automatically. Therefore, it's implementation is not needed"""
Do we really need a `DeepSpeedPlugin` to be initialized if deepspeed isn't available?
@@ -74,10 +84,18 @@ def __init__(
         split_batches: bool = False,
         fp16: bool = None,
         cpu: bool = False,
+        deepspeed_plugin: DeepSpeedPlugin = None,
         rng_types: Optional[List[Union[str, RNGType]]] = None,
         kwargs_handlers: Optional[List[KwargsHandler]] = None,
     ):
-        self.state = AcceleratorState(fp16=fp16, cpu=cpu, _from_accelerator=True)
+        if deepspeed_plugin is None:
+            deepspeed_plugin = DeepSpeedPlugin()  # init from env variables
+        else:
+            assert isinstance(
+                deepspeed_plugin, DeepSpeedPlugin
+            ), "`deepspeed_plugin` must be a DeepSpeedPlugin object."
```suggestion
            "Do you want to use DeepSpeed (yes/no)? [no]",
```
@@ -50,6 +50,39 @@ def get_cluster_input():
                 "What is the port you will use to communicate with the main process? ",
                 lambda x: int(x),
             )
+
+    deepspeed_config = None
+    if distributed_type in [DistributedType.MULTI_GPU, DistributedType.NO]:
+        use_deepspeed = _ask_field(
+            "Do you want to use DeepSpeed (yes/no)? [NO]",
DeepSpeed needs to know the batch size of the data even if we build dataloader by ourself. & we are extracting batch size from dataloader internally. But whenever 2 dataloaders are passed in `accelerator.prepare()` (i.e. both train & eval), we can't directly take the batch size from `train_dataloader`, since internally we don't know which is train & which is  eval (because of the way accelerator.prepare() takes in args). So this argument will help users to decide what is train batch size & what is eval.
@@ -0,0 +1,150 @@
+# 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
+from dataclasses import dataclass, field
+
+from deepspeed import DeepSpeedEngine
+
+from .optimizer import AcceleratedOptimizer
+from .state import is_apex_available
+
+
+if is_apex_available():
+    import amp
+
+
+@dataclass
+class DeepSpeedPlugin:
+
+    gradient_accumulation_steps: int = field(
+        default=None, metadata={"help": "Number of steps to accumulate gradients before updating optimizer states"}
+    )
+    zero_stage: int = field(
+        default=None,
+        metadata={"help": "Possible options are 0,1,2,3; Default will be taken from environment variable"},
+    )
+    is_train_batch_min: str = field(
+        default=True,
+        metadata={"help": "If both train & eval dataloaders are specified, this will decide the train_batch_size"},
+    )
I will try to find other solution.
@@ -0,0 +1,150 @@
+# 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
+from dataclasses import dataclass, field
+
+from deepspeed import DeepSpeedEngine
+
+from .optimizer import AcceleratedOptimizer
+from .state import is_apex_available
+
+
+if is_apex_available():
+    import amp
+
+
+@dataclass
+class DeepSpeedPlugin:
+
+    gradient_accumulation_steps: int = field(
+        default=None, metadata={"help": "Number of steps to accumulate gradients before updating optimizer states"}
+    )
+    zero_stage: int = field(
+        default=None,
+        metadata={"help": "Possible options are 0,1,2,3; Default will be taken from environment variable"},
+    )
+    is_train_batch_min: str = field(
+        default=True,
+        metadata={"help": "If both train & eval dataloaders are specified, this will decide the train_batch_size"},
+    )
+
+    auto_opt_mapping: bool = field(
+        default=True,
+        metadata={"help": "whether to map torch.adam to deepspeed optimizer version of adam based on config"},
+    )
+
+    offload_optimizer_device: bool = field(default=None, metadata={"help": "Possible options are none|cpu|nvme"})
+
+    fp16: bool = field(repr=False, default=None, metadata={"help": "You need not define this here"})
I think you probably want a space at the end of this to be consistent with the other _ask_fields
@@ -50,6 +50,39 @@ def get_cluster_input():
                 "What is the port you will use to communicate with the main process? ",
                 lambda x: int(x),
             )
+
+    deepspeed_config = None
+    if distributed_type in [DistributedType.MULTI_GPU, DistributedType.NO]:
+        use_deepspeed = _ask_field(
+            "Do you want to use DeepSpeed (yes/no)? [NO]",
Should actually be 
```
          "Do you want to use DeepSpeed? [yes/NO]: ",
```
to be compatible with what is actually in the lib for now.
@@ -50,6 +50,39 @@ def get_cluster_input():
                 "What is the port you will use to communicate with the main process? ",
                 lambda x: int(x),
             )
+
+    deepspeed_config = None
+    if distributed_type in [DistributedType.MULTI_GPU, DistributedType.NO]:
+        use_deepspeed = _ask_field(
+            "Do you want to use DeepSpeed (yes/no)? [NO]",
```suggestion
                    "Where to offload optimizer states? [NONE/cpu/nvme]: ",
```

Also to have the same kind of syntax as the others.
@@ -50,6 +50,39 @@ def get_cluster_input():
                 "What is the port you will use to communicate with the main process? ",
                 lambda x: int(x),
             )
+
+    deepspeed_config = None
+    if distributed_type in [DistributedType.MULTI_GPU, DistributedType.NO]:
+        use_deepspeed = _ask_field(
+            "Do you want to use DeepSpeed (yes/no)? [NO]",
+            _convert_yes_no_to_bool,
+            default=False,
+            error_message="Please enter yes or no.",
+        )
+        if use_deepspeed:
+            distributed_type = DistributedType.DEEPSPEED
+
+        deepspeed_config = {}
+        if distributed_type == DistributedType.DEEPSPEED:
+            deepspeed_config["zero_stage"] = _ask_field(
+                "What should be your DeepSpeed's ZeRO optimization stage (0, 1, 2, 3)? [2]: ",
+                lambda x: int(x),
+                default=2,
+            )
+
+            if deepspeed_config["zero_stage"] >= 2:
+                deepspeed_config["offload_optimizer_device"] = _ask_field(
+                    "Where to offload optimizer states (none|cpu|nvme)? [none]: ",
```suggestion
🤗 Accelerate supports training on single/multiple GPUs using DeepSpeed. to use it, you don't need to change anything in your training code; you can set everything using just `accelerate config`. However, if you desire to tweak your DeepSpeed related args from your python script, we provide you the `DeepSpeedPlugin`.
```
@@ -160,6 +160,26 @@ Once you have MPI setup on your cluster, just run:
 mpirun -np 2 python examples/nlp_example.py
 ```
 
+## Launching training using DeepSpeed
+
+🤗 Accelerate supports training on single/multiple GPUs using DeepSpeed. For using deepspeed, you don't need to change anything. You can set everything using just `accelerate config`. However, if you desire to tweak your deepspeed related args from your python script, we provide you the `DeepSpeedPlugin`.
```suggestion
# Remember you still need to do gradient accumulation by yourself, just like you would have done without deepspeed
```
@@ -160,6 +160,26 @@ Once you have MPI setup on your cluster, just run:
 mpirun -np 2 python examples/nlp_example.py
 ```
 
+## Launching training using DeepSpeed
+
+🤗 Accelerate supports training on single/multiple GPUs using DeepSpeed. For using deepspeed, you don't need to change anything. You can set everything using just `accelerate config`. However, if you desire to tweak your deepspeed related args from your python script, we provide you the `DeepSpeedPlugin`.
+
+```python
+from accelerator import Accelerator, DeepSpeedPlugin
+
+# deepspeed needs to know your gradient accumulation steps before hand, so don't forget to pass it
+# Remember you need to do gradient accumulation by yourself, just like you would have done without deepspeed
```suggestion
Note: DeepSpeed support is experimental for now. In case you get into some problem, please open an issue.
```
@@ -160,6 +160,26 @@ Once you have MPI setup on your cluster, just run:
 mpirun -np 2 python examples/nlp_example.py
 ```
 
+## Launching training using DeepSpeed
+
+🤗 Accelerate supports training on single/multiple GPUs using DeepSpeed. For using deepspeed, you don't need to change anything. You can set everything using just `accelerate config`. However, if you desire to tweak your deepspeed related args from your python script, we provide you the `DeepSpeedPlugin`.
+
+```python
+from accelerator import Accelerator, DeepSpeedPlugin
+
+# deepspeed needs to know your gradient accumulation steps before hand, so don't forget to pass it
+# Remember you need to do gradient accumulation by yourself, just like you would have done without deepspeed
+deepspeed_plugin = DeepSpeedPlugin(zero_stage=2, gradient_accumulation_steps=2)
+accelerator = Accelerator(fp16=True, deepspeed_plugin=deepspeed_plugin)
+
+# How to save your 🤗 Transformer?
+accelerator.wait_for_everyone()
+unwrapped_model = accelerator.unwrap_model(model)
+unwrapped_model.save_pretrained(save_dir, save_function=accelerator.save, state_dict=accelerator.get_state_dict(model))
+```
+
+Note: DeepSpeed support is experimental for now. In case you get into some issue, please raise an issue.
Note for before merging this: the PR should be rebased on master and the reference to the model added in the `Accelerator` should be cleaned up in the method introduced in the PR above.
https://api.github.com/repos/huggingface/accelerate/pulls/82
diff --git a/README.md b/README.md
index 9c044ce9e..4f6cf5f40 100644
--- a/README.md
+++ b/README.md
@@ -63,12 +63,12 @@ Here is an example:
 + device = accelerator.device
 
   model = torch.nn.Transformer().to(device)
-  optim = torch.optim.Adam(model.parameters())
+  optimizer = torch.optim.Adam(model.parameters())
 
   dataset = load_dataset('my_dataset')
   data = torch.utils.data.DataLoader(dataset, shuffle=True)
 
-+ model, optim, data = accelerator.prepare(model, optim, data)
++ model, optimizer, data = accelerator.prepare(model, optimizer, data)
 
   model.train()
   for epoch in range(10):
@@ -104,12 +104,12 @@ In particular, the same code can then be run without modification on your local
 
 - model = torch.nn.Transformer().to(device)
 + model = torch.nn.Transformer()
-  optim = torch.optim.Adam(model.parameters())
+  optimizer = torch.optim.Adam(model.parameters())
 
   dataset = load_dataset('my_dataset')
   data = torch.utils.data.DataLoader(dataset, shuffle=True)
 
-+ model, optim, data = accelerator.prepare(model, optim, data)
++ model, optimizer, data = accelerator.prepare(model, optimizer, data)
 
   model.train()
   for epoch in range(10):

https://api.github.com/repos/huggingface/accelerate/pulls/81
diff --git a/src/accelerate/commands/config/config_args.py b/src/accelerate/commands/config/config_args.py
index 0e115bdd4..8cc993b59 100644
--- a/src/accelerate/commands/config/config_args.py
+++ b/src/accelerate/commands/config/config_args.py
@@ -39,7 +39,8 @@
 
 
 def load_config_from_file(config_file):
-    config_file = config_file if config_file is not None else default_config_file
+    config_file_exists = config_file is not None and os.path.isfile(config_file)
+    config_file = config_file if config_file_exists else default_config_file
     with open(config_file, "r", encoding="utf-8") as f:
         if config_file.endswith(".json"):
             if (

https://api.github.com/repos/huggingface/accelerate/pulls/79
diff --git a/src/accelerate/notebook_launcher.py b/src/accelerate/notebook_launcher.py
index 3a9b03e28..4c785b12f 100644
--- a/src/accelerate/notebook_launcher.py
+++ b/src/accelerate/notebook_launcher.py
@@ -106,7 +106,7 @@ def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, use
             launcher = PrepareForLaunch(function, distributed_type="MULTI_GPU")
             try:
                 print(f"Launching a training on {num_processes} GPUs.")
-                start_processes(launcher, nprocs=num_processes, start_method="fork")
+                start_processes(launcher, args=args, nprocs=num_processes, start_method="fork")
             finally:
                 # Clean up the environment variables set.
                 del os.environ["WORLD_SIZE"]

https://api.github.com/repos/huggingface/accelerate/pulls/78
diff --git a/src/accelerate/notebook_launcher.py b/src/accelerate/notebook_launcher.py
index 57f253267..3a9b03e28 100644
--- a/src/accelerate/notebook_launcher.py
+++ b/src/accelerate/notebook_launcher.py
@@ -42,10 +42,10 @@ def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, use
             The port to use to communicate between processes when launching a multi-GPU training.
     """
     # Are we in a google colab or a Kaggle Kernel?
-    if "IPython" in sys.modules:
-        in_colab_or_kaggle = "google.colab" in str(sys.modules["IPython"].get_ipython())
-    elif any(key.startswith("KAGGLE") for key in os.environ.keys()):
+    if any(key.startswith("KAGGLE") for key in os.environ.keys()):
         in_colab_or_kaggle = True
+    elif "IPython" in sys.modules:
+        in_colab_or_kaggle = "google.colab" in str(sys.modules["IPython"].get_ipython())
     else:
         in_colab_or_kaggle = False
 

https://api.github.com/repos/huggingface/accelerate/pulls/73
diff --git a/examples/README.md b/examples/README.md
index 434a230ea..4d941b5d3 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -110,7 +110,7 @@ The same script can be run in any of the following configurations:
 Prior to running it you should install timm and torchvision:
 
 ```bash
-pip install timm, torchvision
+pip install timm torchvision
 ```
 
 and you should download the data with the following commands:

https://api.github.com/repos/huggingface/accelerate/pulls/70
diff --git a/src/accelerate/optimizer.py b/src/accelerate/optimizer.py
index 63b53eccb..1dd9f506a 100644
--- a/src/accelerate/optimizer.py
+++ b/src/accelerate/optimizer.py
@@ -17,6 +17,7 @@
 from packaging import version
 
 from .state import AcceleratorState, DistributedType, is_tpu_available
+from .utils import honor_type
 
 
 if is_tpu_available():
@@ -25,7 +26,7 @@
 
 def move_to_device(state, device):
     if isinstance(state, (list, tuple)):
-        return type(state)(move_to_device(t, device) for t in state)
+        return honor_type(state, (move_to_device(t, device) for t in state))
     elif isinstance(state, dict):
         return type(state)({k: move_to_device(v, device) for k, v in state.items()})
     elif isinstance(state, torch.Tensor):
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 9e38cbad9..7302ead15 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -99,6 +99,17 @@ def synchronize_rng_states(rng_types: List[Union[str, RNGType]], generator: Opti
         synchronize_rng_state(RNGType(rng_type), generator=generator)
 
 
+def honor_type(obj, generator):
+    """
+    Cast a generator to the same type as obj (list, tuple or namedtuple)
+    """
+    # There is no direct check whether an object if of type namedtuple sadly, this is a workaround.
+    if isinstance(obj, tuple) and hasattr(obj, "_fields"):
+        # Can instantiate a namedtuple from a generator directly, contrary to a tuple/list.
+        return type(obj)(*list(generator))
+    return type(obj)(generator)
+
+
 def send_to_device(tensor, device):
     """
     Recursively sends the elements in a nested list/tuple/dictionary of tensors to a given device.
@@ -113,7 +124,7 @@ def send_to_device(tensor, device):
         The same data structure as :obj:`tensor` with all tensors sent to the proper device.
     """
     if isinstance(tensor, (list, tuple)):
-        return type(tensor)(send_to_device(t, device) for t in tensor)
+        return honor_type(tensor, (send_to_device(t, device) for t in tensor))
     elif isinstance(tensor, dict):
         return type(tensor)({k: send_to_device(v, device) for k, v in tensor.items()})
     elif not hasattr(tensor, "to"):
@@ -138,7 +149,7 @@ def extract_model_from_parallel(model):
 
 def _tpu_gather(tensor, name="tensor"):
     if isinstance(tensor, (list, tuple)):
-        return type(tensor)(_tpu_gather(t, name=f"{name}_{i}") for i, t in enumerate(tensor))
+        return honor_type(tensor, (_tpu_gather(t, name=f"{name}_{i}") for i, t in enumerate(tensor)))
     elif isinstance(tensor, dict):
         return type(tensor)({k: _tpu_gather(v, name=f"{name}_{k}") for k, v in tensor.items()})
     elif not isinstance(tensor, torch.Tensor):
@@ -148,7 +159,7 @@ def _tpu_gather(tensor, name="tensor"):
 
 def _gpu_gather(tensor):
     if isinstance(tensor, (list, tuple)):
-        return type(tensor)(_gpu_gather(t) for t in tensor)
+        return honor_type(tensor, (_gpu_gather(t) for t in tensor))
     elif isinstance(tensor, dict):
         return type(tensor)({k: _gpu_gather(v) for k, v in tensor.items()})
     elif not isinstance(tensor, torch.Tensor):
@@ -198,7 +209,7 @@ def pad_across_processes(tensor, dim=0, pad_index=0, pad_first=False):
             Whether to pad at the beginning or the end.
     """
     if isinstance(tensor, (list, tuple)):
-        return type(tensor)(pad_across_processes(t, dim=dim, pad_index=pad_index) for t in tensor)
+        return honor_type(tensor, (pad_across_processes(t, dim=dim, pad_index=pad_index) for t in tensor))
     elif isinstance(tensor, dict):
         return type(tensor)({k: pad_across_processes(v, dim=dim, pad_index=pad_index) for k, v in tensor.items()})
     elif not isinstance(tensor, torch.Tensor):
diff --git a/tests/test_utils.py b/tests/test_utils.py
new file mode 100644
index 000000000..ca617634d
--- /dev/null
+++ b/tests/test_utils.py
@@ -0,0 +1,53 @@
+# 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 unittest
+from collections import namedtuple
+
+import torch
+
+from accelerate.utils import send_to_device
+
+
+TestNamedTuple = namedtuple("TestNamedTuple", "a b")
+
+
+class UtilsTester(unittest.TestCase):
+    def test_send_to_device(self):
+        tensor = torch.randn(5, 2)
+        device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+        result1 = send_to_device(tensor, device)
+        self.assertTrue(torch.equal(result1.cpu(), tensor))
+
+        result2 = send_to_device((tensor, [tensor, tensor]), device)
+        self.assertIsInstance(result2, tuple)
+        self.assertTrue(torch.equal(result2[0].cpu(), tensor))
+        self.assertIsInstance(result2[1], list)
+        self.assertTrue(torch.equal(result2[1][0].cpu(), tensor))
+        self.assertTrue(torch.equal(result2[1][1].cpu(), tensor))
+
+        result2 = send_to_device({"a": tensor, "b": [tensor, tensor]}, device)
+        self.assertIsInstance(result2, dict)
+        self.assertTrue(torch.equal(result2["a"].cpu(), tensor))
+        self.assertIsInstance(result2["b"], list)
+        self.assertTrue(torch.equal(result2["b"][0].cpu(), tensor))
+        self.assertTrue(torch.equal(result2["b"][1].cpu(), tensor))
+
+        result3 = send_to_device(TestNamedTuple(a=tensor, b=[tensor, tensor]), device)
+        self.assertIsInstance(result3, TestNamedTuple)
+        self.assertTrue(torch.equal(result3.a.cpu(), tensor))
+        self.assertIsInstance(result3.b, list)
+        self.assertTrue(torch.equal(result3.b[0].cpu(), tensor))
+        self.assertTrue(torch.equal(result3.b[1].cpu(), tensor))

Great tests :)
@@ -0,0 +1,53 @@
+# 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 unittest
+from collections import namedtuple
+
+import torch
+
+from accelerate.utils import send_to_device
+
+
+TestNamedTuple = namedtuple("TestNamedTuple", "a b")
+
+
+class UtilsTester(unittest.TestCase):
https://api.github.com/repos/huggingface/accelerate/pulls/67
diff --git a/README.md b/README.md
index 6e2211143..9c044ce9e 100644
--- a/README.md
+++ b/README.md
@@ -151,6 +151,15 @@ For instance, here is how you would run the GLUE example on the MRPC task (from
 accelerate launch examples/nlp_example.py
 ```
 
+## Launching multi-CPU run using MPI
+
+🤗 Here is another way to launch multi-CPU run using MPI. You can learn how to install Open MPI on [this page](https://www.open-mpi.org/faq/?category=building#easy-build). You can use Intel MPI or MVAPICH as well.
+Once you have MPI setup on your cluster, just run:
+
+```bash
+mpirun -np 2 python examples/nlp_example.py
+```
+
 ## Launching your training from a notebook
 
 🤗 Accelerate also provides a `notebook_launcher` function you can use in a notebook to launch a distributed training. This is especially useful for Colab or Kaggle notebooks with a TPU backend. Just define your training loop in a `training_function` then in your last cell, add:
@@ -188,6 +197,8 @@ pip install accelerate
 ## Supported integrations
 
 - CPU only
+- multi-CPU on one node (machine)
+- multi-CPU on several nodes (machines)
 - single GPU
 - multi-GPU on one node (machine)
 - multi-GPU on several nodes (machines)
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 0143ac949..c7706d390 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -223,6 +223,9 @@ def prepare_model(self, model):
                 output_device=self.local_process_index,
                 **kwargs,
             )
+        elif self.distributed_type == DistributedType.MULTI_CPU:
+            kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {}
+            model = torch.nn.parallel.DistributedDataParallel(model, **kwargs)
         if self.native_amp:
             model.forward = torch.cuda.amp.autocast()(model.forward)
         return model
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index c8bf1bcd7..ebc176089 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -22,16 +22,16 @@
 
 def get_cluster_input():
     distributed_type = _ask_field(
-        "Which type of machine are you using? ([0] No distributed training, [1] multi-GPU, [2] TPU): ",
+        "Which type of machine are you using? ([0] No distributed training, [1] multi-CPU, [2] multi-GPU, [3] TPU): ",
         _convert_distributed_mode,
-        error_message="Please enter 0, 1 or 2.",
+        error_message="Please enter 0, 1, 2 or 3.",
     )
 
     machine_rank = 0
     num_machines = 1
     main_process_ip = None
     main_process_port = None
-    if distributed_type == DistributedType.MULTI_GPU:
+    if distributed_type == DistributedType.MULTI_GPU or distributed_type == DistributedType.MULTI_CPU:
         num_machines = _ask_field(
             "How many different machines will you use (use more than 1 for multi-node training)? [1]: ",
             lambda x: int(x),
diff --git a/src/accelerate/commands/config/config_utils.py b/src/accelerate/commands/config/config_utils.py
index dd14965f7..2bc0e7ec9 100644
--- a/src/accelerate/commands/config/config_utils.py
+++ b/src/accelerate/commands/config/config_utils.py
@@ -37,7 +37,7 @@ def _convert_compute_environment(value):
 
 def _convert_distributed_mode(value):
     value = int(value)
-    return DistributedType(["NO", "MULTI_GPU", "TPU"][value])
+    return DistributedType(["NO", "MULTI_CPU", "MULTI_GPU", "TPU"][value])
 
 
 def _convert_sagemaker_distributed_mode(value):
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index 042a6bdde..76376038b 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -19,6 +19,14 @@
 import torch
 
 
+try:
+    import torch_ccl  # noqa: F401
+
+    _ccl_available = True
+except ImportError:
+    _ccl_available = False
+
+
 try:
     import torch_xla.core.xla_model as xm
 
@@ -27,6 +35,19 @@
     _tpu_available = False
 
 
+def get_int_from_env(env_keys, default):
+    """Returns the first positive env value found in the `env_keys` list or the default."""
+    for e in env_keys:
+        val = int(os.environ.get(e, -1))
+        if val >= 0:
+            return val
+    return default
+
+
+def is_ccl_available():
+    return _ccl_available
+
+
 def is_tpu_available():
     return _tpu_available
 
@@ -43,12 +64,14 @@ class DistributedType(str, Enum):
     Values:
 
         - **NO** -- Not a distributed environment, just a single process.
+        - **MULTI_CPU** -- Distributed on multiple CPU nodes.
         - **MULTI_GPU** -- Distributed on multiple GPUs.
         - **TPU** -- Distributed on TPUs.
     """
 
     # Subclassing str as well as Enum allows the `DistributedType` to be JSON-serializable out of the box.
     NO = "NO"
+    MULTI_CPU = "MULTI_CPU"
     MULTI_GPU = "MULTI_GPU"
     TPU = "TPU"
 
@@ -107,6 +130,7 @@ class AcceleratorState:
     def __init__(self, fp16: bool = None, cpu: bool = False, _from_accelerator: bool = False):
         self.__dict__ = self._shared_state
         if not getattr(self, "initialized", False):
+            self.backend = None
             if not _from_accelerator:
                 raise ValueError(
                     "Please make sure to properly initialize your accelerator via `accelerator = Accelerator()` "
@@ -123,12 +147,50 @@ def __init__(self, fp16: bool = None, cpu: bool = False, _from_accelerator: bool
                 self.distributed_type = DistributedType.MULTI_GPU
                 if not torch.distributed.is_initialized():
                     torch.distributed.init_process_group(backend="nccl")
+                    self.backend = "nccl"
                 self.num_processes = torch.distributed.get_world_size()
                 self.process_index = torch.distributed.get_rank()
                 self.local_process_index = int(os.environ.get("LOCAL_RANK", -1))
                 self.device = torch.device("cuda", self.local_process_index)
                 torch.cuda.set_device(self.device)
                 self.use_fp16 = parse_flag_from_env("USE_FP16", False) if fp16 is None else fp16
+            elif get_int_from_env(["PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE", "WORLD_SIZE"], 1) > 1:
+                self.distributed_type = DistributedType.MULTI_CPU
+                if is_ccl_available() and get_int_from_env(["CCL_WORKER_COUNT"], 0) > 0:
+                    backend = "ccl"
+                elif torch.distributed.is_mpi_available():
+                    backend = "mpi"
+                else:
+                    backend = "gloo"
+                # Try to get launch configuration from environment variables set by MPI launcher - works for Intel MPI, OpenMPI and MVAPICH
+                rank = get_int_from_env(["RANK", "PMI_RANK", "OMPI_COMM_WORLD_RANK", "MV2_COMM_WORLD_RANK"], 0)
+                size = get_int_from_env(["WORLD_SIZE", "PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE"], 1)
+                local_rank = get_int_from_env(
+                    ["LOCAL_RANK", "MPI_LOCALRANKID", "OMPI_COMM_WORLD_LOCAL_RANK", "MV2_COMM_WORLD_LOCAL_RANK"], 0
+                )
+                local_size = get_int_from_env(
+                    ["MPI_LOCALNRANKS", "OMPI_COMM_WORLD_LOCAL_SIZE", "MV2_COMM_WORLD_LOCAL_SIZE"], 1
+                )
+                self.local_process_index = local_rank
+                os.environ["RANK"] = str(rank)
+                os.environ["WORLD_SIZE"] = str(size)
+                os.environ["LOCAL_RANK"] = str(local_rank)
+                if not os.environ.get("MASTER_PORT", None):
+                    os.environ["MASTER_PORT"] = "29500"
+                if not os.environ.get("MASTER_ADDR", None):
+                    if local_size != size and backend != "mpi":
+                        raise ValueError(
+                            "Looks like distributed multinode run but MASTER_ADDR env not set, "
+                            "please try exporting rank 0's hostname as MASTER_ADDR"
+                        )
+                if not torch.distributed.is_initialized():
+                    torch.distributed.init_process_group(backend, rank=rank, world_size=size)
+                    self.backend = backend
+                self.num_processes = torch.distributed.get_world_size()
+                self.process_index = torch.distributed.get_rank()
+                self.local_process_index = local_rank
+                self.device = torch.device("cpu")
+                self.use_fp16 = False
             else:
                 self.distributed_type = DistributedType.NO
                 self.num_processes = 1
@@ -139,7 +201,7 @@ def __init__(self, fp16: bool = None, cpu: bool = False, _from_accelerator: bool
 
     def __repr__(self):
         return (
-            f"Distributed environment: {self.distributed_type}\n"
+            f"Distributed environment: {self.distributed_type}{('  Backend: ' + self.backend) if self.backend else ''}\n"
             f"Num processes: {self.num_processes}\n"
             f"Process index: {self.process_index}\n"
             f"Local process index: {self.local_process_index}\n"
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 5dfe85e54..9e38cbad9 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -80,6 +80,8 @@ def synchronize_rng_state(rng_type: Optional[RNGType] = None, generator: Optiona
         rng_state = rng_state.to(state.device)
         torch.distributed.broadcast(rng_state, 0)
         rng_state = rng_state.cpu()
+    elif state.distributed_type == DistributedType.MULTI_CPU:
+        torch.distributed.broadcast(rng_state, 0)
 
     # Set the broadcast rng state
     if rng_type == RNGType.TORCH:
@@ -156,6 +158,9 @@ def _gpu_gather(tensor):
     return torch.cat(output_tensors, dim=0)
 
 
+_cpu_gather = _gpu_gather
+
+
 def gather(tensor):
     """
     Recursively gather tensor in a nested list/tuple/dictionary of tensors from all devices.
@@ -171,6 +176,8 @@ def gather(tensor):
         return _tpu_gather(tensor, name="accelerate.utils.gather")
     elif AcceleratorState().distributed_type == DistributedType.MULTI_GPU:
         return _gpu_gather(tensor)
+    elif AcceleratorState().distributed_type == DistributedType.MULTI_CPU:
+        return _cpu_gather(tensor)
     else:
         return tensor
 
@@ -230,7 +237,10 @@ def wait_for_everyone():
 
         Make sure all processes will reach this instruction otherwise one of your processes will hang forever.
     """
-    if AcceleratorState().distributed_type == DistributedType.MULTI_GPU:
+    if (
+        AcceleratorState().distributed_type == DistributedType.MULTI_GPU
+        or AcceleratorState().distributed_type == DistributedType.MULTI_CPU
+    ):
         torch.distributed.barrier()
     elif AcceleratorState().distributed_type == DistributedType.TPU:
         xm.rendezvous("accelerate.utils.wait_for_everyone")
@@ -266,7 +276,7 @@ def __init__(self, launcher, distributed_type="NO"):
         self.distributed_type = DistributedType(distributed_type)
 
     def __call__(self, index, *args):
-        if self.distributed_type == DistributedType.MULTI_GPU:
+        if self.distributed_type == DistributedType.MULTI_GPU or self.distributed_type == DistributedType.MULTI_CPU:
             # Prepare the environment for torch.distributed
             os.environ["LOCAL_RANK"] = str(index)
             os.environ["RANK"] = str(index)

```suggestion
            model = torch.nn.parallel.DistributedDataParallel(model, **kwargs)
```
This can probably fit in one line.
@@ -223,6 +223,12 @@ def prepare_model(self, model):
                 output_device=self.local_process_index,
                 **kwargs,
             )
+        elif self.distributed_type == DistributedType.MULTI_CPU:
+            kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {}
+            model = torch.nn.parallel.DistributedDataParallel(
+                model,
+                **kwargs,
+            )
```suggestion
        "Which type of machine are you using? ([0] No distributed training, [1] multinode-CPU, [2] multi-GPU, [3] TPU): ",
```
This is not multi-CPU on one machine, it's multiple machines for CPU, so I think we should be more explicit than multi-CPU
@@ -22,16 +22,16 @@
 
 def get_cluster_input():
     distributed_type = _ask_field(
-        "Which type of machine are you using? ([0] No distributed training, [1] multi-GPU, [2] TPU): ",
+        "Which type of machine are you using? ([0] No distributed training, [1] multi-CPU, [2] multi-GPU, [3] TPU): ",
The name is a bit too short, and we don't use 2 like this in this repo (for instance we have `send_to_device` not `send2device`). It also returns the first match What about `get_int_from_env`?

`env_list` is also not the right name for the argument as we don't pass a list of envs, but a lsit of names to query in the env, so I would use `env_keys` instead of `env_list`. `default` should not have a default value since there is not a canonical one (the function is used with -1 once, 0 twice and 1 twice below).

Lastly, a short docstring that explains what this function does (returning the first positive env found in the list or the default) would be awesome!
@@ -27,6 +35,18 @@
     _tpu_available = False
 
 
+def env2int(env_list, default=-1):
I think the value should also be `MULTI_NODE_CPU`.
@@ -37,7 +37,7 @@ def _convert_compute_environment(value):
 
 def _convert_distributed_mode(value):
     value = int(value)
-    return DistributedType(["NO", "MULTI_GPU", "TPU"][value])
+    return DistributedType(["NO", "MULTI_CPU", "MULTI_GPU", "TPU"][value])
```suggestion
                if not os.environ.get("MASTER_PORT", None):
                    os.environ["MASTER_PORT"] = "29500"
```
@@ -123,12 +146,43 @@ def __init__(self, fp16: bool = None, cpu: bool = False, _from_accelerator: bool
                 self.distributed_type = DistributedType.MULTI_GPU
                 if not torch.distributed.is_initialized():
                     torch.distributed.init_process_group(backend="nccl")
+                    self.backend = "nccl"
                 self.num_processes = torch.distributed.get_world_size()
                 self.process_index = torch.distributed.get_rank()
                 self.local_process_index = int(os.environ.get("LOCAL_RANK", -1))
                 self.device = torch.device("cuda", self.local_process_index)
                 torch.cuda.set_device(self.device)
                 self.use_fp16 = parse_flag_from_env("USE_FP16", False) if fp16 is None else fp16
+            elif env2int(["PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE", "WORLD_SIZE"]) > 1:
+                self.distributed_type = DistributedType.MULTI_CPU
+                if is_ccl_available() and env2int(["CCL_WORKER_COUNT"]) > 0:
+                    backend = "ccl"
+                elif torch.distributed.is_mpi_available():
+                    backend = "mpi"
+                else:
+                    backend = "gloo"
+                rank = env2int(["RANK", "PMI_RANK", "OMPI_COMM_WORLD_RANK", "MV2_COMM_WORLD_RANK"], 0)
+                size = env2int(["WORLD_SIZE", "PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE"], 1)
+                local_rank = env2int(["LOCAL_RANK", "MPI_LOCALRANKID", "OMPI_COMM_WORLD_LOCAL_RANK", "MV2_COMM_WORLD_LOCAL_RANK"], 0)
+                local_size = env2int(["MPI_LOCALNRANKS", "OMPI_COMM_WORLD_LOCAL_SIZE", "MV2_COMM_WORLD_LOCAL_SIZE"], 1)
+                self.local_process_index = local_rank
+                os.environ["RANK"] = str(rank)
+                os.environ["WORLD_SIZE"] = str(size)
+                if not os.environ.get("MASTER_PORT", None): os.environ["MASTER_PORT"] = "29500"
I thin we should raise a ValueError here telling the user to pass a `MASTER_ADDR` instead of relying on a default.
@@ -123,12 +146,43 @@ def __init__(self, fp16: bool = None, cpu: bool = False, _from_accelerator: bool
                 self.distributed_type = DistributedType.MULTI_GPU
                 if not torch.distributed.is_initialized():
                     torch.distributed.init_process_group(backend="nccl")
+                    self.backend = "nccl"
                 self.num_processes = torch.distributed.get_world_size()
                 self.process_index = torch.distributed.get_rank()
                 self.local_process_index = int(os.environ.get("LOCAL_RANK", -1))
                 self.device = torch.device("cuda", self.local_process_index)
                 torch.cuda.set_device(self.device)
                 self.use_fp16 = parse_flag_from_env("USE_FP16", False) if fp16 is None else fp16
+            elif env2int(["PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE", "WORLD_SIZE"]) > 1:
+                self.distributed_type = DistributedType.MULTI_CPU
+                if is_ccl_available() and env2int(["CCL_WORKER_COUNT"]) > 0:
+                    backend = "ccl"
+                elif torch.distributed.is_mpi_available():
+                    backend = "mpi"
+                else:
+                    backend = "gloo"
+                rank = env2int(["RANK", "PMI_RANK", "OMPI_COMM_WORLD_RANK", "MV2_COMM_WORLD_RANK"], 0)
+                size = env2int(["WORLD_SIZE", "PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE"], 1)
+                local_rank = env2int(["LOCAL_RANK", "MPI_LOCALRANKID", "OMPI_COMM_WORLD_LOCAL_RANK", "MV2_COMM_WORLD_LOCAL_RANK"], 0)
+                local_size = env2int(["MPI_LOCALNRANKS", "OMPI_COMM_WORLD_LOCAL_SIZE", "MV2_COMM_WORLD_LOCAL_SIZE"], 1)
+                self.local_process_index = local_rank
+                os.environ["RANK"] = str(rank)
+                os.environ["WORLD_SIZE"] = str(size)
+                if not os.environ.get("MASTER_PORT", None): os.environ["MASTER_PORT"] = "29500"
+                if not os.environ.get("MASTER_ADDR", None):
+                    if local_size != size and backend != "mpi":
+                        print("Warning: Looks like distributed multinode run but MASTER_ADDR env not set, using '127.0.0.1' as default")
+                        print("If this run hangs, try exporting rank 0's hostname as MASTER_ADDR")
This one is not used anywhere, so I would remove it.
@@ -123,12 +146,43 @@ def __init__(self, fp16: bool = None, cpu: bool = False, _from_accelerator: bool
                 self.distributed_type = DistributedType.MULTI_GPU
                 if not torch.distributed.is_initialized():
                     torch.distributed.init_process_group(backend="nccl")
+                    self.backend = "nccl"
                 self.num_processes = torch.distributed.get_world_size()
                 self.process_index = torch.distributed.get_rank()
                 self.local_process_index = int(os.environ.get("LOCAL_RANK", -1))
                 self.device = torch.device("cuda", self.local_process_index)
                 torch.cuda.set_device(self.device)
                 self.use_fp16 = parse_flag_from_env("USE_FP16", False) if fp16 is None else fp16
+            elif env2int(["PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE", "WORLD_SIZE"]) > 1:
+                self.distributed_type = DistributedType.MULTI_CPU
+                if is_ccl_available() and env2int(["CCL_WORKER_COUNT"]) > 0:
+                    backend = "ccl"
+                elif torch.distributed.is_mpi_available():
+                    backend = "mpi"
+                else:
+                    backend = "gloo"
+                rank = env2int(["RANK", "PMI_RANK", "OMPI_COMM_WORLD_RANK", "MV2_COMM_WORLD_RANK"], 0)
+                size = env2int(["WORLD_SIZE", "PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE"], 1)
+                local_rank = env2int(["LOCAL_RANK", "MPI_LOCALRANKID", "OMPI_COMM_WORLD_LOCAL_RANK", "MV2_COMM_WORLD_LOCAL_RANK"], 0)
+                local_size = env2int(["MPI_LOCALNRANKS", "OMPI_COMM_WORLD_LOCAL_SIZE", "MV2_COMM_WORLD_LOCAL_SIZE"], 1)
+                self.local_process_index = local_rank
+                os.environ["RANK"] = str(rank)
+                os.environ["WORLD_SIZE"] = str(size)
+                if not os.environ.get("MASTER_PORT", None): os.environ["MASTER_PORT"] = "29500"
+                if not os.environ.get("MASTER_ADDR", None):
+                    if local_size != size and backend != "mpi":
+                        print("Warning: Looks like distributed multinode run but MASTER_ADDR env not set, using '127.0.0.1' as default")
+                        print("If this run hangs, try exporting rank 0's hostname as MASTER_ADDR")
+                        os.environ["MASTER_ADDR"] = "127.0.0.1"
+                if not torch.distributed.is_initialized():
+                    torch.distributed.init_process_group(backend, rank=rank, world_size=size)
+                    self.backend = backend
+                self.num_processes = torch.distributed.get_world_size()
+                self.process_index = torch.distributed.get_rank()
+                self.local_process_index = local_rank
+                self.local_num_processes = local_size
We should only display the backend if it's not None.
@@ -139,7 +193,7 @@ def __init__(self, fp16: bool = None, cpu: bool = False, _from_accelerator: bool
 
     def __repr__(self):
         return (
-            f"Distributed environment: {self.distributed_type}\n"
+            f"Distributed environment: {self.distributed_type}  Backend: {self.backend}\n"
Well, actually we can have multi-cpu on a single machine. The typical use case is one process (or rank as it is called in MPI world) per socket. So, on a dual socket machine, we can have two ranks per node. So, I suggest keeping it as multi-CPU.
@@ -22,16 +22,16 @@
 
 def get_cluster_input():
     distributed_type = _ask_field(
-        "Which type of machine are you using? ([0] No distributed training, [1] multi-GPU, [2] TPU): ",
+        "Which type of machine are you using? ([0] No distributed training, [1] multi-CPU, [2] multi-GPU, [3] TPU): ",
Some MPI-like custom backends can initialize without MASTER_ADDR so I just put a warning. I can change to ValueError if you think that's a better choice.
@@ -123,12 +146,43 @@ def __init__(self, fp16: bool = None, cpu: bool = False, _from_accelerator: bool
                 self.distributed_type = DistributedType.MULTI_GPU
                 if not torch.distributed.is_initialized():
                     torch.distributed.init_process_group(backend="nccl")
+                    self.backend = "nccl"
                 self.num_processes = torch.distributed.get_world_size()
                 self.process_index = torch.distributed.get_rank()
                 self.local_process_index = int(os.environ.get("LOCAL_RANK", -1))
                 self.device = torch.device("cuda", self.local_process_index)
                 torch.cuda.set_device(self.device)
                 self.use_fp16 = parse_flag_from_env("USE_FP16", False) if fp16 is None else fp16
+            elif env2int(["PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE", "WORLD_SIZE"]) > 1:
+                self.distributed_type = DistributedType.MULTI_CPU
+                if is_ccl_available() and env2int(["CCL_WORKER_COUNT"]) > 0:
+                    backend = "ccl"
+                elif torch.distributed.is_mpi_available():
+                    backend = "mpi"
+                else:
+                    backend = "gloo"
+                rank = env2int(["RANK", "PMI_RANK", "OMPI_COMM_WORLD_RANK", "MV2_COMM_WORLD_RANK"], 0)
+                size = env2int(["WORLD_SIZE", "PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE"], 1)
+                local_rank = env2int(["LOCAL_RANK", "MPI_LOCALRANKID", "OMPI_COMM_WORLD_LOCAL_RANK", "MV2_COMM_WORLD_LOCAL_RANK"], 0)
+                local_size = env2int(["MPI_LOCALNRANKS", "OMPI_COMM_WORLD_LOCAL_SIZE", "MV2_COMM_WORLD_LOCAL_SIZE"], 1)
+                self.local_process_index = local_rank
+                os.environ["RANK"] = str(rank)
+                os.environ["WORLD_SIZE"] = str(size)
+                if not os.environ.get("MASTER_PORT", None): os.environ["MASTER_PORT"] = "29500"
+                if not os.environ.get("MASTER_ADDR", None):
+                    if local_size != size and backend != "mpi":
+                        print("Warning: Looks like distributed multinode run but MASTER_ADDR env not set, using '127.0.0.1' as default")
+                        print("If this run hangs, try exporting rank 0's hostname as MASTER_ADDR")
We are in a test where `backend != "mpi"`, so I would adapt the error message (to say the backend needs a MASTER_ADDR) and raise an error. For MPI, it will still work without the MASTER_ADDR.
@@ -123,12 +146,43 @@ def __init__(self, fp16: bool = None, cpu: bool = False, _from_accelerator: bool
                 self.distributed_type = DistributedType.MULTI_GPU
                 if not torch.distributed.is_initialized():
                     torch.distributed.init_process_group(backend="nccl")
+                    self.backend = "nccl"
                 self.num_processes = torch.distributed.get_world_size()
                 self.process_index = torch.distributed.get_rank()
                 self.local_process_index = int(os.environ.get("LOCAL_RANK", -1))
                 self.device = torch.device("cuda", self.local_process_index)
                 torch.cuda.set_device(self.device)
                 self.use_fp16 = parse_flag_from_env("USE_FP16", False) if fp16 is None else fp16
+            elif env2int(["PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE", "WORLD_SIZE"]) > 1:
+                self.distributed_type = DistributedType.MULTI_CPU
+                if is_ccl_available() and env2int(["CCL_WORKER_COUNT"]) > 0:
+                    backend = "ccl"
+                elif torch.distributed.is_mpi_available():
+                    backend = "mpi"
+                else:
+                    backend = "gloo"
+                rank = env2int(["RANK", "PMI_RANK", "OMPI_COMM_WORLD_RANK", "MV2_COMM_WORLD_RANK"], 0)
+                size = env2int(["WORLD_SIZE", "PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE"], 1)
+                local_rank = env2int(["LOCAL_RANK", "MPI_LOCALRANKID", "OMPI_COMM_WORLD_LOCAL_RANK", "MV2_COMM_WORLD_LOCAL_RANK"], 0)
+                local_size = env2int(["MPI_LOCALNRANKS", "OMPI_COMM_WORLD_LOCAL_SIZE", "MV2_COMM_WORLD_LOCAL_SIZE"], 1)
+                self.local_process_index = local_rank
+                os.environ["RANK"] = str(rank)
+                os.environ["WORLD_SIZE"] = str(size)
+                if not os.environ.get("MASTER_PORT", None): os.environ["MASTER_PORT"] = "29500"
+                if not os.environ.get("MASTER_ADDR", None):
+                    if local_size != size and backend != "mpi":
+                        print("Warning: Looks like distributed multinode run but MASTER_ADDR env not set, using '127.0.0.1' as default")
+                        print("If this run hangs, try exporting rank 0's hostname as MASTER_ADDR")
Made the change.
@@ -123,12 +146,43 @@ def __init__(self, fp16: bool = None, cpu: bool = False, _from_accelerator: bool
                 self.distributed_type = DistributedType.MULTI_GPU
                 if not torch.distributed.is_initialized():
                     torch.distributed.init_process_group(backend="nccl")
+                    self.backend = "nccl"
                 self.num_processes = torch.distributed.get_world_size()
                 self.process_index = torch.distributed.get_rank()
                 self.local_process_index = int(os.environ.get("LOCAL_RANK", -1))
                 self.device = torch.device("cuda", self.local_process_index)
                 torch.cuda.set_device(self.device)
                 self.use_fp16 = parse_flag_from_env("USE_FP16", False) if fp16 is None else fp16
+            elif env2int(["PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE", "WORLD_SIZE"]) > 1:
+                self.distributed_type = DistributedType.MULTI_CPU
+                if is_ccl_available() and env2int(["CCL_WORKER_COUNT"]) > 0:
+                    backend = "ccl"
+                elif torch.distributed.is_mpi_available():
+                    backend = "mpi"
+                else:
+                    backend = "gloo"
+                rank = env2int(["RANK", "PMI_RANK", "OMPI_COMM_WORLD_RANK", "MV2_COMM_WORLD_RANK"], 0)
+                size = env2int(["WORLD_SIZE", "PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE"], 1)
+                local_rank = env2int(["LOCAL_RANK", "MPI_LOCALRANKID", "OMPI_COMM_WORLD_LOCAL_RANK", "MV2_COMM_WORLD_LOCAL_RANK"], 0)
+                local_size = env2int(["MPI_LOCALNRANKS", "OMPI_COMM_WORLD_LOCAL_SIZE", "MV2_COMM_WORLD_LOCAL_SIZE"], 1)
+                self.local_process_index = local_rank
+                os.environ["RANK"] = str(rank)
+                os.environ["WORLD_SIZE"] = str(size)
+                if not os.environ.get("MASTER_PORT", None): os.environ["MASTER_PORT"] = "29500"
+                if not os.environ.get("MASTER_ADDR", None):
+                    if local_size != size and backend != "mpi":
+                        print("Warning: Looks like distributed multinode run but MASTER_ADDR env not set, using '127.0.0.1' as default")
+                        print("If this run hangs, try exporting rank 0's hostname as MASTER_ADDR")
Does this require an install of some library? Let's add how to!
@@ -151,6 +151,15 @@ For instance, here is how you would run the GLUE example on the MRPC task (from
 accelerate launch examples/nlp_example.py
 ```
 
+## Launching multi-CPU run using MPI
+
+🤗 Here is another way to launch multi-CPU run using MPI
+On your cluster just run:
+
+```bash
+mpirun -np 2 python examples/nlp_example.py
```suggestion
    """Returns the first positive env value found in the `env_keys` list or the default."""
```
@@ -27,6 +35,19 @@
     _tpu_available = False
 
 
+def get_int_from_env(env_keys, default):
+    """Returns the first positive env found in the `env_keys` list or the default."""
There are multiple options e.g. [Intel MPI](https://software.intel.com/content/www/us/en/develop/tools/oneapi/components/mpi-library.html#gs.0hn0hn), [OpenMPI](https://www.open-mpi.org/) or MVAPICH. I assume, typical HPC clusters would have some version available or users can install them in userspace easily. example instructions for Open-MPI are [here](https://www.open-mpi.org/faq/?category=building#easy-build).
@@ -151,6 +151,15 @@ For instance, here is how you would run the GLUE example on the MRPC task (from
 accelerate launch examples/nlp_example.py
 ```
 
+## Launching multi-CPU run using MPI
+
+🤗 Here is another way to launch multi-CPU run using MPI
+On your cluster just run:
+
+```bash
+mpirun -np 2 python examples/nlp_example.py
Ok, added a line on how to get OpenMPI...
@@ -151,6 +151,15 @@ For instance, here is how you would run the GLUE example on the MRPC task (from
 accelerate launch examples/nlp_example.py
 ```
 
+## Launching multi-CPU run using MPI
+
+🤗 Here is another way to launch multi-CPU run using MPI
+On your cluster just run:
+
+```bash
+mpirun -np 2 python examples/nlp_example.py
```suggestion
    import torch_ccl
```
I think this should take care of the problem
@@ -19,6 +19,14 @@
 import torch
 
 
+try:
+    import torch_ccl as ccl
I have made the suggested changes. Please take a look.
Thanks.
Last step, could you run `make style` on your branch? This should take care of the failing test.
Let me know if run into any trouble and I can push it on the PR instead.
Sorry, my system doesn't have black. Make gave error.

```
black tests src examples
make: black: Command not found
make: *** [style] Error 127
```
Can you please go ahead and push it to PR?
Thanks.
Just installed black and fixed formatting... It is still giving an error about torch_ccl import but I don't know how to fix it...
Fixed issue with flake8 and squashed to single commit... Should be good to merge now.
Thanks again for all your work on this!
https://api.github.com/repos/huggingface/accelerate/pulls/63
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index 7a8422a05..44909ab3c 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -364,7 +364,8 @@ def prepare_data_loader(
         )
 
     new_dataset = dataloader.dataset
-    new_batch_sampler = dataloader.batch_sampler
+    # Iterable dataset doesn't like batch_sampler, but data_loader creates a default one for it
+    new_batch_sampler = dataloader.batch_sampler if not isinstance(new_dataset, IterableDataset) else None
     generator = getattr(dataloader, "generator", None)
     # No change if no multiprocess
     if num_processes != 1:
@@ -413,6 +414,11 @@ def prepare_data_loader(
         for k in _PYTORCH_DATALOADER_KWARGS
         if k not in ignore_kwargs
     }
+
+    # Need to provide batch_size as batch_sampler is None for Iterable dataset
+    if new_batch_sampler is None:
+        kwargs["batch_size"] = dataloader.batch_size // num_processes if split_batches else dataloader.batch_size
+
     return DataLoaderShard(
         new_dataset,
         device=device if put_on_device else None,

```suggestion
    # Iterable dataset doesn't like batch_sampler, but data_loader creates a default one for it
    new_batch_sampler = dataloader.batch_sampler if not isinstance(new_dataset, IterableDataset) else None
```
Let's group the two lines together.
@@ -365,6 +365,9 @@ def prepare_data_loader(
 
     new_dataset = dataloader.dataset
     new_batch_sampler = dataloader.batch_sampler
+    if isinstance(new_dataset, IterableDataset):
+        # Iterable dataset doesn't like batch_sampler, data_loader creates a default one for it
+        new_batch_sampler = None
```suggestion
    # Need to provide batch_size as batch_sampler is None for Iterable dataset
    if new_batch_sampler is None:
        kwargs["batch_size"] = dataloader.batch_size // num_processes if split_batches else dataloader.batch_size
```
@@ -413,6 +416,10 @@ def prepare_data_loader(
         for k in _PYTORCH_DATALOADER_KWARGS
         if k not in ignore_kwargs
     }
+
+    # Need to provide batch_size as batch_sampler is None for Iterable dataset
+    if new_batch_sampler is None: kwargs["batch_size"] = dataloader.batch_size // num_processes if split_batches else dataloader.batch_size
https://api.github.com/repos/huggingface/accelerate/pulls/62
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index f1242d3fd..44cff4d55 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -137,6 +137,7 @@ def multi_gpu_launcher(args):
         )
     else:
         cmd.extend(["--nproc_per_node", str(args.num_processes)])
+        cmd.extend(["--master_port", str(args.main_process_port)])
     cmd.append(args.training_script)
     cmd.extend(args.training_script_args)
 

https://api.github.com/repos/huggingface/accelerate/pulls/58
diff --git a/README.md b/README.md
index 280473ee5..6e2211143 100644
--- a/README.md
+++ b/README.md
@@ -163,8 +163,6 @@ notebook_launcher(training_function)
 
 An example can be found in [this notebook](https://github.com/huggingface/notebooks/blob/master/examples/accelerate/simple_nlp_example.ipynb). [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/notebooks/blob/master/examples/accelerate/simple_nlp_example.ipynb)
 
-Note that this launcher does not work in Jupyter Notebook on a machine with multiple GPUs (yet). This feature will be released in a later version.
-
 ## Why should I use 🤗 Accelerate?
 
 You should use 🤗 Accelerate when you want to easily run your training scripts in a distributed environment without having to renounce full control over your training loop. This is not a high-level framework above PyTorch, just a thin wrapper so you don't have to learn a new library, In fact the whole API of 🤗 Accelerate is in one class, the `Accelerator` object.
diff --git a/docs/source/launcher.rst b/docs/source/launcher.rst
index da28ad920..1c2279946 100644
--- a/docs/source/launcher.rst
+++ b/docs/source/launcher.rst
@@ -15,16 +15,15 @@ Notebook Launcher
 =======================================================================================================================
 
 Launch your training function inside a notebook. Currently supports launching a training with TPUs on [Google
-Colab](https://colab.research.google.com/) and [Kaggle kernels](https://www.kaggle.com/code), or training on one GPU,
-but support for training on several GPUs (if the machine on which you are running your notebook has them) is planned
-for a future release.
+Colab](https://colab.research.google.com/) and [Kaggle kernels](https://www.kaggle.com/code), as well as training on
+several GPUs (if the machine on which you are running your notebook has them).
 
 An example can be found in `this notebook
 <https://github.com/huggingface/notebooks/blob/master/examples/accelerate/simple_nlp_example.ipynb>`__.
 
 .. warning::
 
-    If you are training on Colab or a Kaggle kernel with TPUs, your :obj:`Accelerator` object should only be defined
-    inside the training function. This is because the initialization should be done inside the launcher only.
+    Your :obj:`Accelerator` object should only be defined inside the training function. This is because the
+    initialization should be done inside the launcher only.
 
 .. autofunction:: accelerate.notebook_launcher
diff --git a/docs/source/quicktour.rst b/docs/source/quicktour.rst
index d9846b777..1e28fed85 100644
--- a/docs/source/quicktour.rst
+++ b/docs/source/quicktour.rst
@@ -199,9 +199,8 @@ Launching training from a notebook
 -----------------------------------------------------------------------------------------------------------------------
 
 In Accelerate 0.3.0, a new :class:`~accelerate.notebook_launcher` has been introduced to help you launch your training
-function from a notebook. Currently supports launching a training with TPUs on Colab, or training on one GPU, but
-support for training on several GPUs (if the machine on which you are running your notebook has them) is planned for a
-future release.
+function from a notebook. This launcher supports launching a training with TPUs on Colab or Kaggle, as well as training
+on several GPUs (if the machine on which you are running your notebook has them).
 
 Just define a function responsible for your whole training and/or evaluation in a cell of the notebook, then execute a
 cell with the following code:
@@ -214,8 +213,8 @@ cell with the following code:
 
 .. warning::
 
-    If you are training on Colab with TPUs, your :obj:`Accelerator` object should only be defined inside the training
-    function. This is because the initialization should be done inside the launcher only.
+    Your :obj:`Accelerator` object should only be defined inside the training function. This is because the
+    initialization should be done inside the launcher only.
 
 
 Training on TPU
diff --git a/src/accelerate/notebook_launcher.py b/src/accelerate/notebook_launcher.py
index e9dce7677..57f253267 100644
--- a/src/accelerate/notebook_launcher.py
+++ b/src/accelerate/notebook_launcher.py
@@ -14,15 +14,15 @@
 
 import os
 import sys
-import warnings
 
 import torch
+from torch.multiprocessing import start_processes
 
 from .state import AcceleratorState
 from .utils import PrepareForLaunch
 
 
-def notebook_launcher(function, args=(), num_processes=None, start_method="fork", **kwargs):
+def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, use_port="29500"):
     """
     Launches a training function, using several processes if it's possible in the current environment (TPU with
     multiple cores for instance).
@@ -32,17 +32,15 @@ def notebook_launcher(function, args=(), num_processes=None, start_method="fork"
             The training function to execute. If it accepts arguments, the first argument should be the index of the
             process run.
         args (:obj:`Tuple`):
-            Tuple of arguments to pass to the function (it will receive :obj:`(index, *args)`).
+            Tuple of arguments to pass to the function (it will receive :obj:`*args`).
         num_processes (:obj:`int`, `optional`):
             The number of processes to use for training. Will default to 8 in Colab/Kaggle if a TPU is available, to
             the number of GPUs available otherwise.
-
-    .. warning::
-
-        Multiple GPUs is not yet supported.
+        use_fp16 (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            If :obj:`True`, will use mixed precision training on multi-GPU.
+        use_port (:obj:`str`, `optional`, defaults to :obj:`"29500"`):
+            The port to use to communicate between processes when launching a multi-GPU training.
     """
-    launcher = PrepareForLaunch(function)
-
     # Are we in a google colab or a Kaggle Kernel?
     if "IPython" in sys.modules:
         in_colab_or_kaggle = "google.colab" in str(sys.modules["IPython"].get_ipython())
@@ -65,15 +63,60 @@ def notebook_launcher(function, args=(), num_processes=None, start_method="fork"
             if num_processes is None:
                 num_processes = 8
 
+            launcher = PrepareForLaunch(function, distributed_type="TPU")
+            print(f"Launching a training on {num_processes} TPU cores.")
             xmp.spawn(launcher, args=args, nprocs=num_processes, start_method="fork")
         else:
             # No need for a distributed launch otherwise as it's either CPU or one GPU.
-            launcher(0, *args)
+            if torch.cuda.is_available():
+                print("Launching training on one GPU.")
+            else:
+                print("Launching training on CPU.")
+            function(*args)
 
     else:
         if num_processes is None:
-            num_processes = torch.cuda.device_count() if torch.cuda.is_available() else 1
+            raise ValueError(
+                "You have to specify the number of GPUs you would like to use, add `num_process=...` to your call."
+            )
 
         if num_processes > 1:
-            warnings.warn("`notebook_launcher` does not support multiple GPUs yet, launching the training on one GPU.")
-        launcher(0, *args)
+            # Multi-GPU launch
+            if len(AcceleratorState._shared_state) > 0:
+                raise ValueError(
+                    "To launch a multi-GPU training from your notebook, the `Accelerator` should only be initialized "
+                    "inside your training function. Restart your notebook and make sure no cells initializes an "
+                    "`Accelerator`."
+                )
+
+            if torch.cuda.is_initialized():
+                raise ValueError(
+                    "To launch a multi-GPU training from your notebook, you need to avoid running any instruction "
+                    "using `torch.cuda` in any cell. Restart your notebook and make sure no cells use any CUDA "
+                    "function."
+                )
+
+            # torch.distributed will expect a few environment variable to be here. We set the ones common to each
+            # process here (the other ones will be set be the launcher).
+            os.environ["WORLD_SIZE"] = str(num_processes)
+            os.environ["MASTER_ADDR"] = "127.0.0.1"
+            os.environ["MASTER_PORT"] = str(use_port)
+            os.environ["USE_FP16"] = str(use_fp16)
+
+            launcher = PrepareForLaunch(function, distributed_type="MULTI_GPU")
+            try:
+                print(f"Launching a training on {num_processes} GPUs.")
+                start_processes(launcher, nprocs=num_processes, start_method="fork")
+            finally:
+                # Clean up the environment variables set.
+                del os.environ["WORLD_SIZE"]
+                del os.environ["MASTER_ADDR"]
+                del os.environ["MASTER_PORT"]
+
+        else:
+            # No need for a distributed launch otherwise as it's either CPU or one GPU.
+            if torch.cuda.is_available():
+                print("Launching training on one GPU.")
+            else:
+                print("Launching training on CPU.")
+            function(*args)
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 0c2282b2c..5dfe85e54 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import importlib
-import inspect
+import os
 import random
 from enum import Enum
 from typing import List, Optional, Union
@@ -252,15 +252,23 @@ def save(obj, f):
 
 class PrepareForLaunch:
     """
-    Prepare a function that launches a script.
+    Prepare a function that will launched in a distributed setup.
+
+    Args:
+        launcher (:obj:`Callable`):
+            The function to launch.
+        distributed_type (:class:`~accelerate.state.DistributedType`):
+            The distributed type to prepare for.
     """
 
-    def __init__(self, launcher):
+    def __init__(self, launcher, distributed_type="NO"):
         self.launcher = launcher
+        self.distributed_type = DistributedType(distributed_type)
+
+    def __call__(self, index, *args):
+        if self.distributed_type == DistributedType.MULTI_GPU:
+            # Prepare the environment for torch.distributed
+            os.environ["LOCAL_RANK"] = str(index)
+            os.environ["RANK"] = str(index)
 
-    def __call__(self, index):
-        launcher_sig = inspect.signature(self.launcher)
-        if len(launcher_sig.parameters) == 0:
-            self.launcher()
-        else:
-            self.launcher(index)
+        self.launcher(*args)

I am not that experienced with `torch.distributed`, but is it good to have them hardcoded and not overrideable? We could add parameters to adjust them. 

@@ -65,15 +61,60 @@ def notebook_launcher(function, args=(), num_processes=None, start_method="fork"
             if num_processes is None:
                 num_processes = 8
 
+            launcher = PrepareForLaunch(function, distributed_type="TPU")
+            print(f"Launching a training on {num_processes} TPU cores.")
             xmp.spawn(launcher, args=args, nprocs=num_processes, start_method="fork")
         else:
             # No need for a distributed launch otherwise as it's either CPU or one GPU.
-            launcher(0, *args)
+            if torch.cuda.is_available():
+                print("Launching training on one GPU.")
+            else:
+                print("Launching training on CPU.")
+            function(*args)
 
     else:
         if num_processes is None:
-            num_processes = torch.cuda.device_count() if torch.cuda.is_available() else 1
+            raise ValueError(
+                "You have to specify the number of GPUs you would like to use, add `num_process=...` to your call."
+            )
 
         if num_processes > 1:
-            warnings.warn("`notebook_launcher` does not support multiple GPUs yet, launching the training on one GPU.")
-        launcher(0, *args)
+            # Multi-GPU launch
+            if len(AcceleratorState._shared_state) > 0:
+                raise ValueError(
+                    "To launch a multi-GPU training from your notebook, the `Accelerator` should only be initialized "
+                    "inside your training function. Restart your notebook and make sure no cells initializes an "
+                    "`Accelerator`."
+                )
+
+            if torch.cuda.is_initialized():
+                raise ValueError(
+                    "To launch a multi-GPU training from your notebook, you need to avoid running any instruction "
+                    "using `torch.cuda` in any cell. Restart your notebook and make sure no cells use any CUDA "
+                    "function."
+                )
+
+            # torch.distributed will expect a few environment variable to be here. We set the ones common to each
+            # process here (the other ones will be set be the launcher).
+            os.environ["WORLD_SIZE"] = str(num_processes)
+            os.environ["MASTER_ADDR"] = "127.0.0.1"
+            os.environ["MASTER_PORT"] = "29500"
Those are the default values in `torch.distributed.launch`. The master address should not change since this does not support multi-node training. I can add a kwarg for the master port.
@@ -65,15 +61,60 @@ def notebook_launcher(function, args=(), num_processes=None, start_method="fork"
             if num_processes is None:
                 num_processes = 8
 
+            launcher = PrepareForLaunch(function, distributed_type="TPU")
+            print(f"Launching a training on {num_processes} TPU cores.")
             xmp.spawn(launcher, args=args, nprocs=num_processes, start_method="fork")
         else:
             # No need for a distributed launch otherwise as it's either CPU or one GPU.
-            launcher(0, *args)
+            if torch.cuda.is_available():
+                print("Launching training on one GPU.")
+            else:
+                print("Launching training on CPU.")
+            function(*args)
 
     else:
         if num_processes is None:
-            num_processes = torch.cuda.device_count() if torch.cuda.is_available() else 1
+            raise ValueError(
+                "You have to specify the number of GPUs you would like to use, add `num_process=...` to your call."
+            )
 
         if num_processes > 1:
-            warnings.warn("`notebook_launcher` does not support multiple GPUs yet, launching the training on one GPU.")
-        launcher(0, *args)
+            # Multi-GPU launch
+            if len(AcceleratorState._shared_state) > 0:
+                raise ValueError(
+                    "To launch a multi-GPU training from your notebook, the `Accelerator` should only be initialized "
+                    "inside your training function. Restart your notebook and make sure no cells initializes an "
+                    "`Accelerator`."
+                )
+
+            if torch.cuda.is_initialized():
+                raise ValueError(
+                    "To launch a multi-GPU training from your notebook, you need to avoid running any instruction "
+                    "using `torch.cuda` in any cell. Restart your notebook and make sure no cells use any CUDA "
+                    "function."
+                )
+
+            # torch.distributed will expect a few environment variable to be here. We set the ones common to each
+            # process here (the other ones will be set be the launcher).
+            os.environ["WORLD_SIZE"] = str(num_processes)
+            os.environ["MASTER_ADDR"] = "127.0.0.1"
+            os.environ["MASTER_PORT"] = "29500"
Yeah saw the default [here](https://github.com/pytorch/pytorch/blob/master/torch/distributed/elastic_launch.py#L394). But who knows what setup customers are using, e.g. IPV6, custom cnames etc. 
@@ -65,15 +61,60 @@ def notebook_launcher(function, args=(), num_processes=None, start_method="fork"
             if num_processes is None:
                 num_processes = 8
 
+            launcher = PrepareForLaunch(function, distributed_type="TPU")
+            print(f"Launching a training on {num_processes} TPU cores.")
             xmp.spawn(launcher, args=args, nprocs=num_processes, start_method="fork")
         else:
             # No need for a distributed launch otherwise as it's either CPU or one GPU.
-            launcher(0, *args)
+            if torch.cuda.is_available():
+                print("Launching training on one GPU.")
+            else:
+                print("Launching training on CPU.")
+            function(*args)
 
     else:
         if num_processes is None:
-            num_processes = torch.cuda.device_count() if torch.cuda.is_available() else 1
+            raise ValueError(
+                "You have to specify the number of GPUs you would like to use, add `num_process=...` to your call."
+            )
 
         if num_processes > 1:
-            warnings.warn("`notebook_launcher` does not support multiple GPUs yet, launching the training on one GPU.")
-        launcher(0, *args)
+            # Multi-GPU launch
+            if len(AcceleratorState._shared_state) > 0:
+                raise ValueError(
+                    "To launch a multi-GPU training from your notebook, the `Accelerator` should only be initialized "
+                    "inside your training function. Restart your notebook and make sure no cells initializes an "
+                    "`Accelerator`."
+                )
+
+            if torch.cuda.is_initialized():
+                raise ValueError(
+                    "To launch a multi-GPU training from your notebook, you need to avoid running any instruction "
+                    "using `torch.cuda` in any cell. Restart your notebook and make sure no cells use any CUDA "
+                    "function."
+                )
+
+            # torch.distributed will expect a few environment variable to be here. We set the ones common to each
+            # process here (the other ones will be set be the launcher).
+            os.environ["WORLD_SIZE"] = str(num_processes)
+            os.environ["MASTER_ADDR"] = "127.0.0.1"
+            os.environ["MASTER_PORT"] = "29500"
```suggestion
            os.environ["MASTER_PORT"] = str(use_port)
```
Would add the extra conversion, might be else error prone
@@ -65,15 +63,60 @@ def notebook_launcher(function, args=(), num_processes=None, start_method="fork"
             if num_processes is None:
                 num_processes = 8
 
+            launcher = PrepareForLaunch(function, distributed_type="TPU")
+            print(f"Launching a training on {num_processes} TPU cores.")
             xmp.spawn(launcher, args=args, nprocs=num_processes, start_method="fork")
         else:
             # No need for a distributed launch otherwise as it's either CPU or one GPU.
-            launcher(0, *args)
+            if torch.cuda.is_available():
+                print("Launching training on one GPU.")
+            else:
+                print("Launching training on CPU.")
+            function(*args)
 
     else:
         if num_processes is None:
-            num_processes = torch.cuda.device_count() if torch.cuda.is_available() else 1
+            raise ValueError(
+                "You have to specify the number of GPUs you would like to use, add `num_process=...` to your call."
+            )
 
         if num_processes > 1:
-            warnings.warn("`notebook_launcher` does not support multiple GPUs yet, launching the training on one GPU.")
-        launcher(0, *args)
+            # Multi-GPU launch
+            if len(AcceleratorState._shared_state) > 0:
+                raise ValueError(
+                    "To launch a multi-GPU training from your notebook, the `Accelerator` should only be initialized "
+                    "inside your training function. Restart your notebook and make sure no cells initializes an "
+                    "`Accelerator`."
+                )
+
+            if torch.cuda.is_initialized():
+                raise ValueError(
+                    "To launch a multi-GPU training from your notebook, you need to avoid running any instruction "
+                    "using `torch.cuda` in any cell. Restart your notebook and make sure no cells use any CUDA "
+                    "function."
+                )
+
+            # torch.distributed will expect a few environment variable to be here. We set the ones common to each
+            # process here (the other ones will be set be the launcher).
+            os.environ["WORLD_SIZE"] = str(num_processes)
+            os.environ["MASTER_ADDR"] = "127.0.0.1"
+            os.environ["MASTER_PORT"] = use_port
https://api.github.com/repos/huggingface/accelerate/pulls/56
diff --git a/src/accelerate/commands/test.py b/src/accelerate/commands/test.py
index 1940c5039..ec6d2524b 100644
--- a/src/accelerate/commands/test.py
+++ b/src/accelerate/commands/test.py
@@ -46,7 +46,7 @@ def test_command(args):
     script_name = os.path.sep.join(__file__.split(os.path.sep)[:-2] + ["test_utils", "test_script.py"])
 
     test_args = f"""
-        {script_name} --config_file={args.config_file}
+        --config_file={args.config_file} {script_name}
     """.split()
     cmd = ["accelerate-launch"] + test_args
     result = execute_subprocess_async(cmd, env=os.environ.copy())

https://api.github.com/repos/huggingface/accelerate/pulls/54
diff --git a/README.md b/README.md
index 25bc92c13..280473ee5 100644
--- a/README.md
+++ b/README.md
@@ -151,6 +151,20 @@ For instance, here is how you would run the GLUE example on the MRPC task (from
 accelerate launch examples/nlp_example.py
 ```
 
+## Launching your training from a notebook
+
+🤗 Accelerate also provides a `notebook_launcher` function you can use in a notebook to launch a distributed training. This is especially useful for Colab or Kaggle notebooks with a TPU backend. Just define your training loop in a `training_function` then in your last cell, add:
+
+```python
+from accelerate import notebook_launcher
+
+notebook_launcher(training_function)
+```
+
+An example can be found in [this notebook](https://github.com/huggingface/notebooks/blob/master/examples/accelerate/simple_nlp_example.ipynb). [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/notebooks/blob/master/examples/accelerate/simple_nlp_example.ipynb)
+
+Note that this launcher does not work in Jupyter Notebook on a machine with multiple GPUs (yet). This feature will be released in a later version.
+
 ## Why should I use 🤗 Accelerate?
 
 You should use 🤗 Accelerate when you want to easily run your training scripts in a distributed environment without having to renounce full control over your training loop. This is not a high-level framework above PyTorch, just a thin wrapper so you don't have to learn a new library, In fact the whole API of 🤗 Accelerate is in one class, the `Accelerator` object.
diff --git a/docs/source/launcher.rst b/docs/source/launcher.rst
index 1b0520d1b..da28ad920 100644
--- a/docs/source/launcher.rst
+++ b/docs/source/launcher.rst
@@ -19,6 +19,9 @@ Colab](https://colab.research.google.com/) and [Kaggle kernels](https://www.kagg
 but support for training on several GPUs (if the machine on which you are running your notebook has them) is planned
 for a future release.
 
+An example can be found in `this notebook
+<https://github.com/huggingface/notebooks/blob/master/examples/accelerate/simple_nlp_example.ipynb>`__.
+
 .. warning::
 
     If you are training on Colab or a Kaggle kernel with TPUs, your :obj:`Accelerator` object should only be defined

https://api.github.com/repos/huggingface/accelerate/pulls/52
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index 56ba51f64..c8bf1bcd7 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -46,7 +46,7 @@ def get_cluster_input():
             main_process_ip = _ask_field(
                 "What is the IP address of the machine that will host the main process? ",
             )
-            main_process_ip = _ask_field(
+            main_process_port = _ask_field(
                 "What is the port you will use to communicate with the main process? ",
                 lambda x: int(x),
             )

https://api.github.com/repos/huggingface/accelerate/pulls/50
diff --git a/docs/source/index.rst b/docs/source/index.rst
index ceef4fc91..fb58d1e04 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -155,5 +155,6 @@ Supported integrations
     :caption: API reference
 
     accelerator
+    launcher
     kwargs
     internal
diff --git a/docs/source/launcher.rst b/docs/source/launcher.rst
new file mode 100644
index 000000000..1b0520d1b
--- /dev/null
+++ b/docs/source/launcher.rst
@@ -0,0 +1,27 @@
+.. 
+    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.
+
+
+Notebook Launcher
+=======================================================================================================================
+
+Launch your training function inside a notebook. Currently supports launching a training with TPUs on [Google
+Colab](https://colab.research.google.com/) and [Kaggle kernels](https://www.kaggle.com/code), or training on one GPU,
+but support for training on several GPUs (if the machine on which you are running your notebook has them) is planned
+for a future release.
+
+.. warning::
+
+    If you are training on Colab or a Kaggle kernel with TPUs, your :obj:`Accelerator` object should only be defined
+    inside the training function. This is because the initialization should be done inside the launcher only.
+
+.. autofunction:: accelerate.notebook_launcher
diff --git a/docs/source/quicktour.rst b/docs/source/quicktour.rst
index 5e3d9ac17..d9846b777 100644
--- a/docs/source/quicktour.rst
+++ b/docs/source/quicktour.rst
@@ -195,6 +195,29 @@ If you stored the config file in a non-default location, you can indicate it to
 You can also override any of the arguments determined by your config file, see TODO: insert ref here.
 
 
+Launching training from a notebook
+-----------------------------------------------------------------------------------------------------------------------
+
+In Accelerate 0.3.0, a new :class:`~accelerate.notebook_launcher` has been introduced to help you launch your training
+function from a notebook. Currently supports launching a training with TPUs on Colab, or training on one GPU, but
+support for training on several GPUs (if the machine on which you are running your notebook has them) is planned for a
+future release.
+
+Just define a function responsible for your whole training and/or evaluation in a cell of the notebook, then execute a
+cell with the following code:
+
+.. code-block::
+
+    from accelerate import notebook_launcher
+
+    notebook_launcher(training_function)
+
+.. warning::
+
+    If you are training on Colab with TPUs, your :obj:`Accelerator` object should only be defined inside the training
+    function. This is because the initialization should be done inside the launcher only.
+
+
 Training on TPU
 -----------------------------------------------------------------------------------------------------------------------
 
diff --git a/src/accelerate/__init__.py b/src/accelerate/__init__.py
index d005e1235..8edc050c3 100644
--- a/src/accelerate/__init__.py
+++ b/src/accelerate/__init__.py
@@ -6,5 +6,6 @@
 
 from .accelerator import Accelerator
 from .kwargs_handlers import DistributedDataParallelKwargs, GradScalerKwargs
+from .notebook_launcher import notebook_launcher
 from .state import DistributedType
 from .utils import synchronize_rng_states
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 9e3e7a566..f1242d3fd 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -16,7 +16,6 @@
 
 import argparse
 import importlib
-import inspect
 import os
 import subprocess
 import sys
@@ -27,15 +26,7 @@
 from accelerate.commands.config import default_config_file, load_config_from_file
 from accelerate.commands.config.config_args import SageMakerConfig
 from accelerate.state import ComputeEnvironment, DistributedType
-from accelerate.utils import is_sagemaker_available
-
-
-class _AddOneArg:
-    def __init__(self, launcher):
-        self.launcher = launcher
-
-    def __call__(self, index):
-        self.launcher()
+from accelerate.utils import PrepareForLaunch, is_sagemaker_available
 
 
 def launch_command_parser(subparsers=None):
@@ -176,12 +167,7 @@ def tpu_launcher(args):
     # Patch sys.argv
     sys.argv = [args.training_script] + args.training_script_args
 
-    # If the function does not take one argument, launch will fail
-    launcher_sig = inspect.signature(main_function)
-    if len(launcher_sig.parameters) == 0:
-        xmp.spawn(_AddOneArg(main_function), args=(), nprocs=args.num_processes)
-    else:
-        xmp.spawn(main_function, args=(), nprocs=args.num_processes)
+    xmp.spawn(PrepareForLaunch(main_function), args=(), nprocs=args.num_processes)
 
 
 def _convert_nargs_to_dict(nargs: List[str]) -> Dict[str, str]:
diff --git a/src/accelerate/notebook_launcher.py b/src/accelerate/notebook_launcher.py
new file mode 100644
index 000000000..e9dce7677
--- /dev/null
+++ b/src/accelerate/notebook_launcher.py
@@ -0,0 +1,79 @@
+# 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 sys
+import warnings
+
+import torch
+
+from .state import AcceleratorState
+from .utils import PrepareForLaunch
+
+
+def notebook_launcher(function, args=(), num_processes=None, start_method="fork", **kwargs):
+    """
+    Launches a training function, using several processes if it's possible in the current environment (TPU with
+    multiple cores for instance).
+
+    Args:
+        function (:obj:`Callable`):
+            The training function to execute. If it accepts arguments, the first argument should be the index of the
+            process run.
+        args (:obj:`Tuple`):
+            Tuple of arguments to pass to the function (it will receive :obj:`(index, *args)`).
+        num_processes (:obj:`int`, `optional`):
+            The number of processes to use for training. Will default to 8 in Colab/Kaggle if a TPU is available, to
+            the number of GPUs available otherwise.
+
+    .. warning::
+
+        Multiple GPUs is not yet supported.
+    """
+    launcher = PrepareForLaunch(function)
+
+    # Are we in a google colab or a Kaggle Kernel?
+    if "IPython" in sys.modules:
+        in_colab_or_kaggle = "google.colab" in str(sys.modules["IPython"].get_ipython())
+    elif any(key.startswith("KAGGLE") for key in os.environ.keys()):
+        in_colab_or_kaggle = True
+    else:
+        in_colab_or_kaggle = False
+
+    if in_colab_or_kaggle:
+        if os.environ.get("TPU_NAME", None) is not None:
+            # TPU launch
+            import torch_xla.distributed.xla_multiprocessing as xmp
+
+            if len(AcceleratorState._shared_state) > 0:
+                raise ValueError(
+                    "To train on TPU in Colab or Kaggle Kernel, the `Accelerator` should only be initialized inside "
+                    "your training function. Restart your notebook and make sure no cells initializes an "
+                    "`Accelerator`."
+                )
+            if num_processes is None:
+                num_processes = 8
+
+            xmp.spawn(launcher, args=args, nprocs=num_processes, start_method="fork")
+        else:
+            # No need for a distributed launch otherwise as it's either CPU or one GPU.
+            launcher(0, *args)
+
+    else:
+        if num_processes is None:
+            num_processes = torch.cuda.device_count() if torch.cuda.is_available() else 1
+
+        if num_processes > 1:
+            warnings.warn("`notebook_launcher` does not support multiple GPUs yet, launching the training on one GPU.")
+        launcher(0, *args)
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 0ac7d04e5..0c2282b2c 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import importlib
+import inspect
 import random
 from enum import Enum
 from typing import List, Optional, Union
@@ -247,3 +248,19 @@ def save(obj, f):
         xm.save(obj, f)
     elif AcceleratorState().local_process_index == 0:
         torch.save(obj, f)
+
+
+class PrepareForLaunch:
+    """
+    Prepare a function that launches a script.
+    """
+
+    def __init__(self, launcher):
+        self.launcher = launcher
+
+    def __call__(self, index):
+        launcher_sig = inspect.signature(self.launcher)
+        if len(launcher_sig.parameters) == 0:
+            self.launcher()
+        else:
+            self.launcher(index)

```suggestion
Launch your training function inside a notebook. Currently supports launching a training with TPUs on [Google Colab](https://colab.research.google.com/), or training
```
```suggestion
Launch your training function from a notebook. Currently supports launching a training with TPUs on Colab, or training
```
@@ -0,0 +1,26 @@
+.. 
+    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.
+
+
+Notebook Launcher
+=======================================================================================================================
+
+Launch your training function from a notebook. Currently supports launching a training with TPUs on Colab, or training
We could add a check for Kaggle here too. I guess users expect the same behavior there. I checked a Kaggle TPU notebook, sadly `sys.modules["IPython"].get_ipython()` is not hinting that the notebook is run on Kaggle. But we could add a check like `if any(key.startswith('KAGGLE') for key in os.environ.keys())`. Additionally here is a dump of the env vars on kaggle

```python
environ{'LC_ALL': 'C.UTF-8',
        'LD_LIBRARY_PATH': '/opt/conda/lib',
        'KAGGLE_DATA_PROXY_TOKEN': 'eyJhbGciOiJBMTI4S1ciLCJlbmMiOiJBMTI4Q0JDLUhTMjU2In0.B5s-CblnZQjaC_KXAizqBneeEiZmD-pS7Asgb2h-jS6uEkmkIIHLYg.ddiwW2QpGlaxa5Q584WQOg.E8mrT0sr5vZsIb5v4zdJf6uvrzk2wuZCaM6g-xLBTkZo_Nx-lLLqcZqUrZixTOznQqeidkc5tJibiyYBeBGcRnwSuFXx0w4V0EJPfKVqp_B5TMbOROvjnLIpX3eoINPjQd7I38nmCeyrlGUnC-aTHLFqDAxaaYmcESykEixt_Hlc7Q_DfiD_uPkEzJs8buJBvcYGFq8FWbJB7n89rkDY-P9xvZnSsFjbiQYtAbfwcKGdxGKYKAoVdqRhsRHjVjj6HDVQU38Ldnuhzjg2kdt8Xhu__YxSXq0Gh-RFLZxP90M.ZPazCS7UZW9RyI3-jniqPA',
        'LANG': 'C.UTF-8',
        'HOSTNAME': '6497510c329c',
        'TESSERACT_PATH': '/usr/bin/tesseract',
        'XRT_TPU_CONFIG': 'tpu_worker;0;10.0.0.2:8470',
        'KAGGLE_KERNEL_INTEGRATIONS': '',
        'PROJ_LIB': '/opt/conda/share/proj',
        'KAGGLE_KERNEL_RUN_TYPE': 'Interactive',
        'MKL_THREADING_LAYER': 'GNU',
        'NO_GCE_CHECK': 'True',
        'PWD': '/kaggle/working',
        'HOME': '/root',
        'DL_ANACONDA_HOME': '/opt/conda',
        'ANACONDA_PYTHON_VERSION': '3.7',
        'KAGGLE_URL_BASE': 'https://www.kaggle.com',
        'KAGGLE_USER_SECRETS_TOKEN': 'eyJhbGciOiJkaXIiLCJlbmMiOiJBMTI4Q0JDLUhTMjU2In0..kvSYrK_dnCLGYSW8GDnQdg.xlHt52wzWG2SlyvAN6zUq9M2S3m6maqeK_nGH6zCEL-VDZE-CYAC19Ecal8oy7zJQCFYcTDdFGi-Qq6c3tdXinuUy-kG1UQ8MnTSDJJT1AvIWAeQf2MNacS8TMWa7IYG.DLWrZFE6rlAGB7rKowpPKQ',
        'KAGGLE_DOCKER_IMAGE': 'gcr.io/kaggle-images/python:tpu',
        'CONTAINER_NAME': 'base-cpu',
        'KAGGLE_DATA_PROXY_PROJECT': 'kaggle-161607',
        'SHELL': '/bin/bash',
        'MPLBACKEND': 'agg',
        'SHLVL': '1',
        'PYTHONPATH': '/kaggle/lib/kagglegym:/kaggle/lib',
        'CONTAINER_URL': 'gcr.io/deeplearning-platform-release/base-cpu:nightly-2021-03-28',
        'KAGGLE_CONTAINER_NAME': 'kaggle_gq54lOn7oztYNC879IY4P2ItqdI6zTyAKVlolN6KbYE-60764729-webtier',
        'PYTHONUSERBASE': '/root/.local',
        'PATH': '/opt/conda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin',
        'KAGGLE_DATA_PROXY_URL': 'https://dp.kaggle.net',
        'TPU_NAME': 'grpc://10.0.0.2:8470',
        '_': '/opt/conda/bin/jupyter',
        'GIT_PYTHON_REFRESH': 'quiet',
        'JPY_PARENT_PID': '9',
        'TERM': 'xterm-color',
        'CLICOLOR': '1',
        'PAGER': 'cat',
        'GIT_PAGER': 'cat'}
``` 

Additionally, if we add Kaggle we should rename `in_colab` to maybe `in_colab_or_kaggle`.
@@ -0,0 +1,80 @@
+# 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 sys
+import warnings
+
+import torch
+
+from .state import AcceleratorState
+from .utils import PrepareForLaunch
+
+
+def notebook_launcher(function, args=(), num_processes=None, start_method="fork", **kwargs):
+    """
+    Launches a training function, using several processes if it's possible in the current environment (TPU with
+    multiple cores for instance).
+
+    Args:
+        function (:obj:`Callable`):
+            The training function to execute. If it accepts arguments, the first argument should be the index of the
+            process run.
+        args (:obj:`Tuple`):
+            Tuple of arguments to pass to the function (it will receive :obj:`(index, *args)`).
+        num_processes (:obj:`int`, `optional`):
+            The number of processes to use for training. Will default to 8 in Colab if a TPU is available, to the
+            number of GPUs available otherwise.
+        start_method (:obj:`str`, `optional`, defaults to :obj:`"fork"`):
+            The method to use to start the various processes
+        kwargs:
+            Additional keyword arguments passed along to the multiprocess :obj:`spawn` method.
+
+    .. warning::
+
+        Multiple GPUs is not yet supported.
+    """
+    launcher = PrepareForLaunch(function)
+
+    # Are we in a google colab?
+    if "IPython" in sys.modules:
+        in_colab = "google.colab" in str(sys.modules["IPython"].get_ipython())
+    else:
Yes, notebook examples will follow! Working on it this afternoon and hoping to have most NLP tasks by the end of the week!
https://api.github.com/repos/huggingface/accelerate/pulls/44
diff --git a/src/accelerate/optimizer.py b/src/accelerate/optimizer.py
index 55d00ab5d..63b53eccb 100644
--- a/src/accelerate/optimizer.py
+++ b/src/accelerate/optimizer.py
@@ -14,6 +14,8 @@
 
 import torch
 
+from packaging import version
+
 from .state import AcceleratorState, DistributedType, is_tpu_available
 
 
@@ -86,8 +88,18 @@ def load_state_dict(self, state_dict):
     def state_dict(self):
         return self.optimizer.state_dict()
 
-    def zero_grad(self):
-        self.optimizer.zero_grad()
+    def zero_grad(self, set_to_none=None):
+        if version.parse(torch.__version__) < version.parse("1.7.0"):
+            if set_to_none is not None:
+                raise ValueError(
+                    "`set_to_none` for Optimizer.zero_grad` was introduced in PyTorch 1.7.0 and can't be used for "
+                    f"earlier versions (found version {torch.__version__})."
+                )
+            self.optimizer.zero_grad()
+        else:
+            if set_to_none is not None:
+                set_to_none = False
+            self.optimizer.zero_grad(set_to_none=set_to_none)
 
     def step(self):
         if self.state.distributed_type == DistributedType.TPU:

if set_to_none is None:
@@ -86,8 +88,18 @@ def load_state_dict(self, state_dict):
     def state_dict(self):
         return self.optimizer.state_dict()
 
-    def zero_grad(self):
-        self.optimizer.zero_grad()
+    def zero_grad(self, set_to_none=None):
+        if version.parse(torch.__version__) < version.parse("1.7.0"):
+            if set_to_none is not None:
+                raise ValueError(
+                    "`set_to_none` for Optimizer.zero_grad` was introduced in PyTorch 1.7.0 and can't be used for "
+                    f"earlier versions (found version {torch.__version__})."
+                )
+            self.optimizer.zero_grad()
+        else:
+            if set_to_none is not None:
Good catch!
@@ -86,8 +88,18 @@ def load_state_dict(self, state_dict):
     def state_dict(self):
         return self.optimizer.state_dict()
 
-    def zero_grad(self):
-        self.optimizer.zero_grad()
+    def zero_grad(self, set_to_none=None):
+        if version.parse(torch.__version__) < version.parse("1.7.0"):
+            if set_to_none is not None:
+                raise ValueError(
+                    "`set_to_none` for Optimizer.zero_grad` was introduced in PyTorch 1.7.0 and can't be used for "
+                    f"earlier versions (found version {torch.__version__})."
+                )
+            self.optimizer.zero_grad()
+        else:
+            if set_to_none is not None:
https://api.github.com/repos/huggingface/accelerate/pulls/43
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 382636253..903f2b0e0 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -66,10 +66,10 @@ def launch_command_parser(subparsers=None):
         "--num_processes", type=int, default=None, help="The total number of processes to be launched in parallel."
     )
     parser.add_argument(
-        "--num_machines", type=int, default=1, help="The total number of machines used in this training."
+        "--num_machines", type=int, default=None, help="The total number of machines used in this training."
     )
     parser.add_argument(
-        "--machine_rank", type=int, default=0, help="The rank of the machine on which this script is launched."
+        "--machine_rank", type=int, default=None, help="The rank of the machine on which this script is launched."
     )
     parser.add_argument("--main_process_ip", type=str, default=None, help="The IP address of the machine of rank 0.")
     parser.add_argument(
@@ -298,6 +298,16 @@ def launch_command(args):
         if not args.multi_gpu and not args.tpu:
             args.multi_gpu = defaults.distributed_type == DistributedType.MULTI_GPU
             args.tpu = defaults.distributed_type == DistributedType.TPU
+        if defaults.compute_environment == ComputeEnvironment.LOCAL_MACHINE:
+            # Update args with the defaults
+            for name, attr in defaults.__dict__.items():
+                # Those args are handled separately
+                if (
+                    name not in ["compute_environment", "fp16", "distributed_type"]
+                    and getattr(args, name, None) is None
+                ):
+                    setattr(args, name, attr)
+
         if args.num_processes is None and defaults.compute_environment == ComputeEnvironment.LOCAL_MACHINE:
             args.num_processes = defaults.num_processes
         if not args.fp16:

https://api.github.com/repos/huggingface/accelerate/pulls/38
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index a9e412f7d..56ba51f64 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -58,22 +58,22 @@ def get_cluster_input():
     else:
         main_training_function = "main"
 
-        num_processes = _ask_field(
-            "How many processes in total will you use? [1]: ",
-            lambda x: int(x),
-            default=1,
-            error_message="Please enter an integer.",
-        )
+    num_processes = _ask_field(
+        "How many processes in total will you use? [1]: ",
+        lambda x: int(x),
+        default=1,
+        error_message="Please enter an integer.",
+    )
 
-        if distributed_type != DistributedType.TPU:
-            fp16 = _ask_field(
-                "Do you wish to use FP16 (mixed precision)? [yes/NO]: ",
-                _convert_yes_no_to_bool,
-                default=False,
-                error_message="Please enter yes or no.",
-            )
-        else:
-            fp16 = False
+    if distributed_type != DistributedType.TPU:
+        fp16 = _ask_field(
+            "Do you wish to use FP16 (mixed precision)? [yes/NO]: ",
+            _convert_yes_no_to_bool,
+            default=False,
+            error_message="Please enter yes or no.",
+        )
+    else:
+        fp16 = False
 
     return ClusterConfig(
         compute_environment=ComputeEnvironment.LOCAL_MACHINE,

https://api.github.com/repos/huggingface/accelerate/pulls/35
diff --git a/docs/source/quicktour.rst b/docs/source/quicktour.rst
index 1d56f9e4b..5e3d9ac17 100644
--- a/docs/source/quicktour.rst
+++ b/docs/source/quicktour.rst
@@ -199,18 +199,18 @@ Training on TPU
 -----------------------------------------------------------------------------------------------------------------------
 
 If you want to launch your script on TPUs, there are a few caveats you should be aware of. Behind the scenes, the TPUs
-will create a graph of all the operations happening im your training step (forward pass, backward pass and optimizer
+will create a graph of all the operations happening in your training step (forward pass, backward pass and optimizer
 step). This is why your first step of training will always be very long as building and compiling this graph for
 optimizations takes some time.
 
 The good news is that this compilation will be cached so the second step and all the following will be much faster. The
 bas news is that it only applies if all of your steps do exactly the same operations, which implies:
 
-- having all tensors of the same length in all your lenghts
-- having static code (i.e., not a foor loop of length that could change from step to step)
+- having all tensors of the same length in all your lengths
+- having static code (i.e., not a for loop of length that could change from step to step)
 
 Having any of the things above change between two steps will trigger a new compilation which will, once again, take a
-lof of time. In practice, that means you must take special care to have all your tensors in your inputs of the same
+lot of time. In practice, that means you must take special care to have all your tensors in your inputs of the same
 shape (so no dynamic padding for instance if you are in an NLP problem) and should not use layer with for loops that
 have different lengths depending on the inputs (such as an LSTM) or the training will be excruciatingly slow.
 
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 75539559e..0143ac949 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -280,7 +280,7 @@ def gather(self, tensor):
 
         Args:
             tensor (:obj:`torch.Tensor`, or a nested tuple/list/dictionary of :obj:`torch.Tensor`):
-                The tensors to gather accross all processes.
+                The tensors to gather across all processes.
 
         Returns:
             :obj:`torch.Tensor`, or a nested tuple/list/dictionary of :obj:`torch.Tensor`: The gathered tensor(s). Note
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index 24a7ad08a..7a8422a05 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -271,7 +271,7 @@ class DataLoaderShard(DataLoader):
             - :obj:`"xla"`: the XLA random number generator (TPU only)
             - :obj:`"generator"`: an optional :obj:`torch.Generator`
         generator (:obj:`torch.Generator`, `optional`):
-            A random number generator to keep synchronized accross processes.
+            A random number generator to keep synchronized across processes.
         kwargs:
             All other keyword arguments to pass to the regular :obj:`DataLoader` initialization.
     """

https://api.github.com/repos/huggingface/accelerate/pulls/33
diff --git a/src/accelerate/commands/config/config_args.py b/src/accelerate/commands/config/config_args.py
index a05debfda..0e115bdd4 100644
--- a/src/accelerate/commands/config/config_args.py
+++ b/src/accelerate/commands/config/config_args.py
@@ -44,7 +44,7 @@ def load_config_from_file(config_file):
         if config_file.endswith(".json"):
             if (
                 json.load(f).get("compute_environment", ComputeEnvironment.LOCAL_MACHINE)
-                is ComputeEnvironment.LOCAL_MACHINE
+                == ComputeEnvironment.LOCAL_MACHINE
             ):
                 config_class = ClusterConfig
             else:
@@ -53,7 +53,7 @@ def load_config_from_file(config_file):
         else:
             if (
                 yaml.safe_load(f).get("compute_environment", ComputeEnvironment.LOCAL_MACHINE)
-                is ComputeEnvironment.LOCAL_MACHINE
+                == ComputeEnvironment.LOCAL_MACHINE
             ):
                 config_class = ClusterConfig
             else:

https://api.github.com/repos/huggingface/accelerate/pulls/31
diff --git a/examples/README.md b/examples/README.md
index be636e63c..434a230ea 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -28,7 +28,7 @@ pip install datasets, transformers
 
 The same script can be run in any of the following configurations:
 - single CPU or single GPU
-- multi GPUS (using PyTorch distributed mode)
+- multi GPUs (using PyTorch distributed mode)
 - (multi) TPUs
 - fp16 (mixed-precision) or fp32 (normal precision)
 
@@ -57,8 +57,8 @@ To run it in each of these various modes, use the following commands:
         ```
     * from any server with Accelerate launcher
         ```bash
-        accelerate launch --fb16 ./nlp_example.py
-- multi GPUS (using PyTorch distributed mode)
+        accelerate launch --fp16 ./nlp_example.py
+- multi GPUs (using PyTorch distributed mode)
     * With Accelerate config and launcher
         ```bash
         accelerate config  # This will create a config file on your server
@@ -103,7 +103,7 @@ The [cv_example.py](./cv_example.py) script is a simple example to fine-tune a R
 
 The same script can be run in any of the following configurations:
 - single CPU or single GPU
-- multi GPUS (using PyTorch distributed mode)
+- multi GPUs (using PyTorch distributed mode)
 - (multi) TPUs
 - fp16 (mixed-precision) or fp32 (normal precision)
 
@@ -145,8 +145,8 @@ To run it in each of these various modes, use the following commands:
         ```
     * from any server with Accelerate launcher
         ```bash
-        accelerate launch --fb16 ./cv_example.py --data_dir path_to_data
-- multi GPUS (using PyTorch distributed mode)
+        accelerate launch --fp16 ./cv_example.py --data_dir path_to_data
+- multi GPUs (using PyTorch distributed mode)
     * With Accelerate config and launcher
         ```bash
         accelerate config  # This will create a config file on your server

https://api.github.com/repos/huggingface/accelerate/pulls/28
diff --git a/docs/source/index.rst b/docs/source/index.rst
index efdeb8f88..ceef4fc91 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -144,6 +144,11 @@ Supported integrations
     quicktour
     installation
 
+.. toctree::
+    :maxdepth: 2
+    :caption: Guides
+
+    sagemaker
 
 .. toctree::
     :maxdepth: 2
diff --git a/docs/source/sagemaker.rst b/docs/source/sagemaker.rst
new file mode 100644
index 000000000..b889c99b3
--- /dev/null
+++ b/docs/source/sagemaker.rst
@@ -0,0 +1,169 @@
+.. 
+    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.
+
+Amazon SageMaker
+=======================================================================================================================
+
+Hugging Face and Amazon introduced new `Hugging Face Deep Learning Containers (DLCs)
+<https://github.com/aws/deep-learning-containers/blob/master/available_images.md#huggingface-training-containers>`_ to
+make it easier than ever to train Hugging Face Transformer models in `Amazon SageMaker
+<https://aws.amazon.com/sagemaker/>`_.
+
+To learn how to use the new 🤗 DLCs with the Amazon SageMaker to run your 🤗 Accelerate scripts and raw training loops.0
+
+
+
+Getting Started
+-----------------------------------------------------------------------------------------------------------------------
+
+Setup & Installation
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+
+Before you can run your 🤗 Accelerate scripts on Amazon SageMaker you need to sign up for an AWS account. If you do not
+have an AWS account yet learn more `here <https://docs.aws.amazon.com/sagemaker/latest/dg/gs-set-up.html>`__.
+
+After you have your AWS Account you need to install the ``sagemaker`` sdk for 🤗 Accelerate with.
+
+.. code-block:: 
+
+    pip install "accelerate[sagemaker]" --upgrade
+
+
+🤗 Accelerate currently uses the 🤗 DLCs, with ``transformers``, ``datasets`` and ``tokenizers`` pre-installed. 🤗
+Accelerate is not in the DLC yet (will soon be added!) so to use it within Amazon SageMaker you need to create a
+``requirements.txt`` in the same directory where your training script is located and add it as dependency.
+
+.. code-block:: 
+
+    accelerate
+
+You should also add any other dependencies you have to this ``requirements.txt``.
+
+
+Configure 🤗 Accelerate
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+You can configure the launch configuration for Amazon SageMaker the same as you do for non SageMaker training jobs with
+the 🤗 Accelerate CLI.
+
+.. code-block:: 
+
+    accelerate config
+    # In which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)): 1
+
+
+🤗 Accelerate will go through a questionnaire about your Amazon SageMaker setup and create a config file you can edit.
+
+.. note::
+    🤗 Accelerate is not saving any of your credentials.
+
+
+Prepare a 🤗 Accelerate fine-tuning script
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+The training script is very similar to a training script you might run outside of SageMaker, but to save your model
+after training you need to specify either ``/opt/ml/model`` or use ``os.environ["SM_MODEL_DIR"]`` as your save
+directory. After training, artifacts in this directory are uploaded to S3.
+
+
+.. code-block:: diff
+
+    - torch.save('/opt/ml/model`)
+    + accelerator.save('/opt/ml/model')
+
+
+.. warning::
+    SageMaker doesn’t support argparse actions. If you want to use, for example, boolean hyperparameters, you need to
+    specify type as bool in your script and provide an explicit True or False value for this hyperparameter. `[REF]
+    <https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/using_pytorch.html#prepare-a-pytorch-training-script>`__.
+
+
+Launch Training
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+You can launch your training with 🤗 Accelerate CLI with
+
+.. code-block:: 
+
+    accelerate launch path_to_script.py --args_to_the_script
+
+
+This will launch your training script using your configuration. The only thing you have to do is provide all the
+arguments needed by your training script as named arguments.
+
+**Examples**
+
+.. note::
+    If you run one of the example scripts, don't forget to add ``accelerator.save('/opt/ml/model')`` to it.
+
+.. code-block:: 
+
+    accelerate launch ./examples/sagemaker_example.py 
+
+
+Outputs:
+
+.. code-block:: 
+
+    Configuring Amazon SageMaker environment
+    Converting Arguments to Hyperparameters
+    Creating Estimator
+    2021-04-08 11:56:50 Starting - Starting the training job...
+    2021-04-08 11:57:13 Starting - Launching requested ML instancesProfilerReport-1617883008: InProgress
+    .........
+    2021-04-08 11:58:54 Starting - Preparing the instances for training.........
+    2021-04-08 12:00:24 Downloading - Downloading input data
+    2021-04-08 12:00:24 Training - Downloading the training image..................
+    2021-04-08 12:03:39 Training - Training image download completed. Training in progress..
+    ........
+    epoch 0: {'accuracy': 0.7598039215686274, 'f1': 0.8178438661710037}
+    epoch 1: {'accuracy': 0.8357843137254902, 'f1': 0.882249560632689}
+    epoch 2: {'accuracy': 0.8406862745098039, 'f1': 0.8869565217391304}
+    ........
+    2021-04-08 12:05:40 Uploading - Uploading generated training model
+    2021-04-08 12:05:40 Completed - Training job completed
+    Training seconds: 331
+    Billable seconds: 331
+    You can find your model data at: s3://your-bucket/accelerate-sagemaker-1-2021-04-08-11-56-47-108/output/model.tar.gz
+
+
+
+Advanced Features
+-----------------------------------------------------------------------------------------------------------------------
+
+Distributed Training: Data Parallelism
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+*currently in development, will be supported soon.*
+
+Distributed Training: Model Parallelism
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+*currently in development, will be supported soon.*
+
+Python packages and dependencies
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+🤗 Accelerate currently uses the 🤗 DLCs, with ``transformers``, ``datasets`` and ``tokenizers`` pre-installed. If you
+want to use different/other Python packages you can do this by adding them to the ``requirements.txt``. These packages
+will be installed before your training script is started.
+
+Remote scripts: Use scripts located on Github
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+*undecided if feature is needed. Contact us if you would like this feature.*
+
+Use Spot Instances
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+*undecided if feature is needed. Contact us if you would like this feature.*
diff --git a/examples/nlp_example.py b/examples/nlp_example.py
index aad2b8b65..220c6ddcc 100644
--- a/examples/nlp_example.py
+++ b/examples/nlp_example.py
@@ -162,8 +162,8 @@ def collate_fn(examples):
 
 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("--cpu", action="store_true", help="If passed, will train on the CPU.")
+    parser.add_argument("--fp16", type=bool, default=False, help="If passed, will use FP16 training.")
+    parser.add_argument("--cpu", type=bool, default=False, help="If passed, will train on the CPU.")
     args = parser.parse_args()
     config = {"lr": 2e-5, "num_epochs": 3, "correct_bias": True, "seed": 42, "batch_size": 16}
     training_function(config, args)
diff --git a/examples/requirements.txt b/examples/requirements.txt
new file mode 100644
index 000000000..65b6b04c3
--- /dev/null
+++ b/examples/requirements.txt
@@ -0,0 +1 @@
+accelerate # used to be installed in Amazon SageMaker environment
\ No newline at end of file
diff --git a/setup.py b/setup.py
index e2ce2cfbc..b2700c015 100644
--- a/setup.py
+++ b/setup.py
@@ -56,7 +56,7 @@
         ]
     },
     python_requires=">=3.6.0",
-    install_requires=["torch>=1.4.0"],
+    install_requires=["torch>=1.4.0", "pyaml>=20.4.0"],
     extras_require=extras,
     classifiers=[
         "Development Status :: 5 - Production/Stable",
diff --git a/src/accelerate/commands/config/config_args.py b/src/accelerate/commands/config/config_args.py
index fe2cb122a..a05debfda 100644
--- a/src/accelerate/commands/config/config_args.py
+++ b/src/accelerate/commands/config/config_args.py
@@ -42,13 +42,19 @@ def load_config_from_file(config_file):
     config_file = config_file if config_file is not None else default_config_file
     with open(config_file, "r", encoding="utf-8") as f:
         if config_file.endswith(".json"):
-            if json.load(f).get("compute_environment", ComputeEnvironment.LOCAL_MACHINE):
+            if (
+                json.load(f).get("compute_environment", ComputeEnvironment.LOCAL_MACHINE)
+                is ComputeEnvironment.LOCAL_MACHINE
+            ):
                 config_class = ClusterConfig
             else:
                 config_class = SageMakerConfig
             return config_class.from_json_file(json_file=config_file)
         else:
-            if yaml.safe_load(f).get("compute_environment", ComputeEnvironment.LOCAL_MACHINE):
+            if (
+                yaml.safe_load(f).get("compute_environment", ComputeEnvironment.LOCAL_MACHINE)
+                is ComputeEnvironment.LOCAL_MACHINE
+            ):
                 config_class = ClusterConfig
             else:
                 config_class = SageMakerConfig
@@ -121,5 +127,5 @@ class SageMakerConfig(BaseConfig):
     region: str = "us-east-1"
     num_machines: int = 1
     base_job_name: str = f"accelerate-sagemaker-{num_machines}"
-    pytroch_version: str = "1.6"
+    pytorch_version: str = "1.6"
     transformers_version: str = "4.4"
diff --git a/src/accelerate/commands/config/sagemaker.py b/src/accelerate/commands/config/sagemaker.py
index 25fd3e6cb..feff0aecf 100644
--- a/src/accelerate/commands/config/sagemaker.py
+++ b/src/accelerate/commands/config/sagemaker.py
@@ -97,7 +97,8 @@ def get_sagemaker_input():
         os.environ["AWS_PROFILE"] = aws_profile
     else:
         print(
-            "Accelerate will expose AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment variables on the machine"
+            "Note you will need to provide AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY when you launch you training script with,"
+            "`accelerate launch --aws_access_key_id XXX --aws_secret_access_key YYY`"
         )
         aws_access_key_id = _ask_field("AWS Access Key ID: ")
         os.environ["AWS_ACCESS_KEY_ID"] = aws_access_key_id
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index cd3ae65cc..382636253 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -20,10 +20,14 @@
 import os
 import subprocess
 import sys
+from ast import literal_eval
 from pathlib import Path
+from typing import Dict, List
 
 from accelerate.commands.config import default_config_file, load_config_from_file
+from accelerate.commands.config.config_args import SageMakerConfig
 from accelerate.state import ComputeEnvironment, DistributedType
+from accelerate.utils import is_sagemaker_available
 
 
 class _AddOneArg:
@@ -80,6 +84,18 @@ def launch_command_parser(subparsers=None):
         default=None,
         help="The name of the main function to be executed in your script (only for TPU training).",
     )
+    parser.add_argument(
+        "--aws_access_key_id",
+        type=str,
+        default=None,
+        help="The AWS_ACCESS_KEY_ID used to launch the Amazon SageMaker training job",
+    )
+    parser.add_argument(
+        "--aws_secret_access_key",
+        type=str,
+        default=None,
+        help="The AWS_SECRET_ACCESS_KEY used to launch the Amazon SageMaker training job",
+    )
     parser.add_argument(
         "training_script",
         type=str,
@@ -168,11 +184,107 @@ def tpu_launcher(args):
         xmp.spawn(main_function, args=(), nprocs=args.num_processes)
 
 
-def sagemaker_launcher(sagemaker_config, args):
-    raise NotImplementedError(
-        "Support for starting SageMaker training is not yet implemented. But you can create configs for it."
+def _convert_nargs_to_dict(nargs: List[str]) -> Dict[str, str]:
+    if len(nargs) < 0:
+        return {}
+    # helper function to infer type for argsparser
+
+    def _infer_type(s):
+        try:
+            s = float(s)
+
+            if s // 1 == s:
+                return int(s)
+            return s
+        except ValueError:
+            return s
+
+    parser = argparse.ArgumentParser()
+    _, unknown = parser.parse_known_args(nargs)
+    for index, argument in enumerate(unknown):
+        if argument.startswith(("-", "--")):
+            action = None
+            if index + 1 < len(unknown):  # checks if next index would be in list
+                if unknown[index + 1].startswith(("-", "--")):  # checks if next element is an key
+                    # raise an error if element is store_true or store_false
+                    raise ValueError(
+                        "SageMaker doesn’t support argparse actions for `store_true` or `store_false`. Please define explicit types"
+                    )
+            else:  # raise an error if last element is store_true or store_false
+                raise ValueError(
+                    "SageMaker doesn’t support argparse actions for `store_true` or `store_false`. Please define explicit types"
+                )
+            # adds argument to parser based on action_store true
+            if action is None:
+                parser.add_argument(argument, type=_infer_type)
+            else:
+                parser.add_argument(argument, action=action)
+
+    return {
+        key: (literal_eval(value) if value == "True" or value == "False" else value)
+        for key, value in parser.parse_args(nargs).__dict__.items()
+    }
+
+
+def sagemaker_launcher(sagemaker_config: SageMakerConfig, args):
+    if not is_sagemaker_available():
+        raise ImportError(
+            "Please install sagemaker to be able to launch training on Amazon SageMaker with `pip install accelerate[sagemaker]`"
+        )
+    from sagemaker.huggingface import HuggingFace
+
+    # configure environment
+    print("Configuring Amazon SageMaker environment")
+    os.environ["AWS_DEFAULT_REGION"] = sagemaker_config.region
+
+    # configure credentials
+    if sagemaker_config.profile is not None:
+        os.environ["AWS_PROFILE"] = sagemaker_config.profile
+    elif args.aws_access_key_id is not None and args.aws_secret_access_key is not None:
+        os.environ["AWS_ACCESS_KEY_ID"] = args.aws_access_key_id
+        os.environ["AWS_SECRET_ACCESS_KEY"] = args.aws_secret_access_key
+    else:
+        raise EnvironmentError(
+            "You need to provide an aws_access_key_id and aws_secret_access_key when not using aws_profile"
+        )
+
+    # extract needed arguments
+    source_dir = os.path.dirname(args.training_script)
+    if not source_dir:  # checks if string is empty
+        source_dir = "."
+    entry_point = os.path.basename(args.training_script)
+    if not entry_point.endswith(".py"):
+        raise ValueError(f'Your training script should be a python script and not "{entry_point}"')
+
+    print("Converting Arguments to Hyperparameters")
+    hyperparameters = _convert_nargs_to_dict(args.training_script_args)
+
+    environment = {"USE_FP16": args.fp16}  # Environment variables to be set for use during training job
+
+    # configure distribution set up
+    distribution = None  # TODO: not yet implemented
+
+    # configure session
+    print("Creating Estimator")
+    huggingface_estimator = HuggingFace(
+        entry_point=entry_point,
+        source_dir=source_dir,
+        role=sagemaker_config.iam_role_name,
+        transformers_version="4.4",
+        pytorch_version="1.6",
+        py_version="py36",
+        base_job_name=sagemaker_config.base_job_name,
+        instance_count=sagemaker_config.num_machines,
+        instance_type=sagemaker_config.ec2_instance_type,
+        debugger_hook_config=False,
+        distribution=distribution,
+        hyperparameters=hyperparameters,
+        environment=environment,
     )
 
+    huggingface_estimator.fit()
+    print(f"You can find your model data at: {huggingface_estimator.model_data}")
+
 
 def launch_command(args):
     # Sanity checks
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 21b635996..68299c790 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -31,6 +31,10 @@ def is_boto3_available():
     return importlib.util.find_spec("boto3") is not None
 
 
+def is_sagemaker_available():
+    return importlib.util.find_spec("sagemaker") is not None
+
+
 class RNGType(Enum):
     TORCH = "torch"
     CUDA = "cuda"
diff --git a/tests/test_sagemaker.py b/tests/test_sagemaker.py
new file mode 100644
index 000000000..f2d587a39
--- /dev/null
+++ b/tests/test_sagemaker.py
@@ -0,0 +1,62 @@
+import unittest
+from dataclasses import dataclass
+
+import pytest
+from accelerate.commands.config.config_args import SageMakerConfig
+from accelerate.commands.launch import _convert_nargs_to_dict
+from accelerate.state import ComputeEnvironment
+
+
+@dataclass
+class MockLaunchConfig(SageMakerConfig):
+    compute_environment = ComputeEnvironment.AMAZON_SAGEMAKER
+    fp16 = True
+    ec2_instance_type = "ml.p3.2xlarge"
+    iam_role_name = "accelerate_sagemaker_execution_role"
+    profile = "hf-sm"
+    region = "us-east-1"
+    num_machines = 1
+    base_job_name = "accelerate-sagemaker-1"
+    pytorch_version = "1.6"
+    transformers_version = "4.4"
+    training_script = "train.py"
+    success_training_script_args = [
+        "--model_name_or_path",
+        "bert",
+        "--do_train",
+        "False",
+        "--epochs",
+        "3",
+        "--learning_rate",
+        "5e-5",
+        "--max_steps",
+        "50.5",
+    ]
+    fail_training_script_args = [
+        "--model_name_or_path",
+        "bert",
+        "--do_train",
+        "--do_test",
+        "False",
+        "--do_predict",
+        "--epochs",
+        "3",
+        "--learning_rate",
+        "5e-5",
+        "--max_steps",
+        "50.5",
+    ]
+
+
+class SageMakerLaunch(unittest.TestCase):
+    def test_args_convert(self):
+        # If no defaults are changed, `to_kwargs` returns an empty dict.
+        converted_args = _convert_nargs_to_dict(MockLaunchConfig.success_training_script_args)
+        assert isinstance(converted_args["model_name_or_path"], str)
+        assert isinstance(converted_args["do_train"], bool)
+        assert isinstance(converted_args["epochs"], int)
+        assert isinstance(converted_args["learning_rate"], float)
+        assert isinstance(converted_args["max_steps"], float)
+
+        with pytest.raises(ValueError):
+            _convert_nargs_to_dict(MockLaunchConfig.fail_training_script_args)

```suggestion
## Getting Started
```
@@ -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.
+-->
+
+# Amazon SageMaker
+
+Hugging Face and Amazon introduced new [Hugging Face Deep Learning Containers (DLCs)](https://github.com/aws/deep-learning-containers/blob/master/available_images.md#huggingface-training-containers) to make it easier than ever to train Hugging Face Transformer models in [Amazon SageMaker](https://aws.amazon.com/sagemaker/).
+
+To learn how to use the new 🤗 DLCs with the Amazon SageMaker to run your 🤗 Accelerate scripts and raw training loops.0
+
+---
+
+## Getting Started: Train a 🤗 Transformers Model
```suggestion
🤗 Accelerate currently uses the 🤗 DLCs, with `transformers`, `datasets` and `tokenizers` pre-installed. 🤗 Accelerate is not in the DLC yet (swill soon be added!) so to use it within Amazon SageMaker you need to create a `requirements.txt` in the same directory where your training script is located and add it as dependency.
```
@@ -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.
+-->
+
+# Amazon SageMaker
+
+Hugging Face and Amazon introduced new [Hugging Face Deep Learning Containers (DLCs)](https://github.com/aws/deep-learning-containers/blob/master/available_images.md#huggingface-training-containers) to make it easier than ever to train Hugging Face Transformer models in [Amazon SageMaker](https://aws.amazon.com/sagemaker/).
+
+To learn how to use the new 🤗 DLCs with the Amazon SageMaker to run your 🤗 Accelerate scripts and raw training loops.0
+
+---
+
+## Getting Started: Train a 🤗 Transformers Model
+
+### Setup & Installation
+
+Before you can run your 🤗 Accelerate scripts on Amazon SageMaker you need to sign up for an AWS account. If you do not have an AWS account yet learn more [here](https://docs.aws.amazon.com/sagemaker/latest/dg/gs-set-up.html). 
+
+After you have your AWS Account you need to install the `sagemaker` sdk for 🤗 Accelerate with. 
+
+```python
+pip install "accelerate[sagemaker]" --upgrade
+```
+
+🤗 Accelerate currently uses the 🤗 DLCs, with `transformers`, `datasets` and `tokenizers` pre-installed. To be able to use 🤗 Accelerate within Amazon SageMaker you need to create a `requirements.txt` in the same directory where your training script is located and add it as dependency.
```suggestion
You should also add any other dependencies you have to this `requirements.txt`. 
```
@@ -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.
+-->
+
+# Amazon SageMaker
+
+Hugging Face and Amazon introduced new [Hugging Face Deep Learning Containers (DLCs)](https://github.com/aws/deep-learning-containers/blob/master/available_images.md#huggingface-training-containers) to make it easier than ever to train Hugging Face Transformer models in [Amazon SageMaker](https://aws.amazon.com/sagemaker/).
+
+To learn how to use the new 🤗 DLCs with the Amazon SageMaker to run your 🤗 Accelerate scripts and raw training loops.0
+
+---
+
+## Getting Started: Train a 🤗 Transformers Model
+
+### Setup & Installation
+
+Before you can run your 🤗 Accelerate scripts on Amazon SageMaker you need to sign up for an AWS account. If you do not have an AWS account yet learn more [here](https://docs.aws.amazon.com/sagemaker/latest/dg/gs-set-up.html). 
+
+After you have your AWS Account you need to install the `sagemaker` sdk for 🤗 Accelerate with. 
+
+```python
+pip install "accelerate[sagemaker]" --upgrade
+```
+
+🤗 Accelerate currently uses the 🤗 DLCs, with `transformers`, `datasets` and `tokenizers` pre-installed. To be able to use 🤗 Accelerate within Amazon SageMaker you need to create a `requirements.txt` in the same directory where your training script is located and add it as dependency.
+
+```python
+accelerate
+```
+
+You can also add other dependencies to this `requirements.txt`. 
What do you mean here? Shouldn't it be in the example?
@@ -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.
+-->
+
+# Amazon SageMaker
+
+Hugging Face and Amazon introduced new [Hugging Face Deep Learning Containers (DLCs)](https://github.com/aws/deep-learning-containers/blob/master/available_images.md#huggingface-training-containers) to make it easier than ever to train Hugging Face Transformer models in [Amazon SageMaker](https://aws.amazon.com/sagemaker/).
+
+To learn how to use the new 🤗 DLCs with the Amazon SageMaker to run your 🤗 Accelerate scripts and raw training loops.0
+
+---
+
+## Getting Started: Train a 🤗 Transformers Model
+
+### Setup & Installation
+
+Before you can run your 🤗 Accelerate scripts on Amazon SageMaker you need to sign up for an AWS account. If you do not have an AWS account yet learn more [here](https://docs.aws.amazon.com/sagemaker/latest/dg/gs-set-up.html). 
+
+After you have your AWS Account you need to install the `sagemaker` sdk for 🤗 Accelerate with. 
+
+```python
+pip install "accelerate[sagemaker]" --upgrade
+```
+
+🤗 Accelerate currently uses the 🤗 DLCs, with `transformers`, `datasets` and `tokenizers` pre-installed. To be able to use 🤗 Accelerate within Amazon SageMaker you need to create a `requirements.txt` in the same directory where your training script is located and add it as dependency.
+
+```python
+accelerate
+```
+
+You can also add other dependencies to this `requirements.txt`. 
+
+
+### Configure 🤗 Accelerate 
+
+You can configure the launch configuration for Amazon SageMaker the same as you do for non SageMaker training jobs with the 🤗 Accelerate CLI. 
+
+```bash
+accelerate config
+# In which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)): 1
+```
+
+🤗 Accelerate will go through a questionnaire about your Amazon SageMaker setup and create a config file you can edit. 
+
+_**Note**: 🤗 Accelerate is not saving any of your credentials._
+
+
+### Prepare a 🤗 Accelerate fine-tuning script.
+
+The training script is very similar to a training script you might run outside of SageMaker, but to save your model after training you need to specify either `/opt/ml/model` or use `os.environ["SM_MODEL_DIR"]` as your save directory. After training, artifacts in this directory are uploaded to S3.
+
+```diff
+- torch.save()
++ accelerator.save('/opt/ml/model')
+```
+
+### Launch Training
+
+You can launch your training with 🤗 Accelerate CLI with
+
+```bash
+accelerate launch path_to_script.py --args_to_the_script
+```
+
+This will launch your training script using your configuration. The only thing you have to do is provide all the arguments needed by your training script as named arguments.
+
+**Examples**
+
+_**Note**: If you run one of the example scripts, don't for get to add `accelerator.save('/opt/ml/model')`to it._
I see you did something for those later on, is it still necessary? We should make sure we support those as they are very common, or if they are not, it should be made clear in the docs.
@@ -162,8 +162,8 @@ def collate_fn(examples):
 
 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("--cpu", action="store_true", help="If passed, will train on the CPU.")
+    parser.add_argument("--fp16", type=bool, default=False, help="If passed, will use FP16 training.")
+    parser.add_argument("--cpu", type=bool, default=False, help="If passed, will train on the CPU.")
Should we amke one folder for each example? The CV example would require more.
@@ -0,0 +1 @@
+accelerate # used to be installed in Amazon SageMaker environment
There should be something for `store_false` too.
@@ -168,11 +184,100 @@ def tpu_launcher(args):
         xmp.spawn(main_function, args=(), nprocs=args.num_processes)
 
 
-def sagemaker_launcher(sagemaker_config, args):
-    raise NotImplementedError(
-        "Support for starting SageMaker training is not yet implemented. But you can create configs for it."
+def _convert_nargs_to_dict(nargs: List[str]) -> Dict[str, str]:
+    if len(nargs) < 0:
+        return {}
+    # helper function to infer type for argsparser
+
+    def _infer_type(s):
+        try:
+            s = float(s)
+
+            if s // 1 == s:
+                return int(s)
+            return s
+        except ValueError:
+            return s
+
+    parser = argparse.ArgumentParser()
+    _, unknown = parser.parse_known_args(nargs)
+    for index, argument in enumerate(unknown):
+        if argument.startswith(("-", "--")):
+            action = None
+            if index + 1 < len(unknown):  # checks if next index would be in list
+                if unknown[index + 1].startswith(("-", "--")):  # checks if next element is an key
+                    action = "store_true"
+            else:  # action_true parameter is last element
+                action = "store_true"
The current example doesn't show how to save models at all. So that's why I added it here.
@@ -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.
+-->
+
+# Amazon SageMaker
+
+Hugging Face and Amazon introduced new [Hugging Face Deep Learning Containers (DLCs)](https://github.com/aws/deep-learning-containers/blob/master/available_images.md#huggingface-training-containers) to make it easier than ever to train Hugging Face Transformer models in [Amazon SageMaker](https://aws.amazon.com/sagemaker/).
+
+To learn how to use the new 🤗 DLCs with the Amazon SageMaker to run your 🤗 Accelerate scripts and raw training loops.0
+
+---
+
+## Getting Started: Train a 🤗 Transformers Model
+
+### Setup & Installation
+
+Before you can run your 🤗 Accelerate scripts on Amazon SageMaker you need to sign up for an AWS account. If you do not have an AWS account yet learn more [here](https://docs.aws.amazon.com/sagemaker/latest/dg/gs-set-up.html). 
+
+After you have your AWS Account you need to install the `sagemaker` sdk for 🤗 Accelerate with. 
+
+```python
+pip install "accelerate[sagemaker]" --upgrade
+```
+
+🤗 Accelerate currently uses the 🤗 DLCs, with `transformers`, `datasets` and `tokenizers` pre-installed. To be able to use 🤗 Accelerate within Amazon SageMaker you need to create a `requirements.txt` in the same directory where your training script is located and add it as dependency.
+
+```python
+accelerate
+```
+
+You can also add other dependencies to this `requirements.txt`. 
+
+
+### Configure 🤗 Accelerate 
+
+You can configure the launch configuration for Amazon SageMaker the same as you do for non SageMaker training jobs with the 🤗 Accelerate CLI. 
+
+```bash
+accelerate config
+# In which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)): 1
+```
+
+🤗 Accelerate will go through a questionnaire about your Amazon SageMaker setup and create a config file you can edit. 
+
+_**Note**: 🤗 Accelerate is not saving any of your credentials._
+
+
+### Prepare a 🤗 Accelerate fine-tuning script.
+
+The training script is very similar to a training script you might run outside of SageMaker, but to save your model after training you need to specify either `/opt/ml/model` or use `os.environ["SM_MODEL_DIR"]` as your save directory. After training, artifacts in this directory are uploaded to S3.
+
+```diff
+- torch.save()
++ accelerator.save('/opt/ml/model')
+```
+
+### Launch Training
+
+You can launch your training with 🤗 Accelerate CLI with
+
+```bash
+accelerate launch path_to_script.py --args_to_the_script
+```
+
+This will launch your training script using your configuration. The only thing you have to do is provide all the arguments needed by your training script as named arguments.
+
+**Examples**
+
+_**Note**: If you run one of the example scripts, don't for get to add `accelerator.save('/opt/ml/model')`to it._
SageMaker cannot pass `store_true` or `store_false` to a script executed. 

> Note that SageMaker doesn’t support argparse actions. If you want to use, for example, boolean hyperparameters, you need to specify type as bool in your script and provide an explicit True or False value for this hyperparameter when instantiating PyTorch Estimator.

https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/using_pytorch.html#prepare-a-pytorch-training-script

should we add it to the documentation then? 
@@ -162,8 +162,8 @@ def collate_fn(examples):
 
 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("--cpu", action="store_true", help="If passed, will train on the CPU.")
+    parser.add_argument("--fp16", type=bool, default=False, help="If passed, will use FP16 training.")
+    parser.add_argument("--cpu", type=bool, default=False, help="If passed, will train on the CPU.")
for the CV example to work we would only need to add `timm`. `torchvision` should be inside the DLC since it uses the AWS PyTorch DLC as base

@@ -0,0 +1 @@
+accelerate # used to be installed in Amazon SageMaker environment
puh. I that even possible? like how can someone cloud tell the difference between `store_true` and `store_false` without knowing the script... 
@@ -168,11 +184,100 @@ def tpu_launcher(args):
         xmp.spawn(main_function, args=(), nprocs=args.num_processes)
 
 
-def sagemaker_launcher(sagemaker_config, args):
-    raise NotImplementedError(
-        "Support for starting SageMaker training is not yet implemented. But you can create configs for it."
+def _convert_nargs_to_dict(nargs: List[str]) -> Dict[str, str]:
+    if len(nargs) < 0:
+        return {}
+    # helper function to infer type for argsparser
+
+    def _infer_type(s):
+        try:
+            s = float(s)
+
+            if s // 1 == s:
+                return int(s)
+            return s
+        except ValueError:
+            return s
+
+    parser = argparse.ArgumentParser()
+    _, unknown = parser.parse_known_args(nargs)
+    for index, argument in enumerate(unknown):
+        if argument.startswith(("-", "--")):
+            action = None
+            if index + 1 < len(unknown):  # checks if next index would be in list
+                if unknown[index + 1].startswith(("-", "--")):  # checks if next element is an key
+                    action = "store_true"
+            else:  # action_true parameter is last element
+                action = "store_true"
Yes, it should be properly documented.
@@ -162,8 +162,8 @@ def collate_fn(examples):
 
 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("--cpu", action="store_true", help="If passed, will train on the CPU.")
+    parser.add_argument("--fp16", type=bool, default=False, help="If passed, will use FP16 training.")
+    parser.add_argument("--cpu", type=bool, default=False, help="If passed, will train on the CPU.")
I have no idea what you're doing it here then. Also I though the script could not have `--store_true`/`--store_false`

On top of the documentation, we should add a nice error message here if the parser has arguments with `--store_true`/`--store_false`.
@@ -168,11 +184,100 @@ def tpu_launcher(args):
         xmp.spawn(main_function, args=(), nprocs=args.num_processes)
 
 
-def sagemaker_launcher(sagemaker_config, args):
-    raise NotImplementedError(
-        "Support for starting SageMaker training is not yet implemented. But you can create configs for it."
+def _convert_nargs_to_dict(nargs: List[str]) -> Dict[str, str]:
+    if len(nargs) < 0:
+        return {}
+    # helper function to infer type for argsparser
+
+    def _infer_type(s):
+        try:
+            s = float(s)
+
+            if s // 1 == s:
+                return int(s)
+            return s
+        except ValueError:
+            return s
+
+    parser = argparse.ArgumentParser()
+    _, unknown = parser.parse_known_args(nargs)
+    for index, argument in enumerate(unknown):
+        if argument.startswith(("-", "--")):
+            action = None
+            if index + 1 < len(unknown):  # checks if next index would be in list
+                if unknown[index + 1].startswith(("-", "--")):  # checks if next element is an key
+                    action = "store_true"
+            else:  # action_true parameter is last element
+                action = "store_true"
I reworked the documentation to use `.rst` to highlight it as a warning. ([see PR comment](https://github.com/huggingface/accelerate/pull/26#issuecomment-818525901))
@@ -162,8 +162,8 @@ def collate_fn(examples):
 
 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("--cpu", action="store_true", help="If passed, will train on the CPU.")
+    parser.add_argument("--fp16", type=bool, default=False, help="If passed, will use FP16 training.")
+    parser.add_argument("--cpu", type=bool, default=False, help="If passed, will train on the CPU.")
For SageMaker to work I needed to convert the `args.training_script_args` from `--foo bar --john 2 --store_true --another key` to
 ```python
{
"foo":"bar",
"john": 2,
"store_true": True,
"another": "key"
}
```
To achieve this I built this dynamic args parser, which takes the `args.training_script_args` and creates an `argsparser` out of it to apply afterward on `args.training_script_args` to create this `dict`. 

While creating this I haven't thought of `store_false`, since I never have seen it in real use, that's why I expected that there is only `store_true` and implemented it like that. 

@@ -168,11 +184,100 @@ def tpu_launcher(args):
         xmp.spawn(main_function, args=(), nprocs=args.num_processes)
 
 
-def sagemaker_launcher(sagemaker_config, args):
-    raise NotImplementedError(
-        "Support for starting SageMaker training is not yet implemented. But you can create configs for it."
+def _convert_nargs_to_dict(nargs: List[str]) -> Dict[str, str]:
+    if len(nargs) < 0:
+        return {}
+    # helper function to infer type for argsparser
+
+    def _infer_type(s):
+        try:
+            s = float(s)
+
+            if s // 1 == s:
+                return int(s)
+            return s
+        except ValueError:
+            return s
+
+    parser = argparse.ArgumentParser()
+    _, unknown = parser.parse_known_args(nargs)
+    for index, argument in enumerate(unknown):
+        if argument.startswith(("-", "--")):
+            action = None
+            if index + 1 < len(unknown):  # checks if next index would be in list
+                if unknown[index + 1].startswith(("-", "--")):  # checks if next element is an key
+                    action = "store_true"
+            else:  # action_true parameter is last element
+                action = "store_true"
Should we add `timm` to the `requirements.txt` to make both examples runnable out of the box?
@@ -0,0 +1 @@
+accelerate # used to be installed in Amazon SageMaker environment
As discussed I adjusted the convert script to raise an error if someone uses `store` actions and added a case into the test. 
@@ -168,11 +184,100 @@ def tpu_launcher(args):
         xmp.spawn(main_function, args=(), nprocs=args.num_processes)
 
 
-def sagemaker_launcher(sagemaker_config, args):
-    raise NotImplementedError(
-        "Support for starting SageMaker training is not yet implemented. But you can create configs for it."
+def _convert_nargs_to_dict(nargs: List[str]) -> Dict[str, str]:
+    if len(nargs) < 0:
+        return {}
+    # helper function to infer type for argsparser
+
+    def _infer_type(s):
+        try:
+            s = float(s)
+
+            if s // 1 == s:
+                return int(s)
+            return s
+        except ValueError:
+            return s
+
+    parser = argparse.ArgumentParser()
+    _, unknown = parser.parse_known_args(nargs)
+    for index, argument in enumerate(unknown):
+        if argument.startswith(("-", "--")):
+            action = None
+            if index + 1 < len(unknown):  # checks if next index would be in list
+                if unknown[index + 1].startswith(("-", "--")):  # checks if next element is an key
+                    action = "store_true"
+            else:  # action_true parameter is last element
+                action = "store_true"
```suggestion
have an AWS account yet learn more `here <https://docs.aws.amazon.com/sagemaker/latest/dg/gs-set-up.html>`__.
```
@@ -0,0 +1,169 @@
+.. 
+    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.
+
+Amazon SageMaker
+=======================================================================================================================
+
+Hugging Face and Amazon introduced new `Hugging Face Deep Learning Containers (DLCs)
+<https://github.com/aws/deep-learning-containers/blob/master/available_images.md#huggingface-training-containers>`_ to
+make it easier than ever to train Hugging Face Transformer models in `Amazon SageMaker
+<https://aws.amazon.com/sagemaker/>`_.
+
+To learn how to use the new 🤗 DLCs with the Amazon SageMaker to run your 🤗 Accelerate scripts and raw training loops.0
+
+
+
+Getting Started
+-----------------------------------------------------------------------------------------------------------------------
+
+Setup & Installation
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+
+Before you can run your 🤗 Accelerate scripts on Amazon SageMaker you need to sign up for an AWS account. If you do not
+have an AWS account yet learn more `here <https://docs.aws.amazon.com/sagemaker/latest/dg/gs-set-up.html>`_.
```suggestion
    <https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/using_pytorch.html#prepare-a-pytorch-training-script>`__.
```
@@ -0,0 +1,169 @@
+.. 
+    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.
+
+Amazon SageMaker
+=======================================================================================================================
+
+Hugging Face and Amazon introduced new `Hugging Face Deep Learning Containers (DLCs)
+<https://github.com/aws/deep-learning-containers/blob/master/available_images.md#huggingface-training-containers>`_ to
+make it easier than ever to train Hugging Face Transformer models in `Amazon SageMaker
+<https://aws.amazon.com/sagemaker/>`_.
+
+To learn how to use the new 🤗 DLCs with the Amazon SageMaker to run your 🤗 Accelerate scripts and raw training loops.0
+
+
+
+Getting Started
+-----------------------------------------------------------------------------------------------------------------------
+
+Setup & Installation
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+
+Before you can run your 🤗 Accelerate scripts on Amazon SageMaker you need to sign up for an AWS account. If you do not
+have an AWS account yet learn more `here <https://docs.aws.amazon.com/sagemaker/latest/dg/gs-set-up.html>`_.
+
+After you have your AWS Account you need to install the ``sagemaker`` sdk for 🤗 Accelerate with.
+
+.. code-block:: 
+
+    pip install "accelerate[sagemaker]" --upgrade
+
+
+🤗 Accelerate currently uses the 🤗 DLCs, with ``transformers``, ``datasets`` and ``tokenizers`` pre-installed. 🤗
+Accelerate is not in the DLC yet (swill soon be added!) so to use it within Amazon SageMaker you need to create a
+``requirements.txt`` in the same directory where your training script is located and add it as dependency.
+
+.. code-block:: 
+
+    accelerate
+
+You should also add any other dependencies you have to this ``requirements.txt``.
+
+
+Configure 🤗 Accelerate
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+You can configure the launch configuration for Amazon SageMaker the same as you do for non SageMaker training jobs with
+the 🤗 Accelerate CLI.
+
+.. code-block:: 
+
+    accelerate config
+    # In which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)): 1
+
+
+🤗 Accelerate will go through a questionnaire about your Amazon SageMaker setup and create a config file you can edit.
+
+.. note::
+    🤗 Accelerate is not saving any of your credentials.
+
+
+Prepare a 🤗 Accelerate fine-tuning script
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+The training script is very similar to a training script you might run outside of SageMaker, but to save your model
+after training you need to specify either ``/opt/ml/model`` or use ``os.environ["SM_MODEL_DIR"]`` as your save
+directory. After training, artifacts in this directory are uploaded to S3.
+
+
+.. code-block:: diff
+
+    - torch.save()
+    + accelerator.save('/opt/ml/model')
+
+
+.. warning::
+    SageMaker doesn’t support argparse actions. If you want to use, for example, boolean hyperparameters, you need to
+    specify type as bool in your script and provide an explicit True or False value for this hyperparameter. `[REF]
+    <https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/using_pytorch.html#prepare-a-pytorch-training-script>`_
fp16 might not be accepted by the script so it should not be added like this. It should be added in the env instead, in the variable `"USE_FP16"` (which is what accelerate will look for).
@@ -168,11 +184,105 @@ def tpu_launcher(args):
         xmp.spawn(main_function, args=(), nprocs=args.num_processes)
 
 
-def sagemaker_launcher(sagemaker_config, args):
-    raise NotImplementedError(
-        "Support for starting SageMaker training is not yet implemented. But you can create configs for it."
+def _convert_nargs_to_dict(nargs: List[str]) -> Dict[str, str]:
+    if len(nargs) < 0:
+        return {}
+    # helper function to infer type for argsparser
+
+    def _infer_type(s):
+        try:
+            s = float(s)
+
+            if s // 1 == s:
+                return int(s)
+            return s
+        except ValueError:
+            return s
+
+    parser = argparse.ArgumentParser()
+    _, unknown = parser.parse_known_args(nargs)
+    for index, argument in enumerate(unknown):
+        if argument.startswith(("-", "--")):
+            action = None
+            if index + 1 < len(unknown):  # checks if next index would be in list
+                if unknown[index + 1].startswith(("-", "--")):  # checks if next element is an key
+                    # raise an error if element is store_true or store_false
+                    raise ValueError(
+                        "SageMaker doesn’t support argparse actions for `store_true` or `store_false`. Please define explicit types"
+                    )
+            else:  # raise an error if last element is store_true or store_false
+                raise ValueError(
+                    "SageMaker doesn’t support argparse actions for `store_true` or `store_false`. Please define explicit types"
+                )
+            # adds argument to parser based on action_store true
+            if action is None:
+                parser.add_argument(argument, type=_infer_type)
+            else:
+                parser.add_argument(argument, action=action)
+
+    return {
+        key: (literal_eval(value) if value == "True" or value == "False" else value)
+        for key, value in parser.parse_args(nargs).__dict__.items()
+    }
+
+
+def sagemaker_launcher(sagemaker_config: SageMakerConfig, args):
+    if not is_sagemaker_available():
+        raise ImportError(
+            "Please install sagemaker to be able to launch training on Amazon SageMaker with `pip install accelerate[sagemaker]`"
+        )
+    from sagemaker.huggingface import HuggingFace
+
+    # configure environment
+    print("Configuring Amazon SageMaker environment")
+    os.environ["AWS_DEFAULT_REGION"] = sagemaker_config.region
+
+    # configure credentials
+    if sagemaker_config.profile is not None:
+        os.environ["AWS_PROFILE"] = sagemaker_config.profile
+    elif args.aws_access_key_id is not None and args.aws_secret_access_key is not None:
+        os.environ["AWS_ACCESS_KEY_ID"] = args.aws_access_key_id
+        os.environ["AWS_SECRET_ACCESS_KEY"] = args.aws_secret_access_key
+    else:
+        raise EnvironmentError(
+            "You need to provide an aws_access_key_id and aws_secret_access_key when not using aws_profile"
+        )
+
+    # extract needed arguments
+    source_dir = os.path.dirname(args.training_script)
+    if not source_dir:  # checks if string is empty
+        source_dir = "."
+    entry_point = os.path.basename(args.training_script)
+    if not entry_point.endswith(".py"):
+        raise ValueError(f'Your training script should be a python script and not "{entry_point}"')
+
+    print("Converting Arguments to Hyperparameters")
+    converted_hyperparameters = _convert_nargs_to_dict(args.training_script_args)
+    hyperparameters = {"fp16": args.fp16, **converted_hyperparameters}
```suggestion
Accelerate is not in the DLC yet (will soon be added!) so to use it within Amazon SageMaker you need to create a
```
@@ -0,0 +1,169 @@
+.. 
+    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.
+
+Amazon SageMaker
+=======================================================================================================================
+
+Hugging Face and Amazon introduced new `Hugging Face Deep Learning Containers (DLCs)
+<https://github.com/aws/deep-learning-containers/blob/master/available_images.md#huggingface-training-containers>`_ to
+make it easier than ever to train Hugging Face Transformer models in `Amazon SageMaker
+<https://aws.amazon.com/sagemaker/>`_.
+
+To learn how to use the new 🤗 DLCs with the Amazon SageMaker to run your 🤗 Accelerate scripts and raw training loops.0
+
+
+
+Getting Started
+-----------------------------------------------------------------------------------------------------------------------
+
+Setup & Installation
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+
+Before you can run your 🤗 Accelerate scripts on Amazon SageMaker you need to sign up for an AWS account. If you do not
+have an AWS account yet learn more `here <https://docs.aws.amazon.com/sagemaker/latest/dg/gs-set-up.html>`__.
+
+After you have your AWS Account you need to install the ``sagemaker`` sdk for 🤗 Accelerate with.
+
+.. code-block:: 
+
+    pip install "accelerate[sagemaker]" --upgrade
+
+
+🤗 Accelerate currently uses the 🤗 DLCs, with ``transformers``, ``datasets`` and ``tokenizers`` pre-installed. 🤗
+Accelerate is not in the DLC yet (swill soon be added!) so to use it within Amazon SageMaker you need to create a
I would also add the `/opt/ml/model` here for consistency
```suggestion
    - torch.save('/opt/ml/model`)
```
@@ -0,0 +1,169 @@
+.. 
+    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.
+
+Amazon SageMaker
+=======================================================================================================================
+
+Hugging Face and Amazon introduced new `Hugging Face Deep Learning Containers (DLCs)
+<https://github.com/aws/deep-learning-containers/blob/master/available_images.md#huggingface-training-containers>`_ to
+make it easier than ever to train Hugging Face Transformer models in `Amazon SageMaker
+<https://aws.amazon.com/sagemaker/>`_.
+
+To learn how to use the new 🤗 DLCs with the Amazon SageMaker to run your 🤗 Accelerate scripts and raw training loops.0
+
+
+
+Getting Started
+-----------------------------------------------------------------------------------------------------------------------
+
+Setup & Installation
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+
+Before you can run your 🤗 Accelerate scripts on Amazon SageMaker you need to sign up for an AWS account. If you do not
+have an AWS account yet learn more `here <https://docs.aws.amazon.com/sagemaker/latest/dg/gs-set-up.html>`__.
+
+After you have your AWS Account you need to install the ``sagemaker`` sdk for 🤗 Accelerate with.
+
+.. code-block:: 
+
+    pip install "accelerate[sagemaker]" --upgrade
+
+
+🤗 Accelerate currently uses the 🤗 DLCs, with ``transformers``, ``datasets`` and ``tokenizers`` pre-installed. 🤗
+Accelerate is not in the DLC yet (swill soon be added!) so to use it within Amazon SageMaker you need to create a
+``requirements.txt`` in the same directory where your training script is located and add it as dependency.
+
+.. code-block:: 
+
+    accelerate
+
+You should also add any other dependencies you have to this ``requirements.txt``.
+
+
+Configure 🤗 Accelerate
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+You can configure the launch configuration for Amazon SageMaker the same as you do for non SageMaker training jobs with
+the 🤗 Accelerate CLI.
+
+.. code-block:: 
+
+    accelerate config
+    # In which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)): 1
+
+
+🤗 Accelerate will go through a questionnaire about your Amazon SageMaker setup and create a config file you can edit.
+
+.. note::
+    🤗 Accelerate is not saving any of your credentials.
+
+
+Prepare a 🤗 Accelerate fine-tuning script
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+The training script is very similar to a training script you might run outside of SageMaker, but to save your model
+after training you need to specify either ``/opt/ml/model`` or use ``os.environ["SM_MODEL_DIR"]`` as your save
+directory. After training, artifacts in this directory are uploaded to S3.
+
+
+.. code-block:: diff
+
+    - torch.save()
```suggestion
    If you run one of the example scripts, don't forget to add ``accelerator.save('/opt/ml/model')`` to it.
```
@@ -0,0 +1,169 @@
+.. 
+    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.
+
+Amazon SageMaker
+=======================================================================================================================
+
+Hugging Face and Amazon introduced new `Hugging Face Deep Learning Containers (DLCs)
+<https://github.com/aws/deep-learning-containers/blob/master/available_images.md#huggingface-training-containers>`_ to
+make it easier than ever to train Hugging Face Transformer models in `Amazon SageMaker
+<https://aws.amazon.com/sagemaker/>`_.
+
+To learn how to use the new 🤗 DLCs with the Amazon SageMaker to run your 🤗 Accelerate scripts and raw training loops.0
+
+
+
+Getting Started
+-----------------------------------------------------------------------------------------------------------------------
+
+Setup & Installation
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+
+Before you can run your 🤗 Accelerate scripts on Amazon SageMaker you need to sign up for an AWS account. If you do not
+have an AWS account yet learn more `here <https://docs.aws.amazon.com/sagemaker/latest/dg/gs-set-up.html>`__.
+
+After you have your AWS Account you need to install the ``sagemaker`` sdk for 🤗 Accelerate with.
+
+.. code-block:: 
+
+    pip install "accelerate[sagemaker]" --upgrade
+
+
+🤗 Accelerate currently uses the 🤗 DLCs, with ``transformers``, ``datasets`` and ``tokenizers`` pre-installed. 🤗
+Accelerate is not in the DLC yet (swill soon be added!) so to use it within Amazon SageMaker you need to create a
+``requirements.txt`` in the same directory where your training script is located and add it as dependency.
+
+.. code-block:: 
+
+    accelerate
+
+You should also add any other dependencies you have to this ``requirements.txt``.
+
+
+Configure 🤗 Accelerate
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+You can configure the launch configuration for Amazon SageMaker the same as you do for non SageMaker training jobs with
+the 🤗 Accelerate CLI.
+
+.. code-block:: 
+
+    accelerate config
+    # In which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)): 1
+
+
+🤗 Accelerate will go through a questionnaire about your Amazon SageMaker setup and create a config file you can edit.
+
+.. note::
+    🤗 Accelerate is not saving any of your credentials.
+
+
+Prepare a 🤗 Accelerate fine-tuning script
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+The training script is very similar to a training script you might run outside of SageMaker, but to save your model
+after training you need to specify either ``/opt/ml/model`` or use ``os.environ["SM_MODEL_DIR"]`` as your save
+directory. After training, artifacts in this directory are uploaded to S3.
+
+
+.. code-block:: diff
+
+    - torch.save()
+    + accelerator.save('/opt/ml/model')
+
+
+.. warning::
+    SageMaker doesn’t support argparse actions. If you want to use, for example, boolean hyperparameters, you need to
+    specify type as bool in your script and provide an explicit True or False value for this hyperparameter. `[REF]
+    <https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/using_pytorch.html#prepare-a-pytorch-training-script>`__.
+
+
+Launch Training
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+You can launch your training with 🤗 Accelerate CLI with
+
+.. code-block:: 
+
+    accelerate launch path_to_script.py --args_to_the_script
+
+
+This will launch your training script using your configuration. The only thing you have to do is provide all the
+arguments needed by your training script as named arguments.
+
+**Examples**
+
+.. note::
+    If you run one of the example scripts, don't for get to add ``accelerator.save('/opt/ml/model')`` to it.
![screencapture-127-0-0-1-5500-docs-build-html-sagemaker-html-2021-04-13-09_51_47](https://user-images.githubusercontent.com/32632186/114516657-e68efa00-9c3d-11eb-9706-b5c76d21cbc0.png)

https://api.github.com/repos/huggingface/accelerate/pulls/26
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 7d6210327..016feb77f 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -14,17 +14,18 @@
 #
 import os
 import sys
-sys.path.insert(0, os.path.abspath('../../src'))
+
+sys.path.insert(0, os.path.abspath("../../src"))
 
 
 # -- Project information -----------------------------------------------------
 
-project = u'accelerate'
-copyright = u'2020, The Hugging Face Team, Licenced under the Apache License, Version 2.0'
-author = u'huggingface'
+project = "accelerate"
+copyright = "2020, The Hugging Face Team, Licenced under the Apache License, Version 2.0"
+author = "huggingface"
 
 # The short X.Y version
-version = u'0.1.0'
+version = "0.1.0"
 
 # -- General configuration ---------------------------------------------------
 
@@ -36,27 +37,28 @@
 # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
 # ones.
 extensions = [
-    'sphinx.ext.autodoc',
-    'sphinx.ext.extlinks',
-    'sphinx.ext.coverage',
-    'sphinx.ext.napoleon',
-    'recommonmark',
-    'sphinx.ext.viewcode',
-    'sphinx_markdown_tables',
-    'sphinx_copybutton'
+    "sphinx.ext.autodoc",
+    "sphinx.ext.extlinks",
+    "sphinx.ext.coverage",
+    "sphinx.ext.napoleon",
+    "recommonmark",
+    "sphinx.ext.viewcode",
+    "sphinx_markdown_tables",
+    "sphinx_copybutton",
+    "sphinxext.opengraph",
 ]
 
 # Add any paths that contain templates here, relative to this directory.
-templates_path = ['_templates']
+templates_path = ["_templates"]
 
 # The suffix(es) of source filenames.
 # You can specify multiple suffix as a list of string:
 #
-source_suffix = ['.rst', '.md']
+source_suffix = [".rst", ".md"]
 # source_suffix = '.rst'
 
 # The master toctree document.
-master_doc = 'index'
+master_doc = "index"
 
 # The language for content autogenerated by Sphinx. Refer to documentation
 # for a list of supported languages.
@@ -68,7 +70,7 @@
 # List of patterns, relative to source directory, that match files and
 # directories to ignore when looking for source files.
 # This pattern also affects html_static_path and html_extra_path.
-exclude_patterns = [u'_build', 'Thumbs.db', '.DS_Store']
+exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]
 
 # The name of the Pygments (syntax highlighting) style to use.
 pygments_style = None
@@ -82,20 +84,31 @@
 # The theme to use for HTML and HTML Help pages.  See the documentation for
 # a list of builtin themes.
 #
-html_theme = 'sphinx_rtd_theme'
+html_theme = "sphinx_rtd_theme"
 
 # Theme options are theme-specific and customize the look and feel of a theme
 # further.  For a list of options available for each theme, see the
 # documentation.
 #
-html_theme_options = {
-    'analytics_id': 'UA-83738774-2'
-}
+html_theme_options = {"analytics_id": "UA-83738774-2"}
+
+#  Configuration for OpenGraph and Twitter Card Tags.
+# These are responsible for creating nice shareable social images https://ahrefs.com/blog/open-graph-meta-tags/
+# https://ogp.me/#type_website
+ogp_image = "https://huggingface.co/front/thumbnails/docs/accelerate.png"
+ogp_description = "Run your raw PyTorch training script on any kind of device. 🤗 Accelerate provides an easy API to make your scripts run with mixed precision and on any kind of distributed setting (multi-GPUs, TPUs etc.)"
+ogp_description_length = 160
+
+ogp_custom_meta_tags = [
+    f'<meta name="twitter:image" content="{ogp_image}">',
+    f'<meta name="twitter:description" content="{ogp_description}">',
+]
+
 
 # Add any paths that contain custom static files (such as style sheets) here,
 # relative to this directory. They are copied after the builtin static files,
 # so a file named "default.css" will overwrite the builtin "default.css".
-html_static_path = ['_static']
+html_static_path = ["_static"]
 
 # Custom sidebar templates, must be a dictionary that maps document names
 # to template names.
@@ -107,17 +120,17 @@
 #
 # html_sidebars = {}
 
-# This must be the name of an image file (path relative to the configuration 
-# directory) that is the favicon of the docs. Modern browsers use this as 
-# the icon for tabs, windows and bookmarks. It should be a Windows-style 
+# This must be the name of an image file (path relative to the configuration
+# directory) that is the favicon of the docs. Modern browsers use this as
+# the icon for tabs, windows and bookmarks. It should be a Windows-style
 # icon file (.ico).
-html_favicon = 'favicon.ico'
+html_favicon = "favicon.ico"
 
 
 # -- Options for HTMLHelp output ---------------------------------------------
 
 # Output file base name for HTML help builder.
-htmlhelp_basename = 'acceleratedoc'
+htmlhelp_basename = "acceleratedoc"
 
 
 # -- Options for LaTeX output ------------------------------------------------
@@ -126,15 +139,12 @@
     # The paper size ('letterpaper' or 'a4paper').
     #
     # 'papersize': 'letterpaper',
-
     # The font size ('10pt', '11pt' or '12pt').
     #
     # 'pointsize': '10pt',
-
     # Additional stuff for the LaTeX preamble.
     #
     # 'preamble': '',
-
     # Latex figure (float) alignment
     #
     # 'figure_align': 'htbp',
@@ -144,8 +154,7 @@
 # (source start file, target name, title,
 #  author, documentclass [howto, manual, or own class]).
 latex_documents = [
-    (master_doc, 'accelerate.tex', u'accelerate Documentation',
-     u'huggingface', 'manual'),
+    (master_doc, "accelerate.tex", "accelerate Documentation", "huggingface", "manual"),
 ]
 
 
@@ -153,10 +162,7 @@
 
 # One entry per manual page. List of tuples
 # (source start file, name, description, authors, manual section).
-man_pages = [
-    (master_doc, 'accelerate', u'accelerate Documentation',
-     [author], 1)
-]
+man_pages = [(master_doc, "accelerate", "accelerate Documentation", [author], 1)]
 
 
 # -- Options for Texinfo output ----------------------------------------------
@@ -165,9 +171,15 @@
 # (source start file, target name, title, author,
 #  dir menu entry, description, category)
 texinfo_documents = [
-    (master_doc, 'accelerate', u'accelerate Documentation',
-     author, 'accelerate', 'One line description of project.',
-     'Miscellaneous'),
+    (
+        master_doc,
+        "accelerate",
+        "accelerate Documentation",
+        author,
+        "accelerate",
+        "One line description of project.",
+        "Miscellaneous",
+    ),
 ]
 
 
@@ -186,11 +198,13 @@
 # epub_uid = ''
 
 # A list of files that should not be packed into the epub file.
-epub_exclude_files = ['search.html']
+epub_exclude_files = ["search.html"]
+
 
 def setup(app):
-    app.add_css_file('css/huggingface.css')
-    app.add_css_file('css/code-snippets.css')
-    app.add_js_file('js/custom.js')
+    app.add_css_file("css/huggingface.css")
+    app.add_css_file("css/code-snippets.css")
+    app.add_js_file("js/custom.js")
+
 
 # -- Extension configuration -------------------------------------------------
diff --git a/setup.py b/setup.py
index b6722e124..e2ce2cfbc 100644
--- a/setup.py
+++ b/setup.py
@@ -24,6 +24,7 @@
     "sphinx-markdown-tables",
     "sphinx-rtd-theme==0.4.3",
     "sphinx-copybutton",
+    "sphinxext-opengraph==0.4.1",
 ]
 extras["test"] = [
     "pytest",

FYI the moon-landing PR has been merged and deployed :)
https://api.github.com/repos/huggingface/accelerate/pulls/25
diff --git a/src/accelerate/commands/config/config_args.py b/src/accelerate/commands/config/config_args.py
index 5b0dd00ba..fe2cb122a 100644
--- a/src/accelerate/commands/config/config_args.py
+++ b/src/accelerate/commands/config/config_args.py
@@ -42,13 +42,13 @@ def load_config_from_file(config_file):
     config_file = config_file if config_file is not None else default_config_file
     with open(config_file, "r", encoding="utf-8") as f:
         if config_file.endswith(".json"):
-            if json.load(f)["compute_environment"] == ComputeEnvironment.LOCAL_MACHINE:
+            if json.load(f).get("compute_environment", ComputeEnvironment.LOCAL_MACHINE):
                 config_class = ClusterConfig
             else:
                 config_class = SageMakerConfig
             return config_class.from_json_file(json_file=config_file)
         else:
-            if yaml.safe_load(f)["compute_environment"] == ComputeEnvironment.LOCAL_MACHINE:
+            if yaml.safe_load(f).get("compute_environment", ComputeEnvironment.LOCAL_MACHINE):
                 config_class = ClusterConfig
             else:
                 config_class = SageMakerConfig
@@ -73,7 +73,10 @@ def to_dict(self):
     def from_json_file(cls, json_file=None):
         json_file = default_json_config_file if json_file is None else json_file
         with open(json_file, "r", encoding="utf-8") as f:
-            return cls(**json.load(f))
+            config_dict = json.load(f)
+        if "compute_environment" not in config_dict:
+            config_dict["compute_environment"] = ComputeEnvironment.LOCAL_MACHINE
+        return cls(**config_dict)
 
     def to_json_file(self, json_file):
         with open(json_file, "w", encoding="utf-8") as f:
@@ -84,7 +87,10 @@ def to_json_file(self, json_file):
     def from_yaml_file(cls, yaml_file=None):
         yaml_file = default_yaml_config_file if yaml_file is None else yaml_file
         with open(yaml_file, "r", encoding="utf-8") as f:
-            return cls(**yaml.safe_load(f))
+            config_dict = yaml.safe_load(f)
+        if "compute_environment" not in config_dict:
+            config_dict["compute_environment"] = ComputeEnvironment.LOCAL_MACHINE
+        return cls(**config_dict)
 
     def to_yaml_file(self, yaml_file):
         with open(yaml_file, "w", encoding="utf-8") as f:

No problem @philschmid, I should have caught it during the review :-) 
https://api.github.com/repos/huggingface/accelerate/pulls/23
diff --git a/README.md b/README.md
index 68d6866b5..77c6b6ebf 100644
--- a/README.md
+++ b/README.md
@@ -53,38 +53,38 @@ limitations under the License.
 Here is an example:
 
 ```diff
-import torch
-import torch.nn.functional as F
-from datasets import load_dataset
+  import torch
+  import torch.nn.functional as F
+  from datasets import load_dataset
 
 + from accelerate import Accelerator
 + accelerator = Accelerator()
 - device = 'cpu'
 + device = accelerator.device
 
-model = torch.nn.Transformer().to(device)
-optim = torch.optim.Adam(model.parameters())
+  model = torch.nn.Transformer().to(device)
+  optim = torch.optim.Adam(model.parameters())
 
-dataset = load_dataset('my_dataset')
-data = torch.utils.data.Dataloader(dataset)
+  dataset = load_dataset('my_dataset')
+  data = torch.utils.data.Dataloader(dataset)
 
 + model, optim, data = accelerator.prepare(model, optim, data)
 
-model.train()
-for epoch in range(10):
-    for source, targets in data:
-        source = source.to(device)
-        targets = targets.to(device)
+  model.train()
+  for epoch in range(10):
+      for source, targets in data:
+          source = source.to(device)
+          targets = targets.to(device)
 
-        optimizer.zero_grad()
+          optimizer.zero_grad()
 
-        output = model(source, targets)
-        loss = F.cross_entropy(output, targets)
+          output = model(source, targets)
+          loss = F.cross_entropy(output, targets)
 
 +         accelerator.backward(loss)
 -         loss.backward()
 
-        optimizer.step()
+          optimizer.step()
 ```
 
 As you can see in this example, by adding 5-lines to any standard PyTorch training script you can now run on any kind of single or distributed node setting (single CPU, single GPU, multi-GPUs and TPUs) as well as with or without mixed precision (fp16).
@@ -94,9 +94,9 @@ In particular, the same code can then be run without modification on your local
 🤗 Accelerate even handles the device placement for you (which requires a few more changes to your code, but is safer in general), so you can even simplify your training loop further:
 
 ```diff
-import torch
-import torch.nn.functional as F
-from datasets import load_dataset
+  import torch
+  import torch.nn.functional as F
+  from datasets import load_dataset
 
 + from accelerate import Accelerator
 + accelerator = Accelerator()
@@ -104,28 +104,28 @@ from datasets import load_dataset
 
 + model = torch.nn.Transformer()
 - model = torch.nn.Transformer().to(device)
-optim = torch.optim.Adam(model.parameters())
+  optim = torch.optim.Adam(model.parameters())
 
-dataset = load_dataset('my_dataset')
-data = torch.utils.data.Dataloader(dataset)
+  dataset = load_dataset('my_dataset')
+  data = torch.utils.data.Dataloader(dataset)
 
 + model, optim, data = accelerator.prepare(model, optim, data)
 
-model.train()
-for epoch in range(10):
-    for source, targets in data:
+  model.train()
+  for epoch in range(10):
+      for source, targets in data:
 -         source = source.to(device)
 -         targets = targets.to(device)
 
-        optimizer.zero_grad()
+          optimizer.zero_grad()
 
-        output = model(source, targets)
-        loss = F.cross_entropy(output, targets)
+          output = model(source, targets)
+          loss = F.cross_entropy(output, targets)
 
 +         accelerator.backward(loss)
 -         loss.backward()
 
-        optimizer.step()
+          optimizer.step()
 ```
 
 ## Launching script
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 57933888d..efdeb8f88 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -47,53 +47,58 @@ A traditional training loop in PyTorch looks like this:
 
 Changing it to work with accelerate is really easy and only adds a few lines of code:
 
-.. code-block:: python
-
-    from accelerate import Accelerator
-
-    accelerator = Accelerator()
-    # Use the device given by the `accelerator` object.
-    device = accelerator.device
-    my_model.to(device)
-    # Pass every important object (model, optimizer, dataloader) to `accelerator.prepare`
-    my_model, my_optimizer, my_training_dataloader = accelerate.prepare(
-        my_model, my_optimizer, my_training_dataloader
-    )
-
-    for batch in my_training_dataloader:
-        my_optimizer.zero_grad()
-        inputs, targets = batch
-        inputs = inputs.to(device)
-        targets = targets.to(device)
-        outputs = my_model(inputs)
-        loss = my_loss_function(outputs, targets)
-        # Just a small change for the backward instruction
-        accelerate.backward(loss)
-        my_optimizer.step()
+.. code-block:: diff
+
+    + from accelerate import Accelerator
+
+    + accelerator = Accelerator()
+      # Use the device given by the `accelerator` object.
+    + device = accelerator.device
+      my_model.to(device)
+      # Pass every important object (model, optimizer, dataloader) to `accelerator.prepare`
+    + my_model, my_optimizer, my_training_dataloader = accelerate.prepare(
+    +     my_model, my_optimizer, my_training_dataloader
+    + )
+
+      for batch in my_training_dataloader:
+          my_optimizer.zero_grad()
+          inputs, targets = batch
+          inputs = inputs.to(device)
+          targets = targets.to(device)
+          outputs = my_model(inputs)
+          loss = my_loss_function(outputs, targets)
+          # Just a small change for the backward instruction
+    -     loss.backward()
+    +     accelerate.backward(loss)
+          my_optimizer.step()
 
 and with this, your script can now run in a distributed environment (multi-GPU, TPU).
 
 You can even simplify your script a bit by letting 🤗 Accelerate handle the device placement for you (which is safer,
 especially for TPU training):
 
-.. code-block:: python
-
-    from accelerate import Accelerator
-
-    accelerator = Accelerator()
-    # Pass every important object (model, optimizer, dataloader) to `accelerator.prepare`
-    my_model, my_optimizer, my_training_dataloader = accelerate.prepare(
-        my_model, my_optimizer, my_training_dataloader
-    )
-
-    for batch in my_training_dataloader:
-        my_optimizer.zero_grad()
-        inputs, targets = batch
-        outputs = my_model(inputs)
-        loss = my_loss_function(outputs, targets)
-        # Just a small change for the backward instruction
-        accelerate.backward(loss)
-        my_optimizer.step()
+.. code-block:: diff
+
+    + from accelerate import Accelerator
+
+    + accelerator = Accelerator()
+    - my_model.to(device)
+      # Pass every important object (model, optimizer, dataloader) to `accelerator.prepare`
+    + my_model, my_optimizer, my_training_dataloader = accelerate.prepare(
+    +     my_model, my_optimizer, my_training_dataloader
+    + )
+
+      for batch in my_training_dataloader:
+          my_optimizer.zero_grad()
+          inputs, targets = batch
+    -     inputs = inputs.to(device)
+    -     targets = targets.to(device)
+          outputs = my_model(inputs)
+          loss = my_loss_function(outputs, targets)
+          # Just a small change for the backward instruction
+    -     loss.backward()
+    +     accelerate.backward(loss)
+          my_optimizer.step()
 
 
 Script launcher
diff --git a/docs/source/quicktour.rst b/docs/source/quicktour.rst
index 31c94ef81..1d56f9e4b 100644
--- a/docs/source/quicktour.rst
+++ b/docs/source/quicktour.rst
@@ -359,7 +359,7 @@ The random number generator synchronization will by default synchronize:
 - the main random number generator in PyTorch <=1.5.1
 
 You can choose which random number generator(s) to synchronize with the :obj:`rng_types` argument of the main
-:class:`~accelerate.Accelerator`. In PyTorch >= 1.6, it is recommended to rely on local :obj:`generator`s to avoid
+:class:`~accelerate.Accelerator`. In PyTorch >= 1.6, it is recommended to rely on local :obj:`generator` to avoid
 setting the same seed in the main random number generator in all processes.
 
 .. Warning::

https://api.github.com/repos/huggingface/accelerate/pulls/22
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index e67a4cee1..a8c11b887 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -198,10 +198,10 @@ def pad_across_processes(tensor, dim=0, pad_index=0, pad_first=False):
     new_tensor = tensor.new_zeros(tuple(new_size)) + pad_index
     if pad_first:
         indices = tuple(
-            slice(max_size - old_size[dim], max_size) if i == dim else slice(None) for i in range(len(size))
+            slice(max_size - old_size[dim], max_size) if i == dim else slice(None) for i in range(len(new_size))
         )
     else:
-        indices = tuple(slice(0, old_size[dim]) if i == dim else slice(None) for i in range(len(size)))
+        indices = tuple(slice(0, old_size[dim]) if i == dim else slice(None) for i in range(len(new_size)))
     new_tensor[indices] = tensor
     return new_tensor
 

https://api.github.com/repos/huggingface/accelerate/pulls/21
diff --git a/README.md b/README.md
index 219b93248..68d6866b5 100644
--- a/README.md
+++ b/README.md
@@ -52,36 +52,23 @@ limitations under the License.
 
 Here is an example:
 
-<table>
-<tr>
-<th> Original training code <br> (CPU or mono-GPU only)</th>
-<th> With Accelerate <br> (CPU/GPU/multi-GPUs/TPUs/fp16) </th>
-</tr>
-<tr>
-<td>
-
-```python
+```diff
 import torch
 import torch.nn.functional as F
 from datasets import load_dataset
 
-
-
-device = 'cpu'
++ from accelerate import Accelerator
++ accelerator = Accelerator()
+- device = 'cpu'
++ device = accelerator.device
 
 model = torch.nn.Transformer().to(device)
-optim = torch.optim.Adam(
-    model.parameters()
-)
+optim = torch.optim.Adam(model.parameters())
 
 dataset = load_dataset('my_dataset')
-data = torch.utils.data.Dataloader(
-    dataset
-)
-
-
-
+data = torch.utils.data.Dataloader(dataset)
 
++ model, optim, data = accelerator.prepare(model, optim, data)
 
 model.train()
 for epoch in range(10):
@@ -92,166 +79,55 @@ for epoch in range(10):
         optimizer.zero_grad()
 
         output = model(source, targets)
-        loss = F.cross_entropy(
-            output, targets
-        )
-
-        loss.backward()
-
-        optimizer.step()
-```
-
-</td>
-<td>
-
-```python
-  import torch
-  import torch.nn.functional as F
-  from datasets import load_dataset
-
-+ from accelerate import Accelerator
-+ accelerator = Accelerator()
-+ device = accelerator.device
-
-  model = torch.nn.Transformer().to(device)
-  optim = torch.optim.Adam(
-      model.parameters()
-  )
-
-  dataset = load_dataset('my_dataset')
-  data = torch.utils.data.Dataloader(
-      dataset
-  )
-
-+ model, optim, data = accelerator.prepare(
-+     model, optim, data
-+ )
-
-  model.train()
-  for epoch in range(10):
-      for source, targets in data:
-          source = source.to(device)
-          targets = targets.to(device)
-
-          optimizer.zero_grad()
-
-          output = model(source, targets)
-          loss = F.cross_entropy(
-              output, targets
-          )
+        loss = F.cross_entropy(output, targets)
 
 +         accelerator.backward(loss)
+-         loss.backward()
 
-          optimizer.step()
+        optimizer.step()
 ```
 
-</td>
-</tr>
-</table>
-
 As you can see in this example, by adding 5-lines to any standard PyTorch training script you can now run on any kind of single or distributed node setting (single CPU, single GPU, multi-GPUs and TPUs) as well as with or without mixed precision (fp16).
 
 In particular, the same code can then be run without modification on your local machine for debugging or your training environment.
 
 🤗 Accelerate even handles the device placement for you (which requires a few more changes to your code, but is safer in general), so you can even simplify your training loop further:
 
-<table>
-<tr>
-<th> Original training code <br> (CPU or mono-GPU only)</th>
-<th> With Accelerate <br> (CPU/GPU/multi-GPUs/TPUs/fp16) </th>
-</tr>
-<tr>
-<td>
-
-```python
+```diff
 import torch
 import torch.nn.functional as F
 from datasets import load_dataset
 
++ from accelerate import Accelerator
++ accelerator = Accelerator()
+- device = 'cpu'
 
-
-device = 'cpu'
-
-model = torch.nn.Transformer().to(device)
-optim = torch.optim.Adam(
-    model.parameters()
-)
++ model = torch.nn.Transformer()
+- model = torch.nn.Transformer().to(device)
+optim = torch.optim.Adam(model.parameters())
 
 dataset = load_dataset('my_dataset')
-data = torch.utils.data.Dataloader(
-    dataset
-)
-
-
-
+data = torch.utils.data.Dataloader(dataset)
 
++ model, optim, data = accelerator.prepare(model, optim, data)
 
 model.train()
 for epoch in range(10):
     for source, targets in data:
-        source = source.to(device)
-        targets = targets.to(device)
+-         source = source.to(device)
+-         targets = targets.to(device)
 
         optimizer.zero_grad()
 
         output = model(source, targets)
-        loss = F.cross_entropy(
-            output, targets
-        )
-
-        loss.backward()
-
-        optimizer.step()
-```
-
-</td>
-<td>
-
-```python
-  import torch
-  import torch.nn.functional as F
-  from datasets import load_dataset
-
-+ from accelerate import Accelerator
-+ accelerator = Accelerator()
--
-
-+ model = torch.nn.Transformer()
-  optim = torch.optim.Adam(
-      model.parameters()
-  )
-
-  dataset = load_dataset('my_dataset')
-  data = torch.utils.data.Dataloader(
-      dataset
-  )
-
-+ model, optim, data = accelerator.prepare(
-+     model, optim, data
-+ )
-
-  model.train()
-  for epoch in range(10):
-      for source, targets in data:
--
--
-
-          optimizer.zero_grad()
-
-          output = model(source, targets)
-          loss = F.cross_entropy(
-              output, targets
-          )
+        loss = F.cross_entropy(output, targets)
 
 +         accelerator.backward(loss)
+-         loss.backward()
 
-          optimizer.step()
+        optimizer.step()
 ```
 
-</td>
-</tr>
-</table>
-
 ## Launching script
 
 🤗 Accelerate also provides an optional CLI tool that allows you to quickly configure and test your training environment before launching the scripts. No need to remember how to use `torch.distributed.launch` or to write a specific launcher for TPU training!

https://api.github.com/repos/huggingface/accelerate/pulls/20
diff --git a/docs/source/quicktour.rst b/docs/source/quicktour.rst
index 1e3ac4e89..31c94ef81 100644
--- a/docs/source/quicktour.rst
+++ b/docs/source/quicktour.rst
@@ -130,6 +130,13 @@ do with the :meth:`~accelerate.Accelerator.gather` method.
     Any instruction using your training dataloader length (for instance if you need the number of total training steps
     to create a learning rate scheduler) should go after the call to :meth:`~accelerate.Accelerator.prepare`.
 
+.. Warning::
+
+    The :meth:`~accelerate.Accelerator.gather` method requires the tensors to be all the same size on each process. If
+    you have tensors of different sizes on each process (for instance when dynamically padding to the maximum length in
+    a batch), you should use the :meth:`~accelerate.Accelerator.pad_across_processes` method to pad you tensor to the
+    biggest size across processes.
+
 
 Launching your distributed script
 -----------------------------------------------------------------------------------------------------------------------
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index d9f7cf1bc..75539559e 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -22,7 +22,7 @@
 from .kwargs_handlers import DistributedDataParallelKwargs, GradScalerKwargs, KwargsHandler
 from .optimizer import AcceleratedOptimizer
 from .state import AcceleratorState, DistributedType
-from .utils import RNGType, extract_model_from_parallel, gather, save, wait_for_everyone
+from .utils import RNGType, extract_model_from_parallel, gather, pad_across_processes, save, wait_for_everyone
 
 
 class Accelerator:
@@ -289,6 +289,23 @@ def gather(self, tensor):
         """
         return gather(tensor)
 
+    def pad_across_processes(self, tensor, dim=0, pad_index=0, pad_first=False):
+        """
+        Recursively pad the tensors in a nested list/tuple/dictionary of tensors from all devices to the same size so
+        they can safely be gathered.
+
+        Args:
+            tensor (nested list/tuple/dictionary of :obj:`torch.Tensor`):
+                The data to gather.
+            dim (:obj:`int`, `optional`, defaults to 0):
+                The dimension on which to pad.
+            pad_index (:obj:`int`, `optional`, defaults to 0):
+                The value with which to pad.
+            pad_first (:obj:`bool`, `optional`, defaults to :obj:`False`):
+                Whether to pad at the beginning or the end.
+        """
+        return pad_across_processes(tensor, dim=dim, pad_index=pad_index, pad_first=pad_first)
+
     def unwrap_model(self, model):
         """
         Unwraps the :obj:`model` from the additional layer possible added by :meth:`~accelerate.Accelerator.prepare`.
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index f3a9ae92a..96368ba7b 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -21,7 +21,6 @@
 import subprocess
 import sys
 from pathlib import Path
-from typing import Optional
 
 from accelerate.commands.config import LaunchConfig, default_config_file
 from accelerate.state import DistributedType
@@ -68,9 +67,7 @@ def launch_command_parser(subparsers=None):
     parser.add_argument(
         "--machine_rank", type=int, default=0, help="The rank of the machine on which this script is launched."
     )
-    parser.add_argument(
-        "--main_process_ip", type=str, default=None, help="The IP address of the machine of rank 0."
-    )
+    parser.add_argument("--main_process_ip", type=str, default=None, help="The IP address of the machine of rank 0.")
     parser.add_argument(
         "--main_process_port",
         type=int,
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 61a39fc95..e67a4cee1 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -148,7 +148,7 @@ def _gpu_gather(tensor):
 
 def gather(tensor):
     """
-    Recusrively gather tensor in a nested list/tuple/dictionary of tensors from all devices.
+    Recursively gather tensor in a nested list/tuple/dictionary of tensors from all devices.
 
     Args:
         tensor (nested list/tuple/dictionary of :obj:`torch.Tensor`):
@@ -165,6 +165,47 @@ def gather(tensor):
         return tensor
 
 
+def pad_across_processes(tensor, dim=0, pad_index=0, pad_first=False):
+    """
+    Recursively pad the tensors in a nested list/tuple/dictionary of tensors from all devices to the same size so they
+    can safely be gathered.
+
+    Args:
+        tensor (nested list/tuple/dictionary of :obj:`torch.Tensor`):
+            The data to gather.
+        dim (:obj:`int`, `optional`, defaults to 0):
+            The dimension on which to pad.
+        pad_index (:obj:`int`, `optional`, defaults to 0):
+            The value with which to pad.
+        pad_first (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            Whether to pad at the beginning or the end.
+    """
+    if isinstance(tensor, (list, tuple)):
+        return type(tensor)(pad_across_processes(t, dim=dim, pad_index=pad_index) for t in tensor)
+    elif isinstance(tensor, dict):
+        return type(tensor)({k: pad_across_processes(v, dim=dim, pad_index=pad_index) for k, v in tensor.items()})
+    elif not isinstance(tensor, torch.Tensor):
+        raise TypeError(f"Can't pad the values of type {type(tensor)}, only of nested list/tuple/dicts of tensors.")
+
+    # Gather all sizes
+    size = torch.tensor(tensor.shape, device=tensor.device)[None]
+    sizes = gather(size).cpu()
+    # Then pad to the maximum size
+    max_size = max(s[dim] for s in sizes)
+    old_size = tensor.shape
+    new_size = list(old_size)
+    new_size[dim] = max_size
+    new_tensor = tensor.new_zeros(tuple(new_size)) + pad_index
+    if pad_first:
+        indices = tuple(
+            slice(max_size - old_size[dim], max_size) if i == dim else slice(None) for i in range(len(size))
+        )
+    else:
+        indices = tuple(slice(0, old_size[dim]) if i == dim else slice(None) for i in range(len(size)))
+    new_tensor[indices] = tensor
+    return new_tensor
+
+
 def wait_for_everyone():
     """
     Introduces a blocking point in the script, making sure all processes have reached this point before continuing.
diff --git a/tests/test_multigpu.py b/tests/test_multigpu.py
index ce69a3db3..7ffbd94f8 100644
--- a/tests/test_multigpu.py
+++ b/tests/test_multigpu.py
@@ -20,6 +20,7 @@
 import torch
 
 import accelerate
+from accelerate import Accelerator
 from accelerate.test_utils import execute_subprocess_async, require_multi_gpu
 
 
@@ -39,3 +40,43 @@ def test_multi_gpu(self):
         """.split()
         cmd = [sys.executable] + distributed_args
         execute_subprocess_async(cmd, env=os.environ.copy())
+
+    @require_multi_gpu
+    def test_pad_across_processes(self):
+        distributed_args = f"""
+            -m torch.distributed.launch
+            --nproc_per_node={torch.cuda.device_count()}
+            --use_env
+            {inspect.getfile(self.__class__)}
+        """.split()
+        cmd = [sys.executable] + distributed_args
+        execute_subprocess_async(cmd, env=os.environ.copy())
+
+
+if __name__ == "__main__":
+    accelerator = Accelerator()
+    shape = (accelerator.state.process_index + 2, 10)
+    tensor = torch.randint(0, 10, shape).to(accelerator.device)
+
+    error_msg = ""
+
+    tensor1 = accelerator.pad_across_processes(tensor)
+    if tensor1.shape[0] != accelerator.state.num_processes + 1:
+        error_msg += f"Found shape {tensor1.shape} but should have {accelerator.state.num_processes + 1} at dim 0."
+    if not torch.equal(tensor1[: accelerator.state.process_index + 2], tensor):
+        error_msg += "Tensors have different values."
+    if not torch.all(tensor1[accelerator.state.process_index + 2 :] == 0):
+        error_msg += "Padding was not done with the right value (0)."
+
+    tensor2 = accelerator.pad_across_processes(tensor, pad_first=True)
+    if tensor2.shape[0] != accelerator.state.num_processes + 1:
+        error_msg += f"Found shape {tensor2.shape} but should have {accelerator.state.num_processes + 1} at dim 0."
+    index = accelerator.state.num_processes - accelerator.state.process_index - 1
+    if not torch.equal(tensor2[index:], tensor):
+        error_msg += "Tensors have different values."
+    if not torch.all(tensor2[:index] == 0):
+        error_msg += "Padding was not done with the right value (0)."
+
+    # Raise error at the end to make sure we don't stop at the first failure.
+    if len(error_msg) > 0:
+        raise ValueError(error_msg)

https://api.github.com/repos/huggingface/accelerate/pulls/19
diff --git a/.github/workflows/quality.yml b/.github/workflows/quality.yml
new file mode 100644
index 000000000..9b54a269f
--- /dev/null
+++ b/.github/workflows/quality.yml
@@ -0,0 +1,17 @@
+name: Quality Check
+
+on: [pull_request]
+
+jobs:
+  quality:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v2
+    - name: Set up Python 3.6
+      uses: actions/setup-python@v2
+      with:
+        python-version: 3.6
+    - name: Install Python dependencies
+      run: pip install -e .[quality]
+    - name: Run Quality check
+      run: make quality
\ No newline at end of file
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
new file mode 100644
index 000000000..aa04063d8
--- /dev/null
+++ b/.github/workflows/test.yml
@@ -0,0 +1,17 @@
+name: Run Tests
+
+on: [pull_request]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v2
+    - name: Set up Python 3.6
+      uses: actions/setup-python@v2
+      with:
+        python-version: 3.6
+    - name: Install Python dependencies
+      run: pip install -e .[test]
+    - name: Run Tests
+      run: make test
\ No newline at end of file
diff --git a/setup.py b/setup.py
index 7a7a6a9cd..0fdc5e785 100644
--- a/setup.py
+++ b/setup.py
@@ -17,7 +17,21 @@
 
 extras = {}
 extras["quality"] = ["black >= 20.8b1", "isort >= 5.5.4", "flake8 >= 3.8.3"]
-extras["docs"] = ["recommonmark", "sphinx==3.2.1", "sphinx-markdown-tables", "sphinx-rtd-theme==0.4.3", "sphinx-copybutton"]
+extras["docs"] = [
+    "recommonmark",
+    "sphinx==3.2.1",
+    "sphinx-markdown-tables",
+    "sphinx-rtd-theme==0.4.3",
+    "sphinx-copybutton",
+]
+extras["test"] = [
+    "pytest",
+    "pytest-xdist",
+]
+
+extras["sagemaker"] = [
+    "sagemaker",  # boto3 is a required package in sagemaker
+]
 
 setup(
     name="accelerate",
@@ -32,11 +46,13 @@
     url="https://github.com/huggingface/accelerate",
     package_dir={"": "src"},
     packages=find_packages("src"),
-    entry_points={"console_scripts": [
-        "accelerate=accelerate.commands.accelerate_cli:main",
-        "accelerate-config=accelerate.commands.config:main",
-        "accelerate-launch=accelerate.commands.launch:main",
-    ]},
+    entry_points={
+        "console_scripts": [
+            "accelerate=accelerate.commands.accelerate_cli:main",
+            "accelerate-config=accelerate.commands.config:main",
+            "accelerate-launch=accelerate.commands.launch:main",
+        ]
+    },
     python_requires=">=3.6.0",
     install_requires=["torch>=1.4.0"],
     extras_require=extras,
diff --git a/src/accelerate/commands/config.py b/src/accelerate/commands/config.py
deleted file mode 100644
index c32905932..000000000
--- a/src/accelerate/commands/config.py
+++ /dev/null
@@ -1,218 +0,0 @@
-#!/usr/bin/env python
-
-# 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 argparse
-import json
-import os
-from dataclasses import dataclass
-from enum import Enum
-from typing import Optional
-
-import yaml
-from accelerate.state import DistributedType
-
-
-hf_cache_home = os.path.expanduser(
-    os.getenv("HF_HOME", os.path.join(os.getenv("XDG_CACHE_HOME", "~/.cache"), "huggingface"))
-)
-cache_dir = os.path.join(hf_cache_home, "accelerate")
-default_json_config_file = os.path.join(cache_dir, "default_config.json")
-default_yaml_config_file = os.path.join(cache_dir, "default_config.yaml")
-
-# For backward compatibility: the default config is the json one if it's the only existing file.
-if os.path.isfile(default_yaml_config_file) or not os.path.isfile(default_json_config_file):
-    default_config_file = default_yaml_config_file
-else:
-    default_config_file = default_json_config_file
-
-
-@dataclass
-class LaunchConfig:
-    distributed_type: DistributedType
-    num_processes: int
-    fp16: bool
-    machine_rank: int = 0
-    num_machines: int = 1
-    main_process_ip: Optional[str] = None
-    main_process_port: Optional[int] = None
-    main_training_function: str = "main"
-
-    def __post_init__(self):
-        if isinstance(self.distributed_type, str):
-            self.distributed_type = DistributedType(self.distributed_type)
-
-    def to_dict(self):
-        result = self.__dict__
-        # For serialization, it's best to convert Enums to strings (or their underlying value type).
-        for key, value in result.items():
-            if isinstance(value, Enum):
-                result[key] = value.value
-        return result
-
-    @classmethod
-    def from_json_file(cls, json_file=None):
-        json_file = default_json_config_file if json_file is None else json_file
-        with open(json_file, "r", encoding="utf-8") as f:
-            return cls(**json.load(f))
-
-    def to_json_file(self, json_file):
-        with open(json_file, "w", encoding="utf-8") as f:
-            content = json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n"
-            f.write(content)
-
-    @classmethod
-    def from_yaml_file(cls, yaml_file=None):
-        yaml_file = default_yaml_config_file if yaml_file is None else yaml_file
-        with open(yaml_file, "r", encoding="utf-8") as f:
-            return cls(**yaml.safe_load(f))
-
-    def to_yaml_file(self, yaml_file):
-        with open(yaml_file, "w", encoding="utf-8") as f:
-            yaml.safe_dump(self.to_dict(), f)
-
-
-def config_command_parser(subparsers=None):
-    if subparsers is not None:
-        parser = subparsers.add_parser("config")
-    else:
-        parser = argparse.ArgumentParser("Accelerate config command")
-
-    parser.add_argument(
-        "--config_file",
-        default=None,
-        help=(
-            "The path to use to store the config file. Will default to a file named default_config.json in the cache "
-            "location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have "
-            "such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed "
-            "with 'huggingface'."
-        ),
-    )
-
-    if subparsers is not None:
-        parser.set_defaults(func=config_command)
-    return parser
-
-
-def _ask_field(input_text, convert_value=None, default=None, error_message=None):
-    ask_again = True
-    while ask_again:
-        result = input(input_text)
-        try:
-            if default is not None and len(result) == 0:
-                return default
-            return convert_value(result) if convert_value is not None else result
-        except:
-            if error_message is not None:
-                print(error_message)
-
-
-def get_user_input():
-    def _convert_distributed_mode(value):
-        value = int(value)
-        return DistributedType(["NO", "MULTI_GPU", "TPU"][value])
-
-    def _convert_yes_no_to_bool(value):
-        return {"yes": True, "no": False}[value.lower()]
-
-    distributed_type = _ask_field(
-        "Which type of machine are you using? ([0] No distributed training, [1] multi-GPU, [2] TPU): ",
-        _convert_distributed_mode,
-        error_message="Please enter 0, 1 or 2.",
-    )
-
-    machine_rank = 0
-    num_machines = 1
-    main_process_ip = None
-    main_process_port = None
-    if distributed_type == DistributedType.MULTI_GPU:
-        num_machines = _ask_field(
-            "How many different machines will you use (use more than 1 for multi-node training)? [1]: ",
-            lambda x: int(x),
-            default=1,
-        )
-        if num_machines > 1:
-            machine_rank = _ask_field(
-                "What is the rank of this machine (from 0 to the number of machines - 1 )? [0]: ",
-                lambda x: int(x),
-                default=0,
-            )
-            main_process_ip = _ask_field(
-                "What is the IP address of the machine that will host the main process? ",
-            )
-            main_process_ip = _ask_field(
-                "What is the port you will use to communicate with the main process? ",
-                lambda x: int(x),
-            )
-    if distributed_type == DistributedType.TPU:
-        main_training_function = _ask_field(
-            "What is the name of the function in your script that should be launched in all parallel scripts? [main]: ",
-            default="main",
-        )
-    else:
-        main_training_function = "main"
-
-    num_processes = _ask_field(
-        "How many processes in total will you use? [1]: ",
-        lambda x: int(x),
-        default=1,
-        error_message="Please enter an integer.",
-    )
-
-    if distributed_type != DistributedType.TPU:
-        fp16 = _ask_field(
-            "Do you wish to use FP16 (mixed precision)? [yes/NO]: ",
-            _convert_yes_no_to_bool,
-            default=False,
-            error_message="Please enter yes or no.",
-        )
-    else:
-        fp16 = False
-
-    return LaunchConfig(
-        distributed_type=distributed_type,
-        num_processes=num_processes,
-        fp16=fp16,
-        machine_rank=machine_rank,
-        num_machines=num_machines,
-        main_process_ip=main_process_ip,
-        main_process_port=main_process_port,
-        main_training_function=main_training_function,
-    )
-
-
-def config_command(args):
-    config = get_user_input()
-    if args.config_file is not None:
-        config_file = args.config_file
-    else:
-        if not os.path.isdir(cache_dir):
-            os.makedirs(cache_dir)
-        config_file = default_yaml_config_file
-
-    if config_file.endswith(".json"):
-        config.to_json_file(config_file)
-    else:
-        config.to_yaml_file(config_file)
-
-
-def main():
-    parser = config_command_parser()
-    args = parser.parse_args()
-    config_command(args)
-
-
-if __name__ == "__main__":
-    main()
diff --git a/src/accelerate/commands/config/__init__.py b/src/accelerate/commands/config/__init__.py
new file mode 100644
index 000000000..48ea96ab1
--- /dev/null
+++ b/src/accelerate/commands/config/__init__.py
@@ -0,0 +1,85 @@
+#!/usr/bin/env python
+
+# 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 argparse
+import os
+
+from accelerate.state import ComputeEnvironment
+
+from .cluster import get_cluster_input
+from .config_args import cache_dir, default_config_file, default_yaml_config_file, load_config_from_file  # noqa: F401
+from .config_utils import _ask_field, _convert_compute_environment
+from .sagemaker import get_sagemaker_input
+
+
+def get_user_input():
+    compute_environment = _ask_field(
+        "In which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)): ",
+        _convert_compute_environment,
+        error_message="Please enter 0 or 1",
+    )
+    if compute_environment == ComputeEnvironment.AMAZON_SAGEMAKER:
+        config = get_sagemaker_input()
+    else:
+        config = get_cluster_input()
+    return config
+
+
+def config_command_parser(subparsers=None):
+    if subparsers is not None:
+        parser = subparsers.add_parser("config")
+    else:
+        parser = argparse.ArgumentParser("Accelerate config command")
+
+    parser.add_argument(
+        "--config_file",
+        default=None,
+        help=(
+            "The path to use to store the config file. Will default to a file named default_config.json in the cache "
+            "location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have "
+            "such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed "
+            "with 'huggingface'."
+        ),
+    )
+
+    if subparsers is not None:
+        parser.set_defaults(func=config_command)
+    return parser
+
+
+def config_command(args):
+    config = get_user_input()
+    if args.config_file is not None:
+        config_file = args.config_file
+    else:
+        if not os.path.isdir(cache_dir):
+            os.makedirs(cache_dir)
+        config_file = default_yaml_config_file
+
+    if config_file.endswith(".json"):
+        config.to_json_file(config_file)
+    else:
+        config.to_yaml_file(config_file)
+
+
+def main():
+    parser = config_command_parser()
+    args = parser.parse_args()
+    config_command(args)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
new file mode 100644
index 000000000..a9e412f7d
--- /dev/null
+++ b/src/accelerate/commands/config/cluster.py
@@ -0,0 +1,88 @@
+#!/usr/bin/env python
+
+# 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.
+
+from accelerate.state import ComputeEnvironment, DistributedType
+
+from .config_args import ClusterConfig
+from .config_utils import _ask_field, _convert_distributed_mode, _convert_yes_no_to_bool
+
+
+def get_cluster_input():
+    distributed_type = _ask_field(
+        "Which type of machine are you using? ([0] No distributed training, [1] multi-GPU, [2] TPU): ",
+        _convert_distributed_mode,
+        error_message="Please enter 0, 1 or 2.",
+    )
+
+    machine_rank = 0
+    num_machines = 1
+    main_process_ip = None
+    main_process_port = None
+    if distributed_type == DistributedType.MULTI_GPU:
+        num_machines = _ask_field(
+            "How many different machines will you use (use more than 1 for multi-node training)? [1]: ",
+            lambda x: int(x),
+            default=1,
+        )
+        if num_machines > 1:
+            machine_rank = _ask_field(
+                "What is the rank of this machine (from 0 to the number of machines - 1 )? [0]: ",
+                lambda x: int(x),
+                default=0,
+            )
+            main_process_ip = _ask_field(
+                "What is the IP address of the machine that will host the main process? ",
+            )
+            main_process_ip = _ask_field(
+                "What is the port you will use to communicate with the main process? ",
+                lambda x: int(x),
+            )
+    if distributed_type == DistributedType.TPU:
+        main_training_function = _ask_field(
+            "What is the name of the function in your script that should be launched in all parallel scripts? [main]: ",
+            default="main",
+        )
+    else:
+        main_training_function = "main"
+
+        num_processes = _ask_field(
+            "How many processes in total will you use? [1]: ",
+            lambda x: int(x),
+            default=1,
+            error_message="Please enter an integer.",
+        )
+
+        if distributed_type != DistributedType.TPU:
+            fp16 = _ask_field(
+                "Do you wish to use FP16 (mixed precision)? [yes/NO]: ",
+                _convert_yes_no_to_bool,
+                default=False,
+                error_message="Please enter yes or no.",
+            )
+        else:
+            fp16 = False
+
+    return ClusterConfig(
+        compute_environment=ComputeEnvironment.LOCAL_MACHINE,
+        distributed_type=distributed_type,
+        num_processes=num_processes,
+        fp16=fp16,
+        machine_rank=machine_rank,
+        num_machines=num_machines,
+        main_process_ip=main_process_ip,
+        main_process_port=main_process_port,
+        main_training_function=main_training_function,
+    )
diff --git a/src/accelerate/commands/config/config_args.py b/src/accelerate/commands/config/config_args.py
new file mode 100644
index 000000000..5b0dd00ba
--- /dev/null
+++ b/src/accelerate/commands/config/config_args.py
@@ -0,0 +1,119 @@
+#!/usr/bin/env python
+
+# 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 json
+import os
+from dataclasses import dataclass
+from enum import Enum
+from typing import Optional, Union
+
+import yaml
+from accelerate.state import ComputeEnvironment, DistributedType, SageMakerDistributedType
+
+
+hf_cache_home = os.path.expanduser(
+    os.getenv("HF_HOME", os.path.join(os.getenv("XDG_CACHE_HOME", "~/.cache"), "huggingface"))
+)
+cache_dir = os.path.join(hf_cache_home, "accelerate")
+default_json_config_file = os.path.join(cache_dir, "default_config.json")
+default_yaml_config_file = os.path.join(cache_dir, "default_config.yaml")
+
+# For backward compatibility: the default config is the json one if it's the only existing file.
+if os.path.isfile(default_yaml_config_file) or not os.path.isfile(default_json_config_file):
+    default_config_file = default_yaml_config_file
+else:
+    default_config_file = default_json_config_file
+
+
+def load_config_from_file(config_file):
+    config_file = config_file if config_file is not None else default_config_file
+    with open(config_file, "r", encoding="utf-8") as f:
+        if config_file.endswith(".json"):
+            if json.load(f)["compute_environment"] == ComputeEnvironment.LOCAL_MACHINE:
+                config_class = ClusterConfig
+            else:
+                config_class = SageMakerConfig
+            return config_class.from_json_file(json_file=config_file)
+        else:
+            if yaml.safe_load(f)["compute_environment"] == ComputeEnvironment.LOCAL_MACHINE:
+                config_class = ClusterConfig
+            else:
+                config_class = SageMakerConfig
+            return config_class.from_yaml_file(yaml_file=config_file)
+
+
+@dataclass
+class BaseConfig:
+    compute_environment: ComputeEnvironment
+    distributed_type: Union[DistributedType, SageMakerDistributedType]
+    fp16: bool
+
+    def to_dict(self):
+        result = self.__dict__
+        # For serialization, it's best to convert Enums to strings (or their underlying value type).
+        for key, value in result.items():
+            if isinstance(value, Enum):
+                result[key] = value.value
+        return result
+
+    @classmethod
+    def from_json_file(cls, json_file=None):
+        json_file = default_json_config_file if json_file is None else json_file
+        with open(json_file, "r", encoding="utf-8") as f:
+            return cls(**json.load(f))
+
+    def to_json_file(self, json_file):
+        with open(json_file, "w", encoding="utf-8") as f:
+            content = json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n"
+            f.write(content)
+
+    @classmethod
+    def from_yaml_file(cls, yaml_file=None):
+        yaml_file = default_yaml_config_file if yaml_file is None else yaml_file
+        with open(yaml_file, "r", encoding="utf-8") as f:
+            return cls(**yaml.safe_load(f))
+
+    def to_yaml_file(self, yaml_file):
+        with open(yaml_file, "w", encoding="utf-8") as f:
+            yaml.safe_dump(self.to_dict(), f)
+
+    def __post_init__(self):
+        if isinstance(self.compute_environment, str):
+            self.compute_environment = ComputeEnvironment(self.compute_environment)
+        if isinstance(self.distributed_type, str):
+            self.distributed_type = DistributedType(self.distributed_type)
+
+
+@dataclass
+class ClusterConfig(BaseConfig):
+    num_processes: int
+    machine_rank: int = 0
+    num_machines: int = 1
+    main_process_ip: Optional[str] = None
+    main_process_port: Optional[int] = None
+    main_training_function: str = "main"
+
+
+@dataclass
+class SageMakerConfig(BaseConfig):
+    ec2_instance_type: str
+    iam_role_name: str
+    profile: Optional[str] = None
+    region: str = "us-east-1"
+    num_machines: int = 1
+    base_job_name: str = f"accelerate-sagemaker-{num_machines}"
+    pytroch_version: str = "1.6"
+    transformers_version: str = "4.4"
diff --git a/src/accelerate/commands/config/config_utils.py b/src/accelerate/commands/config/config_utils.py
new file mode 100644
index 000000000..dd14965f7
--- /dev/null
+++ b/src/accelerate/commands/config/config_utils.py
@@ -0,0 +1,49 @@
+#!/usr/bin/env python
+
+# 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.
+
+from accelerate.state import ComputeEnvironment, DistributedType, SageMakerDistributedType
+
+
+def _ask_field(input_text, convert_value=None, default=None, error_message=None):
+    ask_again = True
+    while ask_again:
+        result = input(input_text)
+        try:
+            if default is not None and len(result) == 0:
+                return default
+            return convert_value(result) if convert_value is not None else result
+        except:
+            if error_message is not None:
+                print(error_message)
+
+
+def _convert_compute_environment(value):
+    value = int(value)
+    return ComputeEnvironment(["LOCAL_MACHINE", "AMAZON_SAGEMAKER"][value])
+
+
+def _convert_distributed_mode(value):
+    value = int(value)
+    return DistributedType(["NO", "MULTI_GPU", "TPU"][value])
+
+
+def _convert_sagemaker_distributed_mode(value):
+    value = int(value)
+    return SageMakerDistributedType(["NO", "DATA_PARALLEL", "MODEL_PARALLEL"][value])
+
+
+def _convert_yes_no_to_bool(value):
+    return {"yes": True, "no": False}[value.lower()]
diff --git a/src/accelerate/commands/config/sagemaker.py b/src/accelerate/commands/config/sagemaker.py
new file mode 100644
index 000000000..25fd3e6cb
--- /dev/null
+++ b/src/accelerate/commands/config/sagemaker.py
@@ -0,0 +1,157 @@
+#!/usr/bin/env python
+
+# 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 json
+import os
+
+from accelerate.state import ComputeEnvironment, SageMakerDistributedType
+from accelerate.utils import is_boto3_available
+
+from .config_args import SageMakerConfig
+from .config_utils import _ask_field, _convert_sagemaker_distributed_mode, _convert_yes_no_to_bool
+
+
+if is_boto3_available():
+    import boto3  # noqa: F401
+
+
+def _create_iam_role_for_sagemaker(role_name):
+    iam_client = boto3.client("iam")
+
+    sagemaker_trust_policy = {
+        "Version": "2012-10-17",
+        "Statement": [
+            {"Effect": "Allow", "Principal": {"Service": "sagemaker.amazonaws.com"}, "Action": "sts:AssumeRole"}
+        ],
+    }
+    try:
+        # create the role, associated with the chosen trust policy
+        iam_client.create_role(
+            RoleName=role_name, AssumeRolePolicyDocument=json.dumps(sagemaker_trust_policy, indent=2)
+        )
+        policy_document = {
+            "Version": "2012-10-17",
+            "Statement": [
+                {
+                    "Effect": "Allow",
+                    "Action": [
+                        "sagemaker:*",
+                        "ecr:GetDownloadUrlForLayer",
+                        "ecr:BatchGetImage",
+                        "ecr:BatchCheckLayerAvailability",
+                        "ecr:GetAuthorizationToken",
+                        "cloudwatch:PutMetricData",
+                        "cloudwatch:GetMetricData",
+                        "cloudwatch:GetMetricStatistics",
+                        "cloudwatch:ListMetrics",
+                        "logs:CreateLogGroup",
+                        "logs:CreateLogStream",
+                        "logs:DescribeLogStreams",
+                        "logs:PutLogEvents",
+                        "logs:GetLogEvents",
+                        "s3:CreateBucket",
+                        "s3:ListBucket",
+                        "s3:GetBucketLocation",
+                        "s3:GetObject",
+                        "s3:PutObject",
+                    ],
+                    "Resource": "*",
+                }
+            ],
+        }
+        # attach policy to role
+        iam_client.put_role_policy(
+            RoleName=role_name,
+            PolicyName=f"{role_name}_policy_permission",
+            PolicyDocument=json.dumps(policy_document, indent=2),
+        )
+    except iam_client.exceptions.EntityAlreadyExistsException:
+        print(f"role {role_name} already exists. Using existing one")
+
+
+def _get_iam_role_arn(role_name):
+    iam_client = boto3.client("iam")
+    return iam_client.get_role(RoleName=role_name)["Role"]["Arn"]
+
+
+def get_sagemaker_input():
+    credentials_configuration = _ask_field(
+        "How do you want to authorize? ([0] AWS Profile, [1] Credentials (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY)): ",
+        lambda x: int(x),
+    )
+    aws_profile = None
+    if credentials_configuration == 0:
+        aws_profile = _ask_field("Enter your AWS Profile name: [default] ", default="default")
+        os.environ["AWS_PROFILE"] = aws_profile
+    else:
+        print(
+            "Accelerate will expose AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment variables on the machine"
+        )
+        aws_access_key_id = _ask_field("AWS Access Key ID: ")
+        os.environ["AWS_ACCESS_KEY_ID"] = aws_access_key_id
+
+        aws_secret_access_key = _ask_field("AWS Secret Access Key: ")
+        os.environ["AWS_SECRET_ACCESS_KEY"] = aws_secret_access_key
+
+    aws_region = _ask_field("Enter your AWS Region: [us-east-1]", default="us-east-1")
+    os.environ["AWS_DEFAULT_REGION"] = aws_region
+
+    role_management = _ask_field(
+        "Do you already have an IAM Role for executing Amazon SageMaker Training Jobs? ([0] provide IAM Role name, [1] create new IAM role using credentials: ",
+        lambda x: int(x),
+    )
+    if role_management == 0:
+        iam_role_name = _ask_field("Enter your IAM role name: ")
+    else:
+        iam_role_name = "accelerate_sagemaker_execution_role"
+        print(f'Accelerate will create an iam role "{iam_role_name}" using the provided credentials')
+        _create_iam_role_for_sagemaker(iam_role_name)
+
+    distributed_type = _ask_field(
+        "Which type of machine are you using? ([0] No distributed training, [1] data parallelism, [2] model parallelism): ",
+        _convert_sagemaker_distributed_mode,
+        error_message="Please enter 0, 1 or 2",
+    )
+
+    # using the best two instances for single-gpu training or multi-gpu -> can turn into question to make it more diverse
+    ec2_instance_type = "ml.p3.2xlarge" if distributed_type == SageMakerDistributedType.NO else "ml.p3dn.24xlarge"
+    num_machines = 1
+    if (
+        distributed_type == SageMakerDistributedType.DATA_PARALLEL
+        or distributed_type == SageMakerDistributedType.MODEL_PARALLEL
+    ):
+        raise NotImplementedError("Model or Data Parallelism is not implemented yet. We are working on it")
+        num_machines = _ask_field(
+            "How many machines do you want use? [2]: ",
+            lambda x: int(x),
+            default=2,
+        )
+    fp16 = _ask_field(
+        "Do you wish to use FP16 (mixed precision)? [yes/NO]: ",
+        _convert_yes_no_to_bool,
+        default=False,
+        error_message="Please enter yes or no.",
+    )
+
+    return SageMakerConfig(
+        compute_environment=ComputeEnvironment.AMAZON_SAGEMAKER,
+        distributed_type=distributed_type,
+        ec2_instance_type=ec2_instance_type,
+        profile=aws_profile,
+        region=aws_region,
+        iam_role_name=iam_role_name,
+        fp16=fp16,
+        num_machines=num_machines,
+    )
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 061c289ea..b8e1e8156 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -23,8 +23,8 @@
 from pathlib import Path
 from typing import Optional
 
-from accelerate.commands.config import LaunchConfig, default_config_file
-from accelerate.state import DistributedType
+from accelerate.commands.config import default_config_file, load_config_from_file
+from accelerate.state import ComputeEnvironment, DistributedType
 
 
 class _AddOneArg:
@@ -171,6 +171,12 @@ def tpu_launcher(args):
         xmp.spawn(main_function, args=(), nprocs=args.num_processes)
 
 
+def sagemaker_launcher(sagemaker_config, args):
+    raise NotImplementedError(
+        "Support for starting SageMaker training is not yet implemented. But you can create configs for it."
+    )
+
+
 def launch_command(args):
     # Sanity checks
     if args.multi_gpu and args.tpu:
@@ -178,15 +184,15 @@ def launch_command(args):
 
     # Get the default from the config file.
     if args.config_file is not None or os.path.isfile(default_config_file) and not args.cpu:
-        defaults = LaunchConfig.from_json_file(json_file=args.config_file)
+        defaults = load_config_from_file(args.config_file)
         if not args.multi_gpu and not args.tpu:
             args.multi_gpu = defaults.distributed_type == DistributedType.MULTI_GPU
             args.tpu = defaults.distributed_type == DistributedType.TPU
-        if args.num_processes is None:
+        if args.num_processes is None and defaults.compute_environment == ComputeEnvironment.LOCAL_MACHINE:
             args.num_processes = defaults.num_processes
         if not args.fp16:
             args.fp16 = defaults.fp16
-        if args.main_training_function is None:
+        if args.main_training_function is None and defaults.compute_environment == ComputeEnvironment.LOCAL_MACHINE:
             args.main_training_function = defaults.main_training_function
     else:
         if args.num_processes is None:
@@ -197,6 +203,8 @@ def launch_command(args):
         multi_gpu_launcher(args)
     elif args.tpu and not args.cpu:
         tpu_launcher(args)
+    elif defaults.compute_environment == ComputeEnvironment.AMAZON_SAGEMAKER:
+        sagemaker_launcher(defaults, args)
     else:
         simple_launcher(args)
 
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index 5f150add9..55831be0f 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -53,6 +53,38 @@ class DistributedType(str, Enum):
     TPU = "TPU"
 
 
+class SageMakerDistributedType(str, Enum):
+    """
+    Represents a type of distributed environment.
+
+    Values:
+
+        - **NO** -- Not a distributed environment, just a single process.
+        - **DATA_PARALLEL** -- using sagemaker distributed data parallelism.
+        - **MODEL_PARALLEL** -- using sagemaker distributed model parallelism.
+    """
+
+    # Subclassing str as well as Enum allows the `SageMakerDistributedType` to be JSON-serializable out of the box.
+    NO = "NO"
+    DATA_PARALLEL = "DATA_PARALLEL"
+    MODEL_PARALLEL = "MODEL_PARALLEL"
+
+
+class ComputeEnvironment(str, Enum):
+    """
+    Represents a type of the compute environment.
+
+    Values:
+
+        - **LOCAL_MACHINE** -- private/custom cluster hardware.
+        - **AMAZON_SAGEMAKER** -- Amazon SageMaker as compute environment.
+    """
+
+    # Subclassing str as well as Enum allows the `ComputeEnvironment` to be JSON-serializable out of the box.
+    LOCAL_MACHINE = "LOCAL_MACHINE"
+    AMAZON_SAGEMAKER = "AMAZON_SAGEMAKER"
+
+
 # Inspired by Alex Martelli's 'Borg'.
 class AcceleratorState:
     """
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 61a39fc95..ca2d558b0 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import importlib
 import random
 from enum import Enum
 from typing import List, Optional, Union
@@ -26,6 +27,10 @@
     import torch_xla.core.xla_model as xm
 
 
+def is_boto3_available():
+    return importlib.util.find_spec("boto3") is not None
+
+
 class RNGType(Enum):
     TORCH = "torch"
     CUDA = "cuda"

decoupled input gathering from main config `get_user_input()` into the separate options
@@ -0,0 +1,88 @@
+#!/usr/bin/env python
+
+# 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.
+
+from accelerate.state import ComputeEnvironment, DistributedType
+
+from .config_args import ClusterConfig
+from .config_utils import _ask_field, _convert_distributed_mode, _convert_yes_no_to_bool
+
+
+def get_cluster_input():
Determines which config flow should be used. 
@@ -0,0 +1,85 @@
+#!/usr/bin/env python
+
+# 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 argparse
+import os
+
+from accelerate.state import ComputeEnvironment
+
+from .cluster import get_cluster_input
+from .config_args import cache_dir, default_yaml_config_file
+from .config_utils import _ask_field, _convert_compute_environment
+from .sagemaker import get_sagemaker_input
+
+
+def get_user_input():
+    compute_environment = _ask_field(
create a separate file to merge all configuration `dataclasses`. We can also think about moving them so `state.py`
@@ -0,0 +1,119 @@
+#!/usr/bin/env python
+
+# 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.
+
general loading function for determining if config file is `json` or `yaml` and which configuration is used for loading. Open for suggestions on how to identify the `config_class` easier. 
@@ -0,0 +1,119 @@
+#!/usr/bin/env python
+
+# 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 json
+import os
+from dataclasses import dataclass
+from enum import Enum
+from typing import Optional, Union
+
+import yaml
+from accelerate.state import ComputeEnvironment, DistributedType, SageMakerDistributedType
+
+
+hf_cache_home = os.path.expanduser(
+    os.getenv("HF_HOME", os.path.join(os.getenv("XDG_CACHE_HOME", "~/.cache"), "huggingface"))
+)
+cache_dir = os.path.join(hf_cache_home, "accelerate")
+default_json_config_file = os.path.join(cache_dir, "default_config.json")
+default_yaml_config_file = os.path.join(cache_dir, "default_config.yaml")
+
+# For backward compatibility: the default config is the json one if it's the only existing file.
+if os.path.isfile(default_yaml_config_file) or not os.path.isfile(default_json_config_file):
+    default_config_file = default_yaml_config_file
+else:
+    default_config_file = default_json_config_file
+
+
+def load_config_from_file(config_file):
base class contains overlapping properties and methods. 
@@ -0,0 +1,119 @@
+#!/usr/bin/env python
+
+# 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 json
+import os
+from dataclasses import dataclass
+from enum import Enum
+from typing import Optional, Union
+
+import yaml
+from accelerate.state import ComputeEnvironment, DistributedType, SageMakerDistributedType
+
+
+hf_cache_home = os.path.expanduser(
+    os.getenv("HF_HOME", os.path.join(os.getenv("XDG_CACHE_HOME", "~/.cache"), "huggingface"))
+)
+cache_dir = os.path.join(hf_cache_home, "accelerate")
+default_json_config_file = os.path.join(cache_dir, "default_config.json")
+default_yaml_config_file = os.path.join(cache_dir, "default_config.yaml")
+
+# For backward compatibility: the default config is the json one if it's the only existing file.
+if os.path.isfile(default_yaml_config_file) or not os.path.isfile(default_json_config_file):
+    default_config_file = default_yaml_config_file
+else:
+    default_config_file = default_json_config_file
+
+
+def load_config_from_file(config_file):
+    config_file = config_file if config_file is not None else default_config_file
+    with open(config_file, "r", encoding="utf-8") as f:
+        if config_file.endswith(".json"):
+            if json.load(f)["compute_environment"] == ComputeEnvironment.CUSTOM_CLUSTER:
+                config_class = ClusterConfig
+            else:
+                config_class = SageMakerConfig
+            return config_class.from_json_file(json_file=config_file)
+        else:
+            if yaml.safe_load(f)["compute_environment"] == ComputeEnvironment.CUSTOM_CLUSTER:
+                config_class = ClusterConfig
+            else:
+                config_class = SageMakerConfig
+            return config_class.from_yaml_file(yaml_file=config_file)
+
+
+@dataclass
+class BaseConfig:
kept `num_machines` in specialist class due to `dataclass` -> `Dataclass fields without default value cannot appear after data fields with default values`. 
@@ -0,0 +1,119 @@
+#!/usr/bin/env python
+
+# 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 json
+import os
+from dataclasses import dataclass
+from enum import Enum
+from typing import Optional, Union
+
+import yaml
+from accelerate.state import ComputeEnvironment, DistributedType, SageMakerDistributedType
+
+
+hf_cache_home = os.path.expanduser(
+    os.getenv("HF_HOME", os.path.join(os.getenv("XDG_CACHE_HOME", "~/.cache"), "huggingface"))
+)
+cache_dir = os.path.join(hf_cache_home, "accelerate")
+default_json_config_file = os.path.join(cache_dir, "default_config.json")
+default_yaml_config_file = os.path.join(cache_dir, "default_config.yaml")
+
+# For backward compatibility: the default config is the json one if it's the only existing file.
+if os.path.isfile(default_yaml_config_file) or not os.path.isfile(default_json_config_file):
+    default_config_file = default_yaml_config_file
+else:
+    default_config_file = default_json_config_file
+
+
+def load_config_from_file(config_file):
+    config_file = config_file if config_file is not None else default_config_file
+    with open(config_file, "r", encoding="utf-8") as f:
+        if config_file.endswith(".json"):
+            if json.load(f)["compute_environment"] == ComputeEnvironment.CUSTOM_CLUSTER:
+                config_class = ClusterConfig
+            else:
+                config_class = SageMakerConfig
+            return config_class.from_json_file(json_file=config_file)
+        else:
+            if yaml.safe_load(f)["compute_environment"] == ComputeEnvironment.CUSTOM_CLUSTER:
+                config_class = ClusterConfig
+            else:
+                config_class = SageMakerConfig
+            return config_class.from_yaml_file(yaml_file=config_file)
+
+
+@dataclass
+class BaseConfig:
+    compute_environment: ComputeEnvironment
+    distributed_type: Union[DistributedType, SageMakerDistributedType]
+    fp16: bool
+
+    def to_dict(self):
+        result = self.__dict__
+        # For serialization, it's best to convert Enums to strings (or their underlying value type).
+        for key, value in result.items():
+            if isinstance(value, Enum):
+                result[key] = value.value
+        return result
+
+    @classmethod
+    def from_json_file(cls, json_file=None):
+        json_file = default_json_config_file if json_file is None else json_file
+        with open(json_file, "r", encoding="utf-8") as f:
+            return cls(**json.load(f))
+
+    def to_json_file(self, json_file):
+        with open(json_file, "w", encoding="utf-8") as f:
+            content = json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n"
+            f.write(content)
+
+    @classmethod
+    def from_yaml_file(cls, yaml_file=None):
+        yaml_file = default_yaml_config_file if yaml_file is None else yaml_file
+        with open(yaml_file, "r", encoding="utf-8") as f:
+            return cls(**yaml.safe_load(f))
+
+    def to_yaml_file(self, yaml_file):
+        with open(yaml_file, "w", encoding="utf-8") as f:
+            yaml.safe_dump(self.to_dict(), f)
+
+    def __post_init__(self):
+        if isinstance(self.distributed_type, str):
+            self.compute_environment = ComputeEnvironment(self.compute_environment)
+        if isinstance(self.distributed_type, str):
+            self.distributed_type = DistributedType(self.distributed_type)
+
+
+@dataclass
+class ClusterConfig(BaseConfig):
+    num_processes: int
+    machine_rank: int = 0
+    num_machines: int = 1
open for better naming
@@ -53,6 +53,38 @@ class DistributedType(str, Enum):
     TPU = "TPU"
 
 
+class SageMakerDistributedType(str, Enum):
+    """
+    Represents a type of distributed environment.
+
+    Values:
+
+        - **NO** -- Not a distributed environment, just a single process.
+        - **DATA_PARALLEL** -- using sagemaker distributed data parallelism.
+        - **MODEL_PARALLEL** -- using sagemaker distributed model parallelism.
+    """
+
+    # Subclassing str as well as Enum allows the `SageMakerDistributedType` to be JSON-serializable out of the box.
+    NO = "NO"
+    DATA_PARALLEL = "DATA_PARALLEL"
+    MODEL_PARALLEL = "MODEL_PARALLEL"
+
+
+class ComputeEnvironment(str, Enum):
+    """
+    Represents a type of the compute environment.
+
+    Values:
+
+        - **CUSTOM_CLUSTER** -- private/custom cluster hardware.
+        - **AMAZON_SAGEMAKER** -- Amazon SageMaker as compute environment.
+    """
+
+    # Subclassing str as well as Enum allows the `ComputeEnvironment` to be JSON-serializable out of the box.
+    CUSTOM_CLUSTER = "CUSTOM_CLUSTER"
This is a bit weird so would remove one .config by adding `config_command_parser` in the intermediate init.
@@ -16,7 +16,7 @@
 
 from argparse import ArgumentParser
 
-from accelerate.commands.config import config_command_parser
+from accelerate.commands.config.config import config_command_parser
Custom Cluster does sound a bit too complicated for the base config.
```suggestion
        "In which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)): ",
```
@@ -0,0 +1,85 @@
+#!/usr/bin/env python
+
+# 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 argparse
+import os
+
+from accelerate.state import ComputeEnvironment
+
+from .cluster import get_cluster_input
+from .config_args import cache_dir, default_yaml_config_file
+from .config_utils import _ask_field, _convert_compute_environment
+from .sagemaker import get_sagemaker_input
+
+
+def get_user_input():
+    compute_environment = _ask_field(
+        "In which compute environment are you running? ([0] Custom Cluster, [1] AWS (Amazon SageMaker)): ",
```suggestion
        if isinstance(self.compute_environment, str):
```
@@ -0,0 +1,119 @@
+#!/usr/bin/env python
+
+# 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 json
+import os
+from dataclasses import dataclass
+from enum import Enum
+from typing import Optional, Union
+
+import yaml
+from accelerate.state import ComputeEnvironment, DistributedType, SageMakerDistributedType
+
+
+hf_cache_home = os.path.expanduser(
+    os.getenv("HF_HOME", os.path.join(os.getenv("XDG_CACHE_HOME", "~/.cache"), "huggingface"))
+)
+cache_dir = os.path.join(hf_cache_home, "accelerate")
+default_json_config_file = os.path.join(cache_dir, "default_config.json")
+default_yaml_config_file = os.path.join(cache_dir, "default_config.yaml")
+
+# For backward compatibility: the default config is the json one if it's the only existing file.
+if os.path.isfile(default_yaml_config_file) or not os.path.isfile(default_json_config_file):
+    default_config_file = default_yaml_config_file
+else:
+    default_config_file = default_json_config_file
+
+
+def load_config_from_file(config_file):
+    config_file = config_file if config_file is not None else default_config_file
+    with open(config_file, "r", encoding="utf-8") as f:
+        if config_file.endswith(".json"):
+            if json.load(f)["compute_environment"] == ComputeEnvironment.CUSTOM_CLUSTER:
+                config_class = ClusterConfig
+            else:
+                config_class = SageMakerConfig
+            return config_class.from_json_file(json_file=config_file)
+        else:
+            if yaml.safe_load(f)["compute_environment"] == ComputeEnvironment.CUSTOM_CLUSTER:
+                config_class = ClusterConfig
+            else:
+                config_class = SageMakerConfig
+            return config_class.from_yaml_file(yaml_file=config_file)
+
+
+@dataclass
+class BaseConfig:
+    compute_environment: ComputeEnvironment
+    distributed_type: Union[DistributedType, SageMakerDistributedType]
+    fp16: bool
+
+    def to_dict(self):
+        result = self.__dict__
+        # For serialization, it's best to convert Enums to strings (or their underlying value type).
+        for key, value in result.items():
+            if isinstance(value, Enum):
+                result[key] = value.value
+        return result
+
+    @classmethod
+    def from_json_file(cls, json_file=None):
+        json_file = default_json_config_file if json_file is None else json_file
+        with open(json_file, "r", encoding="utf-8") as f:
+            return cls(**json.load(f))
+
+    def to_json_file(self, json_file):
+        with open(json_file, "w", encoding="utf-8") as f:
+            content = json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n"
+            f.write(content)
+
+    @classmethod
+    def from_yaml_file(cls, yaml_file=None):
+        yaml_file = default_yaml_config_file if yaml_file is None else yaml_file
+        with open(yaml_file, "r", encoding="utf-8") as f:
+            return cls(**yaml.safe_load(f))
+
+    def to_yaml_file(self, yaml_file):
+        with open(yaml_file, "w", encoding="utf-8") as f:
+            yaml.safe_dump(self.to_dict(), f)
+
+    def __post_init__(self):
+        if isinstance(self.distributed_type, str):
Would also put those two functions in the intermediate init `accelerate.commands.config.__init__.py`
@@ -23,8 +23,8 @@
 from pathlib import Path
 from typing import Optional
 
-from accelerate.commands.config import LaunchConfig, default_config_file
-from accelerate.state import DistributedType
+from accelerate.commands.config.config_args import default_config_file, load_config_from_file
I think `LOCAL_MACHINE` is better.
@@ -53,6 +53,38 @@ class DistributedType(str, Enum):
     TPU = "TPU"
 
 
+class SageMakerDistributedType(str, Enum):
+    """
+    Represents a type of distributed environment.
+
+    Values:
+
+        - **NO** -- Not a distributed environment, just a single process.
+        - **DATA_PARALLEL** -- using sagemaker distributed data parallelism.
+        - **MODEL_PARALLEL** -- using sagemaker distributed model parallelism.
+    """
+
+    # Subclassing str as well as Enum allows the `SageMakerDistributedType` to be JSON-serializable out of the box.
+    NO = "NO"
+    DATA_PARALLEL = "DATA_PARALLEL"
+    MODEL_PARALLEL = "MODEL_PARALLEL"
+
+
+class ComputeEnvironment(str, Enum):
+    """
+    Represents a type of the compute environment.
+
+    Values:
+
+        - **CUSTOM_CLUSTER** -- private/custom cluster hardware.
+        - **AMAZON_SAGEMAKER** -- Amazon SageMaker as compute environment.
+    """
+
+    # Subclassing str as well as Enum allows the `ComputeEnvironment` to be JSON-serializable out of the box.
+    CUSTOM_CLUSTER = "CUSTOM_CLUSTER"
Would define a public `is_boto3_available()` function instead of a private variable.
@@ -25,6 +26,8 @@
 if is_tpu_available():
     import torch_xla.core.xla_model as xm
 
+_has_boto3 = importlib.util.find_spec("boto3") is not None
+
I think we can add an extra "sagemaker" with boto3 inside.
@@ -17,8 +17,17 @@
 
 extras = {}
 extras["quality"] = ["black >= 20.8b1", "isort >= 5.5.4", "flake8 >= 3.8.3"]
-extras["docs"] = ["recommonmark", "sphinx==3.2.1", "sphinx-markdown-tables", "sphinx-rtd-theme==0.4.3", "sphinx-copybutton"]
-
+extras["docs"] = [
+    "recommonmark",
+    "sphinx==3.2.1",
+    "sphinx-markdown-tables",
+    "sphinx-rtd-theme==0.4.3",
+    "sphinx-copybutton",
+]
+extras["test"] = [
+    "pytest",
+    "pytest-xdist",
+]
 setup(
I can´t reproduce why `quality` is failing. I created a complete new conda env with the same packages. 
https://api.github.com/repos/huggingface/accelerate/pulls/17
diff --git a/src/accelerate/commands/config.py b/src/accelerate/commands/config.py
index e6b5582a8..c32905932 100644
--- a/src/accelerate/commands/config.py
+++ b/src/accelerate/commands/config.py
@@ -18,8 +18,10 @@
 import json
 import os
 from dataclasses import dataclass
+from enum import Enum
 from typing import Optional
 
+import yaml
 from accelerate.state import DistributedType
 
 
@@ -27,7 +29,14 @@
     os.getenv("HF_HOME", os.path.join(os.getenv("XDG_CACHE_HOME", "~/.cache"), "huggingface"))
 )
 cache_dir = os.path.join(hf_cache_home, "accelerate")
-default_config_file = os.path.join(cache_dir, "default_config.json")
+default_json_config_file = os.path.join(cache_dir, "default_config.json")
+default_yaml_config_file = os.path.join(cache_dir, "default_config.yaml")
+
+# For backward compatibility: the default config is the json one if it's the only existing file.
+if os.path.isfile(default_yaml_config_file) or not os.path.isfile(default_json_config_file):
+    default_config_file = default_yaml_config_file
+else:
+    default_config_file = default_json_config_file
 
 
 @dataclass
@@ -41,17 +50,39 @@ class LaunchConfig:
     main_process_port: Optional[int] = None
     main_training_function: str = "main"
 
+    def __post_init__(self):
+        if isinstance(self.distributed_type, str):
+            self.distributed_type = DistributedType(self.distributed_type)
+
+    def to_dict(self):
+        result = self.__dict__
+        # For serialization, it's best to convert Enums to strings (or their underlying value type).
+        for key, value in result.items():
+            if isinstance(value, Enum):
+                result[key] = value.value
+        return result
+
     @classmethod
     def from_json_file(cls, json_file=None):
-        json_file = default_config_file if json_file is None else json_file
+        json_file = default_json_config_file if json_file is None else json_file
         with open(json_file, "r", encoding="utf-8") as f:
             return cls(**json.load(f))
 
     def to_json_file(self, json_file):
         with open(json_file, "w", encoding="utf-8") as f:
-            content = json.dumps(self.__dict__, indent=2, sort_keys=True) + "\n"
+            content = json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n"
             f.write(content)
 
+    @classmethod
+    def from_yaml_file(cls, yaml_file=None):
+        yaml_file = default_yaml_config_file if yaml_file is None else yaml_file
+        with open(yaml_file, "r", encoding="utf-8") as f:
+            return cls(**yaml.safe_load(f))
+
+    def to_yaml_file(self, yaml_file):
+        with open(yaml_file, "w", encoding="utf-8") as f:
+            yaml.safe_dump(self.to_dict(), f)
+
 
 def config_command_parser(subparsers=None):
     if subparsers is not None:
@@ -169,9 +200,12 @@ def config_command(args):
     else:
         if not os.path.isdir(cache_dir):
             os.makedirs(cache_dir)
-        config_file = default_config_file
+        config_file = default_yaml_config_file
 
-    config.to_json_file(config_file)
+    if config_file.endswith(".json"):
+        config.to_json_file(config_file)
+    else:
+        config.to_yaml_file(config_file)
 
 
 def main():

```suggestion
            yaml.safe_dump(self.to_dict(), f)
```

> safe_dump(data, stream=None) serializes the given Python object into the stream. If stream is None, it returns the produced stream. safe_dump produces only standard YAML tags and cannot represent an arbitrary Python object.
@@ -41,17 +50,39 @@ class LaunchConfig:
     main_process_port: Optional[int] = None
     main_training_function: str = "main"
 
+    def __post_init__(self):
+        if isinstance(self.distributed_type, str):
+            self.distributed_type = DistributedType(self.distributed_type)
+
+    def to_dict(self):
+        result = self.__dict__
+        # For serialization, it's best to convert Enums to strings (or their underlying value type).
+        for key, value in result.items():
+            if isinstance(value, Enum):
+                result[key] = value.value
+        return result
+
     @classmethod
     def from_json_file(cls, json_file=None):
-        json_file = default_config_file if json_file is None else json_file
+        json_file = default_json_config_file if json_file is None else json_file
         with open(json_file, "r", encoding="utf-8") as f:
             return cls(**json.load(f))
 
     def to_json_file(self, json_file):
         with open(json_file, "w", encoding="utf-8") as f:
-            content = json.dumps(self.__dict__, indent=2, sort_keys=True) + "\n"
+            content = json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n"
             f.write(content)
 
+    @classmethod
+    def from_yaml_file(cls, yaml_file=None):
+        yaml_file = default_yaml_config_file if yaml_file is None else yaml_file
+        with open(yaml_file, "r", encoding="utf-8") as f:
+            return cls(**yaml.safe_load(f))
+
+    def to_yaml_file(self, yaml_file):
+        with open(yaml_file, "w", encoding="utf-8") as f:
+            yaml.dump(self.to_dict(), f)
https://api.github.com/repos/huggingface/accelerate/pulls/16
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 065329a1d..57933888d 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -145,4 +145,5 @@ Supported integrations
     :caption: API reference
 
     accelerator
+    kwargs
     internal
diff --git a/docs/source/kwargs.rst b/docs/source/kwargs.rst
new file mode 100644
index 000000000..61f42cd11
--- /dev/null
+++ b/docs/source/kwargs.rst
@@ -0,0 +1,30 @@
+.. 
+    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.
+
+
+Kwargs Handlers
+=======================================================================================================================
+
+The following objects can be passed to the main :class:`~accelerate.Accelerator` to customize how some PyTorch objects
+related to distributed training or mixed precision are created.
+
+
+DistributedDataParallelKwargs
+-----------------------------------------------------------------------------------------------------------------------
+
+.. autoclass:: accelerate.DistributedDataParallelKwargs
+
+
+GradScalerKwargs
+-----------------------------------------------------------------------------------------------------------------------
+
+.. autoclass:: accelerate.GradScalerKwargs
diff --git a/examples/cv_example.py b/examples/cv_example.py
index 862fbce15..f4164613f 100644
--- a/examples/cv_example.py
+++ b/examples/cv_example.py
@@ -22,7 +22,6 @@
 from torch.utils.data import DataLoader, Dataset
 
 import PIL
-import torchvision
 from accelerate import Accelerator
 from timm import create_model
 from torchvision.transforms import Compose, RandomResizedCrop, Resize, ToTensor
diff --git a/examples/nlp_example.py b/examples/nlp_example.py
index 13333c4a4..aad2b8b65 100644
--- a/examples/nlp_example.py
+++ b/examples/nlp_example.py
@@ -14,6 +14,7 @@
 # limitations under the License.
 import argparse
 
+import torch
 from torch.utils.data import DataLoader
 
 from accelerate import Accelerator, DistributedType
diff --git a/src/accelerate/__init__.py b/src/accelerate/__init__.py
index 1cc471903..da4f7402e 100644
--- a/src/accelerate/__init__.py
+++ b/src/accelerate/__init__.py
@@ -5,5 +5,6 @@
 __version__ = "0.1.0"
 
 from .accelerator import Accelerator
+from .kwargs_handlers import DistributedDataParallelKwargs, GradScalerKwargs
 from .state import DistributedType
 from .utils import synchronize_rng_states
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 4707e0afe..d9f7cf1bc 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -19,6 +19,7 @@
 from packaging import version
 
 from .data_loader import prepare_data_loader
+from .kwargs_handlers import DistributedDataParallelKwargs, GradScalerKwargs, KwargsHandler
 from .optimizer import AcceleratedOptimizer
 from .state import AcceleratorState, DistributedType
 from .utils import RNGType, extract_model_from_parallel, gather, save, wait_for_everyone
@@ -56,6 +57,9 @@ class Accelerator:
 
             Will default to :obj:`["torch"]` for PyTorch versions <=1.5.1 and :obj:`["generator"]` for PyTorch versions
             >= 1.6.
+        kwargs_handlers (list of kwargs handlers, `optional`)
+            A list of :obj:`KwargHandler` to customize how the objects related to distributed training or mixed
+            precision are created. See :doc:`kwargs` for more information.
 
     Attributes
 
@@ -70,18 +74,37 @@ def __init__(
         fp16: bool = None,
         cpu: bool = False,
         rng_types: Optional[List[Union[str, RNGType]]] = None,
+        kwargs_handlers: Optional[List[KwargsHandler]] = None,
     ):
         self.state = AcceleratorState(fp16=fp16, cpu=cpu, _from_accelerator=True)
 
         self.device_placement = device_placement
         self.split_batches = split_batches
 
+        # Kwargs handlers
+        self.ddp_handler = None
+        self.scaler_handler = None
+        if kwargs_handlers is not None:
+            for handler in kwargs_handlers:
+                assert isinstance(handler, KwargsHandler), f"Unsupported kwargs handler passed: {handler}."
+                if isinstance(handler, DistributedDataParallelKwargs):
+                    if self.ddp_handler is not None:
+                        raise ValueError("You can only pass one `DistributedDataParallelKwargs` in `kwargs_handler`.")
+                    else:
+                        self.ddp_handler = handler
+                elif isinstance(handler, GradScalerKwargs):
+                    if self.scaler_handler is not None:
+                        raise ValueError("You can only pass one `GradScalerKwargs` in `kwargs_handler`.")
+                    else:
+                        self.scaler_handler = handler
+
         # Mixed precision attributes
         self.scaler = None
         self.native_amp = False
         if self.state.use_fp16:
             self.native_amp = version.parse(torch.__version__) >= version.parse("1.6")
-            self.scaler = torch.cuda.amp.GradScaler()
+            kwargs = self.scaler_handler.to_kwargs() if self.scaler_handler is not None else {}
+            self.scaler = torch.cuda.amp.GradScaler(**kwargs)
 
         # Internal references to the training objects
         self._optimizers = []
@@ -187,16 +210,18 @@ def prepare(self, *args):
                 if isinstance(obj, torch.optim.Optimizer):
                     obj._switch_parameters(mapping)
 
-        return result
+        return result if len(result) > 1 else result[0]
 
     def prepare_model(self, model):
         if self.device_placement:
             model = model.to(self.device)
         if self.distributed_type == DistributedType.MULTI_GPU:
+            kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {}
             model = torch.nn.parallel.DistributedDataParallel(
                 model,
                 device_ids=[self.local_process_index],
                 output_device=self.local_process_index,
+                **kwargs,
             )
         if self.native_amp:
             model.forward = torch.cuda.amp.autocast()(model.forward)
diff --git a/src/accelerate/kwargs_handlers.py b/src/accelerate/kwargs_handlers.py
new file mode 100644
index 000000000..7f52d6639
--- /dev/null
+++ b/src/accelerate/kwargs_handlers.py
@@ -0,0 +1,73 @@
+# 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 copy
+from dataclasses import dataclass
+
+
+class KwargsHandler:
+    """
+    Internal mixin that implements a :obj:`to_kwargs()` method for a dataclass.
+    """
+
+    def to_dict(self):
+        return copy.deepcopy(self.__dict__)
+
+    def to_kwargs(self):
+        """
+        Returns a dictionary containing the attributes with values different from the default of this class.
+        """
+        default_dict = self.__class__().to_dict()
+        this_dict = self.to_dict()
+        return {k: v for k, v in this_dict.items() if default_dict[k] != v}
+
+
+@dataclass
+class DistributedDataParallelKwargs(KwargsHandler):
+    """
+    Use this object in your :class:`~accelerate.Accelerator` to customize how your model is wrapped in a
+    :obj:`torch.nn.parallel.DistributedDataParallel`. Please refer to the documentation of this `wrapper
+    <https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html>`__ for more information
+    on each argument.
+
+    .. warning::
+
+        :obj:`gradient_as_bucket_view` is only available in PyTorch 1.7.0 and later versions.
+    """
+
+    dim: int = 0
+    broadcast_buffers: bool = True
+    bucket_cap_mb: int = 25
+    find_unused_parameters: bool = False
+    check_reduction: bool = False
+    gradient_as_bucket_view: bool = False
+
+
+@dataclass
+class GradScalerKwargs(KwargsHandler):
+    """
+    Use this object in your :class:`~accelerate.Accelerator` to customize the behavior of mixed precision, specifically
+    how the :obj:`torch.cuda.amp.GradScaler` used is created. Please refer to the documentation of this `scaler
+    <https://pytorch.org/docs/stable/amp.html?highlight=gradscaler>`__ for more information on each argument.
+
+    .. warning::
+
+        :obj:`GradScaler` is only available in PyTorch 1.5.0 and later versions.
+    """
+
+    init_scale: float = 65536.0
+    growth_factor: float = 2.0
+    backoff_factor: float = 0.5
+    growth_interval: int = 2000
+    enabled: bool = True
diff --git a/src/accelerate/optimizer.py b/src/accelerate/optimizer.py
index a8b95d2d2..55d00ab5d 100644
--- a/src/accelerate/optimizer.py
+++ b/src/accelerate/optimizer.py
@@ -64,7 +64,7 @@ def param_groups(self):
         return self.optimizer.param_groups
 
     @param_groups.setter
-    def param_group(self, param_groups):
+    def param_groups(self, param_groups):
         self.optimizer.param_groups = param_groups
 
     @property
@@ -72,7 +72,7 @@ def defaults(self):
         return self.optimizer.defaults
 
     @defaults.setter
-    def param_group(self, defaults):
+    def defaults(self, defaults):
         self.optimizer.defaults = defaults
 
     def add_param_group(self, param_group):
diff --git a/src/accelerate/test_utils/__init__.py b/src/accelerate/test_utils/__init__.py
index bb9f0f66e..03aef45b3 100644
--- a/src/accelerate/test_utils/__init__.py
+++ b/src/accelerate/test_utils/__init__.py
@@ -2,5 +2,5 @@
 # There's no way to ignore "F401 '...' imported but unused" warnings in this
 # module, but to preserve other warnings. So, don't check this module at all.
 
-from .testing import are_the_same_tensors, execute_subprocess_async, require_multi_gpu, require_tpu
+from .testing import are_the_same_tensors, execute_subprocess_async, require_cuda, require_multi_gpu, require_tpu
 from .training import RegressionDataset, RegressionModel
diff --git a/src/accelerate/test_utils/testing.py b/src/accelerate/test_utils/testing.py
index d01b5276e..507f16eb8 100644
--- a/src/accelerate/test_utils/testing.py
+++ b/src/accelerate/test_utils/testing.py
@@ -33,10 +33,19 @@ def are_the_same_tensors(tensor):
     return True
 
 
+def require_cuda(test_case):
+    """
+    Decorator marking a test that requires CUDA. These tests are skipped when there are no GPU available.
+    """
+    if not torch.cuda.is_available():
+        return unittest.skip("test requires a GPU")(test_case)
+    else:
+        return test_case
+
+
 def require_tpu(test_case):
     """
     Decorator marking a test that requires TPUs. These tests are skipped when there are no TPUs available.
-
     """
     if not is_tpu_available():
         return unittest.skip("test requires TPU")(test_case)
diff --git a/tests/test_kwargs_handlers.py b/tests/test_kwargs_handlers.py
new file mode 100644
index 000000000..aaf15de0c
--- /dev/null
+++ b/tests/test_kwargs_handlers.py
@@ -0,0 +1,96 @@
+# 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 inspect
+import os
+import sys
+import unittest
+from dataclasses import dataclass
+
+import torch
+
+from accelerate import Accelerator, DistributedDataParallelKwargs, GradScalerKwargs
+from accelerate.kwargs_handlers import KwargsHandler
+from accelerate.test_utils import execute_subprocess_async, require_cuda, require_multi_gpu
+
+
+@dataclass
+class MockClass(KwargsHandler):
+    a: int = 0
+    b: bool = False
+    c: float = 3.0
+
+
+class DataLoaderTester(unittest.TestCase):
+    def test_kwargs_handler(self):
+        # If no defaults are changed, `to_kwargs` returns an empty dict.
+        self.assertDictEqual(MockClass().to_kwargs(), {})
+        self.assertDictEqual(MockClass(a=2).to_kwargs(), {"a": 2})
+        self.assertDictEqual(MockClass(a=2, b=True).to_kwargs(), {"a": 2, "b": True})
+        self.assertDictEqual(MockClass(a=2, c=2.25).to_kwargs(), {"a": 2, "c": 2.25})
+
+    @require_cuda
+    def test_grad_scaler_kwargs(self):
+        # If no defaults are changed, `to_kwargs` returns an empty dict.
+        scaler_handler = GradScalerKwargs(init_scale=1024, growth_factor=2)
+        accelerator = Accelerator(fp16=True, kwargs_handlers=[scaler_handler])
+        print(accelerator.use_fp16)
+        scaler = accelerator.scaler
+
+        # Check the kwargs have been applied
+        self.assertEqual(scaler._init_scale, 1024.0)
+        self.assertEqual(scaler._growth_factor, 2.0)
+
+        # Check the other values are at the default
+        self.assertEqual(scaler._backoff_factor, 0.5)
+        self.assertEqual(scaler._growth_interval, 2000)
+        self.assertEqual(scaler._enabled, True)
+
+    @require_multi_gpu
+    def test_ddp_kwargs(self):
+        distributed_args = f"""
+            -m torch.distributed.launch
+            --nproc_per_node={torch.cuda.device_count()}
+            --use_env
+            {inspect.getfile(self.__class__)}
+        """.split()
+        cmd = [sys.executable] + distributed_args
+        execute_subprocess_async(cmd, env=os.environ.copy())
+
+
+if __name__ == "__main__":
+    ddp_scaler = DistributedDataParallelKwargs(bucket_cap_mb=15, find_unused_parameters=True)
+    accelerator = Accelerator(kwargs_handlers=[ddp_scaler])
+    model = torch.nn.Linear(100, 200)
+    model = accelerator.prepare(model)
+
+    # Check the values changed in kwargs
+    error_msg = ""
+    observed_bucket_cap_map = model.bucket_bytes_cap // (1024 * 1024)
+    if observed_bucket_cap_map != 15:
+        error_msg += f"Kwargs badly passed, should have `15` but found {observed_bucket_cap_map}.\n"
+    if model.find_unused_parameters is not True:
+        error_msg += f"Kwargs badly passed, should have `True` but found {model.find_unused_parameters}.\n"
+
+    # Check the values of the defaults
+    if model.dim != 0:
+        error_msg += f"Default value not respected, should have `0` but found {model.dim}.\n"
+    if model.broadcast_buffers is not True:
+        error_msg += f"Default value not respected, should have `True` but found {model.broadcast_buffers}.\n"
+    if model.gradient_as_bucket_view is not False:
+        error_msg += f"Default value not respected, should have `False` but found {model.gradient_as_bucket_view}.\n"
+
+    # Raise error at the end to make sure we don't stop at the first failure.
+    if len(error_msg) > 0:
+        raise ValueError(error_msg)

Unrelated to this PR, this is to make quality pass.
@@ -22,7 +22,6 @@
 from torch.utils.data import DataLoader, Dataset
 
 import PIL
-import torchvision
Unrelated to this PR, this is to make quality pass.
@@ -14,6 +14,7 @@
 # limitations under the License.
 import argparse
 
+import torch
Unrelated to this PR, this is to make quality pass.
@@ -64,15 +64,15 @@ def param_groups(self):
         return self.optimizer.param_groups
 
     @param_groups.setter
-    def param_group(self, param_groups):
+    def param_groups(self, param_groups):
Unrelated to this PR, this is to make quality pass.
@@ -64,15 +64,15 @@ def param_groups(self):
         return self.optimizer.param_groups
 
     @param_groups.setter
-    def param_group(self, param_groups):
+    def param_groups(self, param_groups):
         self.optimizer.param_groups = param_groups
 
     @property
     def defaults(self):
         return self.optimizer.defaults
 
     @defaults.setter
-    def param_group(self, defaults):
+    def defaults(self, defaults):
https://api.github.com/repos/huggingface/accelerate/pulls/15
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 54ef3c12f..61a39fc95 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -105,10 +105,7 @@ def send_to_device(tensor, device):
     elif isinstance(tensor, dict):
         return type(tensor)({k: send_to_device(v, device) for k, v in tensor.items()})
     elif not hasattr(tensor, "to"):
-        raise TypeError(
-            f"Can't send the values of type {type(tensor)} to device {device}, only of nested list/tuple/dicts "
-            "of tensors or objects having a `to` method."
-        )
+        return tensor
     return tensor.to(device)
 
 

https://api.github.com/repos/huggingface/accelerate/pulls/13
diff --git a/docs/source/quicktour.rst b/docs/source/quicktour.rst
index 0a86f41b9..1e3ac4e89 100644
--- a/docs/source/quicktour.rst
+++ b/docs/source/quicktour.rst
@@ -207,7 +207,8 @@ lof of time. In practice, that means you must take special care to have all your
 shape (so no dynamic padding for instance if you are in an NLP problem) and should not use layer with for loops that
 have different lengths depending on the inputs (such as an LSTM) or the training will be excruciatingly slow.
 
-To introduce special behavior in your script for TPUs you can check the :obj:`distributed_type` of your :obj:`accelerator`:
+To introduce special behavior in your script for TPUs you can check the :obj:`distributed_type` of your
+:obj:`accelerator`:
 
 .. code-block:: python
 
@@ -356,9 +357,10 @@ setting the same seed in the main random number generator in all processes.
 
 .. Warning::
 
-    Synchronization the main torch (or CUDA or XLA) random number generator will affect any other potential random artifacts you could have in your
-    dataset (like random data augmentation) in the sense all processes will get the same random numbers from the torch
-    random modules (so will apply the same random data augmentation if it's controlled by torch).
+    Synchronization the main torch (or CUDA or XLA) random number generator will affect any other potential random
+    artifacts you could have in your dataset (like random data augmentation) in the sense all processes will get the
+    same random numbers from the torch random modules (so will apply the same random data augmentation if it's
+    controlled by torch).
 
 .. Note::
 
diff --git a/examples/README.md b/examples/README.md
index 30b376cef..be636e63c 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -18,7 +18,13 @@ limitations under the License.
 
 ## Simple NLP example
 
-The [nlp_example.py](./nlp_example.py) script is a simple example to train a Bert model on a classification task ([GLUE's MRPC]()).
+The [nlp_example.py](./nlp_example.py) script is a simple example to train a Bert model on a classification task ([GLUE's MRPC](https://www.microsoft.com/en-us/download/details.aspx?id=52398)).
+
+Prior to running it you should install 🤗 Dataset and 🤗 Transformers:
+
+```bash
+pip install datasets, transformers
+```
 
 The same script can be run in any of the following configurations:
 - single CPU or single GPU
@@ -89,3 +95,91 @@ To run it in each of these various modes, use the following commands:
         ```
     * In PyTorch:
         Add an `xmp.spawn` line in your script as you usually do.
+
+
+## Simple vision example
+
+The [cv_example.py](./cv_example.py) script is a simple example to fine-tune a ResNet-50 on a classification task ([Ofxord-IIT Pet Dataset](https://www.robots.ox.ac.uk/~vgg/data/pets/)).
+
+The same script can be run in any of the following configurations:
+- single CPU or single GPU
+- multi GPUS (using PyTorch distributed mode)
+- (multi) TPUs
+- fp16 (mixed-precision) or fp32 (normal precision)
+
+Prior to running it you should install timm and torchvision:
+
+```bash
+pip install timm, torchvision
+```
+
+and you should download the data with the following commands:
+
+```bash
+wget https://www.robots.ox.ac.uk/~vgg/data/pets/data/images.tar.gz
+tar -xzf images.tar.gz
+```
+
+To run it in each of these various modes, use the following commands:
+- single CPU:
+    * from a server without GPU
+        ```bash
+        python ./cv_example.py --data_dir path_to_data
+        ```
+    * from any server by passing `cpu=True` to the `Accelerator`.
+        ```bash
+        python ./cv_example.py --data_dir path_to_data --cpu
+        ```
+    * from any server with Accelerate launcher
+        ```bash
+        accelerate launch --cpu ./cv_example.py --data_dir path_to_data
+        ```
+- single GPU:
+    ```bash
+    python ./nlp_example.py  # from a server with a GPU
+    ```
+- with fp16 (mixed-precision)
+    * from any server by passing `fp16=True` to the `Accelerator`.
+        ```bash
+        python ./cv_example.py --data_dir path_to_data --fp16
+        ```
+    * from any server with Accelerate launcher
+        ```bash
+        accelerate launch --fb16 ./cv_example.py --data_dir path_to_data
+- multi GPUS (using PyTorch distributed mode)
+    * With Accelerate config and launcher
+        ```bash
+        accelerate config  # This will create a config file on your server
+        accelerate launch ./cv_example.py --data_dir path_to_data  # This will run the script on your server
+        ```
+    * With traditional PyTorch launcher
+        ```bash
+        python -m torch.distributed.launch --nproc_per_node 2 --use_env ./cv_example.py --data_dir path_to_data
+        ```
+- multi GPUs, multi node (several machines, using PyTorch distributed mode)
+    * With Accelerate config and launcher, on each machine:
+        ```bash
+        accelerate config  # This will create a config file on each server
+        accelerate launch ./cv_example.py --data_dir path_to_data  # This will run the script on each server
+        ```
+    * With PyTorch launcher only
+        ```bash
+        python -m torch.distributed.launch --nproc_per_node 2 \
+            --use_env \
+            --node_rank 0 \
+            --master_addr master_node_ip_address \
+            ./cv_example.py --data_dir path_to_data  # On the first server
+        python -m torch.distributed.launch --nproc_per_node 2 \
+            --use_env \
+            --node_rank 1 \
+            --master_addr master_node_ip_address \
+            ./cv_example.py --data_dir path_to_data  # On the second server
+        ```
+- (multi) TPUs
+    * With Accelerate config and launcher
+        ```bash
+        accelerate config  # This will create a config file on your TPU server
+        accelerate launch ./cv_example.py --data_dir path_to_data  # This will run the script on each server
+        ```
+    * In PyTorch:
+        Add an `xmp.spawn` line in your script as you usually do.
diff --git a/examples/cv_example.py b/examples/cv_example.py
new file mode 100644
index 000000000..862fbce15
--- /dev/null
+++ b/examples/cv_example.py
@@ -0,0 +1,198 @@
+# coding=utf-8
+# Copyright 2021 The HuggingFace Inc. 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 re
+
+import numpy as np
+import torch
+from torch.optim.lr_scheduler import OneCycleLR
+from torch.utils.data import DataLoader, Dataset
+
+import PIL
+import torchvision
+from accelerate import Accelerator
+from timm import create_model
+from torchvision.transforms import Compose, RandomResizedCrop, Resize, ToTensor
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate
+#
+# This example trains a ResNet50 on the Oxford-IIT Pet Dataset
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+
+# Function to get the label from the filename
+def extract_label(fname):
+    stem = fname.split(os.path.sep)[-1]
+    return re.search(r"^(.*)_\d+\.jpg$", stem).groups()[0]
+
+
+class PetsDataset(Dataset):
+    def __init__(self, file_names, image_transform=None, label_to_id=None):
+        self.file_names = file_names
+        self.image_transform = image_transform
+        self.label_to_id = label_to_id
+
+    def __len__(self):
+        return len(self.file_names)
+
+    def __getitem__(self, idx):
+        fname = self.file_names[idx]
+        raw_image = PIL.Image.open(fname)
+        image = raw_image.convert("RGB")
+        if self.image_transform is not None:
+            image = self.image_transform(image)
+        label = extract_label(fname)
+        if self.label_to_id is not None:
+            label = self.label_to_id[label]
+        return {"image": image, "label": label}
+
+
+def training_function(config, args):
+    # Initialize accelerator
+    accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu)
+
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+    image_size = config["image_size"]
+    if not isinstance(image_size, (list, tuple)):
+        image_size = (image_size, image_size)
+
+    # Grab all the image filenames
+    file_names = [os.path.join(args.data_dir, fname) for fname in os.listdir(args.data_dir) if fname.endswith(".jpg")]
+
+    # Build the label correspondences
+    all_labels = [extract_label(fname) for fname in file_names]
+    id_to_label = list(set(all_labels))
+    id_to_label.sort()
+    label_to_id = {lbl: i for i, lbl in enumerate(id_to_label)}
+
+    # Set the seed before splitting the data.
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+    # Split our filenames between train and validation
+    random_perm = np.random.permutation(len(file_names))
+    cut = int(0.8 * len(file_names))
+    train_split = random_perm[:cut]
+    eval_split = random_perm[cut:]
+
+    # For training we use a simple RandomResizedCrop
+    train_tfm = Compose([RandomResizedCrop(image_size, scale=(0.5, 1.0)), ToTensor()])
+    train_dataset = PetsDataset(
+        [file_names[i] for i in train_split], image_transform=train_tfm, label_to_id=label_to_id
+    )
+
+    # For evaluation, we use a deterministic Resize
+    eval_tfm = Compose([Resize(image_size), ToTensor()])
+    eval_dataset = PetsDataset([file_names[i] for i in eval_split], image_transform=eval_tfm, label_to_id=label_to_id)
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size, num_workers=4)
+    eval_dataloader = DataLoader(eval_dataset, shuffle=False, batch_size=batch_size, num_workers=4)
+
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = create_model("resnet50d", pretrained=True, num_classes=len(label_to_id))
+
+    # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
+    # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
+    # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
+    model = model.to(accelerator.device)
+
+    # Freezing the base model
+    for param in model.parameters():
+        param.requires_grad = False
+    for param in model.get_classifier().parameters():
+        param.requires_grad = True
+
+    # We normalize the batches of images to be a bit faster.
+    mean = torch.tensor(model.default_cfg["mean"])[None, :, None, None].to(accelerator.device)
+    std = torch.tensor(model.default_cfg["std"])[None, :, None, None].to(accelerator.device)
+
+    # Instantiate optimizer
+    optimizer = torch.optim.Adam(params=model.parameters(), lr=lr / 25)
+
+    # Prepare everything
+    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
+    # prepare method.
+    model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader
+    )
+
+    # Instantiate learning rate scheduler after preparing the training dataloader as the prepare method
+    # may change its length.
+    lr_scheduler = OneCycleLR(optimizer=optimizer, max_lr=lr, epochs=num_epochs, steps_per_epoch=len(train_dataloader))
+
+    # Now we train the model
+    for epoch in range(num_epochs):
+        model.train()
+        for step, batch in enumerate(train_dataloader):
+            # We could avoid this line since we set the accelerator with `device_placement=True`.
+            batch = {k: v.to(accelerator.device) for k, v in batch.items()}
+            inputs = (batch["image"] - mean) / std
+            outputs = model(inputs)
+            loss = torch.nn.functional.cross_entropy(outputs, batch["label"])
+            accelerator.backward(loss)
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
+
+        model.eval()
+        accurate = 0
+        num_elems = 0
+        for step, batch in enumerate(eval_dataloader):
+            # We could avoid this line since we set the accelerator with `device_placement=True`.
+            batch = {k: v.to(accelerator.device) for k, v in batch.items()}
+            inputs = (batch["image"] - mean) / std
+            with torch.no_grad():
+                outputs = model(inputs)
+            predictions = outputs.argmax(dim=-1)
+            accurate_preds = accelerator.gather(predictions) == accelerator.gather(batch["label"])
+            num_elems += accurate_preds.shape[0]
+            accurate += accurate_preds.long().sum()
+
+        eval_metric = accurate.item() / num_elems
+        # Use accelerator.print to print only on the main process.
+        accelerator.print(f"epoch {epoch}: {100 * eval_metric:.2f}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Simple example of training script.")
+    parser.add_argument("--data_dir", required=True, help="The data folder on disk.")
+    parser.add_argument("--fp16", action="store_true", help="If passed, will use FP16 training.")
+    parser.add_argument("--cpu", action="store_true", help="If passed, will train on the CPU.")
+    args = parser.parse_args()
+    config = {"lr": 3e-2, "num_epochs": 3, "seed": 42, "batch_size": 64, "image_size": 224}
+    training_function(config, args)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/nlp_example.py b/examples/nlp_example.py
index 859309a7a..13333c4a4 100644
--- a/examples/nlp_example.py
+++ b/examples/nlp_example.py
@@ -30,7 +30,7 @@
 ########################################################################
 # This is a fully working simple example to use Accelerate
 #
-# This example train a Bert base model on GLUE MRPC
+# This example trains a Bert base model on GLUE MRPC
 # in any of the following settings (with the same script):
 #   - single CPU or single GPU
 #   - multi GPUS (using PyTorch distributed mode)
@@ -127,7 +127,7 @@ def collate_fn(examples):
         num_training_steps=len(train_dataloader) * num_epochs,
     )
 
-    # Now we train the model - We prune bad trials after each epoch if needed
+    # Now we train the model
     for epoch in range(num_epochs):
         model.train()
         for step, batch in enumerate(train_dataloader):
@@ -146,7 +146,8 @@ def collate_fn(examples):
         for step, batch in enumerate(eval_dataloader):
             # We could avoid this line since we set the accelerator with `device_placement=True`.
             batch.to(accelerator.device)
-            outputs = model(**batch)
+            with torch.no_grad():
+                outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1)
             metric.add_batch(
                 predictions=accelerator.gather(predictions),
@@ -160,16 +161,8 @@ def collate_fn(examples):
 
 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(
-        "--cpu",
-        action="store_true",
-        help="If passed, will train on the CPU.",
-    )
+    parser.add_argument("--fp16", action="store_true", help="If passed, will use FP16 training.")
+    parser.add_argument("--cpu", action="store_true", help="If passed, will train on the CPU.")
     args = parser.parse_args()
     config = {"lr": 2e-5, "num_epochs": 3, "correct_bias": True, "seed": 42, "batch_size": 16}
     training_function(config, args)
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 4bb404c39..4707e0afe 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -51,9 +51,9 @@ class Accelerator:
             - :obj:`"torch"`: the base torch random number generator
             - :obj:`"cuda"`: the CUDA random number generator (GPU only)
             - :obj:`"xla"`: the XLA random number generator (TPU only)
-            - :obj:`"generator"`: the :obj:`torch.Generator` of the sampler (or batch sampler if there is no sampler
-              in your dataloader) or of the iterable dataset (if it exists) if the underlying dataset is of that type.
-            
+            - :obj:`"generator"`: the :obj:`torch.Generator` of the sampler (or batch sampler if there is no sampler in
+              your dataloader) or of the iterable dataset (if it exists) if the underlying dataset is of that type.
+
             Will default to :obj:`["torch"]` for PyTorch versions <=1.5.1 and :obj:`["generator"]` for PyTorch versions
             >= 1.6.
 
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index fa364ef6e..061c289ea 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -27,10 +27,10 @@
 from accelerate.state import DistributedType
 
 
-class _AddOneArg():
+class _AddOneArg:
     def __init__(self, launcher):
         self.launcher = launcher
-    
+
     def __call__(self, index):
         self.launcher()
 
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index 937e4598d..24a7ad08a 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -339,8 +339,8 @@ def prepare_data_loader(
             - :obj:`"torch"`: the base torch random number generator
             - :obj:`"cuda"`: the CUDA random number generator (GPU only)
             - :obj:`"xla"`: the XLA random number generator (TPU only)
-            - :obj:`"generator"`: the :obj:`torch.Generator` of the sampler (or batch sampler if there is no sampler
-              in your dataloader) or of the iterable dataset (if it exists) if the underlying dataset is of that type.
+            - :obj:`"generator"`: the :obj:`torch.Generator` of the sampler (or batch sampler if there is no sampler in
+              your dataloader) or of the iterable dataset (if it exists) if the underlying dataset is of that type.
 
     Returns:
         :obj:`torch.utils.data.dataloader.DataLoader`: A new data loader that will yield the portion of the batches
diff --git a/src/accelerate/optimizer.py b/src/accelerate/optimizer.py
index 3a6046325..a8b95d2d2 100644
--- a/src/accelerate/optimizer.py
+++ b/src/accelerate/optimizer.py
@@ -63,6 +63,18 @@ def __init__(self, optimizer, device_placement=True, scaler=None):
     def param_groups(self):
         return self.optimizer.param_groups
 
+    @param_groups.setter
+    def param_group(self, param_groups):
+        self.optimizer.param_groups = param_groups
+
+    @property
+    def defaults(self):
+        return self.optimizer.defaults
+
+    @defaults.setter
+    def param_group(self, defaults):
+        self.optimizer.defaults = defaults
+
     def add_param_group(self, param_group):
         self.optimizer.add_param_group(param_group)
 
diff --git a/src/accelerate/test_utils/test_script.py b/src/accelerate/test_utils/test_script.py
index 62543a17b..2e5095533 100644
--- a/src/accelerate/test_utils/test_script.py
+++ b/src/accelerate/test_utils/test_script.py
@@ -14,8 +14,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from packaging import version
-
 import torch
 from torch.utils.data import DataLoader
 
@@ -24,6 +22,7 @@
 from accelerate.state import AcceleratorState, DistributedType
 from accelerate.test_utils import RegressionDataset, RegressionModel, are_the_same_tensors
 from accelerate.utils import gather, set_seed, synchronize_rng_states
+from packaging import version
 
 
 def init_state_check():
@@ -201,6 +200,7 @@ def training_check():
     assert torch.allclose(old_model.a, model.a)
     assert torch.allclose(old_model.b, model.b)
 
+
 def main():
     accelerator = Accelerator()
     state = accelerator.state
@@ -220,5 +220,6 @@ def main():
         print("\n**Training integration test**")
     training_check()
 
+
 if __name__ == "__main__":
     main()
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 4a078b475..54ef3c12f 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -61,7 +61,7 @@ def synchronize_rng_state(rng_type: Optional[RNGType] = None, generator: Optiona
     elif rng_type == RNGType.GENERATOR:
         assert generator is not None, "Need a generator to synchronize its seed."
         rng_state = generator.get_state()
-    
+
     # Broadcast the rng state from device 0 to other devices
     state = AcceleratorState()
     if state.distributed_type == DistributedType.TPU:

PyTorch optimizers are expected to have thos attributes setable, so I'm obliging.
@@ -63,6 +63,18 @@ def __init__(self, optimizer, device_placement=True, scaler=None):
     def param_groups(self):
         return self.optimizer.param_groups
 
+    @param_groups.setter
+    def param_group(self, param_groups):
+        self.optimizer.param_groups = param_groups
+
+    @property
+    def defaults(self):
+        return self.optimizer.defaults
+
+    @defaults.setter
+    def param_group(self, defaults):
+        self.optimizer.defaults = defaults
https://api.github.com/repos/huggingface/accelerate/pulls/10
diff --git a/README.md b/README.md
index 0b28b6a50..219b93248 100644
--- a/README.md
+++ b/README.md
@@ -214,7 +214,7 @@ for epoch in range(10):
 
 + from accelerate import Accelerator
 + accelerator = Accelerator()
-+ device = accelerator.device
+-
 
 + model = torch.nn.Transformer()
   optim = torch.optim.Adam(

https://api.github.com/repos/huggingface/accelerate/pulls/9
diff --git a/docs/source/internal.rst b/docs/source/internal.rst
index 4c6d895b9..2d8f0f453 100644
--- a/docs/source/internal.rst
+++ b/docs/source/internal.rst
@@ -78,6 +78,8 @@ Utilities
 
 .. autofunction:: accelerate.utils.set_seed
 
+.. autofunction:: accelerate.utils.synchronize_rng_state
+
 .. autofunction:: accelerate.utils.synchronize_rng_states
 
 .. autofunction:: accelerate.utils.wait_for_everyone
diff --git a/docs/source/quicktour.rst b/docs/source/quicktour.rst
index f74ca9eea..0a86f41b9 100644
--- a/docs/source/quicktour.rst
+++ b/docs/source/quicktour.rst
@@ -340,20 +340,29 @@ library handles the sharding of your data between processes by changing that :ob
 
 The :class:`~accelerate.data_loader.DataLoaderShard` subclasses :obj:`DataLoader` to add the following functionality:
 
-- it synchronizes the torch random number generators of all processes at each new iteration, to ensure any
+- it synchronizes the appropriate random number generator of all processes at each new iteration, to ensure any
   randomization (like shuffling) is done the exact same way across processes.
 - it puts the batches on the proper device before yielding them (unless you have opted out of
   :obj:`device_placement=True`).
 
+The random number generator synchronization will by default synchronize:
+
+- the :obj:`generator` attribute of a given sampler (like the PyTorch :obj:`RandomSampler`) for PyTorch >= 1.6
+- the main random number generator in PyTorch <=1.5.1
+
+You can choose which random number generator(s) to synchronize with the :obj:`rng_types` argument of the main
+:class:`~accelerate.Accelerator`. In PyTorch >= 1.6, it is recommended to rely on local :obj:`generator`s to avoid
+setting the same seed in the main random number generator in all processes.
+
 .. Warning::
 
-    The random number generator synchronization will affect any other potential random artifacts you could have in your
+    Synchronization the main torch (or CUDA or XLA) random number generator will affect any other potential random artifacts you could have in your
     dataset (like random data augmentation) in the sense all processes will get the same random numbers from the torch
-    random modules (so will apply the same random data augmentation if it's controlled by torch). While this is usually
-    fine, you should use the random number generator from the Python :obj:`random` module or NumPy for your data
-    augmentation if you think this will be a problem.
+    random modules (so will apply the same random data augmentation if it's controlled by torch).
+
+.. Note::
 
-    The randomization part of your sampler on the other hand should absolutely be done using the torch random number
-    generator (like in the traditional :obj:`RandomSampler`).
+    The randomization part of your custom sampler, batch sampler or iterable dataset should be done using a local
+    :obj:`torch.Generator` object (in PyTorch >= 1.6), see the traditional :obj:`RandomSampler`, as an example.
 
 See more details about the internal in the :doc:`Internals page <internal>`.
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index a2f407ba4..4bb404c39 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import List, Optional, Union
+
 import torch
 
 from packaging import version
@@ -19,7 +21,7 @@
 from .data_loader import prepare_data_loader
 from .optimizer import AcceleratedOptimizer
 from .state import AcceleratorState, DistributedType
-from .utils import extract_model_from_parallel, gather, save, wait_for_everyone
+from .utils import RNGType, extract_model_from_parallel, gather, save, wait_for_everyone
 
 
 class Accelerator:
@@ -42,6 +44,18 @@ class Accelerator:
         cpu (:obj:`bool`, `optional`):
             Whether or not to force the script to execute on CPU. Will ignore GPU available if set to :obj:`True` and
             force the execution on one process only.
+        rng_types (list of :obj:`str` or :class:`~accelerate.utils.RNGType`):
+            The list of random number generators to synchronize at the beginning of each iteration in your prepared
+            dataloaders. Should be one or several of:
+
+            - :obj:`"torch"`: the base torch random number generator
+            - :obj:`"cuda"`: the CUDA random number generator (GPU only)
+            - :obj:`"xla"`: the XLA random number generator (TPU only)
+            - :obj:`"generator"`: the :obj:`torch.Generator` of the sampler (or batch sampler if there is no sampler
+              in your dataloader) or of the iterable dataset (if it exists) if the underlying dataset is of that type.
+            
+            Will default to :obj:`["torch"]` for PyTorch versions <=1.5.1 and :obj:`["generator"]` for PyTorch versions
+            >= 1.6.
 
     Attributes
 
@@ -50,7 +64,12 @@ class Accelerator:
     """
 
     def __init__(
-        self, device_placement: bool = True, split_batches: bool = False, fp16: bool = None, cpu: bool = False
+        self,
+        device_placement: bool = True,
+        split_batches: bool = False,
+        fp16: bool = None,
+        cpu: bool = False,
+        rng_types: Optional[List[Union[str, RNGType]]] = None,
     ):
         self.state = AcceleratorState(fp16=fp16, cpu=cpu, _from_accelerator=True)
 
@@ -67,6 +86,10 @@ def __init__(
         # Internal references to the training objects
         self._optimizers = []
 
+        # RNG Types
+        if rng_types is None:
+            self.rng_types = ["torch"] if version.parse(torch.__version__) <= version.parse("1.5.1") else ["generator"]
+
     @property
     def distributed_type(self):
         return self.state.distributed_type
@@ -187,6 +210,7 @@ def prepare_data_loader(self, data_loader):
             process_index=self.process_index,
             split_batches=self.split_batches,
             put_on_device=self.device_placement,
+            rng_types=self.rng_types,
         )
 
     def prepare_optimizer(self, optimizer):
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index 92bfc3c99..937e4598d 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
+from typing import List, Optional, Union
 
 import torch
 from torch.utils.data import BatchSampler, DataLoader, IterableDataset
@@ -20,7 +20,7 @@
 from packaging import version
 
 from .state import AcceleratorState, DistributedType, is_tpu_available
-from .utils import send_to_device, synchronize_rng_states
+from .utils import RNGType, send_to_device, synchronize_rng_states
 
 
 if is_tpu_available():
@@ -262,16 +262,29 @@ class DataLoaderShard(DataLoader):
             The dataset to use to build this datalaoder.
         device (:obj:`torch.device`, `optional`):
             If passed, the device to put all batches on.
+        rng_types (list of :obj:`str` or :class:`~accelerate.utils.RNGType`):
+            The list of random number generators to synchronize at the beginning of each iteration. Should be one or
+            several of:
+
+            - :obj:`"torch"`: the base torch random number generator
+            - :obj:`"cuda"`: the CUDA random number generator (GPU only)
+            - :obj:`"xla"`: the XLA random number generator (TPU only)
+            - :obj:`"generator"`: an optional :obj:`torch.Generator`
+        generator (:obj:`torch.Generator`, `optional`):
+            A random number generator to keep synchronized accross processes.
         kwargs:
             All other keyword arguments to pass to the regular :obj:`DataLoader` initialization.
     """
 
-    def __init__(self, dataset, device=None, **kwargs):
+    def __init__(self, dataset, device=None, rng_types=None, generator=None, **kwargs):
         super().__init__(dataset, **kwargs)
         self.device = device
+        self.rng_types = rng_types
+        self.generator = generator
 
     def __iter__(self):
-        synchronize_rng_states()
+        if self.rng_types is not None:
+            synchronize_rng_states(self.rng_types, self.generator)
         state = AcceleratorState()
         for batch in super().__iter__():
             if state.distributed_type == DistributedType.TPU:
@@ -286,6 +299,7 @@ def prepare_data_loader(
     process_index: Optional[int] = None,
     split_batches: bool = False,
     put_on_device: bool = False,
+    rng_types: Optional[List[Union[str, RNGType]]] = None,
 ) -> DataLoader:
     """
     Wraps a PyTorch :obj:`DataLoader` to generate batches for one of the processes only.
@@ -318,6 +332,15 @@ def prepare_data_loader(
         put_on_device (:obj:`bool`, `optional`, defaults to :obj:`False`):
             Whether or not to put the batches on :obj:`device` (only works if the batches are nested list, tuples or
             dictionaries of tensors).
+        rng_types (list of :obj:`str` or :class:`~accelerate.utils.RNGType`):
+            The list of random number generators to synchronize at the beginning of each iteration. Should be one or
+            several of:
+
+            - :obj:`"torch"`: the base torch random number generator
+            - :obj:`"cuda"`: the CUDA random number generator (GPU only)
+            - :obj:`"xla"`: the XLA random number generator (TPU only)
+            - :obj:`"generator"`: the :obj:`torch.Generator` of the sampler (or batch sampler if there is no sampler
+              in your dataloader) or of the iterable dataset (if it exists) if the underlying dataset is of that type.
 
     Returns:
         :obj:`torch.utils.data.dataloader.DataLoader`: A new data loader that will yield the portion of the batches
@@ -342,9 +365,12 @@ def prepare_data_loader(
 
     new_dataset = dataloader.dataset
     new_batch_sampler = dataloader.batch_sampler
+    generator = getattr(dataloader, "generator", None)
     # No change if no multiprocess
     if num_processes != 1:
         if isinstance(new_dataset, IterableDataset):
+            if getattr(dataloader.dataset, "generator", None) is not None:
+                generator = dataloader.dataset.generator
             new_dataset = IterableDatasetShard(
                 new_dataset,
                 batch_size=dataloader.batch_size,
@@ -355,6 +381,13 @@ def prepare_data_loader(
             )
         else:
             # New batch sampler for the current process.
+            if hasattr(dataloader.sampler, "generator"):
+                if dataloader.sampler.generator is None:
+                    dataloader.sampler.generator = torch.Generator()
+                    generator = dataloader.sampler.generator
+                    generator.manual_seed(int(torch.empty((), dtype=torch.int64).random_().item()))
+            elif getattr(dataloader.batch_sampler, "generator", None) is not None:
+                generator = dataloader.batch_sampler.generator
             new_batch_sampler = BatchSamplerShard(
                 dataloader.batch_sampler,
                 num_processes=num_processes,
@@ -369,8 +402,12 @@ def prepare_data_loader(
         "sampler",
         "batch_sampler",
         "drop_last",
+        "generator",
     ]
 
+    if rng_types is not None and generator is None and "generator" in rng_types:
+        rng_types.remove("generator")
+
     kwargs = {
         k: getattr(dataloader, k, _PYTORCH_DATALOADER_KWARGS[k])
         for k in _PYTORCH_DATALOADER_KWARGS
@@ -380,5 +417,7 @@ def prepare_data_loader(
         new_dataset,
         device=device if put_on_device else None,
         batch_sampler=new_batch_sampler,
+        rng_types=rng_types,
+        generator=generator,
         **kwargs,
     )
diff --git a/src/accelerate/test_utils/test_script.py b/src/accelerate/test_utils/test_script.py
index e8282a319..62543a17b 100644
--- a/src/accelerate/test_utils/test_script.py
+++ b/src/accelerate/test_utils/test_script.py
@@ -14,6 +14,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from packaging import version
+
 import torch
 from torch.utils.data import DataLoader
 
@@ -34,10 +36,16 @@ def init_state_check():
 
 def rng_sync_check():
     state = AcceleratorState()
-    synchronize_rng_states()
+    synchronize_rng_states(["torch"])
     assert are_the_same_tensors(torch.get_rng_state())
     if state.distributed_type == DistributedType.MULTI_GPU:
+        synchronize_rng_states(["cuda"])
         assert are_the_same_tensors(torch.cuda.get_rng_state())
+    if version.parse(torch.__version__) >= version.parse("1.6.0"):
+        generator = torch.Generator()
+        synchronize_rng_states(["generator"], generator=generator)
+        assert are_the_same_tensors(generator.get_state())
+
     if state.local_process_index == 0:
         print("All rng are properly synched.")
 
@@ -101,13 +109,14 @@ def dl_preparation_check():
         print("Shuffled dataloader passing.")
 
 
-def mock_training(length, batch_size):
+def mock_training(length, batch_size, generator):
     set_seed(42)
+    generator.manual_seed(42)
     train_set = RegressionDataset(length=length)
-    train_dl = DataLoader(train_set, batch_size=batch_size, shuffle=True)
+    train_dl = DataLoader(train_set, batch_size=batch_size, shuffle=True, generator=generator)
     model = RegressionModel()
     optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
-    for _ in range(3):
+    for epoch in range(3):
         for batch in train_dl:
             model.zero_grad()
             output = model(batch["x"])
@@ -119,21 +128,23 @@ def mock_training(length, batch_size):
 
 def training_check():
     state = AcceleratorState()
+    generator = torch.Generator()
     batch_size = 8
     length = batch_size * 4 * state.num_processes
 
-    train_set, old_model = mock_training(length, batch_size * state.num_processes)
+    train_set, old_model = mock_training(length, batch_size * state.num_processes, generator)
     assert are_the_same_tensors(old_model.a)
     assert are_the_same_tensors(old_model.b)
 
     accelerator = Accelerator()
-    train_dl = DataLoader(train_set, batch_size=batch_size, shuffle=True)
+    train_dl = DataLoader(train_set, batch_size=batch_size, shuffle=True, generator=generator)
     model = RegressionModel()
     optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
 
     train_dl, model, optimizer = accelerator.prepare(train_dl, model, optimizer)
     set_seed(42)
-    for _ in range(3):
+    generator.manual_seed(42)
+    for epoch in range(3):
         for batch in train_dl:
             model.zero_grad()
             output = model(batch["x"])
@@ -145,15 +156,16 @@ def training_check():
     assert torch.allclose(old_model.a, model.a)
     assert torch.allclose(old_model.b, model.b)
 
-    accelerator.print("Training yielded the same results on one CPU or distributes setup with no batch split.")
+    accelerator.print("Training yielded the same results on one CPU or distributed setup with no batch split.")
 
     accelerator = Accelerator(split_batches=True)
-    train_dl = DataLoader(train_set, batch_size=batch_size * state.num_processes, shuffle=True)
+    train_dl = DataLoader(train_set, batch_size=batch_size * state.num_processes, shuffle=True, generator=generator)
     model = RegressionModel()
     optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
 
     train_dl, model, optimizer = accelerator.prepare(train_dl, model, optimizer)
     set_seed(42)
+    generator.manual_seed(42)
     for _ in range(3):
         for batch in train_dl:
             model.zero_grad()
@@ -170,12 +182,13 @@ def training_check():
 
     # Mostly a test that FP16 doesn't crash as the operation inside the model is not converted to FP16
     accelerator = Accelerator(fp16=True)
-    train_dl = DataLoader(train_set, batch_size=batch_size, shuffle=True)
+    train_dl = DataLoader(train_set, batch_size=batch_size, shuffle=True, generator=generator)
     model = RegressionModel()
     optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
 
     train_dl, model, optimizer = accelerator.prepare(train_dl, model, optimizer)
     set_seed(42)
+    generator.manual_seed(42)
     for _ in range(3):
         for batch in train_dl:
             model.zero_grad()
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 9e99d4c7b..4a078b475 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -13,6 +13,8 @@
 # limitations under the License.
 
 import random
+from enum import Enum
+from typing import List, Optional, Union
 
 import numpy as np
 import torch
@@ -24,6 +26,13 @@
     import torch_xla.core.xla_model as xm
 
 
+class RNGType(Enum):
+    TORCH = "torch"
+    CUDA = "cuda"
+    XLA = "xla"
+    GENERATOR = "generator"
+
+
 def set_seed(seed: int):
     """
     Helper function for reproducible behavior to set the seed in ``random``, ``numpy``, ``torch``.
@@ -36,27 +45,46 @@ def set_seed(seed: int):
     torch.manual_seed(seed)
     torch.cuda.manual_seed_all(seed)
     # ^^ safe to call this function even if cuda is not available
+    if is_tpu_available():
+        xm.set_rng_state(seed)
 
 
-def synchronize_rng_states():
-    """
-    Helper function to synchronize the rng states in distributed training.
-    """
+def synchronize_rng_state(rng_type: Optional[RNGType] = None, generator: Optional[torch.Generator] = None):
+    # Get the proper rng state
+    if rng_type == RNGType.TORCH:
+        rng_state = torch.get_rng_state()
+    elif rng_type == RNGType.CUDA:
+        rng_state = torch.cuda.get_rng_state()
+    elif rng_type == RNGType.XLA:
+        assert is_tpu_available(), "Can't synchronize XLA seeds on an environment without TPUs."
+        rng_state = torch.tensor(xm.get_rng_state())
+    elif rng_type == RNGType.GENERATOR:
+        assert generator is not None, "Need a generator to synchronize its seed."
+        rng_state = generator.get_state()
+    
+    # Broadcast the rng state from device 0 to other devices
     state = AcceleratorState()
     if state.distributed_type == DistributedType.TPU:
-        rng_state = torch.get_rng_state()
         rng_state = xm.mesh_reduce("random_seed", rng_state, lambda x: x[0])
-        torch.set_rng_state(rng_state)
     elif state.distributed_type == DistributedType.MULTI_GPU:
-        rng_state = torch.get_rng_state().to(state.device)
-        # Broadcast the state from process 0 to all the others.
+        rng_state = rng_state.to(state.device)
         torch.distributed.broadcast(rng_state, 0)
-        torch.set_rng_state(rng_state.cpu())
+        rng_state = rng_state.cpu()
 
-        # Broadcast the state from process 0 to all the others.
-        rng_state = torch.cuda.get_rng_state().to(state.device)
-        torch.distributed.broadcast(rng_state, 0)
-        torch.cuda.set_rng_state(rng_state.cpu())
+    # Set the broadcast rng state
+    if rng_type == RNGType.TORCH:
+        torch.set_rng_state(rng_state)
+    elif rng_type == RNGType.CUDA:
+        torch.cuda.set_rng_state(rng_state)
+    elif rng_type == RNGType.XLA:
+        xm.set_rng_state(rng_state.item())
+    elif rng_type == RNGType.GENERATOR:
+        generator.set_state(rng_state)
+
+
+def synchronize_rng_states(rng_types: List[Union[str, RNGType]], generator: Optional[torch.Generator] = None):
+    for rng_type in rng_types:
+        synchronize_rng_state(RNGType(rng_type), generator=generator)
 
 
 def send_to_device(tensor, device):

I like this!
@@ -42,6 +44,18 @@ class Accelerator:
         cpu (:obj:`bool`, `optional`):
             Whether or not to force the script to execute on CPU. Will ignore GPU available if set to :obj:`True` and
             force the execution on one process only.
+        rng_types (list of :obj:`str` or :class:`~accelerate.utils.RNGType`):
Maybe you can move this comment above if you want to keep it
@@ -35,28 +44,47 @@ def set_seed(seed: int):
     np.random.seed(seed)
     torch.manual_seed(seed)
     torch.cuda.manual_seed_all(seed)
+    if is_tpu_available():
+        xm.set_rng_state(seed)
     # ^^ safe to call this function even if cuda is not available
https://api.github.com/repos/huggingface/accelerate/pulls/8
diff --git a/README.md b/README.md
index 728976692..0b28b6a50 100644
--- a/README.md
+++ b/README.md
@@ -46,9 +46,9 @@ limitations under the License.
 
 ## Easy to integrate
 
-🤗 Accelerate was created for PyTorch users who like to write the training loop of PyTorch models but are reluctant to write and maintain the boiler code needed to use multi-GPUs/TPU/fp16.
+🤗 Accelerate was created for PyTorch users who like to write the training loop of PyTorch models but are reluctant to write and maintain the boilerplate code needed to use multi-GPUs/TPU/fp16.
 
-🤗 Accelerate abstracts exactly and only the boiler code related to multi-GPUs/TPU/fp16 and let the rest of your code unchanged.
+🤗 Accelerate abstracts exactly and only the boilerplate code related to multi-GPUs/TPU/fp16 and leaves the rest of your code unchanged.
 
 Here is an example:
 
@@ -149,11 +149,11 @@ for epoch in range(10):
 </tr>
 </table>
 
-As you can see on this example, by adding 5-lines to any standard PyTorch training script you can now run on any kind of single or distributed node setting (single CPU, single GPU, multi-GPUs and TPUs) as well as with or without mixed precision (fp16).
+As you can see in this example, by adding 5-lines to any standard PyTorch training script you can now run on any kind of single or distributed node setting (single CPU, single GPU, multi-GPUs and TPUs) as well as with or without mixed precision (fp16).
 
-The same code can then in particular run without modification on your local machine for debugging or your training environment.
+In particular, the same code can then be run without modification on your local machine for debugging or your training environment.
 
-🤗 Accelerate even handles the device placement for you (a bit more changes to your code but safer in general), so you can even simplify your training loop further:
+🤗 Accelerate even handles the device placement for you (which requires a few more changes to your code, but is safer in general), so you can even simplify your training loop further:
 
 <table>
 <tr>
@@ -254,7 +254,7 @@ for epoch in range(10):
 
 ## Launching script
 
-🤗 Accelerate also provides an optional CLI tool that allows you to quickly configure and test your training environment then launch the scripts. No need to remember how to use `torch.distributed.launch` or to write a specific launcher for TPU training!
+🤗 Accelerate also provides an optional CLI tool that allows you to quickly configure and test your training environment before launching the scripts. No need to remember how to use `torch.distributed.launch` or to write a specific launcher for TPU training!
 On your machine(s) just run:
 
 ```bash
@@ -277,7 +277,7 @@ accelerate launch examples/nlp_example.py
 
 You should use 🤗 Accelerate when you want to easily run your training scripts in a distributed environment without having to renounce full control over your training loop. This is not a high-level framework above PyTorch, just a thin wrapper so you don't have to learn a new library, In fact the whole API of 🤗 Accelerate is in one class, the `Accelerator` object.
 
-## Why shouldn't use 🤗 Accelerate?
+## Why shouldn't I use 🤗 Accelerate?
 
 You shouldn't use 🤗 Accelerate if you don't want to write a training loop yourself. There are plenty of high-level libraries above PyTorch that will offer you that, 🤗 Accelerate is not one of them.
 

https://api.github.com/repos/huggingface/accelerate/pulls/7
diff --git a/README.md b/README.md
index 794176758..728976692 100644
--- a/README.md
+++ b/README.md
@@ -140,7 +140,7 @@ for epoch in range(10):
               output, targets
           )
 
-+         accelerate.backward(loss)
++         accelerator.backward(loss)
 
           optimizer.step()
 ```
@@ -243,7 +243,7 @@ for epoch in range(10):
               output, targets
           )
 
-+         accelerate.backward(loss)
++         accelerator.backward(loss)
 
           optimizer.step()
 ```

Thanks a lot!
https://api.github.com/repos/huggingface/accelerate/pulls/6
diff --git a/examples/nlp_example.py b/examples/nlp_example.py
index e20426857..859309a7a 100644
--- a/examples/nlp_example.py
+++ b/examples/nlp_example.py
@@ -49,6 +49,9 @@
 
 
 def training_function(config, args):
+    # Initialize accelerator
+    accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu)
+
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
     lr = config["lr"]
     num_epochs = int(config["num_epochs"])
@@ -97,14 +100,13 @@ def collate_fn(examples):
     )
 
     set_seed(seed)
-    # Initialize accelerator
-    accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu)
 
     # Instantiate the model (we build the model here so that the seed also control new weights initialization)
     model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
-    # We could avoid this line since we set the accelerator with `device_placement=True`.
-    # If setting devices manually, this line absolutely needs to be before the optimizer creation otherwise training
-    # will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
+
+    # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
+    # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
+    # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
     model = model.to(accelerator.device)
 
     # Instantiate optimizer

Ah thanks a lot for fixing!
https://api.github.com/repos/huggingface/accelerate/pulls/4
diff --git a/README.md b/README.md
index 3776c7cc2..a9d6c5704 100644
--- a/README.md
+++ b/README.md
@@ -270,7 +270,7 @@ accelerate launch my_script.py --args_to_my_script
 For instance, here is how you would run the GLUE example on the MRPC task (from the root of the repo):
 
 ```bash
-accelerate launch examples/glue_example.py --task_name mrpc --model_name_or_path bert-base-cased
+accelerate launch examples/nlp_example.py
 ```
 
 ## Why should I use 🤗 Accelerate?

```suggestion
accelerate launch examples/nlp_example.py
```
This one does not take any arg however.
@@ -270,7 +270,7 @@ accelerate launch my_script.py --args_to_my_script
 For instance, here is how you would run the GLUE example on the MRPC task (from the root of the repo):
 
 ```bash
-accelerate launch examples/glue_example.py --task_name mrpc --model_name_or_path bert-base-cased
+accelerate launch examples/nlp_example.py --task_name mrpc --model_name_or_path bert-base-cased
Thanks a lot!
https://api.github.com/repos/huggingface/accelerate/pulls/3
diff --git a/README.md b/README.md
index 3776c7cc2..3fa28ba40 100644
--- a/README.md
+++ b/README.md
@@ -254,7 +254,7 @@ for epoch in range(10):
 
 ## Launching script
 
-🤗 Accelerate also provides a CLI tool that allows you to quickly configure and test your training environment then launch the scripts. No need to remember how to use `torch.distributed.launch` or to write a specific launcher for TPU training!
+🤗 Accelerate also provides an optional CLI tool that allows you to quickly configure and test your training environment then launch the scripts. No need to remember how to use `torch.distributed.launch` or to write a specific launcher for TPU training!
 On your machine(s) just run:
 
 ```bash
diff --git a/examples/nlp_example.py b/examples/nlp_example.py
index 32ff600a9..e20426857 100644
--- a/examples/nlp_example.py
+++ b/examples/nlp_example.py
@@ -39,7 +39,7 @@
 #
 # To run it in each of these various modes, follow the instructions
 # in the readme for examples:
-# https://github.com/huggingface/accelerate/examples
+# https://github.com/huggingface/accelerate/tree/main/examples
 #
 ########################################################################
 

Thanks!
https://api.github.com/repos/huggingface/accelerate/pulls/2
diff --git a/.gitignore b/.gitignore
index b6e47617d..3b89c957d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -127,3 +127,6 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+# VSCode
+.vscode
diff --git a/README.md b/README.md
index 781eea3c9..fcc20a12a 100644
--- a/README.md
+++ b/README.md
@@ -34,7 +34,98 @@
 <p>Run your *raw* PyTorch training script on any kind of device
 </h3>
 
-🤗 Accelerate provides an easy API to make your scripts run with mixed precision and on any kind of distributed setting (multi-GPUs, TPUs etc.) while still letting you write your own training loop. The same code can then runs seamlessly on your local machine for debugging or your training environment.
+🤗 Accelerate was created for PyTorch users who like to write the training loop of PyTorch models but are reluctant to write and maintain the boiler code needed to use multi-GPUs/TPU/fp16.
+
+🤗 Accelerate abstracts exactly and only the boiler code related to multi-GPUs/TPU/fp16 and let the rest of your code unchanged.
+
+Here is an example:
+
+<table>
+<tr>
+<th> Original training code (CPU or mono-GPU only)</th>
+<th> With Accelerate for CPU/GPU/multi-GPUs/TPUs/fp16 </th>
+</tr>
+<tr>
+<td>
+
+```python
+import torch
+import torch.nn.functional as F
+from datasets import load_dataset
+
+
+
+device = 'cpu'
+
+model = torch.nn.Transformer().to(device)
+optim = torch.optim.Adam(model.parameters())
+
+dataset = load_dataset('my_dataset')
+data = torch.utils.data.Dataloader(dataset)
+
+
+
+
+model.train()
+for epoch in range(10):
+    for source, targets in data:
+        source = source.to(device)
+        targets = targets.to(device)
+
+        optimizer.zero_grad()
+
+        output = model(source, targets)
+        loss = F.cross_entropy(output, targets)
+
+        loss.backward()
+
+        optimizer.step()
+```
+
+</td>
+<td>
+
+```python
+  import torch
+  import torch.nn.functional as F
+  from datasets import load_dataset
+
++ from accelerate import Accelerator
++ accelerator = Accelerator()
++ device = accelerator.device
+
+  model = torch.nn.Transformer().to(device)
+  optim = torch.optim.Adam(model.parameters())
+
+  dataset = load_dataset('my_dataset')
+  data = torch.utils.data.Dataloader(dataset)
+
++ model, optim, data = accelerator.prepare(
+                            model, optim, data)
+
+  model.train()
+  for epoch in range(10):
+      for source, targets in data:
+          source = source.to(device)
+          targets = targets.to(device)
+
+          optimizer.zero_grad()
+
+          output = model(source, targets)
+          loss = F.cross_entropy(output, targets)
+
++         accelerate.backward(loss)
+
+          optimizer.step()
+```
+
+</td>
+</tr>
+</table>
+
+As you can see on this example, by adding 5-lines to any standard PyTorch training script you can now run on any kind of single or distributed node setting (single CPU, single GPU, multi-GPUs and TPUs) as well as with or without mixed precision (fp16).
+
+The same code can then in paticular run without modification on your local machine for debugging or your training environment.
 
 🤗 Accelerate also provides a CLI tool that allows you to quickly configure and test your training environment then launch the scripts.
 
diff --git a/examples/README.md b/examples/README.md
new file mode 100644
index 000000000..cb1d2645b
--- /dev/null
+++ b/examples/README.md
@@ -0,0 +1,75 @@
+# In this folder we showcase various full examples using `Accelerate`
+
+## Simple NLP example
+
+The [simple_example.py](./simple_example.py) script is a simple example to train a Bert model on a classification task ([GLUE's MRPC]()).
+
+The same script can be run in any of the following configurations:
+- single CPU or single GPU
+- multi GPUS (using PyTorch distributed mode)
+- (multi) TPUs
+- fp16 (mixed-precision) or fp32 (normal precision)
+
+To run it in each of these various modes, use the following commands:
+- single CPU:
+    * from a server without GPU
+        ```bash
+        python ./simple_example.py
+        ```
+    * from any server by passing `cpu=True` to the `Accelerator`.
+        ```bash
+        python ./simple_example.py --cpu
+        ```
+    * from any server with Accelerate launcher
+        ```bash
+        accelerate launch --cpu ./simple_example.py
+        ```
+- single GPU:
+    ```bash
+    python ./simple_example.py  # from a server with a GPU
+    ```
+- with fp16 (mixed-precision)
+    * from any server by passing `fp16=True` to the `Accelerator`.
+        ```bash
+        python ./simple_example.py --fp16
+        ```
+    * from any server with Accelerate launcher
+        ```bash
+        accelerate launch --fb16 ./simple_example.py
+- multi GPUS (using PyTorch distributed mode)
+    * With Accelerate config and launcher
+        ```bash
+        accelerate config  # This will create a config file on your server
+        accelerate launch ./simple_example.py  # This will run the script on your server
+        ```
+    * With traditional PyTorch launcher
+        ```bash
+        python -m torch.distributed.launch --nproc_per_node 2 --use_env ./simple_example.py
+        ```
+- multi GPUs, multi node (several machines, using PyTorch distributed mode)
+    * With Accelerate config and launcher, on each machine:
+        ```bash
+        accelerate config  # This will create a config file on each server
+        accelerate launch ./simple_example.py  # This will run the script on each server
+        ```
+    * With PyTorch launcher only
+        ```bash
+        python -m torch.distributed.launch --nproc_per_node 2 \
+            --use_env \
+            --node_rank 0 \
+            --master_addr master_node_ip_address \
+            ./simple_example.py  # On the first server
+        python -m torch.distributed.launch --nproc_per_node 2 \
+            --use_env \
+            --node_rank 1 \
+            --master_addr master_node_ip_address \
+            ./simple_example.py  # On the second server
+        ```
+- (multi) TPUs
+    * With Accelerate config and launcher
+        ```bash
+        accelerate config  # This will create a config file on your TPU server
+        accelerate launch ./simple_example.py  # This will run the script on each server
+        ```
+    * In PyTorch:
+        Add an `xmp.spawn` line in your script as you usually do.
diff --git a/examples/simple_example.py b/examples/simple_example.py
index dbbfd22f8..88c81e3a0 100644
--- a/examples/simple_example.py
+++ b/examples/simple_example.py
@@ -11,11 +11,28 @@
 )
 
 
+########################################################################
+# This is a fully working simple example to use Accelerate
+#
+# This example train a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/examples
+#
+########################################################################
+
+
 MAX_GPU_BATCH_SIZE = 16
 EVAL_BATCH_SIZE = 32
 
 
-def training_function(config):
+def training_function(config, args):
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
     lr = config["lr"]
     num_epochs = int(config["num_epochs"])
@@ -62,7 +79,7 @@ def collate_fn(examples):
 
     set_seed(seed)
     # Initialize accelerator
-    accelerator = Accelerator(device_placement=False)
+    accelerator = Accelerator(device_placement=False, fp16=args.fp16, cpu=args.cpu)
 
     # Instantiate the model (we build the model here so that the seed also control new weights initialization)
     model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
@@ -117,8 +134,20 @@ def collate_fn(examples):
 
 
 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(
+        "--cpu",
+        action="store_true",
+        help="If passed, will train on the CPU.",
+    )
+    args = parser.parse_args()
     config = {"lr": 2e-5, "num_epochs": 3, "correct_bias": True, "seed": 42, "batch_size": 16}
-    training_function(config)
+    training_function(config, args)
 
 
 if __name__ == "__main__":
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 96e3b2c86..a18f31372 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -14,7 +14,7 @@ class Accelerator:
 
     Args:
         device_placement (:obj:`bool`, `optional`, defaults to :obj:`True`):
-            Whether or not the accelerator should put objects on device (tensors yielded by the datalaoder, model,
+            Whether or not the accelerator should put objects on device (tensors yielded by the dataloader, model,
             etc...).
         split_batches (:obj:`bool`, `optional`, defaults to :obj:`False`):
             Whether or not the accelerator should split the batches yielded by the dataloaders across the devices. If
@@ -38,6 +38,7 @@ def __init__(
         self, device_placement: bool = True, split_batches: bool = False, fp16: bool = None, cpu: bool = False
     ):
         self.state = AcceleratorState(fp16=fp16, cpu=cpu, _from_accelerator=True)
+
         self.device_placement = device_placement
         self.split_batches = split_batches
 
@@ -99,6 +100,13 @@ def prepare(self, *args):
         """
         Prepare all objects passed in :obj:`args` for distributed training and mixed precision, then return them in the
         same order.
+
+        Accepts the following type of objects:
+
+            - :obj:`torch.utils.data.DataLoader`: PyTorch Dataloader
+            - :obj:`torch.nn.Module`: PyTorch Module
+            - :obj:`torch.optim.Optimizer`: PyTorch Optimizer
+
         """
         # On TPUs, putting the model on the XLA device will create new parameters, so the corresponding optimizer will
         # have parameters disconnected from the model (so no training :-( ).
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index d559d0b0b..601cba147 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -32,6 +32,9 @@ def launch_command_parser(subparsers=None):
     parser.add_argument(
         "--fp16", default=False, action="store_true", help="Whether or not to use mixed precision training."
     )
+    parser.add_argument(
+        "--cpu", default=False, action="store_true", help="Whether or not to force the training on the CPU."
+    )
     parser.add_argument(
         "--num_processes", type=int, default=None, help="The total number of processes to be launched in parallel."
     )
@@ -77,6 +80,7 @@ def simple_launcher(args):
     cmd.extend(args.training_script_args)
 
     current_env = os.environ.copy()
+    current_env["USE_CPU"] = str(args.cpu)
     current_env["USE_FP16"] = str(args.fp16)
 
     process = subprocess.Popen(cmd, env=current_env)
@@ -143,7 +147,7 @@ def launch_command(args):
         raise ValueError("You can only pick one between `--multi_gpu` and `--tpu`.")
 
     # Get the default from the config file.
-    if args.config_file is not None or os.path.isfile(default_config_file):
+    if args.config_file is not None or os.path.isfile(default_config_file) and not args.cpu:
         defaults = LaunchConfig.from_json_file(json_file=args.config_file)
         if not args.multi_gpu and not args.tpu:
             args.multi_gpu = defaults.distributed_type == DistributedType.MULTI_GPU
@@ -157,9 +161,9 @@ def launch_command(args):
             args.num_processes = 1
 
     # Use the proper launcher
-    if args.multi_gpu:
+    if args.multi_gpu and not args.cpu:
         multi_gpu_launcher(args)
-    elif args.tpu:
+    elif args.tpu and not args.cpu:
         tpu_launcher(args)
     else:
         simple_launcher(args)

```suggestion
+ accelerator = Accelerator(device_placement=False)
```
The default is still `True` as it's safer for the user (and it still takes one line of code to have the default to `False`). If the table has space enough, I'd add the same example with the device placement handled by `accelerate` (like the current example below).
@@ -34,7 +34,98 @@
 <p>Run your *raw* PyTorch training script on any kind of device
 </h3>
 
-🤗 Accelerate provides an easy API to make your scripts run with mixed precision and on any kind of distributed setting (multi-GPUs, TPUs etc.) while still letting you write your own training loop. The same code can then run on your local machine for debugging or your training environment.
+🤗 Accelerate was created for PyTorch users who like to write the training loop of PyTorch models but are reluctant to write and maintain the boiler code needed to use multi-GPUs/TPU/fp16.
+
+🤗 Accelerate abstracts exactly and only the boiler code related to multi-GPUs/TPU/fp16 and let the rest of your code unchanged.
+
+Here is an example:
+
+<table>
+<tr>
+<th> Original training code (CPU or mono-GPU only)</th>
+<th> With Accelerate for CPU/GPU/multi-GPUs/TPUs/fp16 </th>
+</tr>
+<tr>
+<td>
+
+```python
+import torch
+import torch.nn.functional as F
+from datasets import load_dataset
+
+
+
+device = 'cpu'
+
+model = torch.nn.Transformer().to(device)
+optim = torch.optim.Adam(model.parameters())
+
+dataset = load_dataset('my_dataset')
+data = torch.utils.data.Dataloader(dataset)
+
+
+
+
+model.train()
+for epoch in range(10):
+    for source, targets in data:
+        source = source.to(device)
+        targets = targets.to(device)
+
+        optimizer.zero_grad()
+
+        output = model(source, targets)
+        loss = F.cross_entropy(output, targets)
+
+        loss.backward()
+
+        optimizer.step()
+```
+
+</td>
+<td>
+
+```python
+  import torch
+  import torch.nn.functional as F
+  from datasets import load_dataset
+
++ from accelerate import Accelerator
++ accelerator = Accelerator()
https://api.github.com/repos/huggingface/accelerate/pulls/1