fix: CUDA memory leak / release BNB dequantization buffers & stale state in OffloadActivations

[butterwecksolutions]

What does this PR do?

The OffloadActivations context manager leaks VRAM through two independent paths.
Both are fixed by cleaning up stale state in __enter__, where the previous
backward is guaranteed to have completed.

Path 1 — Stale state leaks across steps (MoE + compile)

On MoE architectures with sample packing + torch.compile, dynamic expert routing
leaves saved tensors on subgraphs that never contribute to loss. Their backward
nodes never execute, so tracker, storage_to_tensor_id, and the stashes retain
entries from step to step. ~60 tensors leak per micro-step → OOM by step 2
(e.g. Gemma-4 26B-A4B + ScatterMoE, axolotl-ai-cloud/axolotl#3638).

Path 2 — BNB dequantization buffers (QLoRA)

After #5700 fixed the CUDA stream leak, QLoRA (BNB 4-bit) still leaks because
tracker retains references to tensors sharing allocator blocks with BNB
dequant buffers, and empty_cache() is never called between steps.
~0.6 GiB/step → OOM after 30-40 steps on 24 GB GPUs.

Why __enter__ is the right place

The class example shows backward() running after the with block exits:

with OffloadActivations():
outputs = model(inputs, labels=labels)
loss = outputs.loss
loss.backward()

This means exit may fire BEFORE backward — the tracker still holds
active tensors. Cleaning up there would destroy data the backward needs.

enter, however, is called before the NEXT forward pass. By that point
the previous backward has already completed, so any remaining
state is leaked and safe to drop.

What the fix does

enter now clears:

• tracker, storage_to_tensor_id — release stale tensor references
• tensor_id, is_first_forward_call, is_first_backward_call — state reset
• Stream stashes (when use_streams=True)
• Calls accelerator-aware empty_cache() when bitsandbytes is loaded
  (via sys.modules, to avoid penalizing non-BNB workloads)

exit is unchanged from #5700 — stream sync + super().exit().

Pattern	Trigger	Symptom	Covered by
MoE + packed + compile	axolotl#3638	OOM step 2	__enter__
QLoRA BNB buffer	QLoRA training	OOM step 30-40	__enter__
CUDA stream garbage	#5700	OOM step 19	__exit__ / #5700 (merged)

A complete audit confirms no further leak sources.

Verification

Validated via monkey-patch on QLoRA 9B VL training. Full reproducer available,
will be updated after review feedback.

Capability	#5700 (merged)	#5738	#5730
CUDA stream sync + stash clear	✅	—	—
tracker + storage_to_tensor_id clear	—	✅	✅
stash clear (use_streams guard)	—	✅	✅
is_first_forward/backward reset	—	✅	✅
tensor_id reset	—	—	✅
BNB dequant buffer empty_cache	—	—	✅
CUDA stream garbage leak	✅	—	—
MoE + compile leak path	—	✅	✅
QLoRA buffer leak path	—	—	✅

---

• 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 OffloadActivations lifecycle/memory-management logic; incorrect clearing or cache flushing could impact training correctness or performance, though changes are scoped to context entry.

Overview
Fixes activation-offloading VRAM leaks by resetting OffloadActivations state at context entry: clears tracker/dedup maps, resets forward/backward bookkeeping (including tensor_id), and drops stream stashes when enabled.

When bitsandbytes is loaded, __enter__ now calls the appropriate accelerator empty_cache() (cuda/xpu/npu) to release BNB dequantization buffers between steps. Adds a regression test ensuring tracker size does not grow across repeated forward/backward steps with an unused graph branch.

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

[cursor]

Cursor Bugbot has reviewed your changes and found 2 potential issues.

There are 3 total unresolved issues (including 1 from previous review).

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

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

[butterwecksolutions]

@kashif Can you check this again?

[kashif]

added a test that highlights the issue

[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.

[butterwecksolutions]

@kashif @qgallouedec Thanks again for merging. My new experimental VRAM saving tool fits to this topic: https://github.com/butterwecksolutions/vsqz . Feel free to have a look.