rollback deepspeedai#6726 (deepspeedai#7258)

This PR rollback deepspeedai#6726 which caused
deepspeedai#7116 .

---------

Signed-off-by: Guokai Ma <guokai.ma@gmail.com>
Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com>
Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Signed-off-by: Max Kovalenko <mkovalenko@habana.ai>

Author: 3 people

deepspeed/runtime/zero/parameter_offload.py

@@ -372,8 +372,7 @@ def _run_before_forward_function(input):
 
         @torch.compiler.disable
         def _post_backward_module_hook(module, inputs):
-            if not hasattr(module, "ds_grads_remaining"):
-                module.ds_grads_remaining = 0
+            module.ds_grads_remaining = 0
 
             return apply_to_tensors_only(module.post_bwd_fn.apply,
                                          inputs,

deepspeed/runtime/zero/partitioned_param_coordinator.py

@@ -261,6 +261,7 @@ def reset_step(self) -> None:
         self.__most_recent_step_id_param_fetched_for = collections.defaultdict(lambda: int(-1e10))
         self.__step_id_module_fetched_for = collections.defaultdict(lambda: collections.deque())
         self.__step_id = 0
+        self.__n_available_params = 0
         self.__profiler.reset_events()
 
     def _dump_params(self, tag, sub_module, params, step_id=None):
@@ -451,7 +452,7 @@ def release_and_reset_all(self, module: Module) -> None:
             # there's a hook execution issue
             param.ds_active_sub_modules.clear()
             self.__release_param(param)
-        self.__n_available_params = 0
+
         for param in iter_params(module, recurse=True):
             if param.ds_status != ZeroParamStatus.NOT_AVAILABLE:
                 raise RuntimeError(f"{param.ds_summary()} expected to be released")