Repository: deepspeedai/DeepSpeed
Pull request: Enable tensor fragments for zero 2 & 3
Enable tensor fragmentation in zero stage 2 & 3
Fix #2290
TODO: 
test w/ offload (at least cpu to start with)
test w/o offload

Problem statement:
[REQUEST] individual tensor `.grad` getter/setter for flattened tensors
As we know DS flattens individual tensors by optim groups so that each tensor and its attributes can't be accessed by the user once the ds engine takes over.

As @tjruwase is working on the API used by the universal checkpoint to get/set fragments of data/grads/optim states across multiple gpus, I think the next stage is to override some accessors in those flattened tensors to transparently get/set data on those.

For example `weight.grad` could behind the scenes get/set the sharded fragments - allowing users normal access to important data. e.g. researchers often need to access grads. optim states is another really important data that needs to be accessible. e.g. we wanted this to diagnose spikes during the training, but alas that data wasn't accessible.

And of course should look at what other accessors could be overriden and ideally the flattened tensors would appear to the user as normal tensors.

Thank you!

@tjruwase

Hints from issue discussion:

Commit message: Update deepspeed/utils/tensor_fragment.py

Co-authored-by: Stas Bekman <stas00@users.noreply.github.com>

Patch to review:
diff --git a/deepspeed/runtime/zero/stage_1_and_2.py b/deepspeed/runtime/zero/stage_1_and_2.py
index d1e292caee7b..d4709da7f6b0 100755
--- a/deepspeed/runtime/zero/stage_1_and_2.py
+++ b/deepspeed/runtime/zero/stage_1_and_2.py
@@ -555,6 +555,8 @@ def _link_all_hp_params(self):
             link_hp_params(
                 lp_param_list=self.bit16_groups[i],
                 flat_hp_partition=flat_hp_partition,
+                gradient_dict=self.averaged_gradients,
+                param_group_index=i,
                 partition_start=partition_id * partition_size,
                 partition_size=partition_size,
                 partition_optimizer_state=self.optimizer.state[flat_hp_partition],
diff --git a/deepspeed/utils/__init__.py b/deepspeed/utils/__init__.py
index 6f71897e37a6..3b6f44e47778 100644
--- a/deepspeed/utils/__init__.py
+++ b/deepspeed/utils/__init__.py
@@ -4,6 +4,8 @@
 from .init_on_device import OnDevice
 from .groups import *
 from .nvtx import instrument_w_nvtx
-from .tensor_fragment import tensor_fragment, get_full_hp_param, get_hp_fragment_mapping, fragment_address
+# TODO: Move tensor fragment and mixed precision to zero utils
+from .tensor_fragment import tensor_fragment, get_full_hp_param, get_hp_fragment_mapping, fragment_address, get_full_hp_grad
+from .tensor_fragment import safe_get_full_fp32_param, safe_get_full_grad, safe_get_full_optimizer_state
 from .mixed_precision_linkage import link_hp_params
 from deepspeed.runtime.dataloader import RepeatingLoader
diff --git a/deepspeed/utils/mixed_precision_linkage.py b/deepspeed/utils/mixed_precision_linkage.py
index bfd9932b8d7a..5db63398ad69 100644
--- a/deepspeed/utils/mixed_precision_linkage.py
+++ b/deepspeed/utils/mixed_precision_linkage.py
@@ -2,11 +2,13 @@
 Copyright 2022 The Microsoft DeepSpeed Team
 """
 import types
-from deepspeed.utils import get_full_hp_param, get_hp_fragment_mapping
+from deepspeed.utils import get_full_hp_param, get_full_hp_grad, get_hp_fragment_mapping
 
 
 def link_hp_params(lp_param_list,
                    flat_hp_partition,
+                   gradient_dict,
+                   param_group_index,
                    partition_start,
                    partition_size,
                    partition_optimizer_state,
@@ -20,6 +22,8 @@ def link_hp_params(lp_param_list,
         lp_param._hp_mapping = get_hp_fragment_mapping(lp_param,
                                                        lp_start,
                                                        flat_hp_partition,
+                                                       gradient_dict,
+                                                       param_group_index,
                                                        partition_start,
                                                        partition_size,
                                                        partition_optimizer_state)
@@ -29,10 +33,12 @@ def _init_lp_to_hp_mapping(lp_param_list, partition_start, partition_size, dp_gr
     current_offset = 0
     param_and_offset_list = []
     partition_end = partition_start + partition_size
-    for lp_param in lp_param_list:
+    for i, lp_param in enumerate(lp_param_list):
         lp_param._hp_mapping = None
         lp_param._dp_group = dp_group
         lp_param.get_full_hp_param = types.MethodType(get_full_hp_param, lp_param)
+        lp_param.get_full_hp_grad = types.MethodType(get_full_hp_grad, lp_param)
+        lp_param._index_in_param_group = i
 
         # lp_param overlaps with partition if both are true
         # 1) current_offset < partition_end,
diff --git a/deepspeed/utils/tensor_fragment.py b/deepspeed/utils/tensor_fragment.py
index 37e394b666e5..70f0591ac200 100644
--- a/deepspeed/utils/tensor_fragment.py
+++ b/deepspeed/utils/tensor_fragment.py
@@ -16,10 +16,13 @@ class fragment_address:
 @dataclass
 class tensor_fragment:
     lp_fragment: torch.Tensor
+    # lp_grad_fragment: torch.Tensor
     lp_fragment_address: fragment_address
     hp_fragment: torch.Tensor
     hp_fragment_address: fragment_address
     optim_fragment: {}
+    gradient_dict: {}
+    param_group_index: int
 
     def update_hp(self):
         self.hp_fragment.data.copy_(self.lp_fragment.data)
@@ -58,9 +61,59 @@ def get_full_hp_param(self, optim_state_key=None):
     return reduce_buffer.reshape_as(self)
 
 
+def safe_get_full_fp32_param(param):
+    if hasattr(param, '_hp_mapping'):
+        return param.get_full_hp_param()
+    return None
+
+
+def safe_get_full_optimizer_state(param, optim_state_key):
+    if hasattr(param, '_hp_mapping'):
+        return param.get_full_hp_param(optim_state_key)
+    return None
+
+
+def get_full_hp_grad(self):
+    reduce_buffer = torch.zeros_like(self, dtype=torch.float32).flatten()
+    if self._hp_mapping is not None:
+        lp_frag_address = self._hp_mapping.lp_fragment_address
+        reduce_fragment = torch.narrow(reduce_buffer,
+                                       0,
+                                       lp_frag_address.start,
+                                       lp_frag_address.numel)
+
+        hp_frag_address = self._hp_mapping.hp_fragment_address
+        hp_mapping = self._hp_mapping
+        # lp_grad_fragment = torch.narrow(hp_mapping.gradient_dict[hp_mapping.param_group_index],
+        #                                 0,
+        #                                hp_frag_address.start,
+        #                                hp_frag_address.numel)
+        from deepspeed.utils import logger
+        if hp_mapping.param_group_index not in hp_mapping.gradient_dict or hp_mapping.gradient_dict[hp_mapping.param_group_index] is None:
+            logger.warning("Gradients are only available immediately after backward and before egine step")
+            return None
+        lp_grad_fragment = hp_mapping.gradient_dict[hp_mapping.param_group_index][
+            self._index_in_param_group]
+        hp_grad_fragment = lp_grad_fragment.to(torch.float32).flatten()
+        reduce_fragment.data.copy_(hp_grad_fragment.data)
+    dist.all_reduce(reduce_buffer, group=self._dp_group)
+    return reduce_buffer.reshape_as(self)
+
+
+# TODO: Figure out the correct return dtype
+def safe_get_full_grad(param):
+    if param.grad is not None:
+        return param.grad
+    if hasattr(param, '_hp_mapping'):
+        return param.get_full_hp_grad()
+    return None
+
+
 def get_hp_fragment_mapping(lp_param,
                             lp_start,
                             flat_hp_partition,
+                            gradient_dict,
+                            param_group_index,
                             partition_start,
                             partition_size,
                             optimizer_state_dict):
@@ -101,7 +154,9 @@ def get_hp_fragment_mapping(lp_param,
                            lp_fragment_address=lp_frag_address,
                            hp_fragment=hp_fragment_tensor,
                            hp_fragment_address=hp_frag_address,
-                           optim_fragment=optim_fragment)
+                           optim_fragment=optim_fragment,
+                           gradient_dict=gradient_dict,
+                           param_group_index=param_group_index)
 
 
 '''


Code review comments:
diff_hunk: @@ -58,9 +61,59 @@ def get_full_hp_param(self, optim_state_key=None):
     return reduce_buffer.reshape_as(self)
 
 
+def safe_get_full_fp32_param(param):
+    if hasattr(param, '_hp_mapping'):
+        return param.get_full_hp_param()
+    return None
+
+
+def safe_get_full_optimizer_state(param, optim_state_key):
+    if hasattr(param, '_hp_mapping'):
+        return param.get_full_hp_param(optim_state_key)
+    return None
+
+
+def get_full_hp_grad(self):
+    reduce_buffer = torch.zeros_like(self, dtype=torch.float32).flatten()
+    if self._hp_mapping is not None:
+        lp_frag_address = self._hp_mapping.lp_fragment_address
+        reduce_fragment = torch.narrow(reduce_buffer,
+                                       0,
+                                       lp_frag_address.start,
+                                       lp_frag_address.numel)
+
+        hp_frag_address = self._hp_mapping.hp_fragment_address
+        hp_mapping = self._hp_mapping
+        # lp_grad_fragment = torch.narrow(hp_mapping.gradient_dict[hp_mapping.param_group_index],
+        #                                 0,
+        #                                hp_frag_address.start,
+        #                                hp_frag_address.numel)
+        from deepspeed.utils import logger
+        if hp_mapping.param_group_index not in hp_mapping.gradient_dict or hp_mapping.gradient_dict[hp_mapping.param_group_index] is None:
+            logger.warning("Gradients are only available immediately after backward and before egine step")
+            return None
path: deepspeed/utils/tensor_fragment.py
text: @user1:
```suggestion
        if hp_mapping.param_group_index not in hp_mapping.gradient_dict or hp_mapping.gradient_dict[hp_mapping.param_group_index] is None:
           raise ValueError("Gradients are only available immediately after backward and before engine step")
```

@user1:
changed to assert