Fix inter-iteration activation retention in HyperConnectionTransformerLayer

[htesd]

What does this PR do?

Re-submission of #6017, rebased onto current dev (c76ff61). The code change is identical — one line.

Why the original PR disappeared: #6017 was not withdrawn by me. On Aug 5 my fork left the upstream fork network (the repository was turned private for unrelated internal use), and GitHub auto-closes any open PR coming from a detached fork, attributing the close to the repo owner — which is why the timeline reads "htesd closed this". Apologies for the churn, @yaox12 @hxbai: this happened ~3 hours after @yaox12 merged dev into the branch and re-triggered CI. This PR comes from a fresh fork (htesd/Megatron-LM-oss) that is properly attached to the network.

Prior review status on #6017, for reference:

• /claude strict-review: passed, 0 findings (CRITICAL 0 / IMPORTANT 0 / SUGGESTION 0)
• @hxbai: approved + /ok to test
• No CI job ever failed — the runs were cancelled after sitting in the queue for ~42h.

The bug

HyperConnectionTransformerLayer._forward_post_mlp_with_fused_hyper_connection (transformer_layer.py:2305) never releases self.mlp_norm_manager after calling group_offload(), unlike the base class TransformerLayer._forward_post_mlp (transformer_layer.py:1134-1138), which sets self.mlp_norm_manager = None right after use.

Because FineGrainedActivationOffloadingInterface.__init__ stores the wrapped tensor unconditionally (self.tensor = tensor, even when offload=False), the manager left attached to the layer module keeps a reference to one microbatch's pre-MLP-norm input — assigned unconditionally in _forward_mlp (transformer_layer.py:2151) on every forward — until the next forward pass overwrites it. The memory is retained across the idle window between iterations.

The effect is strongly amplified with recompute_granularity='full': the retained tensor is the one produced during the recomputation pass, and its autograd history keeps the layer's entire recomputation subgraph alive (stream-expanded residual buffers, h_post/BDA outputs, input grads). Holding one (s, b, h) tensor ends up pinning (s, b, n*h)-sized buffers several times over.

Measured impact

Numbers below are from #6017 (measured on that PR's base commit); the affected code path is unchanged on current dev.

2-layer GPT config with hyper connections (num_residual_streams=4), bf16, mbs=1, mock data, recompute_granularity='full', single GPU (H200), measuring torch.cuda.memory_allocated() in the idle window between iterations:

metric	before	after
idle floor between iterations, seq 8192	29.30 GB	28.63 GB (−0.67 GB)
floor growth from seq 4096 to 8192	+0.61 GB	+0.27 GB

Extrapolated to a 9-layer pipeline stage at 16k sequence length, the retention is roughly 7 GB of wasted per-GPU memory.

The one-line fix mirrors the existing cleanup in the base class _forward_post_mlp.

A possible follow-up hardening (not included here, to keep this minimal): make FineGrainedActivationOffloadingInterface not store the tensor when offload=False.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.