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
+        }


Review comments: body: ```suggestion
            assert isinstance(deepspeed_plugin, DeepSpeedPlugin), "`deepspeed_plugin` must be a DeepSpeedPlugin object."
```
diff_hunk: @@ -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"
body: ```suggestion
        assert len(batch_size) > 0, "You must specify a training or evaluation dataloader in `accelerate.prepare()` when using DeepSpeed."
```
diff_hunk: @@ -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"
body: Rename to batch_sizes (since there are several of them)?
diff_hunk: @@ -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")]
body: ```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!
diff_hunk: @@ -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`")
body: Nothing needs to be done for the `eval_batch_size`?
diff_hunk: @@ -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
+
body: I don't understand this part, could you explain?
diff_hunk: @@ -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
body: The error message should be more along the lines of "You need to pass the model along the optimizer when using Deepspeed."
diff_hunk: @@ -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"
body: I'm not sure I understand this error message.
diff_hunk: @@ -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"
body: Great to have this handled here :-) 
diff_hunk: @@ -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)
body: I think deepspeed is a subcase of multi-CPU/multi-GPU no? I would rather ask the question after.
diff_hunk: @@ -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.",
body: 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!
diff_hunk: @@ -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]: ",
body: Why add this default?
diff_hunk: @@ -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"
body: Are those all expected by deepspeed?
diff_hunk: @@ -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)
body: There is something done in Transformers for that, if you need help.
diff_hunk: @@ -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
body: Make sure to use the `make style` command, this is not super readable here :-)
diff_hunk: @@ -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"})
body: This is copied from `deepspeed`? Could we subclass their class then and only change the method that requires it?
diff_hunk: @@ -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):
body: Should this subclass the `AcceleratedOptimizer` for most properties and only overrides the necessary ones?
diff_hunk: @@ -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):
body: I can do this next week!
diff_hunk: @@ -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
body: Apex is not supported yet. Where do you need it?
diff_hunk: @@ -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
+
body: 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).
diff_hunk: @@ -34,6 +35,20 @@
 except ImportError:
     _tpu_available = False
 
+try:
+    import deepspeed
+
+    _deepspeed_available = True
+except:
+    _deepspeed_available = False
body: 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)).
diff_hunk: @@ -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):
body: yes in `deepspeed.initialize()` as `config_params`
diff_hunk: @@ -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)
body: `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).
diff_hunk: @@ -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"
body: 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??
diff_hunk: @@ -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
body: its needed in DeepSpeedEngineWrapper.
diff_hunk: @@ -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
+
body: I need to fix it a bit. Will work on it.
diff_hunk: @@ -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"
body: no
diff_hunk: @@ -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
+
body: 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.
diff_hunk: @@ -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
body: Ok, then please use the same as below and not a direct import.
diff_hunk: @@ -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
+
body: okay. another PR after this then :)
diff_hunk: @@ -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
body: Can we use `batch_sizes` here?
diff_hunk: @@ -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")]
body: 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.")
```
diff_hunk: @@ -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."
body: 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.
diff_hunk: @@ -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)
body: 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.
diff_hunk: @@ -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
body: Is this the only reference we create? It seems so but I want to be certain.
diff_hunk: @@ -357,6 +447,7 @@ def free_memory(self):
         method between two trainings with different models/optimizers.
         """
         self._optimizers = []
+        self.deepspeed_engine = None
body: ```suggestion
                "What should be your DeepSpeed's ZeRO optimization stage (0, 1, 2)? [0]: ",
```
Should include 3?
diff_hunk: @@ -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]: ",
body: Should this be a more general check of whether `attr` is a dict?
diff_hunk: @@ -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":
body: 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 :-)
diff_hunk: @@ -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
body: This file should be removed afterward :-) 
diff_hunk: @@ -0,0 +1,461 @@
+# coding=utf-8
body: 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`.
diff_hunk: @@ -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
body: value of `ds_config` will always be a dict since I am defining it to be `dict` when taking deepspeed arguments from `accelerate config`.
diff_hunk: @@ -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":
body: will remove it in the very end (before merge)
diff_hunk: @@ -0,0 +1,461 @@
+# coding=utf-8
body: I didn't understand the question. Could you please clarify??

I will rename `ds_plugin` to `deepspeed_plugin`.
diff_hunk: @@ -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
body: I would prefer we use `deepspeed_plugin` and `deepspeed_config` instead of `ds_plugin` and `ds_config`
diff_hunk: @@ -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
body: Yes, that's why I'm suggesting we check `if isinstance(attr, dict)` instead of the name.
diff_hunk: @@ -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":
body: yes
diff_hunk: @@ -357,6 +447,7 @@ def free_memory(self):
         method between two trainings with different models/optimizers.
         """
         self._optimizers = []
+        self.deepspeed_engine = None
body: 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 ...."
diff_hunk: @@ -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`.
body: ```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 :)
diff_hunk: @@ -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.
body: The file `deepspeed_utils` has a hard requirement on `deepspeed` being available
diff_hunk: @@ -20,12 +20,22 @@
 from packaging import version
 
 from .data_loader import prepare_data_loader
+from .deepspeed_utils import DeepSpeedEngineWrapper, DeepSpeedOptimizerWrapper, DeepSpeedPlugin
body: This should be conditionally imported if objects are to be imported from this file in others where deepspeed isn't necessarily installed
diff_hunk: @@ -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
body: I don't understand this argument
diff_hunk: @@ -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"},
+    )
body: Is it necessary to have this here if it shouldn't be used?
diff_hunk: @@ -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"})
body: 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
```
diff_hunk: @@ -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"""
body: Do we really need a `DeepSpeedPlugin` to be initialized if deepspeed isn't available?
diff_hunk: @@ -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."
body: ```suggestion
            "Do you want to use DeepSpeed (yes/no)? [no]",
```
diff_hunk: @@ -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]",
body: 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.
diff_hunk: @@ -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"},
+    )
body: I will try to find other solution.
diff_hunk: @@ -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"})
body: I think you probably want a space at the end of this to be consistent with the other _ask_fields
diff_hunk: @@ -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]",
body: Should actually be 
```
          "Do you want to use DeepSpeed? [yes/NO]: ",
```
to be compatible with what is actually in the lib for now.
diff_hunk: @@ -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]",
body: ```suggestion
                    "Where to offload optimizer states? [NONE/cpu/nvme]: ",
```

Also to have the same kind of syntax as the others.
diff_hunk: @@ -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]: ",
body: ```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`.
```
diff_hunk: @@ -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`.
body: ```suggestion
# Remember you still need to do gradient accumulation by yourself, just like you would have done without deepspeed
```
diff_hunk: @@ -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
body: ```suggestion
Note: DeepSpeed support is experimental for now. In case you get into some problem, please open an issue.
```
diff_hunk: @@ -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.

Discussion: body: 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.