From 9e35a29c830d2eb4b99ee0190d6b654046ef85f0 Mon Sep 17 00:00:00 2001 From: iiap <1471127927@qq.com> Date: Fri, 24 Jul 2026 16:15:53 +0800 Subject: [PATCH] Fix inter-iteration activation retention in HyperConnectionTransformerLayer Release mlp_norm_manager after group_offload, mirroring the base class _forward_post_mlp. FineGrainedActivationOffloadingInterface stores the wrapped tensor unconditionally (even with offload=False), so a manager left attached to the layer keeps one microbatch's pre-MLP-norm input alive across iterations; with recompute_granularity='full' its autograd history additionally pins the layer's recomputation subgraph. Signed-off-by: iiap <1471127927@qq.com> --- megatron/core/transformer/transformer_layer.py | 1 + 1 file changed, 1 insertion(+) diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index 34a456bfa69..1b22f0a608b 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -2296,6 +2296,7 @@ def _forward_post_mlp_with_fused_hyper_connection( ) hidden_states = self.mlp_norm_manager.group_offload(hidden_states) + self.mlp_norm_manager = None output = make_viewless_tensor( inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True