From db017fd7316d89497590354de8d2337025df7b1c Mon Sep 17 00:00:00 2001
From: Shaden Smith <shaden.smith@microsoft.com>
Date: Sun, 6 Jun 2021 11:27:13 +0000
Subject: [PATCH 01/38] Minor tweaks to support Megatron 2.4 + DS 3D

---
 .../activation_checkpointing/checkpointing.py | 20 +++++++---
 deepspeed/runtime/engine.py                   |  9 ++++-
 deepspeed/runtime/pipe/engine.py              |  3 +-
 deepspeed/runtime/pipe/module.py              |  5 ++-
 deepspeed/runtime/utils.py                    | 39 +++++++++++++++++++
 5 files changed, 67 insertions(+), 9 deletions(-)

diff --git a/deepspeed/runtime/activation_checkpointing/checkpointing.py b/deepspeed/runtime/activation_checkpointing/checkpointing.py
index f516e4de5203..d1f92146a38f 100644
--- a/deepspeed/runtime/activation_checkpointing/checkpointing.py
+++ b/deepspeed/runtime/activation_checkpointing/checkpointing.py
@@ -24,7 +24,7 @@
 
 from deepspeed.runtime.config import DeepSpeedConfig
 from deepspeed.utils import logger
-from deepspeed.runtime.utils import move_to_device, see_memory_usage
+from deepspeed.runtime.utils import move_to_device, see_memory_usage, bwc_tensor_model_parallel_rank
 from deepspeed.utils.timer import SynchronizedWallClockTimer as Timers
 
 # DeepSpeed Checkpointing Enabled or Disabled
@@ -213,9 +213,12 @@ def model_parallel_cuda_manual_seed(seed):
                               model parallel regions.
     """
     global mpu
+
+    tp_rank = bwc_tensor_model_parallel_rank(mpu)
+
     # 2718 is just for fun and any POSITIVE value will work.
     offset = seed + 2718
-    model_parallel_seed = offset + mpu.get_model_parallel_rank()
+    model_parallel_seed = offset + tp_rank
     # Data parallel gets the original sedd.
     data_parallel_seed = seed
 
@@ -225,7 +228,7 @@ def model_parallel_cuda_manual_seed(seed):
             'model parallel rank {}, and data parallel rank {} with '
             'model parallel seed: {} and data parallel seed: {}'.format(
                 torch.distributed.get_rank(),
-                mpu.get_model_parallel_rank(),
+                tp_rank,
                 mpu.get_data_parallel_rank(),
                 model_parallel_seed,
                 data_parallel_seed),
@@ -384,9 +387,14 @@ def save_args_for_backward(*all_args):
         global data_offsets, size_offsets
         if mp_rank is None:
             if mpu is not None:
-                mp_rank = mpu.get_model_parallel_rank()
-                mp_size = mpu.get_model_parallel_world_size()
-                mp_group = mpu.get_model_parallel_group()
+                if hasattr(mpu, 'get_tensor_model_parallel_rank'):
+                    mp_rank = mpu.get_tensor_model_parallel_rank()
+                    mp_size = mpu.get_tensor_model_parallel_world_size()
+                    mp_group = mpu.get_tensor_model_parallel_group()
+                else:
+                    mp_rank = mpu.get_model_parallel_rank()
+                    mp_size = mpu.get_model_parallel_world_size()
+                    mp_group = mpu.get_model_parallel_group()
             else:
                 mp_rank = 0
                 mp_size = 1
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
index c2aa59471a94..24b64fbb039b 100755
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -1275,7 +1275,14 @@ def _take_model_step(self, lr_kwargs, block_eigenvalue={}):
             self.skipped_steps += 1
         else:
             if self.lr_scheduler is not None:
-                self.lr_scheduler.step(**(lr_kwargs or {}))
+                try:
+                    self.lr_scheduler.step(**(lr_kwargs or {}))
+                except TypeError:
+                    # XXX Hack to work with Megatron 2.0 and DeepSpeed pipelines.
+                    # We don't currently have a way to specify lr_kwargs from
+                    # pipe_engine.train_batch()
+                    self.lr_scheduler.step(increment=self.train_batch_size())
+
 
         if report_progress and (self.global_steps + 1) % self.steps_per_print() == 0:
             self._report_progress(self.global_steps + 1)
diff --git a/deepspeed/runtime/pipe/engine.py b/deepspeed/runtime/pipe/engine.py
index 196cbe8c6217..27bd266074f8 100644
--- a/deepspeed/runtime/pipe/engine.py
+++ b/deepspeed/runtime/pipe/engine.py
@@ -110,7 +110,8 @@ def __init__(self, *super_args, **super_kwargs):
         self.is_model_parallel = self.grid.model_parallel_size > 1
 
         # Partition input/output buffers
-        self.is_pipe_partitioned = self.is_model_parallel
+        # XXX temporarily disable while I revert some partition hacks.
+        self.is_pipe_partitioned = False #self.is_model_parallel
         self.is_grad_partitioned = False
 
         model_parameters = filter(lambda p: p.requires_grad, self.module.parameters())
diff --git a/deepspeed/runtime/pipe/module.py b/deepspeed/runtime/pipe/module.py
index c4e111e47315..31144f7a69b4 100644
--- a/deepspeed/runtime/pipe/module.py
+++ b/deepspeed/runtime/pipe/module.py
@@ -585,7 +585,10 @@ def load_state_dir(self, load_dir, strict=True):
         self._synchronize_tied_weights()
 
     def _is_checkpointable(self, funcs):
-        if self.__class__.__name__ == 'GPT2ModelPipe':
+        # This is an unfortunate hack related to torch and deepspeed activation checkpoint implementations.
+        # Some layers like torch.nn.Embedding will not receive grads if checkpointed, which breaks things.
+        # I presume it's related to the discrete inputs that cannot require_grad? Need to revisit.
+        if self.__class__.__name__ in ('GPTModelPipe', 'GPT2ModelPipe'):
             return all('ParallelTransformerLayerPipe' in f.__class__.__name__
                        for f in funcs)
 
diff --git a/deepspeed/runtime/utils.py b/deepspeed/runtime/utils.py
index c792d2a6d0db..c3ababa43c3e 100755
--- a/deepspeed/runtime/utils.py
+++ b/deepspeed/runtime/utils.py
@@ -49,6 +49,45 @@ def set_random_seed(seed):
     torch.manual_seed(seed)
 
 
+def bwc_tensor_model_parallel_rank(mpu=None):
+    """Backwards-compatible way of querying the tensor model parallel rank from
+    an ``mpu`` object.
+
+    *Tensor* model parallelism means that tensors are physically split across
+    processes. This contrasts with *pipeline* model parallelism, in which the
+    layers are partitioned but tensors left intact.
+
+    The API for tensor model parallelism has changed across versions and this
+    helper provides a best-effort implementation across versions of ``mpu``
+    objects.  The preferred mechanism is
+    ``mpu.get_tensor_model_parallel_rank()``.
+
+    This should "just work" with both Megatron-LM and DeepSpeed's pipeline
+    parallelism.
+
+    Args:
+        mpu (model parallel unit, optional): The tensor model parallel rank.
+            If ``mpu=None``, returns 0. Defaults to ``None``.
+
+    Returns:
+        int: the rank
+    """
+    if mpu is None:
+        # No model parallelism in easy :)
+        return 0
+
+    if hasattr(mpu, 'get_tensor_model_parallel_rank'):
+        # New Megatron and DeepSpeed convention (post pipeline-parallelism release)
+        return mpu.get_tensor_model_parallel_rank()
+    elif hasattr(mpu, 'get_slice_parallel_rank'):
+        # Some DeepSpeed + pipeline parallelism versions
+        return mpu.get_slice_parallel_rank()
+    else:
+        # Deprecated Megatron and DeepSpeed convention
+        return mpu.get_model_parallel_rank()
+
+
+
 def move_to_device(item, device):
     """
     Move tensor onto device. Works on individual tensors, and tensors contained/nested in lists, tuples, and dicts.

From 407ff0f1543be06d52de22abeaf9bf9af49295d9 Mon Sep 17 00:00:00 2001
From: Shaden Smith <shaden.smith@microsoft.com>
Date: Sun, 6 Jun 2021 14:55:18 +0000
Subject: [PATCH 02/38] pipe partitioning

---
 deepspeed/runtime/pipe/engine.py | 48 +++++++++++++++-----------------
 1 file changed, 23 insertions(+), 25 deletions(-)

diff --git a/deepspeed/runtime/pipe/engine.py b/deepspeed/runtime/pipe/engine.py
index 27bd266074f8..1019a358b356 100644
--- a/deepspeed/runtime/pipe/engine.py
+++ b/deepspeed/runtime/pipe/engine.py
@@ -111,8 +111,8 @@ def __init__(self, *super_args, **super_kwargs):
 
         # Partition input/output buffers
         # XXX temporarily disable while I revert some partition hacks.
-        self.is_pipe_partitioned = False #self.is_model_parallel
-        self.is_grad_partitioned = False
+        self.is_pipe_partitioned = self.is_model_parallel
+        self.is_grad_partitioned = False #self.is_model_parallel
 
         model_parameters = filter(lambda p: p.requires_grad, self.module.parameters())
         num_params = sum([p.numel() for p in model_parameters])
@@ -554,12 +554,18 @@ def _exec_forward_pass(self, buffer_id):
                 local_part=inputs[1],
                 group=self.grid.get_slice_parallel_group())
 
+            inputs = part_input.full()
+            inputs.requires_grad = True
+            part_input = None
+            self.pipe_buffers['inputs'][buffer_id] = inputs
+            '''
             inputs = tuple([part_input.full(), inputs[2]])
             inputs[0].requires_grad = True
             # skip mask
             #inputs[1].requires_grad = True
             part_input = None
             self.pipe_buffers['inputs'][buffer_id] = inputs
+            '''
 
         # Zero out the gradients each time we use the tensor because only the data in
         # tensor changes across batches
@@ -569,13 +575,14 @@ def _exec_forward_pass(self, buffer_id):
 
         # Partition the outputs if we are not the last stage
         if self.is_pipe_partitioned and not self.is_last_stage():
-            part = PartitionedTensor(tensor=outputs[0],
+            assert torch.is_tensor(outputs)
+            part = PartitionedTensor(tensor=outputs,
                                      group=self.grid.get_slice_parallel_group())
             # Clear the large output data, but save the computation graph
-            outputs[0].data = torch.zeros(1)
-            self.pipe_buffers['output_tensors'][buffer_id] = outputs[0]
+            outputs.data = torch.zeros(1)
+            self.pipe_buffers['output_tensors'][buffer_id] = outputs
             # Inject the partitioned tensor into the output before sending
-            outputs = tuple([part.to_meta(), part.data(), outputs[1]])
+            outputs = tuple([part.to_meta(), part.data()])
             part = None
 
         self.pipe_buffers['outputs'][buffer_id] = outputs
@@ -633,15 +640,11 @@ def _exec_backward_pass(self, buffer_id):
                     local_part=outputs[1],
                     group=self.grid.get_slice_parallel_group())
                 self.pipe_buffers['output_tensors'][buffer_id].data = part_output.full()
-                outputs = tuple(
-                    [self.pipe_buffers['output_tensors'][buffer_id],
-                     outputs[2]])
+                outputs = self.pipe_buffers['output_tensors'][buffer_id]
             else:
                 # Already restored from partition
-                self.pipe_buffers['output_tensors'][buffer_id].data = outputs[0]
-                outputs = tuple(
-                    [self.pipe_buffers['output_tensors'][buffer_id],
-                     outputs[1]])
+                self.pipe_buffers['output_tensors'][buffer_id].data = outputs
+                outputs = self.pipe_buffers['output_tensors'][buffer_id]
 
         grad_tensors = self.grad_layer
         if self.is_grad_partitioned:
@@ -650,7 +653,7 @@ def _exec_backward_pass(self, buffer_id):
                 meta=self.grad_layer[0],
                 local_part=self.grad_layer[1],
                 group=self.grid.get_slice_parallel_group())
-            grad_tensors = tuple([part_grad.full(), self.grad_layer[2]])
+            grad_tensors = part_grad.full()
             part_grad = None
             #print(f'RANK={self.global_rank} BEFORE-BWD restored grad={self.grad_layer[0].size()} {self.grad_layer[1].size()}')
 
@@ -873,13 +876,10 @@ def _exec_send_grads(self, buffer_id):
 
         # Partition the gradient
         if self.is_grad_partitioned:
-            part = PartitionedTensor(tensor=inputs[0].grad,
+            assert torch.is_tensor(inputs)
+            part = PartitionedTensor(tensor=inputs.grad,
                                      group=self.grid.get_slice_parallel_group())
-            # Clear the large output data, but save the computation graph
-            # Inject the partitoned tensor into the output before sending
-
-            # XXX Hack
-            inputs = tuple([part.to_meta(), part.data(), inputs[1]])
+            inputs = tuple([part.to_meta(), part.data()])
 
         # XXX Terrible hack
         # Drop the attention mask from the input buffer here. It does not have
@@ -900,8 +900,6 @@ def _exec_send_grads(self, buffer_id):
                 # First two sends are partitioned gradient
                 p2p.send(inputs[0], self.prev_stage)
                 p2p.send(inputs[1], self.prev_stage)
-                # XXX hack hack hack
-                #p2p.send(inputs[2].grad, self.prev_stage)
             else:
                 for idx, buffer in enumerate(inputs):
                     # Skip tensors that will not produce a grad
@@ -975,7 +973,7 @@ def _exec_recv_grads(self, buffer_id):
                 local_part=outputs[1],
                 group=self.grid.get_slice_parallel_group())
             outputs[0].data = part_output.full()
-            outputs = tuple([outputs[0], outputs[2]])
+            outputs = outputs[0]
             # save for backward
             self.pipe_buffers['outputs'][buffer_id] = outputs
 
@@ -985,7 +983,7 @@ def _exec_recv_grads(self, buffer_id):
                 s = list(outputs.size())
                 self.grad_layer = self._allocate_buffer(s, num_buffers=1)[0]
             else:
-                sizes = [list(t.size()) for t in outputs if t.is_floating_point()]
+                sizes = [list(t.size()) for t in outputs]# if t.is_floating_point()]
                 self.grad_layer = self._allocate_buffers(sizes, num_buffers=1)[0]
 
         if isinstance(self.grad_layer, torch.Tensor):
@@ -999,7 +997,7 @@ def _exec_recv_grads(self, buffer_id):
                                               dtype=torch.long,
                                               device=self.device)
                 p2p.recv(buffer, self.next_stage)
-
+            
         if self.wall_clock_breakdown():
             self.timers('pipe_recv_grad').stop()
 

From 8e48756c123b037db9c273788f1339d7b2572045 Mon Sep 17 00:00:00 2001
From: Shaden Smith <Shaden.Smith@microsoft.com>
Date: Wed, 9 Jun 2021 09:40:14 -0700
Subject: [PATCH 03/38] unit test for bugfix #1135 (#1148)

* unit test for bugfix #1135

* formatter

* fix test in presence of mpi4py

Co-authored-by: Jeff Rasley <jerasley@microsoft.com>
---
 tests/unit/test_zero_context.py | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/tests/unit/test_zero_context.py b/tests/unit/test_zero_context.py
index 8b8d51131503..98ee0c7ad00b 100644
--- a/tests/unit/test_zero_context.py
+++ b/tests/unit/test_zero_context.py
@@ -281,3 +281,32 @@ def test_stage_3_output_type(output_type):
                 loss = loss['loss']
             engine.backward(loss)
             engine.step()
+
+
+class ConvX(torch.nn.Conv1d):
+    def __init__(self, *args):
+        super().__init__(*args)
+        # This would not be partitioned before bugfix 5ca8167
+        self.param_in = torch.nn.Parameter(torch.FloatTensor(5).uniform_())
+
+    def forward(self, x):
+        return x
+
+
+class ConvNet(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = ConvX(1, 3, 4)
+        self.param = torch.nn.Parameter(torch.FloatTensor(5).uniform_())
+
+    def forward(self, x):
+        return x
+
+
+def test_subclass_param():
+    setup_serial_env()
+    with deepspeed.zero.Init(config=config):
+        model = ConvNet()
+
+    assert model.param.ds_status == ZeroParamStatus.NOT_AVAILABLE
+    assert model.conv1.param_in.ds_status == ZeroParamStatus.NOT_AVAILABLE

From 71ecf7e62583cf8cd3661d30d5f82851df3bb82c Mon Sep 17 00:00:00 2001
From: eltonzheng <eltonz@microsoft.com>
Date: Wed, 9 Jun 2021 17:24:43 -0700
Subject: [PATCH 04/38] Add Windows support in README, use c++17 on Windows  to
 support latest VC & cuda build tool (#1151)

* Add Windows support in README, use c++17 on Windows to support latest vc build tool

* Add detailed cpp build tools version in README

Co-authored-by: Jeff Rasley <jerasley@microsoft.com>
---
 README.md                           |  6 ++++++
 op_builder/builder.py               | 12 ++++++++++++
 op_builder/cpu_adam.py              | 12 ------------
 op_builder/quantizer.py             | 12 ------------
 op_builder/transformer.py           | 12 ------------
 op_builder/transformer_inference.py | 12 ------------
 6 files changed, 18 insertions(+), 48 deletions(-)

diff --git a/README.md b/README.md
index c602f90144bf..4b752f021787 100755
--- a/README.md
+++ b/README.md
@@ -96,6 +96,12 @@ If you would like to pre-install any of the DeepSpeed extensions/ops (instead
 of JIT compiling) or install pre-compiled ops via PyPI please see our [advanced
 installation instructions](https://www.deepspeed.ai/tutorials/advanced-install/).
 
+On Windows you can build wheel with following steps, currently only inference mode is supported.
+1. Install pytorch, such as pytorch 1.8 + cuda 11.1
+2. Install visual cpp build tools, such as VS2019 C++ x64/x86 build tools
+3. Launch cmd console with Administrator privilege for creating required symlink folders
+4. Run `python setup.py bdist_wheel` to build wheel in `dist` folder
+
 # Features
 Below we provide a brief feature list, see our detailed [feature
 overview](https://www.deepspeed.ai/features/) for descriptions and usage.
diff --git a/op_builder/builder.py b/op_builder/builder.py
index 50dfa4f9c550..97ae936c68bf 100644
--- a/op_builder/builder.py
+++ b/op_builder/builder.py
@@ -364,6 +364,18 @@ def cxx_args(self):
         else:
             return ['-O3', '-std=c++14', '-g', '-Wno-reorder']
 
+    def nvcc_args(self):
+        args = [
+            '-O3',
+            '--use_fast_math',
+            '-std=c++17' if sys.platform == "win32" else '-std=c++14',
+            '-U__CUDA_NO_HALF_OPERATORS__',
+            '-U__CUDA_NO_HALF_CONVERSIONS__',
+            '-U__CUDA_NO_HALF2_OPERATORS__'
+        ]
+
+        return args + self.compute_capability_args()
+
     def libraries_args(self):
         if sys.platform == "win32":
             return ['cublas', 'curand']
diff --git a/op_builder/cpu_adam.py b/op_builder/cpu_adam.py
index 75fa042613fc..b93f27b81709 100644
--- a/op_builder/cpu_adam.py
+++ b/op_builder/cpu_adam.py
@@ -62,15 +62,3 @@ def cxx_args(self):
             '-fopenmp',
             SIMD_WIDTH
         ]
-
-    def nvcc_args(self):
-        args = [
-            '-O3',
-            '--use_fast_math',
-            '-std=c++14',
-            '-U__CUDA_NO_HALF_OPERATORS__',
-            '-U__CUDA_NO_HALF_CONVERSIONS__',
-            '-U__CUDA_NO_HALF2_OPERATORS__'
-        ]
-        args += self.compute_capability_args()
-        return args
diff --git a/op_builder/quantizer.py b/op_builder/quantizer.py
index f77e43c62408..b3d6b2fd4509 100644
--- a/op_builder/quantizer.py
+++ b/op_builder/quantizer.py
@@ -21,15 +21,3 @@ def sources(self):
 
     def include_paths(self):
         return ['csrc/includes']
-
-    def nvcc_args(self):
-        args = [
-            '-O3',
-            '--use_fast_math',
-            '-std=c++14',
-            '-U__CUDA_NO_HALF_OPERATORS__',
-            '-U__CUDA_NO_HALF_CONVERSIONS__',
-            '-U__CUDA_NO_HALF2_OPERATORS__'
-        ]
-
-        return args + self.compute_capability_args()
diff --git a/op_builder/transformer.py b/op_builder/transformer.py
index a94c6da62e39..82ab26f7f25c 100644
--- a/op_builder/transformer.py
+++ b/op_builder/transformer.py
@@ -30,15 +30,3 @@ def sources(self):
 
     def include_paths(self):
         return ['csrc/includes']
-
-    def nvcc_args(self):
-        args = [
-            '-O3',
-            '--use_fast_math',
-            '-std=c++14',
-            '-U__CUDA_NO_HALF_OPERATORS__',
-            '-U__CUDA_NO_HALF_CONVERSIONS__',
-            '-U__CUDA_NO_HALF2_OPERATORS__'
-        ]
-
-        return args + self.compute_capability_args()
diff --git a/op_builder/transformer_inference.py b/op_builder/transformer_inference.py
index 76ee22e0b2c6..cf5e870c906c 100755
--- a/op_builder/transformer_inference.py
+++ b/op_builder/transformer_inference.py
@@ -24,15 +24,3 @@ def sources(self):
 
     def include_paths(self):
         return ['csrc/transformer/inference/includes']
-
-    def nvcc_args(self):
-        args = [
-            '-O3',
-            '--use_fast_math',
-            '-std=c++14',
-            '-U__CUDA_NO_HALF_OPERATORS__',
-            '-U__CUDA_NO_HALF_CONVERSIONS__',
-            '-U__CUDA_NO_HALF2_OPERATORS__',
-        ]
-
-        return args + self.compute_capability_args()

From a096d32d7c0f03cf976424a7881b8d1e21db85af Mon Sep 17 00:00:00 2001
From: Shaden Smith <shaden.smith@microsoft.com>
Date: Fri, 11 Jun 2021 00:30:13 +0000
Subject: [PATCH 05/38] re-enable grad buffer partitioning

---
 deepspeed/runtime/pipe/engine.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deepspeed/runtime/pipe/engine.py b/deepspeed/runtime/pipe/engine.py
index 1019a358b356..14601eb0a007 100644
--- a/deepspeed/runtime/pipe/engine.py
+++ b/deepspeed/runtime/pipe/engine.py
@@ -112,7 +112,7 @@ def __init__(self, *super_args, **super_kwargs):
         # Partition input/output buffers
         # XXX temporarily disable while I revert some partition hacks.
         self.is_pipe_partitioned = self.is_model_parallel
-        self.is_grad_partitioned = False #self.is_model_parallel
+        self.is_grad_partitioned = self.is_model_parallel
 
         model_parameters = filter(lambda p: p.requires_grad, self.module.parameters())
         num_params = sum([p.numel() for p in model_parameters])

From 9b4093b37ab0596fef1743c62e4c446bbb1faeab Mon Sep 17 00:00:00 2001
From: Tunji Ruwase <olruwase@microsoft.com>
Date: Fri, 11 Jun 2021 18:48:02 +0000
Subject: [PATCH 06/38] Avoid partitioning small activations

---
 deepspeed/runtime/activation_checkpointing/checkpointing.py | 2 +-
 deepspeed/runtime/engine.py                                 | 1 -
 deepspeed/runtime/pipe/engine.py                            | 4 ++--
 deepspeed/runtime/utils.py                                  | 1 -
 4 files changed, 3 insertions(+), 5 deletions(-)

diff --git a/deepspeed/runtime/activation_checkpointing/checkpointing.py b/deepspeed/runtime/activation_checkpointing/checkpointing.py
index d1f92146a38f..efe95f91bac5 100644
--- a/deepspeed/runtime/activation_checkpointing/checkpointing.py
+++ b/deepspeed/runtime/activation_checkpointing/checkpointing.py
@@ -424,7 +424,7 @@ def save_args_for_backward(*all_args):
 
             inputs = []
             for i, item in enumerate(args[:-1]):
-                if not torch.is_tensor(item):
+                if not torch.is_tensor(item) or mp_size > item.numel():
                     inputs.append(item)
                     continue
 
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
index 24b64fbb039b..146a6c6e931c 100755
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -1283,7 +1283,6 @@ def _take_model_step(self, lr_kwargs, block_eigenvalue={}):
                     # pipe_engine.train_batch()
                     self.lr_scheduler.step(increment=self.train_batch_size())
 
-
         if report_progress and (self.global_steps + 1) % self.steps_per_print() == 0:
             self._report_progress(self.global_steps + 1)
 
diff --git a/deepspeed/runtime/pipe/engine.py b/deepspeed/runtime/pipe/engine.py
index 14601eb0a007..09231bf3e96a 100644
--- a/deepspeed/runtime/pipe/engine.py
+++ b/deepspeed/runtime/pipe/engine.py
@@ -983,7 +983,7 @@ def _exec_recv_grads(self, buffer_id):
                 s = list(outputs.size())
                 self.grad_layer = self._allocate_buffer(s, num_buffers=1)[0]
             else:
-                sizes = [list(t.size()) for t in outputs]# if t.is_floating_point()]
+                sizes = [list(t.size()) for t in outputs]  # if t.is_floating_point()]
                 self.grad_layer = self._allocate_buffers(sizes, num_buffers=1)[0]
 
         if isinstance(self.grad_layer, torch.Tensor):
@@ -997,7 +997,7 @@ def _exec_recv_grads(self, buffer_id):
                                               dtype=torch.long,
                                               device=self.device)
                 p2p.recv(buffer, self.next_stage)
-            
+
         if self.wall_clock_breakdown():
             self.timers('pipe_recv_grad').stop()
 
diff --git a/deepspeed/runtime/utils.py b/deepspeed/runtime/utils.py
index c3ababa43c3e..f3a4c4c0f1b3 100755
--- a/deepspeed/runtime/utils.py
+++ b/deepspeed/runtime/utils.py
@@ -87,7 +87,6 @@ def bwc_tensor_model_parallel_rank(mpu=None):
         return mpu.get_model_parallel_rank()
 
 
-
 def move_to_device(item, device):
     """
     Move tensor onto device. Works on individual tensors, and tensors contained/nested in lists, tuples, and dicts.

From 3e948df134d0a8815f8875432c17b177a03f59c9 Mon Sep 17 00:00:00 2001
From: Shaden Smith <shaden.smith@microsoft.com>
Date: Sun, 13 Jun 2021 21:50:47 +0000
Subject: [PATCH 07/38] send/recv

---
 deepspeed/runtime/pipe/p2p.py | 123 ++++++++++++++++++++--------------
 1 file changed, 74 insertions(+), 49 deletions(-)

diff --git a/deepspeed/runtime/pipe/p2p.py b/deepspeed/runtime/pipe/p2p.py
index 24c0f250a4b9..162f4882b230 100644
--- a/deepspeed/runtime/pipe/p2p.py
+++ b/deepspeed/runtime/pipe/p2p.py
@@ -2,6 +2,10 @@
 Copyright 2019 The Microsoft DeepSpeed Team
 '''
 
+import pickle
+import typing
+
+import torch
 import torch.distributed as dist
 
 _groups = None
@@ -16,8 +20,6 @@ def init_process_groups(grid):
 
     assert _grid.pipe_parallel_size > 1, "There is no pipeline parallelism"
 
-    _groups = [dist.new_group(ranks=group) for group in _grid.p2p_groups]
-
 
 def _is_valid_send_recv(src_stage, dest_stage):
     first_stage = 0
@@ -30,61 +32,84 @@ def _is_valid_send_recv(src_stage, dest_stage):
 
 def send(tensor, dest_stage, async_op=False):
     global _groups
-
-    async_op = False
+    assert async_op == False, "Doesnt support async_op true"
     src_stage = _grid.get_stage_id()
     _is_valid_send_recv(src_stage, dest_stage)
 
-    group = _get_send_recv_group(src_stage, dest_stage)
-    src_rank = _grid.stage_to_global(stage_id=src_stage)
-
-    return dist.broadcast(tensor, src_rank, group=group, async_op=async_op)
+    dest_rank = _grid.stage_to_global(stage_id=dest_stage)
+    return dist.send(tensor, dest_rank)
 
 
 def recv(tensor, src_stage, async_op=False):
-
     global _groups
-
-    async_op = False
+    assert async_op == False, "Doesnt support async_op true"
     dest_stage = _grid.get_stage_id()
     _is_valid_send_recv(src_stage, dest_stage)
 
-    group = _get_send_recv_group(src_stage, dest_stage)
     src_rank = _grid.stage_to_global(stage_id=src_stage)
-
-    return dist.broadcast(tensor, src_rank, group=group, async_op=async_op)
-
-
-def barrier(stage_id):
-    global _groups, _grid
-    group_id = _grid.stage_to_global(stage_id=stage_id)
-    if (dist.get_rank() >= 0):
-        print("Barrier Group ID", group_id)
-        print("Barrier Group", _grid.p2p_groups[group_id])
-    dist.barrier(group=_groups[group_id])
-    if (dist.get_rank() >= 0):
-        print("Exiting Barrier ", group_id)
-
-
-def _get_send_recv_group(src_stage, dest_stage):
-    '''the group id is always the smaller rank unless its a wrap around'''
-
-    stage_id = None
-
-    first_stage = 0
-    last_stage = _grid.pipe_parallel_size - 1
-
-    if (src_stage == first_stage and dest_stage == last_stage
-            or dest_stage == first_stage and src_stage == last_stage):
-        stage_id = last_stage
-    elif src_stage > dest_stage:
-        stage_id = dest_stage
-    else:
-        stage_id = src_stage
-    '''group_id corresponds to group of [group_id, group_id+1]
-     unless group_id is the rank of the last stage
-     in which case group_id correspods to group[group_id-num_stages+1, group_id]
-     '''
-    group_id = _grid.stage_to_global(stage_id=stage_id)
-
-    return _groups[group_id]
+    return dist.recv(tensor, src_rank)
+
+
+def send_obj(msg: typing.Any, dest: int):
+    """Send an arbitrary python object to ``dest``.
+
+    Note: ``msg`` must be pickleable.
+
+    WARN: This incurs a CPU -> GPU transfer and should be used sparingly
+    for performance reasons.
+
+    Args:
+        msg (typing.Any): The object to send.
+        dest (int): Destination rank.
+    """
+    # serialize the message
+    msg = pickle.dumps(msg)
+    # construct a tensor to send
+    msg = torch.ByteTensor(torch.ByteStorage.from_buffer(msg)).cuda()
+
+    # Send meta and message
+    length_tensor = torch.tensor([len(msg)], dtype=torch.long).cuda()
+    dist.send(length_tensor, dst=dest)
+    dist.send(msg, dst=dest)
+
+
+def recv_obj(sender: int) -> typing.Any:
+    """Receive an arbitrary python object from ``sender``.
+
+    WARN: This incur a CPU <-> GPU transfers and should be used sparingly
+    for performance reasons.
+
+    Args:
+        sender (int): The rank sending the message.
+    """
+    # Get message meta
+    length = torch.tensor([0], dtype=torch.long).cuda()
+    dist.recv(length, src=sender)
+
+    # Receive and deserialize
+    msg = torch.empty(length.item(), dtype=torch.uint8).cuda()
+    dist.recv(msg, src=sender)
+
+    msg = pickle.loads(msg.cpu().numpy().tobytes())
+
+    def _to(x):
+        """Recursively move to the current device."""
+        if torch.is_tensor(x):
+            return x.cuda()
+        if isinstance(x, (tuple, list)):
+            ret = [_to(x_) for x_ in x]
+            if isinstance(x, tuple):
+                ret = tuple(ret)
+            return ret
+        # handle kwargs
+        if isinstance(x, dict):
+            ret = dict()
+            for key, val in x.items():
+                ret[_to(key)] = _to(val)
+            return ret
+
+        # Anything else is a no-op
+        return x
+
+    msg = _to(msg)
+    return msg

From b6a2cb379867bd4450f739b6b985306de5b4c3e8 Mon Sep 17 00:00:00 2001
From: Shaden Smith <shaden.smith@microsoft.com>
Date: Sun, 13 Jun 2021 22:15:07 +0000
Subject: [PATCH 08/38] isend/irecv missing wait

---
 deepspeed/runtime/pipe/p2p.py | 33 +++++++++++++++++++++++++++------
 1 file changed, 27 insertions(+), 6 deletions(-)

diff --git a/deepspeed/runtime/pipe/p2p.py b/deepspeed/runtime/pipe/p2p.py
index 162f4882b230..64d9712639e5 100644
--- a/deepspeed/runtime/pipe/p2p.py
+++ b/deepspeed/runtime/pipe/p2p.py
@@ -12,6 +12,8 @@
 _grid = None
 
 
+_async = []
+
 #initializes adjacent process groups
 #run this only after torch.distributed.init_process_group() has been called
 def init_process_groups(grid):
@@ -30,25 +32,44 @@ def _is_valid_send_recv(src_stage, dest_stage):
     "Functionality currently limited to send and receive between adjacent ranks only"
 
 
-def send(tensor, dest_stage, async_op=False):
+def send(tensor, dest_stage, async_op=True):
     global _groups
-    assert async_op == False, "Doesnt support async_op true"
+    #assert async_op == False, "Doesnt support async_op true"
     src_stage = _grid.get_stage_id()
     _is_valid_send_recv(src_stage, dest_stage)
 
     dest_rank = _grid.stage_to_global(stage_id=dest_stage)
-    return dist.send(tensor, dest_rank)
+    if async_op:
+        global _async
+        op = dist.isend(tensor, dest_rank)
+        _async.append(op)
+    else:
+        return dist.send(tensor, dest_rank)
 
 
-def recv(tensor, src_stage, async_op=False):
+def recv(tensor, src_stage, async_op=True):
     global _groups
-    assert async_op == False, "Doesnt support async_op true"
+    #assert async_op == False, "Doesnt support async_op true"
     dest_stage = _grid.get_stage_id()
     _is_valid_send_recv(src_stage, dest_stage)
 
     src_rank = _grid.stage_to_global(stage_id=src_stage)
-    return dist.recv(tensor, src_rank)
 
+    if async_op:
+        global _async
+        op = dist.irecv(tensor, src_rank)
+        _async.append(op)
+    else:
+        return dist.recv(tensor, src_rank)
+
+
+def wait():
+    global _async
+    for op in _async:
+        op.wait()
+    _async = []
+
+    torch.cuda.synchronize()
 
 def send_obj(msg: typing.Any, dest: int):
     """Send an arbitrary python object to ``dest``.

From 6bb63b8560ac04c528fa5075f967d89bf0072447 Mon Sep 17 00:00:00 2001
From: Shaden Smith <shaden.smith@microsoft.com>
Date: Mon, 14 Jun 2021 19:20:21 +0000
Subject: [PATCH 09/38] turn off async ops

---
 deepspeed/runtime/pipe/p2p.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/deepspeed/runtime/pipe/p2p.py b/deepspeed/runtime/pipe/p2p.py
index 64d9712639e5..0fa31b89b984 100644
--- a/deepspeed/runtime/pipe/p2p.py
+++ b/deepspeed/runtime/pipe/p2p.py
@@ -32,7 +32,7 @@ def _is_valid_send_recv(src_stage, dest_stage):
     "Functionality currently limited to send and receive between adjacent ranks only"
 
 
-def send(tensor, dest_stage, async_op=True):
+def send(tensor, dest_stage, async_op=False):
     global _groups
     #assert async_op == False, "Doesnt support async_op true"
     src_stage = _grid.get_stage_id()
@@ -47,7 +47,7 @@ def send(tensor, dest_stage, async_op=True):
         return dist.send(tensor, dest_rank)
 
 
-def recv(tensor, src_stage, async_op=True):
+def recv(tensor, src_stage, async_op=False):
     global _groups
     #assert async_op == False, "Doesnt support async_op true"
     dest_stage = _grid.get_stage_id()

From aa16828c267dc7a790a5cc2f43f93a67617a2364 Mon Sep 17 00:00:00 2001
From: Hyunwoong Ko <gusdnd852@naver.com>
Date: Tue, 15 Jun 2021 09:24:18 +0900
Subject: [PATCH 10/38] Fix bugs in the tutorial documentation (#1157)

* Add `import os` to inference tutorials

* assign deepspeed-initialized model to hf model
---
 docs/_tutorials/inference-tutorial.md | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/docs/_tutorials/inference-tutorial.md b/docs/_tutorials/inference-tutorial.md
index ab13b3dfe264..1686e75be9a2 100644
--- a/docs/_tutorials/inference-tutorial.md
+++ b/docs/_tutorials/inference-tutorial.md
@@ -71,6 +71,7 @@ DeepSpeed inference can be used in conjunction with HuggingFace `pipeline`. Belo
 
 ```python
 # Filename: gpt-neo-2.7b-generation.py
+import os
 import deepspeed
 import torch
 import transformers
@@ -82,10 +83,10 @@ generator = pipeline('text-generation', model='EleutherAI/gpt-neo-2.7B', device=
 
 
 
-deepspeed.init_inference(generator.model,
-                         mp_size=world_size,
-                         dtype=torch.float,
-                         replace_method='auto')
+generator.model = deepspeed.init_inference(generator.model,
+                                           mp_size=world_size,
+                                           dtype=torch.float,
+                                           replace_method='auto')
 
 string = generator("DeepSpeed is", do_sample=True, min_length=50)
 if torch.distributed.get_rank() == 0:

From 4eaf910616bda4be87dc7777b7f3adff792adc49 Mon Sep 17 00:00:00 2001
From: Samyam Rajbhandari <samyamr@microsoft.com>
Date: Wed, 16 Jun 2021 07:43:07 -1000
Subject: [PATCH 11/38] Samyamr/largest partitioned params calculation fix
 (#1150)

* largest_partitioned_params calculation fix

largest partitioned params was getting calculated incorrectly

* Update stage3.py

* Update stage3.py

* formatting fix

* changing sub-group size default to 1e9

Co-authored-by: Olatunji Ruwase <olruwase@microsoft.com>
---
 deepspeed/runtime/zero/constants.py | 2 +-
 deepspeed/runtime/zero/stage3.py    | 5 ++++-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/deepspeed/runtime/zero/constants.py b/deepspeed/runtime/zero/constants.py
index eaeb2a95ccd9..7beebe00e717 100755
--- a/deepspeed/runtime/zero/constants.py
+++ b/deepspeed/runtime/zero/constants.py
@@ -91,7 +91,7 @@
 ZERO_OPTIMIZATION_OFFLOAD_OPTIMIZER_DEFAULT = None
 
 ZERO_OPTIMIZATION_SUB_GROUP_SIZE = 'sub_group_size'
-ZERO_OPTIMIZATION_SUB_GROUP_SIZE_DEFAULT = 1000000000000
+ZERO_OPTIMIZATION_SUB_GROUP_SIZE_DEFAULT = 1000000000
 
 #maximum number of parameters per GPU before releasing them
 ZERO_OPTIMIZATION_MAX_LIVE_PARAMETERS = 'stage3_max_live_parameters'
diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py
index cac206ed759b..548c38a072c3 100755
--- a/deepspeed/runtime/zero/stage3.py
+++ b/deepspeed/runtime/zero/stage3.py
@@ -852,7 +852,10 @@ def __init__(self,
                 count = count + 1
 
         #Largest partitioned param
-        largest_partitioned_param_numel = max(self.fp16_partitioned_groups_flat_numel)
+        largest_partitioned_param_numel = max([
+            max([tensor.numel() for tensor in fp16_partitioned_group])
+            for fp16_partitioned_group in self.fp16_partitioned_groups
+        ])
         print_rank_0(
             f'Largest partitioned param numel = {largest_partitioned_param_numel}',
             force=True)

From fa7921e23f15844373dc3225a27a1b8f40084a2a Mon Sep 17 00:00:00 2001
From: Olatunji Ruwase <olruwase@microsoft.com>
Date: Wed, 16 Jun 2021 11:31:22 -0700
Subject: [PATCH 12/38] [Doc] Fix steps_per_print description (#1163)

* Fix docstring

* Make screenshots clickable for easier viewing

* Navigation menu in alphabetical order; More clicable screenshots

* Rename 1Cycle doc

* Tweak naming

* Remove no longer used flag

* ZeRO3 Offload release

* Single GPU results

* Rearrange figures

* Single GPU text

* tweak intro

* zero3-offload section

* Add asynchronous i/o docs

* Fix print_per_steps doc
---
 docs/_pages/config-json.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/_pages/config-json.md b/docs/_pages/config-json.md
index f49d55a9a2a4..0a847b2fa5e5 100755
--- a/docs/_pages/config-json.md
+++ b/docs/_pages/config-json.md
@@ -557,7 +557,7 @@ Configuring the asynchronous I/O module for offloading parameter and optimizer s
 
 | Description                    | Default |
 | ------------------------------ | ------- |
-| Print train loss every N steps | `10`    |
+| Print progress report every N training steps. The report includes the number of training steps, number of skipped optimizer updates (likely due to overflows in mixed-precision training), current learning rate, and current momentum.  | `10`    |
 
 <i>**wall_clock_breakdown**</i>: [boolean]
 

From da1fe2f82c9787af822af98d3d9372c92d7f0c69 Mon Sep 17 00:00:00 2001
From: Jeff Rasley <jerasley@microsoft.com>
Date: Wed, 16 Jun 2021 14:18:37 -0700
Subject: [PATCH 13/38] Remove hard torch dependency at install (#1166)

---
 docs/_tutorials/advanced-install.md  |  7 +++----
 op_builder/builder.py                |  8 +++++++-
 op_builder/cpu_adam.py               |  3 ++-
 op_builder/fused_adam.py             |  1 -
 op_builder/fused_lamb.py             |  1 -
 op_builder/quantizer.py              |  1 -
 op_builder/sparse_attn.py            |  7 ++++++-
 op_builder/stochastic_transformer.py |  1 -
 op_builder/transformer.py            |  1 -
 op_builder/transformer_inference.py  |  1 -
 setup.py                             | 27 +++++++++++++++++++--------
 11 files changed, 37 insertions(+), 21 deletions(-)

diff --git a/docs/_tutorials/advanced-install.md b/docs/_tutorials/advanced-install.md
index b1c54325c3a7..a1493a9e9e24 100755
--- a/docs/_tutorials/advanced-install.md
+++ b/docs/_tutorials/advanced-install.md
@@ -11,10 +11,6 @@ just-in-time (JIT) using [torch's JIT C++ extension loader that relies on
 ninja](https://pytorch.org/docs/stable/cpp_extension.html) to build and
 dynamically link them at runtime.
 
-**Note:** [PyTorch](https://pytorch.org/) must be installed _before_ installing
-DeepSpeed.
-{: .notice--info}
-
 ```bash
 pip install deepspeed
 ```
@@ -30,6 +26,9 @@ ds_report
 
 ## Pre-install DeepSpeed Ops
 
+**Note:** [PyTorch](https://pytorch.org/) must be installed _before_ pre-compiling any DeepSpeed c++/cuda ops. However, this is not required if using the default mode of JIT compilition of ops.
+{: .notice--info}
+
 Sometimes we have found it useful to pre-install either some or all DeepSpeed
 C++/CUDA ops instead of using the JIT compiled path. In order to support
 pre-installation we introduce build environment flags to turn on/off building
diff --git a/op_builder/builder.py b/op_builder/builder.py
index 97ae936c68bf..21547f896473 100644
--- a/op_builder/builder.py
+++ b/op_builder/builder.py
@@ -4,7 +4,6 @@
 import os
 import sys
 import time
-import torch
 import importlib
 from pathlib import Path
 import subprocess
@@ -17,6 +16,13 @@
 DEFAULT_TORCH_EXTENSION_PATH = "/tmp/torch_extensions"
 DEFAULT_COMPUTE_CAPABILITIES = "6.0;6.1;7.0"
 
+try:
+    import torch
+except ImportError:
+    print(
+        f"{WARNING} unable to import torch, please install it if you want to pre-compile any deepspeed ops."
+    )
+
 
 def installed_cuda_version():
     import torch.utils.cpp_extension
diff --git a/op_builder/cpu_adam.py b/op_builder/cpu_adam.py
index b93f27b81709..129ddeea3a29 100644
--- a/op_builder/cpu_adam.py
+++ b/op_builder/cpu_adam.py
@@ -3,7 +3,6 @@
 """
 import os
 import sys
-import torch
 import subprocess
 from .builder import CUDAOpBuilder
 
@@ -26,6 +25,7 @@ def sources(self):
         return ['csrc/adam/cpu_adam.cpp', 'csrc/adam/custom_cuda_kernel.cu']
 
     def include_paths(self):
+        import torch
         CUDA_INCLUDE = os.path.join(torch.utils.cpp_extension.CUDA_HOME, "include")
         return ['csrc/includes', CUDA_INCLUDE]
 
@@ -47,6 +47,7 @@ def simd_width(self):
         return '-D__SCALAR__'
 
     def cxx_args(self):
+        import torch
         CUDA_LIB64 = os.path.join(torch.utils.cpp_extension.CUDA_HOME, "lib64")
         SIMD_WIDTH = self.simd_width()
 
diff --git a/op_builder/fused_adam.py b/op_builder/fused_adam.py
index 0340ed02a8fb..c9a0d4436d01 100644
--- a/op_builder/fused_adam.py
+++ b/op_builder/fused_adam.py
@@ -1,7 +1,6 @@
 """
 Copyright 2020 The Microsoft DeepSpeed Team
 """
-import torch
 from .builder import CUDAOpBuilder
 
 
diff --git a/op_builder/fused_lamb.py b/op_builder/fused_lamb.py
index 4c73ecb404e3..169654809d06 100644
--- a/op_builder/fused_lamb.py
+++ b/op_builder/fused_lamb.py
@@ -1,7 +1,6 @@
 """
 Copyright 2020 The Microsoft DeepSpeed Team
 """
-import torch
 from .builder import CUDAOpBuilder
 
 
diff --git a/op_builder/quantizer.py b/op_builder/quantizer.py
index b3d6b2fd4509..43bc5778ea20 100644
--- a/op_builder/quantizer.py
+++ b/op_builder/quantizer.py
@@ -1,4 +1,3 @@
-import torch
 from .builder import CUDAOpBuilder
 
 
diff --git a/op_builder/sparse_attn.py b/op_builder/sparse_attn.py
index 58b8ca4f0e6b..6f30cc84da23 100644
--- a/op_builder/sparse_attn.py
+++ b/op_builder/sparse_attn.py
@@ -1,7 +1,6 @@
 """
 Copyright 2020 The Microsoft DeepSpeed Team
 """
-import torch
 import warnings
 from .builder import OpBuilder
 
@@ -28,6 +27,12 @@ def is_compatible(self):
         #command_status = list(map(self.command_exists, required_commands))
         #deps_compatible = all(command_status)
 
+        try:
+            import torch
+        except ImportError:
+            self.warning(f"unable to import torch, please install it first")
+            return False
+
         # torch-cpu will not have a cuda version
         if torch.version.cuda is None:
             cuda_compatible = False
diff --git a/op_builder/stochastic_transformer.py b/op_builder/stochastic_transformer.py
index b7e2f3845117..aa47c13c49e4 100644
--- a/op_builder/stochastic_transformer.py
+++ b/op_builder/stochastic_transformer.py
@@ -1,7 +1,6 @@
 """
 Copyright 2020 The Microsoft DeepSpeed Team
 """
-import torch
 from .transformer import TransformerBuilder
 
 
diff --git a/op_builder/transformer.py b/op_builder/transformer.py
index 82ab26f7f25c..2d48e2421b82 100644
--- a/op_builder/transformer.py
+++ b/op_builder/transformer.py
@@ -1,7 +1,6 @@
 """
 Copyright 2020 The Microsoft DeepSpeed Team
 """
-import torch
 from .builder import CUDAOpBuilder
 
 
diff --git a/op_builder/transformer_inference.py b/op_builder/transformer_inference.py
index cf5e870c906c..94db63711914 100755
--- a/op_builder/transformer_inference.py
+++ b/op_builder/transformer_inference.py
@@ -1,4 +1,3 @@
-import torch
 from .builder import CUDAOpBuilder
 
 
diff --git a/setup.py b/setup.py
index 654b983eec81..2424fdc6c46a 100755
--- a/setup.py
+++ b/setup.py
@@ -22,12 +22,14 @@
 from setuptools import setup, find_packages
 import time
 
+torch_available = True
 try:
     import torch
     from torch.utils.cpp_extension import BuildExtension
 except ImportError:
-    raise ImportError('Unable to import torch, please visit https://pytorch.org/ '
-                      'to see how to properly install torch on your system.')
+    torch_available = False
+    print('[WARNING] Unable to import torch, pre-compiling ops will be disabled. ' \
+        'Please visit https://pytorch.org/ to see how to properly install torch on your system.')
 
 from op_builder import ALL_OPS, get_default_compute_capatabilities
 
@@ -45,7 +47,7 @@ def fetch_requirements(path):
 }
 
 # If MPI is available add 1bit-adam requirements
-if torch.cuda.is_available():
+if torch_available and torch.cuda.is_available():
     if shutil.which('ompi_info') or shutil.which('mpiname'):
         cupy = f"cupy-cuda{torch.version.cuda.replace('.','')[:3]}"
         extras_require['1bit_adam'].append(cupy)
@@ -60,12 +62,17 @@ def fetch_requirements(path):
 cmdclass = {}
 
 # For any pre-installed ops force disable ninja
-cmdclass['build_ext'] = BuildExtension.with_options(use_ninja=False)
+if torch_available:
+    cmdclass['build_ext'] = BuildExtension.with_options(use_ninja=False)
 
-TORCH_MAJOR = torch.__version__.split('.')[0]
-TORCH_MINOR = torch.__version__.split('.')[1]
+if torch_available:
+    TORCH_MAJOR = torch.__version__.split('.')[0]
+    TORCH_MINOR = torch.__version__.split('.')[1]
+else:
+    TORCH_MAJOR = "0"
+    TORCH_MINOR = "0"
 
-if not torch.cuda.is_available():
+if torch_available and not torch.cuda.is_available():
     # Fix to allow docker builds, similar to https://github.com/NVIDIA/apex/issues/486
     print(
         "[WARNING] Torch did not find cuda available, if cross-compiling or running with cpu only "
@@ -81,6 +88,9 @@ def fetch_requirements(path):
 BUILD_OP_DEFAULT = int(os.environ.get('DS_BUILD_OPS', BUILD_OP_PLATFORM))
 print(f"DS_BUILD_OPS={BUILD_OP_DEFAULT}")
 
+if BUILD_OP_DEFAULT:
+    assert torch_available, "Unable to pre-compile ops without torch installed. Please install torch before attempting to pre-compile ops."
+
 
 def command_exists(cmd):
     if sys.platform == "win32":
@@ -109,6 +119,7 @@ def op_enabled(op_name):
 
     # If op install enabled, add builder to extensions
     if op_enabled(op_name) and op_compatible:
+        assert torch_available, f"Unable to pre-compile {op_name}, please first install torch"
         install_ops[op_name] = op_enabled(op_name)
         ext_modules.append(builder.builder())
 
@@ -170,7 +181,7 @@ def create_dir_symlink(src, dest):
 torch_version = ".".join([TORCH_MAJOR, TORCH_MINOR])
 # Set cuda_version to 0.0 if cpu-only
 cuda_version = "0.0"
-if torch.version.cuda is not None:
+if torch_available and torch.version.cuda is not None:
     cuda_version = ".".join(torch.version.cuda.split('.')[:2])
 torch_info = {"version": torch_version, "cuda_version": cuda_version}
 

From 4c114a27a067ecb30a30a918f60e91d3c5e7e254 Mon Sep 17 00:00:00 2001
From: Jeff Rasley <jerasley@microsoft.com>
Date: Wed, 16 Jun 2021 15:46:35 -0700
Subject: [PATCH 14/38] patch in ompi local rank if local_rank isn't set
 (#1164)

---
 deepspeed/runtime/engine.py | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
index 28ad2e977f17..74e91ac98a88 100755
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -562,6 +562,14 @@ def _configure_with_arguments(self, args, mpu):
         # environment variable is set. We must align args.local_rank to this value for
         # backwards compatability with scripts relying on [args|self].local_rank containing
         # the correct local rank info. _do_args_sanity_check will ensure this is the case.
+
+        if "OMPI_COMM_WORLD_LOCAL_RANK" in os.environ:
+            ompi_local_rank = os.environ.get("OMPI_COMM_WORLD_LOCAL_RANK")
+            local_rank = os.environ.get('LOCAL_RANK', ompi_local_rank)
+            assert ompi_local_rank == local_rank, f"LOCAL_RANK ({local_rank}) != OMPI_COMM_WORLD_LOCAL_RANK ({mpi_local_rank}), " \
+                "not sure how to proceed as we're seeing conficting local rank info."
+            os.environ['LOCAL_RANK'] = local_rank
+
         self.local_rank = int(os.environ['LOCAL_RANK'])
         if hasattr(args, 'local_rank'):
             args.local_rank = self.local_rank
@@ -581,8 +589,10 @@ def _do_args_sanity_check(self, args):
                 assert args.deepspeed_config is None, "Not sure how to proceed, we were given both a deepscale_config and deepspeed_config"
             args.deepspeed_config = args.deepscale_config
 
-        assert "LOCAL_RANK" in os.environ, "DeepSpeed requires the LOCAL_RANK environment variable, it is set by the deepspeed launcher, " \
-            "deepspeed.init_distributed, or the torch.distributed launcher. If using a different launcher please ensure LOCAL_RANK is set prior to initializing deepspeed."
+        assert "LOCAL_RANK" in os.environ or "OMPI_COMM_WORLD_LOCAL_RANK" in os.environ, "DeepSpeed requires the LOCAL_RANK environment " \
+            "variable, it is set by the deepspeed launcher, deepspeed.init_distributed, or the torch.distributed launcher. If using a " \
+            "different launcher please ensure LOCAL_RANK is set prior to initializing deepspeed."
+
         if hasattr(args, 'local_rank') and args.local_rank != None:
             assert isinstance(args.local_rank, int), f"args.local_rank of {args.local_rank} is an unknown type {type(args.local_rank)}"
             if args.local_rank >= 0:

From 1ba3e8e3364049ecfc647aae4715ec9441676c51 Mon Sep 17 00:00:00 2001
From: Jeff Rasley <jerasley@microsoft.com>
Date: Thu, 17 Jun 2021 10:53:50 -0700
Subject: [PATCH 15/38] add assert to make sure zero isn't enabled with 1bit-*
 (#1169)

---
 deepspeed/runtime/engine.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
index 74e91ac98a88..d6aa5a0c291b 100755
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -778,6 +778,7 @@ def _configure_basic_optimizer(self, model_parameters):
             from deepspeed.ops.lamb import FusedLamb
             optimizer = FusedLamb(model_parameters, **optimizer_parameters)
         elif self.optimizer_name() == ONEBIT_ADAM_OPTIMIZER:
+            assert not self.zero_optimization(), "1bit-Adam is not compatible with ZeRO"
             from deepspeed.runtime.fp16.onebit.adam import OnebitAdam
             optimizer = OnebitAdam(model_parameters, self, **optimizer_parameters)
             if not self.fp16_enabled():
@@ -785,6 +786,7 @@ def _configure_basic_optimizer(self, model_parameters):
                     f'Currently the convergence of 1-bit Adam is only verified under FP16'
                 )
         elif self.optimizer_name() == ONEBIT_LAMB_OPTIMIZER:
+            assert not self.zero_optimization(), "1bit-Lamb is not compatible with ZeRO"
             from deepspeed.runtime.fp16.onebit.lamb import OnebitLamb
             optimizer = OnebitLamb(model_parameters, self, **optimizer_parameters)
             if not self.fp16_enabled():

From b1669c0d8f6dbe65162b534f2c6967f64d143dc3 Mon Sep 17 00:00:00 2001
From: Olatunji Ruwase <olruwase@microsoft.com>
Date: Thu, 17 Jun 2021 17:01:25 -0700
Subject: [PATCH 16/38] Avoid partitioning small activations (#1154)

---
 deepspeed/runtime/activation_checkpointing/checkpointing.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deepspeed/runtime/activation_checkpointing/checkpointing.py b/deepspeed/runtime/activation_checkpointing/checkpointing.py
index f516e4de5203..f56e4b853313 100644
--- a/deepspeed/runtime/activation_checkpointing/checkpointing.py
+++ b/deepspeed/runtime/activation_checkpointing/checkpointing.py
@@ -416,7 +416,7 @@ def save_args_for_backward(*all_args):
 
             inputs = []
             for i, item in enumerate(args[:-1]):
-                if not torch.is_tensor(item):
+                if not torch.is_tensor(item) or mp_size > item.numel():
                     inputs.append(item)
                     continue
 

From 3b68984498cd10f138f69725597ca83452294ac7 Mon Sep 17 00:00:00 2001
From: Jeff Rasley <jerasley@microsoft.com>
Date: Mon, 21 Jun 2021 10:41:27 -0700
Subject: [PATCH 17/38] remove torchvision dependency (#1178)

---
 requirements/requirements.txt | 1 -
 1 file changed, 1 deletion(-)

diff --git a/requirements/requirements.txt b/requirements/requirements.txt
index afe7d231d4bf..c0e94ec4dabe 100644
--- a/requirements/requirements.txt
+++ b/requirements/requirements.txt
@@ -1,5 +1,4 @@
 torch>=1.2
-torchvision>=0.4.0
 tqdm
 tensorboardX==1.8
 ninja

From 0982ecca24b08ae02a2e87ba61641772b03886fe Mon Sep 17 00:00:00 2001
From: Jeff Rasley <jerasley@microsoft.com>
Date: Wed, 23 Jun 2021 10:18:10 -0700
Subject: [PATCH 18/38] bump to 0.4.2

---
 version.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/version.txt b/version.txt
index 267577d47e49..2b7c5ae01848 100644
--- a/version.txt
+++ b/version.txt
@@ -1 +1 @@
-0.4.1
+0.4.2

From df8b1f884fbc7297bdaaa017a6b58e98c3d4780c Mon Sep 17 00:00:00 2001
From: Stas Bekman <stas00@users.noreply.github.com>
Date: Wed, 23 Jun 2021 12:53:31 -0700
Subject: [PATCH 19/38] zero_to_fp32: restore persistent buffers (#1146)

Co-authored-by: Olatunji Ruwase <olruwase@microsoft.com>
Co-authored-by: Jeff Rasley <jerasley@microsoft.com>
---
 deepspeed/runtime/engine.py     | 22 +++++++++++++++
 deepspeed/utils/zero_to_fp32.py | 50 +++++++++++++++++++++++++++++++--
 2 files changed, 69 insertions(+), 3 deletions(-)

diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
index d6aa5a0c291b..4692bb709079 100755
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -1915,6 +1915,7 @@ def _save_checkpoint(self, save_dir, tag, client_state={}):
         self._curr_ckpt_path = os.path.join(save_dir, tag)
 
         state = dict(module=self.module_state_dict(),
+                     buffer_names=self._get_buffer_names(),
                      optimizer=self.optimizer.state_dict()
                      if self.optimizer and not self.zero_optimization() else None,
                      lr_scheduler=self.lr_scheduler.state_dict()
@@ -1934,6 +1935,27 @@ def _save_checkpoint(self, save_dir, tag, client_state={}):
         torch.save(state, save_path)
         self._curr_save_path = None
 
+    def _get_buffer_names(self):
+        buffer_names = []
+
+        # we save buffer names so that we could extract later the real buffers from the saved
+        # state_dict["module"] in the non-zero checkpoint - the buffers are already there but they
+        # are intermixed with param placeholders
+
+        # have to traverse the tree to be able to skip non-persistent buffers
+        def get_layer_named_buffers(module, prefix=""):
+            for name, buf in module.named_buffers(recurse=False):
+                if buf is not None and name not in module._non_persistent_buffers_set:
+                    buffer_names.append(prefix + name)
+
+            for name, child in module.named_children():
+                if child is not None:
+                    get_layer_named_buffers(child, prefix + name + ".")
+
+        get_layer_named_buffers(self.module, prefix="")
+
+        return buffer_names
+
     def _get_param_shapes(self):
         param_shapes = OrderedDict()
         for name, param in self.module.named_parameters():
diff --git a/deepspeed/utils/zero_to_fp32.py b/deepspeed/utils/zero_to_fp32.py
index 2d98a39e3fc7..1b58c403618c 100644
--- a/deepspeed/utils/zero_to_fp32.py
+++ b/deepspeed/utils/zero_to_fp32.py
@@ -20,13 +20,27 @@
 debug = 0
 
 
+def get_model_state_file(checkpoint_dir):
+
+    if not os.path.isdir(checkpoint_dir):
+        raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")
+
+    # there should be only one file
+    file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt")
+
+    if not os.path.exists(file):
+        raise FileNotFoundError(f"can't find '{file}' in directory '{checkpoint_dir}'")
+
+    return file
+
+
 def get_optim_files(checkpoint_dir):
 
     if not os.path.isdir(checkpoint_dir):
         raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")
 
     # XXX: need to test that this simple glob rule works for multi-node setup too
-    optim_files = sorted(glob.glob(f"{checkpoint_dir}/*_optim_states.pt"))
+    optim_files = sorted(glob.glob(os.path.join(checkpoint_dir, "*_optim_states.pt")))
 
     if len(optim_files) == 0:
         raise FileNotFoundError(
@@ -35,13 +49,34 @@ def get_optim_files(checkpoint_dir):
     return optim_files
 
 
+def parse_model_state(file):
+
+    # load to cpu
+    device = torch.device('cpu')
+    state_dict = torch.load(file, map_location=device)
+
+    if "buffer_names" not in state_dict:
+        raise ValueError(f"{file} is not a model state checkpoint")
+    buffer_names = state_dict["buffer_names"]
+    if debug:
+        print(buffer_names)
+
+    # recover just the buffers while restoring them to fp32 if they were saved in fp16
+    buffers = {
+        k: v.float()
+        for k,
+        v in state_dict["module"].items() if k in buffer_names
+    }
+    return buffers
+
+
 def parse_optim_states(files):
     state_dicts = []
     for f in files:
         state_dicts.append(torch.load(f))
 
     if not "zero_stage" in state_dicts[0]['optimizer_state_dict']:
-        raise ValueError(f"non zero checkpoint")
+        raise ValueError(f"{files[0]} is not a zero checkpoint")
     zero_stage = state_dicts[0]['optimizer_state_dict']["zero_stage"]
     world_size = state_dicts[0]['optimizer_state_dict']["partition_count"]
     param_shapes = state_dicts[0]["param_shapes"]
@@ -86,7 +121,9 @@ def convert_zero_chkpt_to_fp32_consolid_state_dict(checkpoint_dir, output_file):
     """
     print(f"Processing zero checkpoint '{checkpoint_dir}'")
 
+    model_file = get_model_state_file(checkpoint_dir)
     optim_files = get_optim_files(checkpoint_dir)
+    buffers = parse_model_state(model_file)
     zero_stage, world_size, param_shapes, fp32_flat_groups = parse_optim_states(optim_files)
     print(
         f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")
@@ -108,9 +145,16 @@ def convert_zero_chkpt_to_fp32_consolid_state_dict(checkpoint_dir, output_file):
         # XXX: memory usage doubles here (zero2)
         full_single_fp32_vector = torch.cat(fp32_flat_groups, 0)
 
+    state_dict = OrderedDict()
+
+    # buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    # params
     # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
     # out-of-core computing solution
-    state_dict = OrderedDict()
     offset = 0
     total_numel = 0
     for name, shape in param_shapes.items():

From 0c1802cc8b877a510eea61dd562b82c433876802 Mon Sep 17 00:00:00 2001
From: Stas Bekman <stas00@users.noreply.github.com>
Date: Wed, 23 Jun 2021 12:59:42 -0700
Subject: [PATCH 20/38] ZeRO 2+3 memory estimators (#965)

Co-authored-by: Samyam Rajbhandari <samyamr@microsoft.com>
Co-authored-by: Jeff Rasley <jerasley@microsoft.com>
---
 deepspeed/runtime/zero/stage2.py | 103 +++++++++++
 deepspeed/runtime/zero/stage3.py | 152 ++++++++++++++++
 docs/code-docs/source/index.rst  |   8 +
 docs/code-docs/source/memory.rst | 288 +++++++++++++++++++++++++++++++
 4 files changed, 551 insertions(+)
 create mode 100644 docs/code-docs/source/memory.rst

diff --git a/deepspeed/runtime/zero/stage2.py b/deepspeed/runtime/zero/stage2.py
index a9216160f2c5..ad529aa96391 100755
--- a/deepspeed/runtime/zero/stage2.py
+++ b/deepspeed/runtime/zero/stage2.py
@@ -2000,3 +2000,106 @@ def _handle_overflow(cpu_sum, x, i):
         logger.info(
             f"rank {rank} detected overflow {cpu_sum} in tensor {i}:{t_i} shape {x.shape}"
         )
+
+
+def estimate_zero2_model_states_mem_needs(total_params,
+                                          num_gpus_per_node=1,
+                                          num_nodes=1,
+                                          cpu_offload=True,
+                                          additional_buffer_factor=1.5):
+
+    total_gpus = num_nodes * num_gpus_per_node
+
+    if cpu_offload:
+        gpu_mem = 2 * total_params
+        cpu_mem = total_params * max(4 * total_gpus, 16) * additional_buffer_factor
+    else:
+        gpu_mem = 4 * total_params + int(16 * total_params / total_gpus)
+        cpu_mem = total_params * 4 * num_gpus_per_node * additional_buffer_factor
+
+    return int(cpu_mem), int(gpu_mem)
+
+
+def model_to_params(model):
+    # shared params calculated only once
+    total_params = sum(
+        dict((p.data_ptr(),
+              p.numel()) for p in model.parameters()).values())
+    return total_params
+
+
+def estimate_zero2_model_states_mem_needs_all_live(model,
+                                                   num_gpus_per_node=1,
+                                                   num_nodes=1,
+                                                   additional_buffer_factor=1.5):
+    """
+    Print out estimates on memory usage requirements for ZeRO 2 params, optim states and gradients
+    for a given ``model`` and hardware setup.
+
+    If you have an actual model object, use this function and everything will be derived
+    automatically.
+
+    If it's a hypothetical model, use ``estimate_zero2_model_states_mem_needs_all_cold`` where you have to pass
+    the ``total_params`` explicitly.
+
+    Args:
+        - ``model``: ``nn.Module`` object
+        - ``num_gpus_per_node``: how many gpus per node (defaults to 1)
+        - ``num_nodes``: how many nodes (defaults to 1),
+        - ``additional_buffer_factor``: estimation factor (defaults to 1.5):
+
+    """
+
+    total_params = model_to_params(model)
+
+    estimate_zero2_model_states_mem_needs_all_cold(
+        total_params=total_params,
+        num_gpus_per_node=num_gpus_per_node,
+        num_nodes=num_nodes,
+        additional_buffer_factor=additional_buffer_factor)
+
+
+def estimate_zero2_model_states_mem_needs_all_cold(total_params,
+                                                   num_gpus_per_node=1,
+                                                   num_nodes=1,
+                                                   additional_buffer_factor=1.5):
+    """
+    Print out estimates on memory usage requirements for ZeRO 2 params, optim states and gradients
+    for a given ``model`` and hardware setup.
+
+    If it's a hypothetical model, use this function where you have to pass
+    the ``total_params`` and ``largest_layer_params`` explicitly.
+
+    If you have an actual model object, use ``estimate_zero2_model_states_mem_needs_all_live`` and everything
+    will be derived automatically.
+
+    Args:
+        - ``total_params``: total  model params
+        - ``num_gpus_per_node``: how many gpus per node (defaults to 1)
+        - ``num_nodes``: how many nodes (defaults to 1),
+        - ``additional_buffer_factor``: estimation factor (defaults to 1.5):
+
+    """
+    def format_options(cpu_offload):
+        enabled = []
+        enabled.append(f"cpu_offload={1 if cpu_offload else 0}")
+        return ", ".join(enabled)
+
+    nodes_str = "nodes" if num_nodes > 1 else "node"
+    gpus_str = "GPUs" if num_gpus_per_node > 1 else "GPU"
+    print(
+        "Estimated memory needed for params, optim states and gradients for a:\n"
+        f"HW: Setup with {num_nodes} {nodes_str}, {num_gpus_per_node} {gpus_str} per node.\n"
+        f"SW: Model with {int(total_params/1e6)}M total params.")
+    print("  per CPU  |  per GPU |   Options")
+    for cpu_offload in [True, False]:
+        cpu_mem, gpu_mem = estimate_zero2_model_states_mem_needs(
+            total_params=total_params,
+            num_gpus_per_node=num_gpus_per_node,
+            num_nodes=num_nodes,
+            cpu_offload=cpu_offload,
+            additional_buffer_factor=additional_buffer_factor
+        )
+
+        options_str = format_options(cpu_offload=cpu_offload)
+        print(f" {cpu_mem/2**30:7.2f}GB | {gpu_mem/2**30:6.2f}GB | {options_str}")
diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py
index 548c38a072c3..d492df42067d 100755
--- a/deepspeed/runtime/zero/stage3.py
+++ b/deepspeed/runtime/zero/stage3.py
@@ -3258,3 +3258,155 @@ def _handle_overflow(cpu_sum, x, i):
         logger.info(
             f"rank {rank} detected overflow {cpu_sum} in tensor {i}:{t_i} shape {x.shape}"
         )
+
+
+def estimate_zero3_model_states_mem_needs(total_params,
+                                          largest_layer_params,
+                                          num_gpus_per_node=1,
+                                          num_nodes=1,
+                                          cpu_offload=True,
+                                          cpu_offload_params=True,
+                                          zero_init=True,
+                                          additional_buffer_factor=1.5):
+
+    total_gpus = num_nodes * num_gpus_per_node
+    gpus_factor = 1 / num_nodes
+    largest_layer_memory = (4 * largest_layer_params)
+
+    if cpu_offload:
+        if cpu_offload_params:
+            gpu_mem = largest_layer_memory
+
+            if zero_init:
+                cpu_mem = total_params * 18 * gpus_factor * additional_buffer_factor
+            else:
+                cpu_mem = total_params * max(4 * num_gpus_per_node,
+                                             18 * gpus_factor) * additional_buffer_factor
+
+        else:
+            gpu_mem = largest_layer_memory + int(2 * total_params / total_gpus)
+
+            if zero_init:
+                cpu_mem = total_params * 16 * gpus_factor * additional_buffer_factor
+            else:
+                cpu_mem = total_params * max(4 * num_gpus_per_node,
+                                             16 * gpus_factor) * additional_buffer_factor
+    else:
+        gpu_mem = largest_layer_memory + int(18 * total_params / total_gpus)
+        if zero_init:
+            cpu_mem = largest_layer_params * 4 * num_gpus_per_node * additional_buffer_factor
+        else:
+            cpu_mem = total_params * 4 * num_gpus_per_node * additional_buffer_factor
+
+    return int(cpu_mem), int(gpu_mem), largest_layer_memory
+
+
+def model_to_params(model):
+    # shared params calculated only once
+    total_params = sum(
+        dict((p.data_ptr(),
+              p.numel()) for p in model.parameters()).values())
+
+    largest_layer_params = 0
+    for m in model.modules():
+        # assuming no shared params within a single layer
+        layer_params = sum(p.numel() for p in m.parameters(recurse=False))
+        largest_layer_params = max(largest_layer_params, layer_params)
+
+    return total_params, largest_layer_params
+
+
+import math
+
+
+def estimate_zero3_model_states_mem_needs_all_live(model,
+                                                   num_gpus_per_node=1,
+                                                   num_nodes=1,
+                                                   additional_buffer_factor=1.5):
+    """
+    Print out estimates on memory usage requirements for ZeRO 3 params, optim states and gradients
+    for a given ``model`` and hardware setup.
+
+    If you have an actual model object, use this function and everything will be derived
+    automatically.
+
+    If it's a hypothetical model, use ``estimate_zero3_model_states_mem_needs_all_cold`` where you have to pass
+    the ``total_params`` and ``largest_layer_params`` explicitly.
+
+    Args:
+        - ``model``: ``nn.Module`` object
+        - ``num_gpus_per_node``: how many gpus per node (defaults to 1)
+        - ``num_nodes``: how many nodes (defaults to 1),
+        - ``additional_buffer_factor``: estimation factor (defaults to 1.5):
+
+    """
+
+    total_params, largest_layer_params = model_to_params(model)
+
+    estimate_zero3_model_states_mem_needs_all_cold(
+        total_params=total_params,
+        largest_layer_params=largest_layer_params,
+        num_gpus_per_node=num_gpus_per_node,
+        num_nodes=num_nodes,
+        additional_buffer_factor=additional_buffer_factor)
+
+
+def estimate_zero3_model_states_mem_needs_all_cold(total_params,
+                                                   largest_layer_params,
+                                                   num_gpus_per_node=1,
+                                                   num_nodes=1,
+                                                   additional_buffer_factor=1.5):
+    """
+    Print out estimates on memory usage requirements for ZeRO 3 params, optim states and gradients
+    for a given ``model`` and hardware setup.
+
+    If it's a hypothetical model, use this function where you have to pass
+    the ``total_params`` and ``largest_layer_params`` explicitly.
+
+    If you have an actual model object, use ``estimate_zero3_model_states_mem_needs_all_live`` and everything
+    will be derived automatically.
+
+    Args:
+        - ``total_params``: total  model params
+        - ``largest_layer_params``: largest layer's params
+        - ``num_gpus_per_node``: how many gpus per node (defaults to 1)
+        - ``num_nodes``: how many nodes (defaults to 1),
+        - ``additional_buffer_factor``: estimation factor (defaults to 1.5):
+
+    """
+    def format_options(cpu_offload, cpu_offload_params, zero_init):
+        enabled = []
+        enabled.append(f"cpu_offload={1 if cpu_offload else 0}")
+        enabled.append(f"cpu_offload_params={1 if cpu_offload_params else 0}")
+        enabled.append(f"zero_init={1 if zero_init else 0}")
+        return ", ".join(enabled)
+
+    nodes_str = "nodes" if num_nodes > 1 else "node"
+    gpus_str = "GPUs" if num_gpus_per_node > 1 else "GPU"
+    print(
+        "Estimated memory needed for params, optim states and gradients for a:\n"
+        f"HW: Setup with {num_nodes} {nodes_str}, {num_gpus_per_node} {gpus_str} per node.\n"
+        f"SW: Model with {int(total_params/1e6)}M total params, {int(largest_layer_params/1e6)}M largest layer params."
+    )
+    print("  per CPU  |  per GPU |   Options")
+    for cpu_offload in [True, False]:
+        for cpu_offload_params in [True, False]:
+            if not cpu_offload and cpu_offload_params:
+                continue
+            for zero_init in [True, False]:
+                cpu_mem, gpu_mem, largest_layer_memory = estimate_zero3_model_states_mem_needs(
+                    total_params=total_params,
+                    largest_layer_params=largest_layer_params,
+                    num_gpus_per_node=num_gpus_per_node,
+                    num_nodes=num_nodes,
+                    cpu_offload=cpu_offload,
+                    cpu_offload_params=cpu_offload_params,
+                    zero_init=zero_init,
+                    additional_buffer_factor=additional_buffer_factor
+                )
+
+                options_str = format_options(cpu_offload=cpu_offload,
+                                             cpu_offload_params=cpu_offload_params,
+                                             zero_init=zero_init)
+                print(
+                    f" {cpu_mem/2**30:7.2f}GB | {gpu_mem/2**30:6.2f}GB | {options_str}")
diff --git a/docs/code-docs/source/index.rst b/docs/code-docs/source/index.rst
index 2f39da002d5e..5175209cc1c4 100644
--- a/docs/code-docs/source/index.rst
+++ b/docs/code-docs/source/index.rst
@@ -79,6 +79,14 @@ Flops Profiler
 
    flops-profiler
 
+
+Memory Usage
+------------------
+.. toctree::
+   :maxdepth: 2
+
+   memory
+
 Indices and tables
 ------------------
 
diff --git a/docs/code-docs/source/memory.rst b/docs/code-docs/source/memory.rst
new file mode 100644
index 000000000000..7f810309a80b
--- /dev/null
+++ b/docs/code-docs/source/memory.rst
@@ -0,0 +1,288 @@
+Memory Requirements
+-----------------------
+
+
+API To Estimate Memory Usage
+============================
+
+ZeRO2:
+
+.. autofunction:: deepspeed.runtime.zero.stage2.estimate_zero2_model_states_mem_needs_all_live
+
+.. autofunction:: deepspeed.runtime.zero.stage2.estimate_zero2_model_states_mem_needs_all_cold
+
+Examples:
+
+Let's try a 3B model with just 1 node with 8 gpus, using live model:
+
+.. code-block:: bash
+
+    python -c 'from transformers import AutoModel; \
+    from deepspeed.runtime.zero.stage2 import estimate_zero2_model_states_mem_needs_all_live; \
+    model = AutoModel.from_pretrained("t5-3b"); \
+    estimate_zero2_model_states_mem_needs_all_live(model, num_gpus_per_node=8, num_nodes=1)'
+    Estimated memory needed for params, optim states and gradients for a:
+    HW: Setup with 1 node, 8 GPUs per node.
+    SW: Model with 2851M total params.
+      per CPU  |  per GPU |   Options
+      127.48GB |   5.31GB | cpu_offload=1
+      127.48GB |  15.93GB | cpu_offload=0
+
+Now, without the actual model, which requires us to know ``total_params`` and
+``largest_layer_params``, but we got those from the run above, so future estimators are now much
+faster as we don't need to load the model.
+
+.. code-block:: bash
+
+    python -c 'from deepspeed.runtime.zero.stage2 import estimate_zero2_model_states_mem_needs_all_cold; \
+    estimate_zero2_model_states_mem_needs_all_cold(total_params=2851e6, num_gpus_per_node=8, num_nodes=1)'
+    Estimated memory needed for params, optim states and gradients for a:
+    HW: Setup with 1 node, 8 GPUs per node.
+    SW: Model with 2851M total params.
+      per CPU  |  per GPU |   Options
+      127.45GB |   5.31GB | cpu_offload=1
+      127.45GB |  15.93GB | cpu_offload=0
+
+There is a slight difference due to rounding - the actual live model has a few more params
+
+
+ZeRO3:
+
+.. autofunction:: deepspeed.runtime.zero.stage3.estimate_zero3_model_states_mem_needs_all_live
+
+.. autofunction:: deepspeed.runtime.zero.stage3.estimate_zero3_model_states_mem_needs_all_cold
+
+Examples:
+
+Let's try a 3B model with just 1 node with 8 gpus, using live model:
+
+.. code-block:: bash
+
+    python -c 'from transformers import AutoModel; \
+    from deepspeed.runtime.zero.stage3 import estimate_zero3_model_states_mem_needs_all_live; \
+    model = AutoModel.from_pretrained("t5-3b"); \
+    estimate_zero3_model_states_mem_needs_all_live(model, num_gpus_per_node=8, num_nodes=1)'
+
+    Estimated memory needed for params, optim states and gradients for a:
+    HW: Setup with 1 node, 8 GPUs per node.
+    SW: Model with 2851M total params, 32M largest layer params.
+      per CPU  |  per GPU |   Options
+       71.71GB |   0.12GB | cpu_offload=1, cpu_offload_params=1, zero_init=1
+      127.48GB |   0.12GB | cpu_offload=1, cpu_offload_params=1, zero_init=0
+       63.74GB |   0.79GB | cpu_offload=1, cpu_offload_params=0, zero_init=1
+      127.48GB |   0.79GB | cpu_offload=1, cpu_offload_params=0, zero_init=0
+        1.47GB |   6.10GB | cpu_offload=0, cpu_offload_params=0, zero_init=1
+      127.48GB |   6.10GB | cpu_offload=0, cpu_offload_params=0, zero_init=0
+
+Now, without the actual model, which requires us to know ``total_params`` and
+``largest_layer_params``, but we got those from the run above, so future estimators are now much
+faster as we don't need to load the model.
+
+.. code-block:: bash
+
+    python -c 'from deepspeed.runtime.zero.stage3 import estimate_zero3_model_states_mem_needs_all_cold; \
+    estimate_zero3_model_states_mem_needs_all_cold(total_params=2851e6, largest_layer_params=32e6, num_gpus_per_node=8, num_nodes=1)'
+
+    Estimated memory needed for params, optim states and gradients for a:
+    HW: Setup with 1 node, 8 GPUs per node.
+    SW: Model with 2851M total params, 32M largest layer params.
+      per CPU  |  per GPU |   Options
+       71.69GB |   0.12GB | cpu_offload=1, cpu_offload_params=1, zero_init=1
+      127.45GB |   0.12GB | cpu_offload=1, cpu_offload_params=1, zero_init=0
+       63.72GB |   0.78GB | cpu_offload=1, cpu_offload_params=0, zero_init=1
+      127.45GB |   0.78GB | cpu_offload=1, cpu_offload_params=0, zero_init=0
+        1.43GB |   6.09GB | cpu_offload=0, cpu_offload_params=0, zero_init=1
+      127.45GB |   6.09GB | cpu_offload=0, cpu_offload_params=0, zero_init=0
+
+There is a slight difference due to rounding - the actual live model has a few more params
+
+
+
+Discussion
+==========
+
+Let's look in detail how the memory estimator API calculates these numbers and also discuss some additional numbers that aren't covered by the API.
+
+In the following discussion:
+
+- ``params`` - total number of model params, which can be calculated as:
+
+.. code-block:: python
+
+    print(sum(dict((p.data_ptr(), p.numel()) for p in model.parameters()).values()))
+
+Some models already include the number of params in the model name, e.g. t5-11b (11B params), gpt-neo-1.3B (1.3B params), etc.
+
+Also if the model weights are stored in ``fp32`` the other quick way to calculate the size of the model is to simply divide the size of the ``state_dict`` file by 4 (fp32 == 4 bytes). For example, you can see that `t5-11b's pytorch_model.bin <https://huggingface.co/t5-11b/tree/main>`__ is 42.1GB in size, so if we divide it by 4, we can immediately tell it's an 11B model.
+
+The following calculations show how much memory is required by model params, gradients and optimizer states. In addition to those you will need enough memory to fit activation calculations and any temporary memory for intermediate calculations, which for long sequences could be very significant (e.g. could take the same amount of memory as params+grads+optim_states combined).
+
+The optimizer states assume that ``Adam`` is used, where 4 bytes per parameter are used by momentum and another 4 by variance (8 in total).
+
+Gradients at ``fp32`` take 4 bytes, and parameters take 2 bytes at ``fp16` and 4 bytes at ``fp32``.
+
+**GPU RAM**
+
+The big question is how big of a model you can fit on the hardware you have? Or rather what size of a GPU RAM do you need to fit the desired model.
+
+
+* ZeRO-2:
+
+   - ``"cpu_offload": true``: 2 * params
+
+   Example: a 40GB GPU can fit ~11B param model (regardless of how many GPUs are used). Here the model is loaded in ``fp16`` so just the model weights take about 22GB and the remaining 18GB are used by other components. You can barely fit a very small batch size in this scenario.
+
+   - ``"cpu_offload": false``: 4 params + 16 params/ (total number of gpus)
+
+* ZeRO-3:
+
+largest_layer_memory = 4*largest_layer_params - GPU memory needed to gather the largest layer on a single GPU. 2 bytes fp16 params are gathered and 2 bytes fp16 grads are computed (total 4x). The optimizer states and fp32 parameters are updated in partitioned form and copied to fp16 params in partitioned form. This happens during the optimizer step. After that the fp16 params are sufficient.
+
+   - case 1: ``"cpu_offload": false, "cpu_offload_params": false`` - largest_layer_memory + 18 * params / total number of gpus across all nodes
+   - case 2: ``"cpu_offload": true, "cpu_offload_params": true``- largest_layer_memory. The main limit here is general RAM.
+   - case 3: ``"cpu_offload": true, "cpu_offload_params": false``- largest_layer_memory + 2 * params / total number of gpus across all nodes
+
+     Example:
+
+.. code-block:: python
+
+from transformers import AutoModel
+model = AutoModel.from_pretrained("t5-large")
+
+# shared params calculated only ones
+total_params = sum(dict((p.data_ptr(), p.numel()) for p in model.parameters()).values())
+
+largest_layer_params = 0
+for m in model.modules():
+    # assuming no shared params within a single layer
+    layer_params = sum(p.numel() for p in m.parameters(recurse=False))
+    largest_layer_params = max(largest_layer_params, layer_params)
+
+largest_layer_memory = (4*largest_layer_params)
+
+total_gpus = 4
+
+case1 = largest_layer_memory + int(18*total_params/total_gpus)
+case2 = largest_layer_memory
+case3 = largest_layer_memory + int(2*total_params/total_gpus)
+
+print(f"total params:         {total_params/1e6:6.2f}M")
+print(f"largest layer params: {largest_layer_params/1e6:6.2f}M")
+print(f"largest layer memory: {largest_layer_memory>>20:6}MB")
+print(f"case1 gpu memory: {(case1)>>20:6}MB")
+print(f"case2 gpu memory: {(case2)>>20:6}MB")
+print(f"case3 gpu memory: {(case3)>>20:6}MB")
+
+total params:         737.67M
+largest layer params:  32.90M
+largest layer memory:    125MB
+case1 gpu memory:   3291MB
+case2 gpu memory:    125MB
+case3 gpu memory:    477MB
+
+
+**General RAM**:
+
+One of the key features of ZeRO is its CPU offload which can dramatically extend the total memory pool accessible to the project by using general RAM. One can easily expand their general RAM by 10x times, at a significantly lower cost than what it'd take to have the same GPU RAM. And often, it's not even possible to buy GPUs with a lot of RAM (112GB GPU anybody?) since they simply don't yet exist.
+
+In the following calculations we will use:
+
+- ``additional_buffer_factor=1.5`` as an additional buffer factor to be conservative
+- ``n_gpus`` the number of GPUs on a single node (machine)
+- ``total_gpus`` the total number of GPUs across all nodes
+- ``params`` - total number of model params (see above for how to get this number)
+
+* ZeRO-2:
+
+   - ``"cpu_offload": false``:
+
+      params * 4 * n_gpus * additional_buffer_factor - this is the memory needed only at the beginning to initialize the model on CPU memory
+
+   - ``"cpu_offload": true``:
+
+      params * max(4 * n_gpus, 16) * additional_buffer_factor
+
+   Example: xxx
+
+* ZeRO-3:
+
+   gpus_factor = n_gpus / total_gpus
+
+   - case 1: ``"cpu_offload": false``:
+
+      Without ``zero.Init``:
+
+          params * 4 * n_gpus * additional_buffer_factor
+
+          this is the memory needed only at the beginning to initialize the model on CPU memory. Once the model is transferred to GPUs this memory is freed.
+
+      With ``zero.Init``:
+
+          largest_layer_params * 4 * n_gpus * additional_buffer_factor
+
+          assuming Pytorch is deallocating the memory once the tensors are moved to the GPU by ZeRO.Init
+
+   - case 2: ``"cpu_offload": true, cpu_offload_params true``:
+
+      Without ``zero.Init``:
+
+          params * max(4 * n_gpus, 18 * gpus_factor) * additional_buffer_factor
+
+      With ``zero.Init``:
+
+          params * 18 * gpus_factor * additional_buffer_factor
+
+   - case 3: ``"cpu_offload": true, cpu_offload_params false``:
+
+      Without ``zero.Init``:
+
+          params * max(4 * n_gpus, 16 * gpus_factor) * additional_buffer_factor
+
+      With ``zero.Init``:
+
+          params * 16 * gpus_factor * additional_buffer_factor
+
+
+Here is a breakdown for the 16 and 18 multipliers (b = bytes):
+
+4 (in ``4*n_gpus``):
+
+- when pytorch creates a model it creates it in fp32 by default (4 bytes)
+
+16:
+
+- 16b for fp32: 4b params, 4b grads, 4b momentum and 4b variance per parameter
+
+18:
+
+- 16b for fp32: 4b params, 4b grads, 4b momentum and 4b variance per parameter
+- +2b for fp16 params
+
+
+**Pinned Memory**
+
+Pinned general RAM is included in normal general RAM allocations (i.e. this is not extra memory allocations but simply shows how much of the general RAM is pinned)
+
+* ZeRO-2: can't be controlled
+
+* ZeRO-3
+
+To enable add: ``"cpu_offload_use_pin_memory" : true``
+
+Now there are 2 sub-cases:
+
+1. ``"cpu_offload_params": true``:
+
+   - 6 * params (2b for fp16 params + 4b for fp32 gradients)
+   - if ``gradient_accumulation_steps > 1`` an additional 2b for fp16 gradients are pinned
+
+2. ``"cpu_offload_params": false``:
+
+   - 4b for fp32 gradients
+
+
+**Activation Memory**
+
+XXX: For Transformers is probably around (2* seq * attn_heads + 16 * hidden_size) * sequence * batch/gpu
+
+This needs to be completed.

From c0c4ebf143b60bae28c32beaed0a9f526f25264d Mon Sep 17 00:00:00 2001
From: Stas Bekman <stas00@users.noreply.github.com>
Date: Wed, 23 Jun 2021 13:08:29 -0700
Subject: [PATCH 21/38] introduce debug utils (#1136)

Co-authored-by: Olatunji Ruwase <olruwase@microsoft.com>
Co-authored-by: Jeff Rasley <jerasley@microsoft.com>
---
 deepspeed/runtime/engine.py                   |   4 +
 .../runtime/zero/partition_parameters.py      |  21 ++-
 deepspeed/runtime/zero/stage3.py              |  35 +++--
 deepspeed/utils/debug.py                      | 122 ++++++++++++++++++
 4 files changed, 162 insertions(+), 20 deletions(-)
 create mode 100644 deepspeed/utils/debug.py

diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
index 4692bb709079..3c695d921a5a 100755
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -38,6 +38,7 @@
 import deepspeed.runtime.lr_schedules as lr_schedules
 from deepspeed.utils import logger, log_dist, init_distributed
 from deepspeed.utils.timer import ThroughputTimer, SynchronizedWallClockTimer
+from deepspeed.utils.debug import debug_extract_module_and_param_names
 from deepspeed.runtime.progressive_layer_drop import ProgressiveLayerDrop
 from deepspeed.runtime.eigenvalue import Eigenvalue
 
@@ -122,6 +123,9 @@ def __init__(self,
         self.gas_boundary_ctr = 0
         self.dist_backend = "nccl"
 
+        # for debug purposes - can then debug print: debug_get_module_name(module)
+        debug_extract_module_and_param_names(model)
+
         # Set config using config_params for backwards compat
         if self.config is None and config_params is not None:
             self.config = config_params
diff --git a/deepspeed/runtime/zero/partition_parameters.py b/deepspeed/runtime/zero/partition_parameters.py
index 04a82d177611..5904b01cc45b 100755
--- a/deepspeed/runtime/zero/partition_parameters.py
+++ b/deepspeed/runtime/zero/partition_parameters.py
@@ -18,6 +18,7 @@
 
 from ..utils import see_memory_usage
 from deepspeed.utils import log_dist, init_distributed
+from deepspeed.utils.debug import debug_param2name_id_shape, debug_module2name, debug_param2name, debug_param2name_id_shape_status, printflock, log_rank_file
 
 from ..swap_tensor.partitioned_param_swapper import AsyncPartitionedParameterSwapper, PartitionedParamStatus
 from ..config import DeepSpeedConfig
@@ -27,8 +28,14 @@
 
 
 def print_rank_0(message, debug=False, force=False):
-    if torch.distributed.get_rank() == 0 and (debug or force):
+    rank = torch.distributed.get_rank()
+    if rank == 0 and (debug or force):
         print(message)
+    # other variations
+    # - print for all ranks w/o interleaving
+    # printflock(f"[{rank}] {message}")
+    # - print to log file per rank
+    # log_rank_file(rank, message)
 
 
 def is_zero_param(parameter):
@@ -481,12 +488,12 @@ def _post_init_method(self, module):
             force=False)
 
         global param_count
-        for name, param in module.named_parameters(recurse=False):
+        for param in module.parameters(recurse=False):
             param_count += param.numel()
             if not is_zero_param(param):
                 self._convert_to_deepspeed_param(param)
                 print_rank_0(
-                    f"Partitioning param with ds id {param.ds_id} and shape {param.data.shape}"
+                    f"Partitioning param {debug_param2name_id_shape(param)} module={debug_module2name(module)}"
                 )
                 param.partition()
         see_memory_usage(
@@ -797,23 +804,23 @@ def _allgather_param(self, param, async_op=False, hierarchy=0):
         assert tensor_size == aligned_param_size, f'param id {param.ds_id} aligned size {aligned_param_size} does not match tensor size {tensor_size}'
 
         print_rank_0(
-            f"{'--'* hierarchy}---- Before allocating Allgather param with id {param.ds_id} and status {param.ds_status} Partition Size {partition_size} and data shape {param.ds_shape}"
+            f"{'--'* hierarchy}---- Before allocating allgather param {debug_param2name_id_shape_status(param)} partition size={partition_size}"
         )
 
         see_memory_usage(
-            f'Before allocate allgather param {param.ds_id} {param.ds_status} {aligned_param_size} {partition_size} {param.ds_shape}',
+            f'Before allocate allgather param {debug_param2name_id_shape_status(param)} partition_size={partition_size} ',
             force=False)
         flat_tensor = torch.zeros(aligned_param_size,
                                   dtype=param.dtype,
                                   device=param.device).view(-1)
         see_memory_usage(
-            f'After allocate allgather param {param.ds_id} {param.ds_status} {aligned_param_size} {partition_size} {param.ds_shape}',
+            f'After allocate allgather param {debug_param2name_id_shape_status(param)} {aligned_param_size} {partition_size} ',
             force=False)
 
         torch.cuda.synchronize()
 
         print_rank_0(
-            f"{'--'* hierarchy}----Allgather param with id {param.ds_id} and status {param.ds_status} Partition Size {partition_size} and data shape {param.ds_shape}"
+            f"{'--'* hierarchy}----allgather param with {debug_param2name_id_shape_status(param)} partition size={partition_size}"
         )
         #        if not flat_tensor.numel() > 100000:
         #            replicated_tensor = flat_tensor.narrow(0,
diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py
index d492df42067d..5d47edaa44ac 100755
--- a/deepspeed/runtime/zero/stage3.py
+++ b/deepspeed/runtime/zero/stage3.py
@@ -32,11 +32,18 @@
 pg_correctness_test = False
 
 FWD_MODULE_STACK = list()
+from deepspeed.utils.debug import debug_module2name_id, debug_param2name_id_numel, debug_param2name_id_shape_device, debug_module2name_class, printflock, log_rank_file
 
 
-def print_rank_0(message, debug=False, force=False):
-    if torch.distributed.get_rank() == 0 and (debug or force):
-        logger.info(message)
+def print_rank_0(message, debug=False, force=True):
+    rank = torch.distributed.get_rank()
+    if rank == 0 and (debug or force):
+        print(message)
+    # other variations
+    # - print for all ranks w/o interleaving
+    # printflock(f"[{rank}] {message}")
+    # - print to log file per rank
+    # log_rank_file(rank, message)
 
 
 def input(msg):
@@ -211,7 +218,7 @@ def get_params_to_prefetch(self, sub_module, numel=2000000):
         # tracing failed. The sub_module passed at the step_id must match with the sub_module during tracing
         if sub_module.id != self.sub_module_trace[self.step_id]:
             print_rank_0(
-                f"Tracing failed. Prefetching is disabled at sub-module: {sub_module.id}"
+                f"Tracing failed. Prefetching is disabled at sub-module: {debug_module2name_id(sub_module)}"
             )
             return []
 
@@ -390,11 +397,13 @@ def reset_step(self):
     def fetch_sub_module(self, sub_module):
         partitioned_params = []
         params_in_flight = False
-        #print_rank_0(f"{'--' * self.hierarchy}Fetching params in module {sub_module.__class__.__name__}")
+        # print_rank_0(f"{'--' * self.hierarchy}Fetching params in module {sub_module.__class__.__name__}")
         params_to_fetch = [
             param for _,
             param in sub_module.named_parameters(recurse=False)
         ]
+        # print([n for n,p in sub_module.named_parameters(recurse=False)])
+
         if hasattr(sub_module, 'ds_external_parameters'):
             print_rank_0(
                 f"{'--' * self.hierarchy}--Fetching external parameters {sub_module.ds_external_parameters()}"
@@ -407,7 +416,7 @@ def fetch_sub_module(self, sub_module):
         for param in params_to_fetch:
             param.ds_active_sub_modules += 1
             print_rank_0(
-                f"{'--' * self.hierarchy}--Fetching parameters {param.ds_id} with active sub modules {param.ds_active_sub_modules}"
+                f"{'--' * self.hierarchy}--Fetching parameters {param.ds_id} {param.ds_shape} with active sub modules {param.ds_active_sub_modules}"
             )
 
             if param.ds_status == ZeroParamStatus.AVAILABLE:
@@ -441,14 +450,14 @@ def fetch_sub_module(self, sub_module):
         for _, param in sub_module.named_parameters(recurse=False):
             param.ds_status = ZeroParamStatus.AVAILABLE
             print_rank_0(
-                f"Param id {param.ds_id}, Shape {param.shape}, device {param.device} norm {param.norm()}",
+                f"Param {debug_param2name_id_shape_device(param)} norm={param.norm()}",
                 force=False)
         #print_rank_0(f"After fetching (id, shape, device): {[(param.ds_id, param.shape, param.device) for param in sub_module.named_parameters(recurse=False)]}")
 
     def release_sub_module(self, sub_module):
         self.hierarchy -= 1
         print_rank_0(
-            f"{'--' * self.hierarchy}Releasing params in module {sub_module.__class__.__name__}"
+            f"{'--' * self.hierarchy}Releasing params in module {debug_module2name_class(sub_module)}"
         )
         params_to_release = [
             param for _,
@@ -468,31 +477,31 @@ def release_sub_module(self, sub_module):
             if not param.ds_active_sub_modules and not self._keep_for_later(
                     sub_module) and not param.ds_persist:
                 print_rank_0(
-                    f"{'--' * self.hierarchy}--Releasing parameters {param.ds_id} with numel {param.numel()} active sub modules {param.ds_active_sub_modules} and keep for later {self._keep_for_later(sub_module)}",
+                    f"{'--' * self.hierarchy}--Releasing parameter {debug_param2name_id_numel(param)} active sub modules {param.ds_active_sub_modules} and keep for later {self._keep_for_later(sub_module)}",
                     force=False)
 
                 # Keeping track of number of elements that are consumed by available parameters
                 self._decrement_available_parameter_numel(param.ds_numel)
                 see_memory_usage(
-                    f"Before releasing param {param.ds_id} with numel {param.numel()}",
+                    f"Before releasing param {debug_param2name_id_numel(param)}",
                     force=False)
                 param.partition(hierarchy=self.hierarchy)
                 see_memory_usage(
-                    f"After releasing param {param.ds_id} has numel {param.numel()} ",
+                    f"After releasing param {debug_param2name_id_numel(param)}",
                     force=False)
 
                 param.ds_status = ZeroParamStatus.NOT_AVAILABLE
             else:
 
                 print_rank_0(
-                    f"{'--' * self.hierarchy}--Did not release parameters {param.ds_id} with numel {param.numel()} with active sub modules {param.ds_active_sub_modules}, keep for later {self._keep_for_later(sub_module)} and persistence {param.ds_persist}",
+                    f"{'--' * self.hierarchy}--Did not release param {debug_param2name_id_numel(param)} with active sub modules {param.ds_active_sub_modules}, keep for later={self._keep_for_later(sub_module)} and persistence={param.ds_persist}",
                     force=False)
 
     def release_and_reset_parameter(self, param):
         param.ds_active_sub_modules = 0
         if param.ds_status == ZeroParamStatus.AVAILABLE:
             print_rank_0(
-                f"Releasing unpartitioned {param.ds_id} active sub-modules {param.ds_active_sub_modules} size {param.ds_numel} and persisitence {param.ds_persist}"
+                f"Releasing unpartitioned param {debug_param2name_id_numel(param)} active sub-modules {param.ds_active_sub_modules} and persisitence {param.ds_persist}"
             )
             self._decrement_available_parameter_numel(param.ds_numel)
             param.partition()
diff --git a/deepspeed/utils/debug.py b/deepspeed/utils/debug.py
new file mode 100644
index 000000000000..06f966dfa7a5
--- /dev/null
+++ b/deepspeed/utils/debug.py
@@ -0,0 +1,122 @@
+""" debug utils """
+
+import fcntl
+
+# for debug purposes map module and param objects to their fully qualified names
+module_names = {}
+param_names = {}
+
+
+def debug_extract_module_and_param_names(model):
+    # extract the fully qualified names as soon as the model is acquired
+    global module_names
+    global param_names
+    # XXX: can probably make a map of param2module and vice-versa
+    module_names = {module: name for name, module in model.named_modules()}
+    param_names = {param: name for name, param in model.named_parameters()}
+
+
+def debug_module2name(module):
+    if module in module_names:
+        return module_names[module]
+    else:
+        return "unknown"
+
+
+def debug_module2name_id(module):
+    return f"name={debug_module2name(module)} id={module.id}"
+
+
+def debug_module2name_class(module):
+    return f"name={debug_module2name(module)} {module.__class__.__name__}"
+
+
+def debug_param2name(param):
+    if param in param_names:
+        return param_names[param]
+    else:
+        return "unknown"
+
+
+def debug_param2name_id(param):
+    return f"name={debug_param2name(param)} id={param.ds_id}"
+
+
+def debug_param2name_id_shape(param):
+    return f"name={debug_param2name(param)} id={param.ds_id} shape={param.data.shape}"
+
+
+def debug_param2name_id_shape_device(param):
+    return f"name={debug_param2name(param)} id={param.ds_id} shape={param.data.shape} device={param.device}"
+
+
+def debug_param2name_id_numel(param):
+    return f"name={debug_param2name(param)} id={param.ds_id} numel={param.numel()}"
+
+
+def debug_param2name_id_shape_status(param):
+    return f"name={debug_param2name(param)} id={param.ds_id} shape={param.data.shape} status={param.ds_status}"
+
+
+def printflock(*msgs):
+    """
+
+    For printing messages for all concurrent gpus w/o getting interleaved text.
+
+    This is useful when debugging issues where multi-gpus don't sync.
+
+    1. Enable the force debug in say partitioning and zero3 files
+    2. Override the usual versions with ::
+
+        def print_rank_0(message, debug=False, force=True):
+            rank = torch.distributed.get_rank()
+            printflock(f"[{rank}] {message}")
+    3. run the program and you get both logs non-interleaved
+
+    But this makes it very difficult to make sense of the output, so the ``log_rank_file`` helper
+    function might be more useful, as it's easier to send each log stream into a separate file and
+    then compare those.
+
+    """
+
+    with open(__file__, "r") as fh:
+        fcntl.flock(fh, fcntl.LOCK_EX)
+        try:
+            print(*msgs)
+        finally:
+            fcntl.flock(fh, fcntl.LOCK_UN)
+
+
+fh = None
+
+
+def log_rank_file(rank, *msgs):
+    """
+    Print to a log file of the given rank
+
+    This is useful for debugging hanging in sync processes. Here is a possible workflow:
+
+    1. Enable the force debug in say partitioning and zero3 files
+    2. Override the usual versions of print_rank_0 in those files with ::
+
+        def print_rank_0(message, debug=False, force=True):
+            rank = torch.distributed.get_rank()
+            log_rank_file(rank, message)
+
+    3. run the program
+    4. fix up the expected differences, e.g. different cuda numbers ::
+
+        perl -pi -e 's|cuda:1|cuda:0|' log_rank_*
+
+    5. now diff and see where names and ids diverge - you will find where the gpus don't do the same
+    work (e.g. when some layers get conditionally skipped on one gpu but not all)
+
+        diff -u log_rank_0.txt log_rank_1.txt | less
+
+    """
+    global fh
+    if fh is None:
+        fh = open(f"log_rank_{rank}.txt", "w")
+    for m in msgs:
+        fh.write(f"{m}\n")
+    fh.flush()

From 429cbc89afdb1a801dc56af25304dd2e34af3632 Mon Sep 17 00:00:00 2001
From: Hyunwoong Ko <gusdnd852@naver.com>
Date: Thu, 24 Jun 2021 12:52:54 +0900
Subject: [PATCH 22/38] Fix bugs about non-contiguous tensor broadcasting
 (#1168)

* Fix bugs about non-contiguous tensor broadcasting

* Fix typo

Co-authored-by: Jeff Rasley <jerasley@microsoft.com>
Co-authored-by: Reza Yazdani <44502768+RezaYazdaniAminabadi@users.noreply.github.com>
---
 deepspeed/inference/engine.py             | 8 ++++++++
 deepspeed/module_inject/replace_policy.py | 2 +-
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/deepspeed/inference/engine.py b/deepspeed/inference/engine.py
index 42ec654a7101..08cf1fa9bacf 100644
--- a/deepspeed/inference/engine.py
+++ b/deepspeed/inference/engine.py
@@ -189,12 +189,16 @@ def _pre_forward_hook(self, module, *inputs, **kwargs):
             if torch.is_tensor(input):
                 input = input.to(torch.cuda.current_device())
                 if self.mp_world_size > 1:
+                    if not input.is_contiguous():
+                        input = input.contiguous()
                     dist.broadcast(input, 0)
 
         for k in kwargs:
             if torch.is_tensor(kwargs[k]):
                 kwargs[k] = kwargs[k].to(torch.cuda.current_device())
                 if self.mp_world_size > 1:
+                    if not kwargs[k].is_contiguous():
+                        kwargs[k] = kwargs[k].contiguous()
                     dist.broadcast(kwargs[k], 0)
 
     def forward(self, *inputs, **kwargs):
@@ -210,12 +214,16 @@ def forward(self, *inputs, **kwargs):
                     if torch.is_tensor(input):
                         input = input.to(torch.cuda.current_device())
                         if self.mp_world_size > 1:
+                            if not input.is_contiguous():
+                                input = input.contiguous()
                             dist.broadcast(input, 0)
 
                 for k in kwargs:
                     if torch.is_tensor(kwargs[k]):
                         kwargs[k] = kwargs[k].to(torch.cuda.current_device())
                         if self.mp_world_size > 1:
+                            if not kwargs[k].is_contiguous():
+                                kwargs[k] = kwargs[k].contiguous()
                             dist.broadcast(kwargs[k], 0)
 
             return self.model_orig_fwd(*inputs, **kwargs)
diff --git a/deepspeed/module_inject/replace_policy.py b/deepspeed/module_inject/replace_policy.py
index 325bb37efd74..3758ffd9b522 100755
--- a/deepspeed/module_inject/replace_policy.py
+++ b/deepspeed/module_inject/replace_policy.py
@@ -19,7 +19,7 @@ def attention(self):
 
     def get_hidden_heads(self):
         """
-        retun hidden_size and number of heads
+        return hidden_size and number of heads
         """
         raise NotImplementedError
 

From bc019a533971e1dcfde93972ae892b1a4f425041 Mon Sep 17 00:00:00 2001
From: Stas Bekman <stas00@users.noreply.github.com>
Date: Fri, 25 Jun 2021 19:50:38 -0700
Subject: [PATCH 23/38] undo noise (#1191)

* undo noise

* another
---
 deepspeed/runtime/zero/stage3.py | 2 +-
 deepspeed/utils/debug.py         | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py
index 5d47edaa44ac..5350b06e654e 100755
--- a/deepspeed/runtime/zero/stage3.py
+++ b/deepspeed/runtime/zero/stage3.py
@@ -35,7 +35,7 @@
 from deepspeed.utils.debug import debug_module2name_id, debug_param2name_id_numel, debug_param2name_id_shape_device, debug_module2name_class, printflock, log_rank_file
 
 
-def print_rank_0(message, debug=False, force=True):
+def print_rank_0(message, debug=False, force=False):
     rank = torch.distributed.get_rank()
     if rank == 0 and (debug or force):
         print(message)
diff --git a/deepspeed/utils/debug.py b/deepspeed/utils/debug.py
index 06f966dfa7a5..bb2b5f7ef75b 100644
--- a/deepspeed/utils/debug.py
+++ b/deepspeed/utils/debug.py
@@ -68,7 +68,7 @@ def printflock(*msgs):
     1. Enable the force debug in say partitioning and zero3 files
     2. Override the usual versions with ::
 
-        def print_rank_0(message, debug=False, force=True):
+        def print_rank_0(message, debug=False, force=False):
             rank = torch.distributed.get_rank()
             printflock(f"[{rank}] {message}")
     3. run the program and you get both logs non-interleaved
@@ -99,7 +99,7 @@ def log_rank_file(rank, *msgs):
     1. Enable the force debug in say partitioning and zero3 files
     2. Override the usual versions of print_rank_0 in those files with ::
 
-        def print_rank_0(message, debug=False, force=True):
+        def print_rank_0(message, debug=False, force=False):
             rank = torch.distributed.get_rank()
             log_rank_file(rank, message)
 

From bd9e9539159f244540549ba76261de61c3db8b0f Mon Sep 17 00:00:00 2001
From: Shaden Smith <shaden.smith@microsoft.com>
Date: Sat, 26 Jun 2021 05:01:29 +0000
Subject: [PATCH 24/38] less verbose load

---
 deepspeed/runtime/state_dict_factory.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/deepspeed/runtime/state_dict_factory.py b/deepspeed/runtime/state_dict_factory.py
index f6be1dbbe57b..ab7f2c8e562c 100755
--- a/deepspeed/runtime/state_dict_factory.py
+++ b/deepspeed/runtime/state_dict_factory.py
@@ -52,9 +52,6 @@ def load(self,
         num_ckpt = len(self.ckpt_list)
         idx = mp_rank * num_ckpt // mp_world_size
 
-        logger.info(
-            f'mp_world_size: {mp_world_size}, mp_rank: {mp_rank}, module_key: {module_key}'
-        )
         """ We have multiple cases to handle here for both training and inference:
             1. PipeModule loading mp_rank_*.pt files, is_pipe_parallel=True, module_key is not None
                 a. if no mp_size/pp_size resizing occurs, for both training & inference, loading

From a029239812e15cf35334514449ed3127b915780a Mon Sep 17 00:00:00 2001
From: Stas Bekman <stas00@users.noreply.github.com>
Date: Mon, 28 Jun 2021 12:42:39 -0700
Subject: [PATCH 25/38] clean up logging (#1190)

Co-authored-by: Jeff Rasley <jerasley@microsoft.com>
---
 csrc/adam/cpu_adam.cpp                        | 45 +++++++++----------
 deepspeed/ops/adam/cpu_adam.py                |  4 +-
 deepspeed/runtime/engine.py                   |  2 +-
 .../runtime/zero/partition_parameters.py      |  2 +-
 deepspeed/runtime/zero/stage3.py              | 20 ++++-----
 deepspeed/runtime/zero/utils.py               |  2 +-
 deepspeed/utils/logging.py                    | 37 +++++++++++++++
 7 files changed, 73 insertions(+), 39 deletions(-)

diff --git a/csrc/adam/cpu_adam.cpp b/csrc/adam/cpu_adam.cpp
index 6726b895f12c..ef756263b6d6 100644
--- a/csrc/adam/cpu_adam.cpp
+++ b/csrc/adam/cpu_adam.cpp
@@ -322,42 +322,37 @@ int create_adam_optimizer(int optimizer_id,
                           float betta2 = 0.999,
                           float eps = 1e-8,
                           float weight_decay = 0,
-                          bool adamw_mode = true)
+                          bool adamw_mode = true,
+                          bool should_log = false)
 {
     auto opt =
         std::make_shared<Adam_Optimizer>(alpha, betta1, betta2, eps, weight_decay, adamw_mode);
 
     s_optimizers[optimizer_id] = opt;
+
+    if (should_log) {
+        std::string avx_type = "";
 #if defined(__AVX512__)
-    std::cout << "Adam Optimizer #" << optimizer_id
-              << " is created with AVX512 arithmetic capability." << std::endl;
-    printf("Config: alpha=%f, betas=(%f, %f), weight_decay=%f, adam_w=%d\n",
-           alpha,
-           betta1,
-           betta2,
-           weight_decay,
-           (int)adamw_mode);
+        avx_type = "AVX512";
 #else
 #if defined(__AVX256__)
-    std::cout << "Adam Optimizer #" << optimizer_id
-              << " is created with AVX2 arithmetic capability." << std::endl;
-    printf("Config: alpha=%f, betas=(%f, %f), weight_decay=%f, adam_w=%d\n",
-           alpha,
-           betta1,
-           betta2,
-           weight_decay,
-           (int)adamw_mode);
+        avx_type = "AVX2";
 #else
-    std::cout << "Adam Optimizer #" << optimizer_id
-              << " is created with scalar arithmetic capability." << std::endl;
-    printf("Config: alpha=%f, betas=(%f, %f), weight_decay=%f, adam_w=%d\n",
-           alpha,
-           betta1,
-           betta2,
-           weight_decay,
-           (int)adamw_mode);
+        avx_type = "scalar";
 #endif
 #endif
+
+        printf("Adam Optimizer #%d is created with %s arithmetic capability.\n",
+               optimizer_id,
+               avx_type.c_str());
+        printf("Config: alpha=%f, betas=(%f, %f), weight_decay=%f, adam_w=%d\n",
+               alpha,
+               betta1,
+               betta2,
+               weight_decay,
+               (int)adamw_mode);
+    }
+
     return 0;
 }
 
diff --git a/deepspeed/ops/adam/cpu_adam.py b/deepspeed/ops/adam/cpu_adam.py
index 35eeedb86b5d..835742a35604 100755
--- a/deepspeed/ops/adam/cpu_adam.py
+++ b/deepspeed/ops/adam/cpu_adam.py
@@ -7,6 +7,7 @@
 import time
 from pathlib import Path
 from ..op_builder import CPUAdamBuilder
+from deepspeed.utils.logging import should_log_le
 
 
 class DeepSpeedCPUAdam(torch.optim.Optimizer):
@@ -83,7 +84,8 @@ def __init__(self,
                                      betas[1],
                                      eps,
                                      weight_decay,
-                                     adamw_mode)
+                                     adamw_mode,
+                                     should_log_le("info"))
 
     def __del__(self):
         # need to destroy the C++ object explicitly to avoid a memory leak when deepspeed.initialize
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
index 3c695d921a5a..325222db4669 100755
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -932,7 +932,7 @@ def _configure_zero_optimizer(self, optimizer):
                 ignore_unused_parameters=self.zero_ignore_unused_parameters(),
                 partition_grads=zero_stage == ZERO_OPTIMIZATION_GRADIENTS)
         elif zero_stage == ZERO_OPTIMIZATION_WEIGHTS:
-            print("Initializing ZeRO Stage 3") if dist.get_rank() == 0 else None
+            logger.info("Initializing ZeRO Stage 3") if dist.get_rank() == 0 else None
             from deepspeed.runtime.zero.stage3 import FP16_DeepSpeedZeroOptimizer_Stage3
             optimizer = FP16_DeepSpeedZeroOptimizer_Stage3(
                 self.module,
diff --git a/deepspeed/runtime/zero/partition_parameters.py b/deepspeed/runtime/zero/partition_parameters.py
index 5904b01cc45b..6ef87f9e00aa 100755
--- a/deepspeed/runtime/zero/partition_parameters.py
+++ b/deepspeed/runtime/zero/partition_parameters.py
@@ -266,7 +266,7 @@ def _init_subclass(cls, **kwargs):
         if self.mem_efficient_linear:
             print_rank_0(
                 "nn.functional.linear has been overridden with a more memory efficient version. This will persist unless manually reset.",
-                force=True)
+                force=False)
             self.linear_bk = torch.nn.functional.linear
             torch.nn.functional.linear = LinearFunctionForZeroStage3.apply
 
diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py
index 5350b06e654e..2806f5d420da 100755
--- a/deepspeed/runtime/zero/stage3.py
+++ b/deepspeed/runtime/zero/stage3.py
@@ -627,7 +627,7 @@ def __init__(self,
                  elastic_checkpoint=False,
                  aio_config=None):
 
-        see_memory_usage("Stage 3 initialize beginning", force=True)
+        see_memory_usage("Stage 3 initialize beginning", force=False)
 
         if dist.get_rank() == 0:
             logger.info(f"Reduce bucket size {reduce_bucket_size}")
@@ -702,7 +702,7 @@ def __init__(self,
             self.max_params_in_cpu = offload_param_config[OFFLOAD_PARAM_MAX_IN_CPU]
             print_rank_0(
                 f"FP16 params swapping is {self.params_in_nvme_and_cpu}, Max params in CPU is {self.max_params_in_cpu}",
-                force=True)
+                force=False)
 
         self.deepspeed_adam_offload = (self.offload_optimizer
                                        and type(init_optimizer) == DeepSpeedCPUAdam)
@@ -798,7 +798,7 @@ def __init__(self,
         self.sub_group_size = sub_group_size
 
         self.sub_group_to_group_id = {}
-        see_memory_usage("Before creating fp16 partitions", force=True)
+        see_memory_usage("Before creating fp16 partitions", force=False)
         self._create_fp16_partitions_with_defragmentation()
         num_fp16_subgroups = len(self.fp16_partitioned_groups_flat)
         see_memory_usage(f"After creating fp16 partitions: {num_fp16_subgroups}",
@@ -867,7 +867,7 @@ def __init__(self,
         ])
         print_rank_0(
             f'Largest partitioned param numel = {largest_partitioned_param_numel}',
-            force=True)
+            force=False)
 
         see_memory_usage(f"Before Set Grad positions", force=False)
 
@@ -922,7 +922,7 @@ def __init__(self,
         self.debug_fp16_grads = [{} for _ in self.fp16_groups]
 
         if dist.get_rank(group=self.dp_process_group) == 0:
-            see_memory_usage(f"After initializing ZeRO optimizer", force=True)
+            see_memory_usage(f"After initializing ZeRO optimizer", force=False)
 
     def _configure_tensor_swapping(self, offload_optimizer_config, aio_config):
         nvme_swap_folder = os.path.join(
@@ -1096,7 +1096,7 @@ def _create_fp16_partitions_with_defragmentation(self):
         for j, param_group in enumerate(self.optimizer.param_groups):
 
             sub_groups = self._create_fp16_sub_groups(param_group['params'])
-            print_rank_0(f'fp16 group {j} has {len(sub_groups)} subgroups', force=True)
+            print_rank_0(f'fp16 group {j} has {len(sub_groups)} subgroups', force=False)
 
             flat_offset = 0
             for sub_group in sub_groups:
@@ -1331,19 +1331,19 @@ def _create_fp32_partitions(self):
         nvme_gigabytes = nvme_memory_usage / GIGA_BYTES
         print_rank_0(
             f'Swappable FP32 Partitions: count={num_swappable_partitions} size={nvme_gigabytes:5.2f} GB',
-            force=True)
+            force=False)
         if self.params_in_nvme_and_cpu:
             print_rank_0(
                 f'Swap from NVMe Partitions: count = {num_swap_from_nvme_partitions}, size = {swap_from_nvme_memory_usage/GIGA_BYTES:5.2f}GB',
-                force=True)
+                force=False)
             print_rank_0(
                 f'Swap from CPU Partitions: count = {num_swap_from_cpu_partitions}, size = {swap_from_cpu_memory_usage/GIGA_BYTES:5.2f}GB',
-                force=True)
+                force=False)
 
         cpu_memory_gigabytes = cpu_memory_usage / GIGA_BYTES
         print_rank_0(
             f'In-Memory FP32 Partitions: count={cpu_memory_sub_groups} size={cpu_memory_gigabytes:5.2f} GB',
-            force=True)
+            force=False)
 
         # Clear for on-the-fly population before the optimizer step
         for param_group in self.optimizer.param_groups:
diff --git a/deepspeed/runtime/zero/utils.py b/deepspeed/runtime/zero/utils.py
index 8873c8db55d5..f90497069901 100755
--- a/deepspeed/runtime/zero/utils.py
+++ b/deepspeed/runtime/zero/utils.py
@@ -40,7 +40,7 @@ def _initialize_parameter_parallel_groups(parameter_parallel_size=None):
 
 def is_zero_supported_optimizer(optimizer):
     if dist.get_rank() == 0:
-        print(
+        logger.info(
             f'Checking ZeRO support for optimizer={optimizer.__class__.__name__} type={type(optimizer)}'
         )
     return type(optimizer) in ZERO_SUPPORTED_OPTIMIZERS
diff --git a/deepspeed/utils/logging.py b/deepspeed/utils/logging.py
index d16682f902e8..ca694514678e 100644
--- a/deepspeed/utils/logging.py
+++ b/deepspeed/utils/logging.py
@@ -3,6 +3,14 @@
 
 import torch.distributed as dist
 
+log_levels = {
+    "debug": logging.DEBUG,
+    "info": logging.INFO,
+    "warning": logging.WARNING,
+    "error": logging.ERROR,
+    "critical": logging.CRITICAL,
+}
+
 
 class LoggerFactory:
     @staticmethod
@@ -58,3 +66,32 @@ def log_dist(message, ranks=None, level=logging.INFO):
     if should_log:
         final_message = "[Rank {}] {}".format(my_rank, message)
         logger.log(level, final_message)
+
+
+def get_current_level():
+    """
+    Return logger's current log level
+    """
+    return logger.getEffectiveLevel()
+
+
+def should_log_le(max_log_level_str):
+    """
+    Args:
+        max_log_level_str: maximum log level as a string
+
+    Returns ``True`` if the current log_level is less or equal to the specified log level. Otherwise ``False``.
+
+    Example:
+
+        ``should_log_le("info")`` will return ``True`` if the current log level is either ``logging.INFO`` or ``logging.DEBUG``
+    """
+
+    if not isinstance(max_log_level_str, str):
+        raise ValueError(f"{max_log_level_str} is not a string")
+
+    max_log_level_str = max_log_level_str.lower()
+    if max_log_level_str not in log_levels:
+        raise ValueError(f"{max_log_level_str} is not one of the `logging` levels")
+
+    return get_current_level() <= log_levels[max_log_level_str]

From d26c258b08a2a82dd0fe68325a8be98cee87dc9d Mon Sep 17 00:00:00 2001
From: Jeff Rasley <jerasley@microsoft.com>
Date: Tue, 13 Jul 2021 10:38:48 -0700
Subject: [PATCH 26/38] added shaden's set_train_batch_size patches, plus
 formatting

---
 deepspeed/runtime/engine.py             | 21 +++++++++++++++++++++
 deepspeed/runtime/pipe/engine.py        | 13 +++++++++++++
 deepspeed/runtime/pipe/p2p.py           |  3 ++-
 deepspeed/runtime/state_dict_factory.py |  1 -
 4 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
index efc9ac0fa18f..c06e4186e464 100755
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -234,6 +234,27 @@ def get_batch_info(self):
         """
         return self.train_batch_size, self.train_micro_batch_size_per_gpu, self.gradient_accumulation_steps
 
+    def set_train_batch_size(self, train_batch_size):
+        """Adjust the global batch size by increasing or decreasing the number of
+        micro-batches (i.e., gradient accumulation steps). The size of each micro-batch
+        (i.e., ``train_micro_batch_size_per_gpu``) is not changed.
+        Args:
+            train_batch_size (int): The new global batch size for training.
+        Raises:
+            ValueError: if ``train_batch_size`` is not divisible by the
+                configured micro-batch size and data parallelism.
+        """
+        if train_batch_size % (self.train_micro_batch_size_per_gpu() *
+                               self.dp_world_size) != 0:
+            #print(f'{train_batch_size=} {self.train_micro_batch_size_per_gpu()=} {self.dp_world_size=}')
+            raise ValueError(
+                f'Train batch size must be divisible by micro-batch data parallelism')
+        new_gas = train_batch_size // (self.train_micro_batch_size_per_gpu() *
+                                       self.dp_world_size)
+        # overwrite config
+        self._config.train_batch_size = train_batch_size
+        self._config.gradient_accumulation_steps = new_gas
+
     def checkpoint_tag_validation_enabled(self):
         return self._config.checkpoint_tag_validation_enabled
 
diff --git a/deepspeed/runtime/pipe/engine.py b/deepspeed/runtime/pipe/engine.py
index 09231bf3e96a..cadb5b82e36f 100644
--- a/deepspeed/runtime/pipe/engine.py
+++ b/deepspeed/runtime/pipe/engine.py
@@ -394,6 +394,19 @@ def eval_batch(self, data_iter, compute_loss=True, reduce_output='avg'):
 
         return eval_output
 
+    def set_train_batch_size(self, train_batch_size):
+        """Adjust the global batch size by increasing or decreasing the number of
+        micro-batches (i.e., gradient accumulation steps). The size of each micro-batch
+        (i.e., ``train_micro_batch_size_per_gpu``) is not changed.
+        Args:
+            train_batch_size (int): The new global batch size for training.
+        Raises:
+            ValueError: if ``train_batch_size`` is not divisible by the
+                configured micro-batch size and data parallelism.
+        """
+        super().set_train_batch_size(train_batch_size)
+        self.micro_batches = self.gradient_accumulation_steps()
+
     def is_first_stage(self):
         """True if this process is in the first stage in the pipeline."""
         return self.stage_id == 0
diff --git a/deepspeed/runtime/pipe/p2p.py b/deepspeed/runtime/pipe/p2p.py
index 0fa31b89b984..13a448484061 100644
--- a/deepspeed/runtime/pipe/p2p.py
+++ b/deepspeed/runtime/pipe/p2p.py
@@ -11,9 +11,9 @@
 _groups = None
 _grid = None
 
-
 _async = []
 
+
 #initializes adjacent process groups
 #run this only after torch.distributed.init_process_group() has been called
 def init_process_groups(grid):
@@ -71,6 +71,7 @@ def wait():
 
     torch.cuda.synchronize()
 
+
 def send_obj(msg: typing.Any, dest: int):
     """Send an arbitrary python object to ``dest``.
 
diff --git a/deepspeed/runtime/state_dict_factory.py b/deepspeed/runtime/state_dict_factory.py
index ab7f2c8e562c..033b6726ddb3 100755
--- a/deepspeed/runtime/state_dict_factory.py
+++ b/deepspeed/runtime/state_dict_factory.py
@@ -51,7 +51,6 @@ def load(self,
         self.module_key = module_key
         num_ckpt = len(self.ckpt_list)
         idx = mp_rank * num_ckpt // mp_world_size
-
         """ We have multiple cases to handle here for both training and inference:
             1. PipeModule loading mp_rank_*.pt files, is_pipe_parallel=True, module_key is not None
                 a. if no mp_size/pp_size resizing occurs, for both training & inference, loading

From 9dbfdbd4739312f2891122246c88aee46c2128a3 Mon Sep 17 00:00:00 2001
From: Shaden Smith <Shaden.Smith@microsoft.com>
Date: Mon, 26 Jul 2021 08:32:40 -0700
Subject: [PATCH 27/38] Adds engine.was_step_applied() (#1251)

---
 deepspeed/runtime/engine.py | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
index c06e4186e464..013924def623 100755
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -122,6 +122,7 @@ def __init__(self,
         self.block_eigenvalue = None
         self.gas_boundary_ctr = 0
         self.dist_backend = "nccl"
+        self._step_applied = False
 
         # for debug purposes - can then debug print: debug_get_module_name(module)
         debug_extract_module_and_param_names(model)
@@ -1013,6 +1014,18 @@ def is_iterable_style_dataset(obj):
                           torch.utils.data.IterableDataset
                           )  # hasattr(obj, "__iter__") should work as well
 
+    def was_step_applied(self) -> bool:
+        """Returns True if the latest ``step()`` produced in parameter updates.
+
+        Note that a ``False`` return is not an error condition. Steps are frequently
+        no-ops, such as between gradient accumulation boundaries or when overflows
+        occur.
+
+        Returns:
+            bool: Whether the latest ``step()`` modified model parameters.
+        """
+        return self._step_applied
+
     def deepspeed_io(self,
                      dataset,
                      batch_size=None,
@@ -1307,6 +1320,7 @@ def _take_model_step(self, lr_kwargs, block_eigenvalue={}):
         overflow = False
         if hasattr(self.optimizer, 'overflow'):
             overflow = self.optimizer.overflow
+        self._step_applied = not overflow
 
         if overflow:
             self.skipped_steps += 1
@@ -1338,6 +1352,8 @@ def step(self, lr_kwargs=None):
                                            "init in order to use step"
         report_progress = self.global_rank == 0 if self.global_rank else True
 
+        self._step_applied = False  # assume False, will flip to True
+
         # Update the model when we reach gradient accumulation boundaries
         if self.is_gradient_accumulation_boundary():
             self.gas_boundary_ctr += 1

From d6945dea8d3d85636ea388a14dcb00ff8231fda3 Mon Sep 17 00:00:00 2001
From: Shaden Smith <Shaden.Smith@microsoft.com>
Date: Mon, 26 Jul 2021 08:52:56 -0700
Subject: [PATCH 28/38] Cleaning up tensor/pipe parallel accounting. (#1252)

* removes repeated overflow log

* pipe_replicated

* _pipe_replicated -> ds_pipe_replicated

* Adds send/recv fallback to bcast when torch version <= 1.8
---
 deepspeed/runtime/pipe/module.py |  4 +-
 deepspeed/runtime/pipe/p2p.py    | 55 ++++++++++++++++--
 deepspeed/runtime/utils.py       | 95 +++++++++++++++++++++++---------
 deepspeed/runtime/zero/stage2.py | 16 +++---
 4 files changed, 132 insertions(+), 38 deletions(-)

diff --git a/deepspeed/runtime/pipe/module.py b/deepspeed/runtime/pipe/module.py
index 31144f7a69b4..af4df2a10ebf 100644
--- a/deepspeed/runtime/pipe/module.py
+++ b/deepspeed/runtime/pipe/module.py
@@ -253,7 +253,7 @@ def _build(self):
         # All pipeline parameters should be considered as model parallel in the context
         # of our FP16 optimizer
         for p in self.parameters():
-            p.model_parallel = True
+            p.ds_pipe_replicated = False
 
     def _count_layer_params(self):
         """Count the trainable parameters in individual layers.
@@ -472,7 +472,7 @@ def _index_tied_modules(self):
                             # Only count the tied module once in the eyes of the FP16 optimizer
                             if self.global_rank != tied_ranks[0]:
                                 for p in self.tied_modules[key].parameters():
-                                    p.model_parallel = False
+                                    p.ds_pipe_replicated = True
         '''
         if len(tied_comms) > 0:
             print(f'RANK={self.global_rank} tied_comms={tied_comms}')
diff --git a/deepspeed/runtime/pipe/p2p.py b/deepspeed/runtime/pipe/p2p.py
index 13a448484061..e129d3d5b177 100644
--- a/deepspeed/runtime/pipe/p2p.py
+++ b/deepspeed/runtime/pipe/p2p.py
@@ -8,12 +8,22 @@
 import torch
 import torch.distributed as dist
 
+# To query whether we have send/recv support
+from packaging.version import Version
+from deepspeed.git_version_info import torch_info
+
 _groups = None
 _grid = None
 
 _async = []
 
 
+def can_send_recv() -> bool:
+    torch_version = Version(torch_info['version'])
+    sendrecv_min = Version('1.8')
+    return torch_version >= sendrecv_min
+
+
 #initializes adjacent process groups
 #run this only after torch.distributed.init_process_group() has been called
 def init_process_groups(grid):
@@ -22,6 +32,9 @@ def init_process_groups(grid):
 
     assert _grid.pipe_parallel_size > 1, "There is no pipeline parallelism"
 
+    if not can_send_recv():
+        _groups = [dist.new_group(ranks=group) for group in _grid.p2p_groups]
+
 
 def _is_valid_send_recv(src_stage, dest_stage):
     first_stage = 0
@@ -34,7 +47,7 @@ def _is_valid_send_recv(src_stage, dest_stage):
 
 def send(tensor, dest_stage, async_op=False):
     global _groups
-    #assert async_op == False, "Doesnt support async_op true"
+    assert async_op == False, "Doesnt support async_op true"
     src_stage = _grid.get_stage_id()
     _is_valid_send_recv(src_stage, dest_stage)
 
@@ -44,12 +57,18 @@ def send(tensor, dest_stage, async_op=False):
         op = dist.isend(tensor, dest_rank)
         _async.append(op)
     else:
-        return dist.send(tensor, dest_rank)
+
+        if can_send_recv():
+            return dist.send(tensor, dest_rank)
+        else:
+            group = _get_send_recv_group(src_stage, dest_stage)
+            src_rank = _grid.stage_to_global(stage_id=src_stage)
+            return dist.broadcast(tensor, src_rank, group=group, async_op=async_op)
 
 
 def recv(tensor, src_stage, async_op=False):
     global _groups
-    #assert async_op == False, "Doesnt support async_op true"
+    assert async_op == False, "Doesnt support async_op true"
     dest_stage = _grid.get_stage_id()
     _is_valid_send_recv(src_stage, dest_stage)
 
@@ -60,7 +79,11 @@ def recv(tensor, src_stage, async_op=False):
         op = dist.irecv(tensor, src_rank)
         _async.append(op)
     else:
-        return dist.recv(tensor, src_rank)
+        if can_send_recv():
+            return dist.recv(tensor, src_rank)
+        else:
+            group = _get_send_recv_group(src_stage, dest_stage)
+            return dist.broadcast(tensor, src_rank, group=group, async_op=async_op)
 
 
 def wait():
@@ -135,3 +158,27 @@ def _to(x):
 
     msg = _to(msg)
     return msg
+
+
+def _get_send_recv_group(src_stage, dest_stage):
+    '''the group id is always the smaller rank unless its a wrap around'''
+
+    stage_id = None
+
+    first_stage = 0
+    last_stage = _grid.pipe_parallel_size - 1
+
+    if (src_stage == first_stage and dest_stage == last_stage
+            or dest_stage == first_stage and src_stage == last_stage):
+        stage_id = last_stage
+    elif src_stage > dest_stage:
+        stage_id = dest_stage
+    else:
+        stage_id = src_stage
+    '''group_id corresponds to group of [group_id, group_id+1]
+     unless group_id is the rank of the last stage
+     in which case group_id correspods to group[group_id-num_stages+1, group_id]
+     '''
+    group_id = _grid.stage_to_global(stage_id=stage_id)
+
+    return _groups[group_id]
diff --git a/deepspeed/runtime/utils.py b/deepspeed/runtime/utils.py
index f3a4c4c0f1b3..8bf8abe21b33 100755
--- a/deepspeed/runtime/utils.py
+++ b/deepspeed/runtime/utils.py
@@ -49,6 +49,10 @@ def set_random_seed(seed):
     torch.manual_seed(seed)
 
 
+def is_model_parallel_parameter(p) -> bool:
+    return hasattr(p, 'model_parallel') and p.model_parallel
+
+
 def bwc_tensor_model_parallel_rank(mpu=None):
     """Backwards-compatible way of querying the tensor model parallel rank from
     an ``mpu`` object.
@@ -269,15 +273,19 @@ def get_grad_norm(parameters, norm_type=2, mpu=None):
         total_norm = total_norm_cuda[0].item()
     else:
         total_norm = 0.
+        tensor_mp_rank = bwc_tensor_model_parallel_rank(mpu=mpu)
         for p in parameters:
-            if mpu is not None:
-                if (mpu.get_model_parallel_rank() == 0
-                    ) or is_model_parallel_parameter(p):
-                    param_norm = p.grad.data.float().norm(norm_type)
-                    total_norm += param_norm.item()**norm_type
-            else:
-                param_norm = p.grad.data.float().norm(norm_type)
-                total_norm += param_norm.item()**norm_type
+            # Pipeline parallelism may replicate parameters. Avoid multi-counting.
+            if hasattr(p, 'ds_pipe_replicated') and p.ds_pipe_replicated:
+                continue
+
+            # Filter to avoid over-counting replicated tensors from tensor
+            # model parallelism
+            if (tensor_mp_rank > 0) and not is_model_parallel_parameter(p):
+                continue
+
+            param_norm = p.grad.data.float().norm(norm_type)
+            total_norm += param_norm.item()**norm_type
 
         # Sum across all model parallel GPUs.
         total_norm_cuda = torch.cuda.FloatTensor([float(total_norm)])
@@ -294,6 +302,48 @@ def get_grad_norm(parameters, norm_type=2, mpu=None):
     return total_norm
 
 
+def get_grad_zeros(parameters, mpu=None):
+    """Compute the number of grads with zero values.
+
+    This is adapted from get_grad_norm
+
+    Arguments:
+        parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a
+            single Tensor that will have gradients normalized
+
+    Returns:
+        Total number of params with zero values (viewed as a single vector).
+    """
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    parameters = list(filter(lambda p: p.grad is not None, parameters))
+
+    total_zeros = 0.
+    tensor_mp_rank = bwc_tensor_model_parallel_rank(mpu=mpu)
+    for p in parameters:
+        # Pipeline parallelism may replicate parameters. Avoid multi-counting.
+        if hasattr(p, 'ds_pipe_replicated') and p.ds_pipe_replicated:
+            continue
+
+        # Filter to avoid over-counting replicated tensors from tensor
+        # model parallelism
+        if (tensor_mp_rank > 0) and not is_model_parallel_parameter(p):
+            continue
+
+        count_zeros = p.grad.numel() - torch.count_nonzero(p.grad)
+        total_zeros += count_zeros.item()
+
+    # Sum across all model parallel GPUs.
+    total_zeros_cuda = torch.cuda.FloatTensor([float(total_zeros)])
+    if mpu is not None:
+        torch.distributed.all_reduce(total_zeros_cuda,
+                                     op=torch.distributed.ReduceOp.SUM,
+                                     group=mpu.get_model_parallel_group())
+    total_zeros = total_zeros_cuda[0].item()
+
+    return total_zeros
+
+
 def get_weight_norm(parameters, norm_type=2, mpu=None):
     """Clips gradient norm of an iterable of parameters.
 
@@ -326,24 +376,19 @@ def get_weight_norm(parameters, norm_type=2, mpu=None):
         total_norm = total_norm_cuda[0].item()
     else:
         total_norm = 0.
+        tensor_mp_rank = bwc_tensor_model_parallel_rank(mpu=mpu)
         for p in parameters:
-            if mpu is not None:
-                if (mpu.get_model_parallel_rank() == 0
-                    ) or is_model_parallel_parameter(p):
-                    try:
-                        param_norm = float(torch.norm(p, norm_type, dtype=torch.float32))
-                    except TypeError as err:
-                        param_norm = float(torch.norm(p.float(), norm_type))
-
-                    #param_norm = p.data.float().norm(norm_type)
-                    total_norm += param_norm**norm_type
-            else:
-                try:
-                    param_norm = float(torch.norm(p, norm_type, dtype=torch.float32))
-                except TypeError as err:
-                    param_norm = float(torch.norm(p.float(), norm_type))
-                #param_norm = p.data.float().norm(norm_type)
-                total_norm += param_norm**norm_type
+            # Pipeline parallelism may replicate parameters. Avoid multi-counting.
+            if hasattr(p, 'ds_pipe_replicated') and p.ds_pipe_replicated:
+                continue
+
+            # Filter to avoid over-counting replicated tensors from tensor
+            # model parallelism
+            if (tensor_mp_rank > 0) and not is_model_parallel_parameter(p):
+                continue
+
+            param_norm = p.data.float().norm(norm_type)
+            total_norm += param_norm**norm_type
 
         # Sum across all model parallel GPUs.
         total_norm_cuda = torch.cuda.FloatTensor([float(total_norm)])
diff --git a/deepspeed/runtime/zero/stage2.py b/deepspeed/runtime/zero/stage2.py
index ad529aa96391..64955ff400dc 100755
--- a/deepspeed/runtime/zero/stage2.py
+++ b/deepspeed/runtime/zero/stage2.py
@@ -13,7 +13,7 @@
 import collections
 
 from deepspeed.runtime.fp16.loss_scaler import LossScaler, DynamicLossScaler
-from deepspeed.runtime.utils import see_memory_usage, is_model_parallel_parameter
+from deepspeed.runtime.utils import bwc_tensor_model_parallel_rank, see_memory_usage, is_model_parallel_parameter
 from deepspeed.runtime.zero.config import ZERO_OPTIMIZATION_GRADIENTS
 from deepspeed.ops.adam import DeepSpeedCPUAdam
 from deepspeed.ops.op_builder import UtilsBuilder
@@ -149,7 +149,7 @@ def __init__(self,
             self.model_parallel_rank = 0
         else:
             self.model_parallel_group = mpu.get_model_parallel_group()
-            self.model_parallel_rank = mpu.get_model_parallel_rank()
+            self.model_parallel_rank = bwc_tensor_model_parallel_rank(mpu)
 
         self.overflow = False
         self.clip_grad = clip_grad
@@ -970,6 +970,10 @@ def complete_grad_norm_calculation_for_cpu_offload(self, params):
         total_norm = 0.0
         norm_type = 2.0
         for p in params:
+            # Pipeline parallelism may replicate parameters. Avoid multi-counting.
+            if hasattr(p, 'ds_pipe_replicated') and p.ds_pipe_replicated:
+                continue
+
             if is_model_parallel_parameter(p) or (self.model_parallel_rank == 0):
                 param_id = self.get_param_id(p)
                 # as some model have trainable parameters but skipped in training,
@@ -1366,6 +1370,9 @@ def get_grad_norm_direct(self, gradients, params, norm_type=2):
             #if dist.get_rank() == 0:
             #    logger.info(f"Total Norm begining {total_norm}")
             for g, p in zip(gradients, params):
+                # Pipeline parallelism may replicate parameters. Avoid multi-counting.
+                if hasattr(p, 'ds_pipe_replicated') and p.ds_pipe_replicated:
+                    continue
                 if is_model_parallel_parameter(p) or (self.model_parallel_rank == 0):
                     param_norm = g.data.double().norm(2)
                     total_norm += param_norm.item()**2
@@ -1499,11 +1506,6 @@ def step(self, closure=None):
 
             see_memory_usage('After overflow after clearing gradients')
 
-            logger.info(
-                "[deepspeed] fp16 dynamic loss scale overflow! Rank {} Skipping step. Attempted loss scale: {}, "
-                "reducing to {}".format(dist.get_rank(),
-                                        prev_scale,
-                                        self.loss_scale))
             self.start_timers(timer_names)
             self.stop_timers(timer_names)
             return

From f93e22b38c3a90d1f9980ea99b3c37a1d5a4e4f9 Mon Sep 17 00:00:00 2001
From: Jeff Rasley <jerasley@microsoft.com>
Date: Fri, 30 Jul 2021 15:58:44 -0700
Subject: [PATCH 29/38] Correctness fix PP+ZeRO for gradient accumulation +
 updates from master (#1263)

Co-authored-by: Olatunji Ruwase <olruwase@microsoft.com>
---
 deepspeed/runtime/engine.py         | 23 ++++++++++++++++++++---
 deepspeed/runtime/zero/config.py    |  6 ++++++
 deepspeed/runtime/zero/constants.py | 11 +++++++++--
 deepspeed/runtime/zero/stage2.py    | 18 +++++++++++++-----
 4 files changed, 48 insertions(+), 10 deletions(-)

diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
index 013924def623..5791b3fab85d 100755
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -501,6 +501,9 @@ def steps_per_print(self):
     def zero_allgather_partitions(self):
         return self._config.zero_config.allgather_partitions
 
+    def zero_round_robin_gradients(self):
+        return self._config.zero_config.round_robin_gradients
+
     def dump_state(self):
         return self._config.dump_state
 
@@ -592,7 +595,7 @@ def _configure_with_arguments(self, args, mpu):
         if "OMPI_COMM_WORLD_LOCAL_RANK" in os.environ:
             ompi_local_rank = os.environ.get("OMPI_COMM_WORLD_LOCAL_RANK")
             local_rank = os.environ.get('LOCAL_RANK', ompi_local_rank)
-            assert ompi_local_rank == local_rank, f"LOCAL_RANK ({local_rank}) != OMPI_COMM_WORLD_LOCAL_RANK ({mpi_local_rank}), " \
+            assert ompi_local_rank == local_rank, f"LOCAL_RANK ({local_rank}) != OMPI_COMM_WORLD_LOCAL_RANK ({ompi_local_rank}), " \
                 "not sure how to proceed as we're seeing conficting local rank info."
             os.environ['LOCAL_RANK'] = local_rank
 
@@ -926,6 +929,15 @@ def _configure_zero_optimizer(self, optimizer):
                 gradient_predivide=self.gradient_predivide)
         elif zero_stage <= ZERO_OPTIMIZATION_GRADIENTS:
             overlap_comm = self.zero_overlap_comm()
+            contiguous_gradients = self.zero_contiguous_gradients()
+            round_robin_gradients = self.zero_round_robin_gradients()
+
+            # Overlap and contiguous grads are meaningless in stage 1 and are ignored
+            if zero_stage == ZERO_OPTIMIZATION_OPTIMIZER_STATES:
+                overlap_comm = False
+                contiguous_gradients = False
+                round_robin_gradients = False
+
             if isinstance(self.module, PipelineModule):
                 if overlap_comm:
                     logger.warning(
@@ -940,7 +952,7 @@ def _configure_zero_optimizer(self, optimizer):
                 dynamic_loss_scale=self.dynamic_loss_scale(),
                 dynamic_loss_args=self.dynamic_loss_scale_args(),
                 clip_grad=self.gradient_clipping(),
-                contiguous_gradients=self.zero_contiguous_gradients(),
+                contiguous_gradients=contiguous_gradients,
                 reduce_bucket_size=self.zero_reduce_bucket_size(),
                 allgather_bucket_size=self.zero_allgather_bucket_size(),
                 dp_process_group=self.data_parallel_group,
@@ -952,7 +964,8 @@ def _configure_zero_optimizer(self, optimizer):
                 gradient_predivide_factor=self.gradient_predivide_factor(),
                 gradient_accumulation_steps=self.gradient_accumulation_steps(),
                 ignore_unused_parameters=self.zero_ignore_unused_parameters(),
-                partition_grads=zero_stage == ZERO_OPTIMIZATION_GRADIENTS)
+                partition_grads=zero_stage == ZERO_OPTIMIZATION_GRADIENTS,
+                round_robin_gradients=round_robin_gradients)
         elif zero_stage == ZERO_OPTIMIZATION_WEIGHTS:
             logger.info("Initializing ZeRO Stage 3") if dist.get_rank() == 0 else None
             from deepspeed.runtime.zero.stage3 import FP16_DeepSpeedZeroOptimizer_Stage3
@@ -1163,6 +1176,10 @@ def forward(self, *inputs, **kwargs):
         return loss
 
     def allreduce_gradients(self, bucket_size=MEMORY_OPT_ALLREDUCE_SIZE):
+        # Pass (PP) gas boundary flag to optimizer (required for zero)
+        self.optimizer.is_gradient_accumulation_boundary = self.is_gradient_accumulation_boundary(
+        )
+
         # ZeRO stage 2 communicates during non gradient accumulation boundaries as well
         if self.zero_optimization_partition_gradients():
             self.optimizer.overlapping_partition_gradients_reduce_epilogue()
diff --git a/deepspeed/runtime/zero/config.py b/deepspeed/runtime/zero/config.py
index 377ad94549a7..a48dd4e620b4 100755
--- a/deepspeed/runtime/zero/config.py
+++ b/deepspeed/runtime/zero/config.py
@@ -39,6 +39,7 @@ def __init__(self, param_dict):
         self.gather_fp16_weights_on_model_save = None
 
         self.ignore_unused_parameters = None
+        self.round_robin_gradients = None
 
         if ZERO_OPTIMIZATION in param_dict.keys():
             zero_config_dict = param_dict[ZERO_OPTIMIZATION]
@@ -184,3 +185,8 @@ def _initialize(self, zero_config_dict):
         self.legacy_stage1 = get_scalar_param(zero_config_dict,
                                               ZERO_OPTIMIZATION_LEGACY_STAGE1,
                                               ZERO_OPTIMIZATION_LEGACY_STAGE1_DEFAULT)
+
+        self.round_robin_gradients = get_scalar_param(
+            zero_config_dict,
+            ZERO_OPTIMIZATION_ROUND_ROBIN_GRADIENTS,
+            ZERO_OPTIMIZATION_ROUND_ROBIN_GRADIENTS_DEFAULT)
diff --git a/deepspeed/runtime/zero/constants.py b/deepspeed/runtime/zero/constants.py
index 7beebe00e717..e0a26d53609f 100755
--- a/deepspeed/runtime/zero/constants.py
+++ b/deepspeed/runtime/zero/constants.py
@@ -30,7 +30,8 @@
     "sub_group_size" : 1000000000000,
     "offload_param": {...},
     "offload_optimizer": {...},
-    "ignore_unused_parameters": [true|false]
+    "ignore_unused_parameters": [true|false],
+    "round_robin_gradients": [true|false]
     }
 }
 '''
@@ -124,6 +125,10 @@
 ZERO_OPTIMIZATION_LEGACY_STAGE1 = "legacy_stage1"
 ZERO_OPTIMIZATION_LEGACY_STAGE1_DEFAULT = False
 
+# Stage 2 - partition gradients in a round robin fashsion to load-balance reduction and offload copying
+ZERO_OPTIMIZATION_ROUND_ROBIN_GRADIENTS = 'round_robin_gradients'
+ZERO_OPTIMIZATION_ROUND_ROBIN_GRADIENTS_DEFAULT = False
+
 #yapf: disable
 ZERO_OPTIMIZATION_DEFAULT = {
     ZERO_OPTIMIZATION_STAGE:
@@ -161,5 +166,7 @@
     ZERO_OPTIMIZATION_IGNORE_UNUSED_PARAMETERS:
     ZERO_OPTIMIZATION_IGNORE_UNUSED_PARAMETERS_DEFAULT,
     ZERO_OPTIMIZATION_LEGACY_STAGE1:
-    ZERO_OPTIMIZATION_LEGACY_STAGE1_DEFAULT
+    ZERO_OPTIMIZATION_LEGACY_STAGE1_DEFAULT,
+    ZERO_OPTIMIZATION_ROUND_ROBIN_GRADIENTS:
+    ZERO_OPTIMIZATION_ROUND_ROBIN_GRADIENTS_DEFAULT
 }
diff --git a/deepspeed/runtime/zero/stage2.py b/deepspeed/runtime/zero/stage2.py
index 64955ff400dc..a7597773641d 100755
--- a/deepspeed/runtime/zero/stage2.py
+++ b/deepspeed/runtime/zero/stage2.py
@@ -99,12 +99,14 @@ def __init__(self,
                  gradient_predivide_factor=1.0,
                  gradient_accumulation_steps=1,
                  ignore_unused_parameters=True,
-                 partition_grads=True):
+                 partition_grads=True,
+                 round_robin_gradients=False):
 
         if dist.get_rank() == 0:
             logger.info(f"Reduce bucket size {reduce_bucket_size}")
             logger.info(f"Allgather bucket size {allgather_bucket_size}")
             logger.info(f"CPU Offload: {cpu_offload}")
+            logger.info(f'Round robin gradient partitioning: {round_robin_gradients}')
         # The fused optimizer does all the work. We need this layer for two reason:
         # 1. maintain same user API from apex.fp16_utils
         # 2. keep common stuff here in case we need to add ne552w fused optimizer later
@@ -159,6 +161,7 @@ def __init__(self,
         self.gradient_accumulation_steps = gradient_accumulation_steps
         self.micro_step_id = 0
         self.ignore_unused_parameters = ignore_unused_parameters
+        self.round_robin_gradients = round_robin_gradients
 
         self.extra_large_param_to_reduce = None
 
@@ -232,10 +235,15 @@ def __init__(self,
             # This ensures that gradients are reduced in a fashion such that ownership round robins among the ranks.
             # For example, rather than 3 gradients (g_n+2, g_n+1, g_n) that are reduced consecutively belonging
             # to the same rank, instead they will belong to 3 ranks (r_m+2, r_m+1, r_m).
-            round_robin_tensors, round_robin_indices = self._round_robin_reorder(
-                self.fp16_groups[i],
-                dist.get_world_size(group=self.dp_process_group)
-            )
+            if self.round_robin_gradients:
+                round_robin_tensors, round_robin_indices = self._round_robin_reorder(
+                    self.fp16_groups[i],
+                    dist.get_world_size(group=self.dp_process_group)
+                )
+            else:
+                round_robin_tensors = self.fp16_groups[i]
+                round_robin_indices = list(range(len(self.fp16_groups[i])))
+
             self.round_robin_fp16_groups.append(round_robin_tensors)
             self.round_robin_fp6_indices.append(round_robin_indices)
 

From e9b5dffa5e4af7dea4f6e9921014d511e7ada288 Mon Sep 17 00:00:00 2001
From: Jeff Rasley <jerasley@microsoft.com>
Date: Fri, 30 Jul 2021 19:04:05 -0700
Subject: [PATCH 30/38] dont clear grads in stage 1 code path

---
 deepspeed/runtime/zero/stage2.py | 23 ++++++++++++-----------
 1 file changed, 12 insertions(+), 11 deletions(-)

diff --git a/deepspeed/runtime/zero/stage2.py b/deepspeed/runtime/zero/stage2.py
index a7597773641d..c3840087915d 100755
--- a/deepspeed/runtime/zero/stage2.py
+++ b/deepspeed/runtime/zero/stage2.py
@@ -1094,17 +1094,18 @@ def reduce_ipg_grads(self):
 
                 self.params_already_reduced[param_id] = True
 
-                if not self.is_param_in_current_partition[param_id]:
-                    if self.overlap_comm and self.contiguous_gradients is False:
-                        # Clear grads of other partitions during the next reduction
-                        # to avoid clearing them before the reduction is complete.
-                        if self.previous_reduced_grads is None:
-                            self.previous_reduced_grads = []
-                        self.previous_reduced_grads.append(param)
-                    else:
-                        param.grad = None
-                elif self.contiguous_gradients:
-                    self.copy_grads_in_partition(param)
+                if self.partition_gradients:
+                    if not self.is_param_in_current_partition[param_id]:
+                        if self.overlap_comm and self.contiguous_gradients is False:
+                            # Clear grads of other partitions during the next reduction
+                            # to avoid clearing them before the reduction is complete.
+                            if self.previous_reduced_grads is None:
+                                self.previous_reduced_grads = []
+                            self.previous_reduced_grads.append(param)
+                        else:
+                            param.grad = None
+                    elif self.contiguous_gradients:
+                        self.copy_grads_in_partition(param)
 
         self.grads_in_ipg_bucket = []
         self.params_in_ipg_bucket = []

From 4b35409668a6e14a80a2dfc338d200894434cd0c Mon Sep 17 00:00:00 2001
From: Jeff Rasley <jerasley@microsoft.com>
Date: Fri, 30 Jul 2021 19:22:54 -0700
Subject: [PATCH 31/38] prevent none grads from being reduced

---
 deepspeed/runtime/zero/stage2.py | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/deepspeed/runtime/zero/stage2.py b/deepspeed/runtime/zero/stage2.py
index c3840087915d..ed0853ec4054 100755
--- a/deepspeed/runtime/zero/stage2.py
+++ b/deepspeed/runtime/zero/stage2.py
@@ -470,7 +470,7 @@ def initialize_optimizer_states(self):
 
         if not self.cpu_offload:
             for group in self.single_partition_of_fp32_groups:
-                group.grad = None
+                group.grad = None  #class init
 
         return
 
@@ -494,7 +494,8 @@ def reduce_gradients(self, pipeline_parallel=False):
         if not self.overlap_comm:
             for i, group in enumerate(self.fp16_groups):
                 for param in group:
-                    self.reduce_ready_partitions_and_remove_grads(param, i)
+                    if param.grad is not None:
+                        self.reduce_ready_partitions_and_remove_grads(param, i)
 
         # reduce any pending grads in either hook/non-hook case
         self.overlapping_partition_gradients_reduce_epilogue()
@@ -972,7 +973,7 @@ def async_inplace_copy_grad_to_fp32_buffer_from_gpu(self, param):
 
         src_tensor = param.grad.view(-1).narrow(0, source_offset, num_elements).float()
         dest_tensor.copy_(src_tensor, non_blocking=True)
-        param.grad = None
+        param.grad = None  #offload only
 
     def complete_grad_norm_calculation_for_cpu_offload(self, params):
         total_norm = 0.0
@@ -1103,7 +1104,7 @@ def reduce_ipg_grads(self):
                                 self.previous_reduced_grads = []
                             self.previous_reduced_grads.append(param)
                         else:
-                            param.grad = None
+                            param.grad = None  #only if self.partition_gradients
                     elif self.contiguous_gradients:
                         self.copy_grads_in_partition(param)
 
@@ -1125,7 +1126,7 @@ def are_all_related_partitions_reduced(params_id):
 
         for params_id in self.is_grad_computed[i][partition_id]:
             if are_all_related_partitions_reduced(params_id):
-                self.param_dict[params_id].grad = None
+                self.param_dict[params_id].grad = None  # dead code
 
     def flatten_and_print(self, message, tensors, start=0, n=5):
         flatten_tensor = self.flatten(tensors)
@@ -1214,7 +1215,7 @@ def allreduce_bucket(self, bucket, allreduce_always_fp32=False, rank=None, log=N
     def _clear_previous_reduced_grads(self):
         if self.previous_reduced_grads is not None:
             for param in self.previous_reduced_grads:
-                param.grad = None
+                param.grad = None  # overlap enabled
             self.previous_reduced_grads = None
 
     #if rank is specified do a reduction instead of an allreduce
@@ -1329,7 +1330,7 @@ def zero_grad(self, set_grads_to_None=True):
         for group in self.fp16_groups:
             for p in group:
                 if set_grads_to_None:
-                    p.grad = None
+                    p.grad = None  # epilogue and in step
                 else:
                     if p.grad is not None:
                         p.grad.detach_()
@@ -1458,7 +1459,7 @@ def get_flat_partition(self,
 
     def free_grad_in_param_list(self, param_list):
         for p in param_list:
-            p.grad = None
+            p.grad = None  # in step
 
     def reset_cpu_buffers(self):
         self.norm_for_param_grads = {}
@@ -1581,7 +1582,7 @@ def step(self, closure=None):
             #get rid of the fp32 gradients. Not needed anymore
             if not self.cpu_offload:
                 for group in self.single_partition_of_fp32_groups:
-                    group.grad = None
+                    group.grad = None  # in step
 
             for fp16_partitions, fp32_partition in zip(self.parallel_partitioned_fp16_groups, self.single_partition_of_fp32_groups):
                 fp16_partitions[partition_id].data.copy_(fp32_partition.data)

From bc17042a54a154a16573dfda5feba20ae0facc3d Mon Sep 17 00:00:00 2001
From: Jeff Rasley <jerasley@microsoft.com>
Date: Mon, 2 Aug 2021 10:00:11 -0700
Subject: [PATCH 32/38] fix empty grad zero tests

---
 tests/unit/simple_model.py |  5 +----
 tests/unit/test_fp16.py    | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/tests/unit/simple_model.py b/tests/unit/simple_model.py
index 9c6062d79faa..15c40976b6a1 100755
--- a/tests/unit/simple_model.py
+++ b/tests/unit/simple_model.py
@@ -17,10 +17,7 @@ def __init__(self, hidden_dim, empty_grad=False):
 
     def forward(self, x, y):
         hidden_dim = x
-        if self.empty_grad and torch.distributed.get_rank() == 0:
-            hidden_dim = self.linear(hidden_dim) + self.linear2(hidden_dim)
-        else:
-            hidden_dim = self.linear(hidden_dim)
+        hidden_dim = self.linear(hidden_dim)
         return self.cross_entropy_loss(hidden_dim, y)
 
 
diff --git a/tests/unit/test_fp16.py b/tests/unit/test_fp16.py
index b2e76f0b7b82..0c0ef3edd3a8 100755
--- a/tests/unit/test_fp16.py
+++ b/tests/unit/test_fp16.py
@@ -856,3 +856,38 @@ def _go(args):
             model.step()
 
     _go(args=args)
+
+
+@pytest.mark.parametrize('stage', [1, 2, 3])
+def test_zero_empty_grad(tmpdir, stage):
+    config_dict = {
+        "train_batch_size": 1,
+        "steps_per_print": 1,
+        "fp16": {
+            "enabled": True
+        },
+        "zero_optimization": {
+            "stage": stage
+        }
+    }
+    args = args_from_dict(tmpdir, config_dict)
+    hidden_dim = 10
+
+    model = SimpleModel(hidden_dim)
+
+    @distributed_test(world_size=[1])
+    def _go(args, model, hidden_dim):
+        optimizer = torch.optim.Adam(model.parameters())
+        model, _, _, _ = deepspeed.initialize(args=args,
+                                              model=model,
+                                              optimizer=optimizer)
+        data_loader = random_dataloader(model=model,
+                                        total_samples=50,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device)
+        for n, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+
+    _go(args=args, model=model, hidden_dim=hidden_dim)

From 6b428821fd0091fbbf0add3661a378c26e2e5b3f Mon Sep 17 00:00:00 2001
From: Olatunji Ruwase <olruwase@microsoft.com>
Date: Mon, 9 Aug 2021 11:37:56 -0700
Subject: [PATCH 33/38] Use mpu in DeepSpeedConfig() call (#1271)

* Use mpu in DeepSpeedConfig() call

* Improve argument naming
---
 .../runtime/zero/partition_parameters.py      | 67 +++++++++++--------
 1 file changed, 38 insertions(+), 29 deletions(-)

diff --git a/deepspeed/runtime/zero/partition_parameters.py b/deepspeed/runtime/zero/partition_parameters.py
index 6ef87f9e00aa..7dc9076eed3e 100755
--- a/deepspeed/runtime/zero/partition_parameters.py
+++ b/deepspeed/runtime/zero/partition_parameters.py
@@ -215,11 +215,16 @@ def recurse(cl):
 # Inserts _post_init_method at the end of init method
 # for all sub classes of torch.nn.Module
 class InsertPostInitMethodToModuleSubClasses(object):
-    def __init__(self, enabled=True, mem_efficient_linear=True, config=None, dtype=None):
+    def __init__(self,
+                 enabled=True,
+                 mem_efficient_linear=True,
+                 ds_config=None,
+                 dtype=None):
         self.mem_efficient_linear = mem_efficient_linear
         self.enabled = enabled
-        self._set_dtype(config, dtype)
-        assert self.dtype in [torch.half, torch.float], f"Invalid data type {self.dtype}, allowed values are [torch.half, torch.float]"
+        self._set_dtype(ds_config, dtype)
+        assert self.dtype in [
+            torch.half, torch.float], f"Invalid data type {self.dtype}, allowed values are [torch.half, torch.float]"
 
     def __enter__(self):
         if not self.enabled:
@@ -287,8 +292,8 @@ def _disable_class(cls):
         torch.Tensor.__new__ = torch.Tensor.__old_new__
         torch.empty = _orig_torch_empty
 
-        #un doing it here will undo it during training
-        #if self.mem_efficient_linear:
+        # un doing it here will undo it during training
+        # if self.mem_efficient_linear:
         #    torch.nn.functional.linear = self.linear_bk
         #        if self.mem_efficient_linear:
         #            torch.nn.functional.linear = self.linear_bk
@@ -303,8 +308,7 @@ def _post_init_method(self, module):
 
     def _set_dtype(self, ds_config, dtype):
         if ds_config is not None and dtype is None:
-            _ds_config = DeepSpeedConfig(ds_config)
-            self.dtype = torch.half if _ds_config.fp16_enabled else torch.float
+            self.dtype = torch.half if ds_config.fp16_enabled else torch.float
         elif dtype is None:
             self.dtype = torch.half
         else:
@@ -321,9 +325,11 @@ def __init__(self,
                  mem_efficient_linear=True,
                  remote_device=None,
                  pin_memory=False,
+                 config_dict_or_path=None,
                  config=None,
                  enabled=True,
-                 dtype=None):
+                 dtype=None,
+                 mpu=None):
         """A context to enable massive model construction for training with
         ZeRO-3. Models are automatically partitioned (or, sharded) across the
         system and converted to half precision.
@@ -343,12 +349,14 @@ def __init__(self,
             pin_memory (bool, optional): Potentially increase performance by
                 using pinned memory for model weights. ``remote_device`` must be
                 ``"cpu"``. Defaults to ``False``.
-            config (``json file`` or dict, optional): If provided, provides configuration
+            config_dict_or_path (dict or ``json file``, optional): If provided, provides configuration
                 for swapping fp16 params to NVMe.
+            config (dict or ``json file``, optional): Deprecated, use config_dict_or_path instead.
             enabled (bool, optional): If ``False``, this context has no
                 effect. Defaults to ``True``.
             dtype (``dtype``, optional): Can be used to change the data type of the parameters.
                 Supported options are ``torch.half`` and ``torch.float``. Defaults to ``None``
+            mpu (``object``, optional): A model parallelism unit object that implements get_{model,data}_parallel_{rank,group,wolrd_size}
 
         This context accelerates model initialization and enables models that
         are too large to allocate in their entirety in CPU memory. It has the
@@ -420,9 +428,11 @@ def get_model():
                 model = deepspeed.zero.Init(module=model)
         """
 
+        _ds_config = DeepSpeedConfig(config_dict_or_path,
+                                     mpu) if config_dict_or_path is not None else None
         super().__init__(enabled=enabled,
                          mem_efficient_linear=mem_efficient_linear,
-                         config=config,
+                         ds_config=_ds_config,
                          dtype=dtype)
         if not torch.distributed.is_initialized():
             init_distributed()
@@ -435,21 +445,20 @@ def get_model():
         self.rank = torch.distributed.get_rank(group=self.ds_process_group)
         self.world_size = torch.distributed.get_world_size(group=self.ds_process_group)
 
-        #Local device is the device where the parameters are consumed
-        #It is the device where parameters are fully instantiated using allgather
+        # Local device is the device where the parameters are consumed
+        # It is the device where parameters are fully instantiated using allgather
         self.local_device = torch.device('cuda:{}'.format(os.environ["LOCAL_RANK"]))
 
-        self._validate_remote_device(remote_device, config)
+        self._validate_remote_device(remote_device, _ds_config)
 
-        #Remote device is the device where parameter partiitons are stored
-        #It can be same as local_device or it could be CPU or NVMe.
+        # Remote device is the device where parameter partiitons are stored
+        # It can be same as local_device or it could be CPU or NVMe.
         self.remote_device = self.local_device if remote_device is None else remote_device
         self.pin_memory = pin_memory if (
             self.remote_device == OFFLOAD_CPU_DEVICE) else False
 
         # Enable fp16 param swapping to NVMe
         if self.remote_device == OFFLOAD_NVME_DEVICE:
-            _ds_config = DeepSpeedConfig(config)
             self.param_swapper = AsyncPartitionedParameterSwapper(_ds_config)
         else:
             self.param_swapper = None
@@ -463,22 +472,21 @@ def get_model():
                 self._convert_to_deepspeed_param(param)
                 param.partition()
 
-    def _validate_remote_device(self, remote_device, ds_config):
-        if ds_config is not None:
-            _ds_config = DeepSpeedConfig(ds_config)
+    def _validate_remote_device(self, remote_device, _ds_config):
+        if _ds_config is not None:
             if remote_device in [None, OFFLOAD_CPU_DEVICE]:
                 if _ds_config.zero_config.offload_param is not None:
                     offload_param_device = _ds_config.zero_config.offload_param[
                         OFFLOAD_PARAM_DEVICE]
                     assert offload_param_device != OFFLOAD_NVME_DEVICE, \
-                    f"{OFFLOAD_PARAM_DEVICE} in DeepSpeed Config cannot be {offload_param_device} if remote device is {remote_device}."
+                        f"{OFFLOAD_PARAM_DEVICE} in DeepSpeed Config cannot be {offload_param_device} if remote device is {remote_device}."
 
             if remote_device == OFFLOAD_NVME_DEVICE:
                 assert _ds_config.zero_config.offload_param is not None, \
-                f'{OFFLOAD_PARAM} must be defined in DeepSpeed Config if remote device is {OFFLOAD_NVME_DEVICE}.'
+                    f'{OFFLOAD_PARAM} must be defined in DeepSpeed Config if remote device is {OFFLOAD_NVME_DEVICE}.'
 
                 assert _ds_config.zero_config.offload_param[OFFLOAD_PARAM_NVME_PATH] is not None, \
-                f'{OFFLOAD_PARAM_NVME_PATH} in DeepSpeed Config cannot be None if remote device is {OFFLOAD_NVME_DEVICE}'
+                    f'{OFFLOAD_PARAM_NVME_PATH} in DeepSpeed Config cannot be None if remote device is {OFFLOAD_NVME_DEVICE}'
 
     def _post_init_method(self, module):
         #see_memory_usage(f"Before converting parmas in {module.__class__.__name__}", force=False)
@@ -624,7 +632,7 @@ def _ensure_availability_of_partitioned_params(self, params):
 
     def _all_gather(self, param_list, async_op=False, hierarchy=None):
 
-        #fetches from nvme if the partition is not available and in nvme
+        # fetches from nvme if the partition is not available and in nvme
         self._ensure_availability_of_partitioned_params(param_list)
 
         handles = []
@@ -651,10 +659,10 @@ def _all_gather(self, param_list, async_op=False, hierarchy=None):
     def _partition(self, param_list, force=False, has_been_updated=False):
         for param in param_list:
             #print_rank_0(f"Before Partitioning Param {param.ds_id}")
-            #self._param_status(param)
+            # self._param_status(param)
             self._partition_param(param, has_been_updated=has_been_updated)
             param.ds_status = ZeroParamStatus.NOT_AVAILABLE
-            #if param.ds_tensor is not None:
+            # if param.ds_tensor is not None:
             #    assert id(param.data) == id(param.ds_tensor.data), \
             #    "After the parameters are initially partitioned, make sure we are not recreating the partition."
             #print_rank_0(f"After Partitioning Param {param.ds_id}")
@@ -678,7 +686,7 @@ def _partition_param(self, param, buffer=None, has_been_updated=False):
             #     if numel in empty_buffers:
             #         empty_buffers[numel].append(buffer)
 
-            #if torch.distributed.get_rank():
+            # if torch.distributed.get_rank():
             #    print(f"Releasing {param.data.numel()}")
             if param.ds_tensor is not None and not has_been_updated:
 
@@ -687,7 +695,7 @@ def _partition_param(self, param, buffer=None, has_been_updated=False):
                 see_memory_usage(
                     f'Before partitioning param {param.ds_id} {param.shape}',
                     force=False)
-                #param.data does not store anything meaningful in partitioned state
+                # param.data does not store anything meaningful in partitioned state
                 param.data = torch.ones(1, dtype=self.dtype).to(param.device)
                 see_memory_usage(f'After partitioning param {param.ds_id} {param.shape}',
                                  force=False)
@@ -765,7 +773,7 @@ def _partition_param(self, param, buffer=None, has_been_updated=False):
 
             #param.data = param.ds_tensor.data
 
-            #param.data does not store anything meaningful in partitioned state
+            # param.data does not store anything meaningful in partitioned state
 
             see_memory_usage(f'Before partitioning param {param.ds_id} {param.shape}',
                              force=False)
@@ -1002,7 +1010,8 @@ def _partition_gradient(self, param, partition_buffer=None, accumulate=False):
                                            dtype=param.dtype,
                                            device=param.device)
         else:
-            assert partition_buffer.numel() >= partition_size, f"The partition buffer size {partition_buffer.numel()} should match the size of param.ds_tensor {partition_size}"
+            assert partition_buffer.numel(
+            ) >= partition_size, f"The partition buffer size {partition_buffer.numel()} should match the size of param.ds_tensor {partition_size}"
 
         rank = torch.distributed.get_rank(group=self.ds_process_group)
         start = partition_size * rank

From cce85b893b8a2490979d8e4b5d1e1f980b24bd71 Mon Sep 17 00:00:00 2001
From: Olatunji Ruwase <olruwase@microsoft.com>
Date: Mon, 9 Aug 2021 16:42:24 -0700
Subject: [PATCH 34/38] API for obtaining global gradient norm (#1292)

* FP16 fused and unfused grad norm query.

* API for obtaining global unclipped gradient norm across parameter groups

* Use global norm not group norms

Co-authored-by: Shaden Smith <shaden.smith@microsoft.com>
---
 deepspeed/runtime/engine.py                 | 18 ++++++++++++++++++
 deepspeed/runtime/fp16/fused_optimizer.py   | 20 +++++++++++---------
 deepspeed/runtime/fp16/unfused_optimizer.py | 17 ++++++++---------
 deepspeed/runtime/utils.py                  | 11 ++++++++++-
 deepspeed/runtime/zero/stage1.py            | 15 +++++++--------
 deepspeed/runtime/zero/stage2.py            | 14 ++++++--------
 deepspeed/runtime/zero/stage3.py            | 14 +++++---------
 7 files changed, 65 insertions(+), 44 deletions(-)

diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
index 5791b3fab85d..770f961908f1 100755
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -36,6 +36,7 @@
     ZERO_OPTIMIZATION_OPTIMIZER_STATES, ZERO_OPTIMIZATION_GRADIENTS, ZERO_OPTIMIZATION_WEIGHTS
 from deepspeed.runtime.csr_tensor import CSRTensor
 import deepspeed.runtime.lr_schedules as lr_schedules
+from deepspeed.runtime.utils import get_grad_norm
 from deepspeed.utils import logger, log_dist, init_distributed
 from deepspeed.utils.timer import ThroughputTimer, SynchronizedWallClockTimer
 from deepspeed.utils.debug import debug_extract_module_and_param_names
@@ -123,6 +124,7 @@ def __init__(self,
         self.gas_boundary_ctr = 0
         self.dist_backend = "nccl"
         self._step_applied = False
+        self._global_grad_norm = None
 
         # for debug purposes - can then debug print: debug_get_module_name(module)
         debug_extract_module_and_param_names(model)
@@ -256,6 +258,19 @@ def set_train_batch_size(self, train_batch_size):
         self._config.train_batch_size = train_batch_size
         self._config.gradient_accumulation_steps = new_gas
 
+    def get_global_grad_norm(self) -> float:
+        """Return the 2-norm of all gradients. If there is model parallelism,
+        the norm will be global.
+
+        The computed norm will be cached and reused until the next step() pass.
+        .. note::
+            In the presence of model parallelism, this is a collective call
+            and acts as a barrier among ``mpu.get_model_parallel_group()``.
+        Returns:
+            float: norm
+        """
+        return self._global_grad_norm
+
     def checkpoint_tag_validation_enabled(self):
         return self._config.checkpoint_tag_validation_enabled
 
@@ -1315,6 +1330,9 @@ def _take_model_step(self, lr_kwargs, block_eigenvalue={}):
 
         self.optimizer.step()
 
+        if hasattr(self.optimizer, '_global_grad_norm'):
+            self._global_grad_norm = self.optimizer._global_grad_norm
+
         # Quantize the updated parameter if there no overflow
         if self.quantizer:
             self.quantizer.quantize(
diff --git a/deepspeed/runtime/fp16/fused_optimizer.py b/deepspeed/runtime/fp16/fused_optimizer.py
index fba0d6b1fd59..72dd2c161845 100755
--- a/deepspeed/runtime/fp16/fused_optimizer.py
+++ b/deepspeed/runtime/fp16/fused_optimizer.py
@@ -9,7 +9,7 @@
 import math
 from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
 
-from deepspeed.runtime.utils import get_grad_norm, CheckOverflow, get_weight_norm
+from deepspeed.runtime.utils import get_global_norm, get_grad_norm, CheckOverflow, get_weight_norm
 from deepspeed.runtime.fp16.loss_scaler import INITIAL_LOSS_SCALE, SCALE_WINDOW, MIN_LOSS_SCALE
 from deepspeed.utils import logger, log_dist
 
@@ -45,6 +45,8 @@ def __init__(self,
         self.fp16_groups_flat = []
         self.fp32_groups_flat = []
 
+        self._global_grad_norm = 0.
+
         # loop to deal with groups
         for i, param_group in enumerate(self.optimizer.param_groups):
             # push this group to list before modify
@@ -161,8 +163,11 @@ def step_fused_adam(self, closure=None):
                     "scale: {}, reducing to {}".format(prev_scale,
                                                        self.cur_scale))
             return self.overflow
+
+        self._global_grad_norm = get_global_norm(norm_list=norm_groups)
+
         combined_scale = self.unscale_and_clip_grads(grads_groups_flat,
-                                                     norm_groups,
+                                                     self._global_grad_norm,
                                                      apply_scale=False)
         # norm is in fact norm*cur_scale
         self.optimizer.step(grads=[[g] for g in grads_groups_flat],
@@ -251,8 +256,10 @@ def step(self, closure=None):
         all_groups_norm = get_grad_norm(self.fp32_groups_flat, mpu=self.mpu)
         self.stop_timers([COMPUTE_NORM])
 
+        self._global_grad_norm = get_global_norm(norm_list=[all_groups_norm])
+
         self.start_timers([UNSCALE_AND_CLIP])
-        self.unscale_and_clip_grads(grads_groups_flat, [all_groups_norm])
+        self.unscale_and_clip_grads(grads_groups_flat, self._global_grad_norm)
         self.stop_timers([UNSCALE_AND_CLIP])
 
         self.start_timers([BASIC_STEP])
@@ -277,12 +284,7 @@ def step(self, closure=None):
 
         return self.overflow
 
-    def unscale_and_clip_grads(self, grad_groups_flat, norm_groups, apply_scale=True):
-        total_norm = 0.0
-        for norm in norm_groups:
-            total_norm += norm**2.0
-        total_norm = math.sqrt(total_norm)
-
+    def unscale_and_clip_grads(self, grad_groups_flat, total_norm, apply_scale=True):
         # compute combined scale factor for this group
         combined_scale = self.cur_scale
         if self.clip_grad > 0.:
diff --git a/deepspeed/runtime/fp16/unfused_optimizer.py b/deepspeed/runtime/fp16/unfused_optimizer.py
index c30df0bef1d0..08a4a7c41cb2 100755
--- a/deepspeed/runtime/fp16/unfused_optimizer.py
+++ b/deepspeed/runtime/fp16/unfused_optimizer.py
@@ -9,7 +9,7 @@
 from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
 import math
 
-from deepspeed.runtime.utils import get_grad_norm, CheckOverflow, get_weight_norm
+from deepspeed.runtime.utils import get_global_norm, get_grad_norm, CheckOverflow, get_weight_norm
 from deepspeed.runtime.fp16.loss_scaler import INITIAL_LOSS_SCALE, SCALE_WINDOW, MIN_LOSS_SCALE
 from deepspeed.utils import logger
 
@@ -32,6 +32,7 @@ def __init__(self,
                  fused_lamb_legacy=False):
 
         self.fused_lamb_legacy = fused_lamb_legacy
+        self._global_grad_norm = 0.
 
         if torch.distributed.get_rank() == 0:
             logger.info(f'Fused Lamb Legacy : {self.fused_lamb_legacy} ')
@@ -148,7 +149,9 @@ def step_fused_lamb(self, closure=None):
                                                        self.cur_scale))
             return self.overflow
 
-        combined_scale = self.unscale_and_clip_grads(norm_groups, apply_scale=False)
+        self._global_grad_norm = get_global_norm(norm_list=norm_groups)
+        combined_scale = self.unscale_and_clip_grads(self._global_grad_norm,
+                                                     apply_scale=False)
         self.optimizer.step(grads=grads_groups,
                             output_params=self.fp16_groups,
                             scale=combined_scale)
@@ -197,7 +200,8 @@ def step(self, closure=None):
                 else:
                     fp32_param.grad = fp16_param.grad.to(fp32_param.dtype)
 
-        self.unscale_and_clip_grads(norm_groups)
+        self._global_grad_norm = get_global_norm(norm_list=norm_groups)
+        self.unscale_and_clip_grads(self._global_grad_norm)
 
         self.optimizer.step()
 
@@ -212,12 +216,7 @@ def step(self, closure=None):
 
         return self.overflow
 
-    def unscale_and_clip_grads(self, norm_groups, apply_scale=True):
-        total_norm = 0.0
-        for norm in norm_groups:
-            total_norm += norm**2.0
-        total_norm = math.sqrt(total_norm)
-
+    def unscale_and_clip_grads(self, total_norm, apply_scale=True):
         # compute combined scale factor for this group
         combined_scale = self.cur_scale
         if self.clip_grad > 0.:
diff --git a/deepspeed/runtime/utils.py b/deepspeed/runtime/utils.py
index 8bf8abe21b33..da9727284232 100755
--- a/deepspeed/runtime/utils.py
+++ b/deepspeed/runtime/utils.py
@@ -9,7 +9,7 @@
 import os
 import psutil
 import gc
-from math import ceil
+from math import ceil, sqrt
 from math import floor
 from bisect import bisect_left, bisect_right
 
@@ -240,6 +240,15 @@ def _handle_overflow(cpu_sum, x, i):
         )
 
 
+def get_global_norm(norm_list):
+    """ Compute total from a list of norms
+    """
+    total_norm = 0.0
+    for norm in norm_list:
+        total_norm += norm**2.0
+    return sqrt(total_norm)
+
+
 def get_grad_norm(parameters, norm_type=2, mpu=None):
     """Clips gradient norm of an iterable of parameters.
 
diff --git a/deepspeed/runtime/zero/stage1.py b/deepspeed/runtime/zero/stage1.py
index 7660c9917b84..20a6c5a21944 100755
--- a/deepspeed/runtime/zero/stage1.py
+++ b/deepspeed/runtime/zero/stage1.py
@@ -5,7 +5,7 @@
 
 from deepspeed.runtime.zero.utils import _initialize_parameter_parallel_groups
 from deepspeed.runtime.fp16.loss_scaler import LossScaler, DynamicLossScaler
-from deepspeed.runtime.utils import get_grad_norm, CheckOverflow
+from deepspeed.runtime.utils import get_global_norm, get_grad_norm, CheckOverflow
 from deepspeed.runtime.zero.config import ZERO_OPTIMIZATION_OPTIMIZER_STATES
 from deepspeed.utils import logger, log_dist
 from deepspeed.ops.op_builder import UtilsBuilder
@@ -104,6 +104,7 @@ def __init__(self,
         self.postscale_gradients = postscale_gradients
         self.gradient_predivide_factor = gradient_predivide_factor
         self.gradient_average = gradient_average
+        self._global_grad_norm = 0.
 
         # TODO: automatically turn off if #params > some_limit
         self.all_gather_partitions = all_gather_partitions
@@ -683,8 +684,11 @@ def step(self, closure=None):
 
             local_sub_partitions_grad_groups.append(local_grad_sub_partitions)
 
+        self._global_grad_norm = get_global_norm(norm_list=norm_groups)
+
         #RS: update unscale/clip with sub partitions
-        self.unscale_and_clip_grads(local_sub_partitions_grad_groups, norm_groups)
+        self.unscale_and_clip_grads(local_sub_partitions_grad_groups,
+                                    self._global_grad_norm)
 
         self.optimizer.step()
 
@@ -720,12 +724,7 @@ def step(self, closure=None):
 
         return self.overflow
 
-    def unscale_and_clip_grads(self, grad_groups_flat, norm_groups):
-        total_norm = 0.0
-        for norm in norm_groups:
-            total_norm += norm**2.0
-        total_norm = math.sqrt(total_norm)
-
+    def unscale_and_clip_grads(self, grad_groups_flat, total_norm):
         # compute combined scale factor for this group
         combined_scale = self.loss_scale
         if self.clip_grad > 0.:
diff --git a/deepspeed/runtime/zero/stage2.py b/deepspeed/runtime/zero/stage2.py
index ed0853ec4054..09cdfe80b59b 100755
--- a/deepspeed/runtime/zero/stage2.py
+++ b/deepspeed/runtime/zero/stage2.py
@@ -13,7 +13,7 @@
 import collections
 
 from deepspeed.runtime.fp16.loss_scaler import LossScaler, DynamicLossScaler
-from deepspeed.runtime.utils import bwc_tensor_model_parallel_rank, see_memory_usage, is_model_parallel_parameter
+from deepspeed.runtime.utils import bwc_tensor_model_parallel_rank, get_global_norm, see_memory_usage, is_model_parallel_parameter
 from deepspeed.runtime.zero.config import ZERO_OPTIMIZATION_GRADIENTS
 from deepspeed.ops.adam import DeepSpeedCPUAdam
 from deepspeed.ops.op_builder import UtilsBuilder
@@ -146,6 +146,8 @@ def __init__(self,
 
         self.is_gradient_accumulation_boundary = True
 
+        self._global_grad_norm = 0.
+
         if mpu is None:
             self.model_parallel_group = None
             self.model_parallel_rank = 0
@@ -1560,7 +1562,8 @@ def step(self, closure=None):
 
             single_partition_grad_groups.append(single_grad_partition)
 
-        self.unscale_and_clip_grads(single_partition_grad_groups, norm_groups)
+        self._global_grad_norm = get_global_norm(norm_list=norm_groups)
+        self.unscale_and_clip_grads(single_partition_grad_groups, self._global_grad_norm)
         self.stop_timers([OPTIMIZER_GRADIENTS])
 
         self.start_timers([OPTIMIZER_STEP])
@@ -1636,12 +1639,7 @@ def step(self, closure=None):
 
         return
 
-    def unscale_and_clip_grads(self, grad_groups_flat, norm_groups):
-        total_norm = 0.0
-        for norm in norm_groups:
-            total_norm += norm**2.0
-        total_norm = math.sqrt(total_norm)
-
+    def unscale_and_clip_grads(self, grad_groups_flat, total_norm):
         # compute combined scale factor for this group
         combined_scale = self.loss_scale
         if self.clip_grad > 0.:
diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py
index 2806f5d420da..9b06d177fefd 100755
--- a/deepspeed/runtime/zero/stage3.py
+++ b/deepspeed/runtime/zero/stage3.py
@@ -16,7 +16,7 @@
 
 from deepspeed.utils.logging import logger
 from deepspeed.runtime.fp16.loss_scaler import LossScaler, DynamicLossScaler
-from deepspeed.runtime.utils import see_memory_usage, is_model_parallel_parameter
+from deepspeed.runtime.utils import get_global_norm, see_memory_usage, is_model_parallel_parameter
 from deepspeed.runtime.zero.partition_parameters import *
 from deepspeed.runtime.zero.partition_parameters import _init_external_params
 from deepspeed.runtime.zero.constants import ZERO_OPTIMIZATION_WEIGHTS
@@ -650,6 +650,7 @@ def __init__(self,
         self.flatten = util_ops.flatten
         self.unflatten = util_ops.unflatten
         self.dtype = self.optimizer.param_groups[0]['params'][0].dtype
+        self._global_grad_norm = 0.
 
         if not all(is_zero_param(p) for p in module.parameters()):
             group = None
@@ -2764,6 +2765,7 @@ def step(self, closure=None):
             return
 
         norm_groups = self._get_norm_groups()
+        self._global_grad_norm = get_global_norm(norm_list=norm_groups)
 
         timer_names = set()
 
@@ -2777,7 +2779,7 @@ def step(self, closure=None):
             self._prepare_sub_group(sub_group_id, timer_names)
 
             #scale the fp32 gradients
-            self.unscale_and_clip_grads(sub_group_id, norm_groups)
+            self.unscale_and_clip_grads(sub_group_id, self._global_grad_norm)
 
             #apply the optimizer step on the sub group and copy fp32 parameters to fp16
             self._optimizer_step(sub_group_id)
@@ -2826,15 +2828,9 @@ def dump_post_step_gradients(self):
                 norm_list = [param_norm, ds_norm] + unflat_norm
                 print(f'Post-Step Norms {i} {param_id} = {norm_list}')
 
-    def unscale_and_clip_grads(self, sub_group_id, norm_groups):
-
+    def unscale_and_clip_grads(self, sub_group_id, total_norm):
         grad_groups_flat = [self.fp32_partitioned_groups_flat[sub_group_id].grad]
 
-        total_norm = 0.0
-        for norm in norm_groups:
-            total_norm += norm**2.0
-        total_norm = math.sqrt(total_norm)
-
         # compute combined scale factor for this group
         combined_scale = self.loss_scale
         if self.clip_grad > 0.:

From e65e511b48129906bf15ca3f870c128d393d4bf5 Mon Sep 17 00:00:00 2001
From: Stas Bekman <stas00@users.noreply.github.com>
Date: Wed, 11 Aug 2021 05:20:38 -0700
Subject: [PATCH 35/38] turn excessive noise off (#1293)

---
 deepspeed/runtime/engine.py             | 4 ++--
 deepspeed/runtime/pipe/module.py        | 8 ++++----
 deepspeed/runtime/state_dict_factory.py | 4 ++--
 3 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
index 770f961908f1..3437998e534e 100755
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -2051,7 +2051,7 @@ def _copy_recovery_script(self, save_path):
         script = "zero_to_fp32.py"
         src = os.path.join(base_dir, "utils", script)
         dst = os.path.join(save_path, script)
-        logger.info(f"creating recovery script {dst}")
+        #logger.info(f"creating recovery script {dst}")
         copyfile(src, dst)
         # make executable
         os.chmod(dst, os.stat(dst).st_mode | stat.S_IEXEC)
@@ -2064,7 +2064,7 @@ def _save_zero_checkpoint(self, save_path, tag):
                        ds_version=version)
         torch.save(zero_sd, zero_checkpoint_name)
         self._copy_recovery_script(save_path)
-        logger.info('zero checkpoint saved {}'.format(zero_checkpoint_name))
+        #logger.info('zero checkpoint saved {}'.format(zero_checkpoint_name))
 
     def _zero3_consolidated_fp16_state_dict(self):
         """
diff --git a/deepspeed/runtime/pipe/module.py b/deepspeed/runtime/pipe/module.py
index af4df2a10ebf..977ce599622f 100644
--- a/deepspeed/runtime/pipe/module.py
+++ b/deepspeed/runtime/pipe/module.py
@@ -577,10 +577,10 @@ def load_state_dir(self, load_dir, strict=True):
 
             layer.load_state_dict(checkpoint)
 
-            if self._grid.data_parallel_id == 0:
-                logger.info(
-                    f'RANK={self.global_rank} Loaded layer={idx+self._local_start} file={load_path}'
-                )
+            # if self._grid.data_parallel_id == 0:
+            #     logger.info(
+            #         f'RANK={self.global_rank} Loaded layer={idx+self._local_start} file={load_path}'
+            #     )
 
         self._synchronize_tied_weights()
 
diff --git a/deepspeed/runtime/state_dict_factory.py b/deepspeed/runtime/state_dict_factory.py
index 033b6726ddb3..f5562e141c3f 100755
--- a/deepspeed/runtime/state_dict_factory.py
+++ b/deepspeed/runtime/state_dict_factory.py
@@ -78,7 +78,7 @@ def load(self,
         merge_count = 1
         if num_ckpt == mp_world_size:
             assert os.path.exists(load_path)
-            logger.info(f'rank: {mp_rank} loading checkpoint: {load_path}')
+            #logger.info(f'rank: {mp_rank} loading checkpoint: {load_path}')
             sd = torch.load(load_path, map_location=lambda storage, loc: storage)
 
             if quantize:
@@ -158,7 +158,7 @@ def set_module(self, sd, module):
         return sd
 
     def check_ckpt_list(self):
-        logger.info(f'checkpoint file list: {self.ckpt_list}')
+        #logger.info(f'checkpoint file list: {self.ckpt_list}')
         assert len(self.ckpt_list) > 0
 
         sd = torch.load(self.ckpt_list[0], map_location=lambda storage, loc: storage)

From db2f8a033be0e4a20bd1a25c922a567729445645 Mon Sep 17 00:00:00 2001
From: Jeff Rasley <jerasley@microsoft.com>
Date: Wed, 25 Aug 2021 07:11:55 -0700
Subject: [PATCH 36/38] [zero] restore fp16 params if no zero ckpts available
 (#1322)

* restore fp16 params if no zero ckpts available

* formatting
---
 deepspeed/runtime/engine.py | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
index 3437998e534e..a782581dde8d 100755
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -1752,9 +1752,12 @@ def load_checkpoint(self,
                                                          load_lr_scheduler_states=load_lr_scheduler_states)
 
         if self.zero_optimization() and load_path is not None:
-            self._load_zero_checkpoint(load_dir,
-                                       tag,
-                                       load_optimizer_states=load_optimizer_states)
+            success = self._load_zero_checkpoint(
+                load_dir,
+                tag,
+                load_optimizer_states=load_optimizer_states)
+            if not success:
+                self.optimizer._restore_from_fp16_weights()
 
         return load_path, client_states
 
@@ -1824,7 +1827,7 @@ def _load_checkpoint(self,
     def _load_zero_checkpoint(self, load_dir, tag, load_optimizer_states=True):
         zero_sd_list = self._get_all_zero_checkpoints(load_dir, tag)
         if zero_sd_list is None:
-            return
+            return False
 
         self.optimizer.load_state_dict(
             state_dict_list=zero_sd_list,
@@ -1833,6 +1836,7 @@ def _load_zero_checkpoint(self, load_dir, tag, load_optimizer_states=True):
         print(
             f'loading {len(zero_sd_list)} zero partition checkpoints for rank {self.global_rank}'
         )
+        return True
 
     def _get_mp_rank_zero_checkpoint_names(self, load_dir, tag, mp_rank, dp_world_size):
         zero_ckpt_names = []

From 72ce55ab59cdc2a2048b4a270b5e16455cbcaefe Mon Sep 17 00:00:00 2001
From: Olatunji Ruwase <olruwase@microsoft.com>
Date: Wed, 25 Aug 2021 10:18:07 -0700
Subject: [PATCH 37/38] Fix PP checkpoint bloat (#1324)

---
 deepspeed/runtime/pipe/module.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/deepspeed/runtime/pipe/module.py b/deepspeed/runtime/pipe/module.py
index 977ce599622f..e61dd1c72878 100644
--- a/deepspeed/runtime/pipe/module.py
+++ b/deepspeed/runtime/pipe/module.py
@@ -559,7 +559,18 @@ def save_state_dict(self, save_dir):
             model_ckpt_path = self.ckpt_layer_path(save_dir, idx)
             if not hasattr(layer, 'state_dict'):
                 continue
-            torch.save(layer.state_dict(), model_ckpt_path)
+            # We pass cloned tensors to torch.save() to avoid checkpoint bloat which occurs because torch.save()
+            # saves the underlying storage rather than the slice of the storage corresponding to individual tensors.
+            # This is a problem in DeepSpeed because we often allocate tensors using slices of large flattened buffers.
+            # Tensor cloning helps to avoid this problem because the storage of cloned tensors are closer to the true size.
+            # It is expected that the garbage collector will reclaim the cloned tensor storage to avoid memory bloat.
+            # See https://pytorch.org/docs/stable/notes/serialization.html#preserve-storage-sharing
+            orig_state_dict = layer.state_dict()
+            final_state_dict = type(orig_state_dict)(
+                {k: v.clone()
+                 for k,
+                 v in orig_state_dict.items()})
+            torch.save(final_state_dict, model_ckpt_path)
 
     def load_state_dir(self, load_dir, strict=True):
         for idx, layer in enumerate(self.forward_funcs):

From c7f3bc51c27884ad80dcafe4aa60f070c1dfa26e Mon Sep 17 00:00:00 2001
From: Stas Bekman <stas00@users.noreply.github.com>
Date: Mon, 30 Aug 2021 09:48:09 -0700
Subject: [PATCH 38/38] update for cuda-11.4 (#1329)

Co-authored-by: Jeff Rasley <jerasley@microsoft.com>
---
 op_builder/builder.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/op_builder/builder.py b/op_builder/builder.py
index 21547f896473..70f0fd3d3e55 100644
--- a/op_builder/builder.py
+++ b/op_builder/builder.py
@@ -63,7 +63,8 @@ def get_default_compute_capatabilities():
     11: ["11.0",
          "11.1",
          "11.2",
-         "11.3"],
+         "11.3",
+         "11.4"],
 }