Fix inter-iteration activation retention in HyperConnectionTransforme… - #6017
Fix inter-iteration activation retention in HyperConnectionTransforme…#6017htesd wants to merge 2 commits into
Conversation
…rLayer 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>
|
/claude strict-review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Strict review passed — no significant issues found. LGTM
The one-line fix (self.mlp_norm_manager = None after group_offload()) correctly mirrors the existing cleanup in the base class TransformerLayer._forward_post_mlp (line 1138). mlp_norm_manager is set unconditionally earlier in the same forward path (line 2144), so it is always non-None at the group_offload call; nulling it afterward only releases the retained reference and has no effect on the returned hidden_states or subsequent computation. No API, config, checkpoint, or process-group concerns.
Findings — CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 0. Risk level: low.
|
/ok to test 9e35a29 |
|
Hi @hxbai, thanks for the approval! It looks like the CI run was cancelled due to infrastructure issues (unit test jobs were queued for ~42h, then stuck for 28h+ before the whole run was cancelled on Jul 27) — no job actually failed. Since the head commit is unchanged (9e35a29), could you re-trigger with |
hi there is something wrong in auto CI test could you re-trigger with /ok to test or re-run the cancelled workflow when you get a chance? Thanks! |
|
/ok to test 6a1b608 |
|
Heads-up @yaox12 @hxbai: this PR was not closed by me on purpose. My fork left the upstream fork network on Aug 5 (repo turned private for unrelated internal use), and GitHub auto-closes open PRs from a detached fork, attributing the close to the repo owner — hence the misleading "htesd closed this". Re-opened as #6338 from a clean fork, rebased onto current |
4 participants
What does this PR do?
HyperConnectionTransformerLayer._forward_post_mlp_with_fused_hyper_connectionneverreleases
self.mlp_norm_managerafter callinggroup_offload(), unlike the base classTransformerLayer._forward_post_mlp, which setsself.mlp_norm_manager = Nonerightafter use.
Because
FineGrainedActivationOffloadingInterface.__init__stores the wrapped tensorunconditionally (
self.tensor = tensor, even whenoffload=False), the manager leftattached to the layer module keeps a reference to one microbatch's pre-MLP-norm input
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 tensoris 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)-sizedbuffers several times over.
Measured impact
2-layer GPT config with hyper connections (
num_residual_streams=4), bf16, mbs=1, mockdata,
recompute_granularity='full', single GPU (H200), measuringtorch.cuda.memory_allocated()in the idle window between iterations: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
FineGrainedActivationOffloadingInterfacenot store the tensor whenoffload=False.