fix: prevent 5 GB+ CUDA memory leak in activation offloading by syncing streams and clear stashes in OffloadActivations.__exit__

[butterwecksolutions]

Problem

OffloadActivations.__exit__ defers entirely to
torch.autograd.graph.saved_tensors_hooks.__exit__, which does not
synchronize the offload CUDA stream (s1) or the compute stream (s0).
Async copies still in-flight when the context manager exits can leave
stashed tensors referencing a stream whose lifetime may already be
over, leaking CUDA blocks with garbage stream IDs.

Observed via torch.cuda.memory._dump_snapshot during QLoRA vision
training on a 9B model (RTX 3090 24 GB):

Stream 93898984661968 (garbage ID): 5.70 GB ← likely orphaned offload tensors
Stream 0 (compute): 14.69 GB ← normal model
Fragmented/unreleased: 5.83 GB

The garbage stream blocks (~8.4 MB each) match the size of activation
offload buffers. After adding stream sync + stash clear, garbage-stream
blocks no longer appear in post-training snapshots.

Fix

Add an __exit__ method that runs BEFORE the parent's cleanup:

1. Synchronize s0 (compute stream) and s1 (offload stream)
2. Clear stashed tensors (bwd_tensor_stash, bwd_ev_stash, fwd_stash)
   — only when use_streams=True (stashes don't exist otherwise)

Tracker is intentionally NOT cleared — the backward pass unpacks
tensors via tracker AFTER __exit__ (see class docstring example).

Zero performance impact — called once per context manager, sync+clear
takes ~0.1ms.

Verification

Applied as a monkey-patch since May 2026. Post-training snapshots show
zero garbage-stream blocks. Without the fix, garbage-stream blocks
accumulate monotonically per step at ~0.2 GiB/step.

---

• This PR fixes a typo or improves the docs

• Read the contributor guideline

• Discussed via GitHub issue (N/A — found during QLoRA debugging)

• Documentation update needed? No — internal behavior change only

• New tests? Not easily testable — requires CUDA memory snapshot assertions

• AI-assisted: PR text and structure refined with AI; root cause found
  through human debugging

---

Note

Medium Risk
Touches CUDA stream/autograd hook lifecycle in OffloadActivations, so mistakes could cause hangs or perf regressions, but the change is small and gated to use_streams=True.

Overview
Prevents activation-offloading CUDA memory from being leaked after leaving the OffloadActivations context by adding a custom __exit__.

When use_streams=True, __exit__ now synchronizes the compute/offload streams (s0, s1) and clears the forward/backward stash dictionaries before delegating to saved_tensors_hooks.__exit__, using a try/finally to ensure the parent hook cleanup always runs.

^{Reviewed by Cursor Bugbot for commit f1377a4. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit e259a88. Configure here.}

[kashif]

Suggested change

	def __exit__(self, *args, **kwargs):
	"""Sync streams and clear stashes before parent cleanup.
	Prevents leaked CUDA blocks with garbage stream IDs when the
	context manager exits with in-flight async copies or stashed
	tensors still referencing the offload stream.
	NOTE: tracker is NOT cleared — the backward pass unpacks tensors
	via tracker AFTER __exit__ (see class docstring example).
	"""
	import torch
	# 1. Sync both streams so async copies finish before cleanup
	if self.use_streams and self.s0 is not None:
	self.s0.synchronize()
	if self.use_streams and self.s1 is not None:
	self.s1.synchronize()
	# 2. Clear stashed tensors (only exist when use_streams=True)
	if self.use_streams:
	self.bwd_tensor_stash.clear()
	self.bwd_ev_stash.clear()
	self.fwd_stash.clear()
	return super().__exit__(*args, **kwargs)
	def __exit__(self, *args, **kwargs):
	"""Sync streams and clear stashes before parent cleanup."""
	try:
	if self.use_streams:
	self.s0.synchronize()
	self.s1.synchronize()
	self.bwd_tensor_stash.clear()
	self.bwd_ev_stash.clear()
	self.fwd_stash.clear()
	finally:
	result = super().__exit__(*args, **kwargs)
	return result

[kashif]

Lets make the parent __exit__ run in a finally so the saved tensor hooks are always popped. I’d also drop the is not None checks unless there is a real code path where self.s0 or self.s1 can be None. In this class, self.s0 is always initialized, and self.s1 exists whenever use_streams=True.

[kashif]

thanks @butterwecksolutions if you can try my suggestion then we should be good!

[butterwecksolutions]

Fix Verified — 7.0 GB Saved via Systematic Isolation

Updated to @kashif's suggested pattern (try/finally, no is-not-None guards,
no tracker.clear).

Then verified against a 14-test reproducer (7 patches × 2 modes, fresh
GPU isolation per test, Qwopus3.5-9B, RTX 3090, PyTorch 2.10, TRL 0.29.0):

Mode	Baseline	P3 (#5700)	Saved
VL (seq_len=512)	20.3 GB	13.3 GB	7.0 GB
TEXT (seq_len=4096)	8.5 GB	8.5 GB	— (leak VL-triggered)

Sole root cause among 5 candidate fixes:

• P1 (bitsandbytes-foundation/bitsandbytes 1935) = 0.0 GB → closed
• P2 (transformers #45769) = 0.0 GB → closed
• P5 (naive backward hook) = OOM at 24.2 GB

Thanks @kashif for the review and the suggestion — the try/finally is the
right call. Thanks @qgallouedec for merging #5694 and maintaining TRL.
Happy I could contribute again!

Full reproducer, rendered HTML report, raw training logs:

• https://butterwecksolutions.github.io/Testrepo/vram-fix-reproducer/vram_test_results.html
• https://github.com/butterwecksolutions/Testrepo/tree/main/vram-fix-reproducer

Viele Grüße aus Arnsberg! 🇩🇪

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.