From 43e037ee2e3d3430dc1482e6514bded6654b8816 Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Wed, 29 Apr 2026 09:41:09 +0800 Subject: [PATCH] remove dead manual dgrad release path in 1F1B overlap schedule MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `ScheduleNode.manual_release_grads` was hardcoded to False at construction and never reassigned anywhere in the codebase, so the two `if self.manual_release_grads:` branches in `ScheduleNode._backward` and `TransformerLayerNode.backward_dw` were unreachable. The supporting bookkeeping — `self.delay_grads_release`, `self.output_grads`, the `assert self.delay_grads_release` in `backward_dw`, and the `if self.delay_wgrad_compute: self.output_grads = grads; ...` block in `backward_impl` — existed only to feed the unreachable branches. PyTorch's caching allocator already reclaims the dgrad tensors when the last Python reference drops at the end of `_backward` / `backward_dw`, so this whole machinery was a no-op. A/B comparison with the flag temporarily exposed as a TransformerConfig field (qwen3_moe_proxy, PP=1 EP=8, overlap_moe_expert_parallel_comm + delay_wgrad_compute, 100 iters on real 8xH100): peak max-allocated memory was identical at 12479.13 MB with the flag both ON and OFF. --- megatron/core/models/gpt/fine_grained_callables.py | 13 ------------- megatron/core/pipeline_parallel/utils.py | 8 -------- 2 files changed, 21 deletions(-) diff --git a/megatron/core/models/gpt/fine_grained_callables.py b/megatron/core/models/gpt/fine_grained_callables.py index 1a16cbf6434..e78e4121a8b 100644 --- a/megatron/core/models/gpt/fine_grained_callables.py +++ b/megatron/core/models/gpt/fine_grained_callables.py @@ -316,12 +316,6 @@ def backward_impl(self, outputs, output_grad): detached_grad = tuple([e.grad for e in self.detached]) grads = output_grad + detached_grad self.default_backward_func(outputs + self.before_detached, grads) - # release the output grad memory after backward finishes, - # except when delay_wgrad_comptue is enabled, the grad should be - # kept until all modules' backward_dw has been invoked. - if self.delay_wgrad_compute: - self.output_grads = grads - self.delay_grads_release = len(self.bwd_dw_callables) > 0 # return grads for record stream return grads @@ -339,13 +333,6 @@ def backward_dw(self): module.backward_dw() nvtx_range_pop(nvtx_msg) - # 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." - 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 def __del__(self): diff --git a/megatron/core/pipeline_parallel/utils.py b/megatron/core/pipeline_parallel/utils.py index 48ce3d34a3c..0593693501c 100644 --- a/megatron/core/pipeline_parallel/utils.py +++ b/megatron/core/pipeline_parallel/utils.py @@ -182,8 +182,6 @@ def __init__( self.free_input = free_input 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""" @@ -263,12 +261,6 @@ def _backward(self, *output_grad): for g in output_grad: if g is not None: g.record_stream(self.stream) - # Manually trigger the memory release of dgrad tensor - # to avoid delayed garbage collection. If - # delay_grads_release is True, dgrad is last used in - # wgrad compute and skip the release here. - if self.manual_release_grads and not self.delay_grads_release: - g.untyped_storage().resize_(0) grads = self.get_grad() self._release_state()