Add comment clarifying refresh-before-reenter ordering

Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

Author: tohtana

deepspeed/runtime/zero/stage3.py

@@ -1279,6 +1279,7 @@ def wrapper(param):
 
                         @instrument_w_nvtx
                         def reduce_partition_and_remove_grads(*notneeded):
+                            # Evaluate refresh condition before reenter_backward_if_needed()
                             refresh_expected = self.should_refresh_expected_hook_count()
                             # Re-enter backward for subsequent phases in reentrant checkpointing
                             self.reenter_backward_if_needed()
@@ -1307,6 +1308,7 @@ def reduce_partition_and_remove_grads(*notneeded):
             def make_hook(params):
 
                 def reduce_leaf_module_grads(module, grad_input, grad_output):
+                    # Evaluate refresh condition before reenter_backward_if_needed()
                     refresh_expected = self.should_refresh_expected_hook_count()
                     self.reenter_backward_if_needed()
 

deepspeed/runtime/zero/stage_1_and_2.py

@@ -1046,6 +1046,7 @@ def create_gradient_handling_hooks(self):
                     def wrapper(param, i):
 
                         def grad_handling_hook(*notneeded):
+                            # Evaluate refresh condition before reenter_backward_if_needed()
                             refresh_expected = self.should_refresh_expected_hook_count()
                             self.reenter_backward_if_needed()
                             self.process_gradients(param, i)