[Dev] Disable ep overlap memory optimization (NVIDIA#2750)

Author: Wohox

megatron/core/models/gpt/fine_grained_callables.py

@@ -304,8 +304,9 @@ def backward_dw(self):
 
         # the output grad memory is last used in wgrad compute, should be safe to release.
         assert self.delay_grads_release, "output grad memory should be valid before wgrad."
-        for tensor in self.output_grads:
-            tensor.untyped_storage().resize_(0)
+        if self.manual_release_grads:
+            for tensor in self.output_grads:
+                tensor.untyped_storage().resize_(0)
         self.output_grads = None
 
         self.bwd_dw_callables = None

megatron/core/pipeline_parallel/utils.py

@@ -183,6 +183,7 @@ def __init__(
         self.inputs = None
         self.outputs = None
         self.delay_grads_release = False
+        self.manual_release_grads = False
 
     def default_backward_func(self, outputs, output_grad):
         """Default backward function"""
@@ -268,7 +269,7 @@ def _backward(self, *output_grad):
                     # to avoid delayed garbage collection. If
                     # delay_grads_release is True, dgrad is last used in
                     # wgrad compute and skip the release here.
-                    if not self.delay_grads_release:
+                    if self.manual_release_grads and not self.delay_grads_release:
                         g.untyped_storage().resize_(0)
 
         grads = self.get_grad()