diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index ea714552464..a9b7a9cfd1a 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -2303,6 +2303,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