[Bug] ShadowPEFT hook-lifecycle leak: per-forward GPU growth + UAF on unload_shadow(copy=False)→delete_adapter

[Sravanjangam]

Summary

ShadowPEFT (#3354, peft@9c16ee66) leaks per-forward state and can hit a use-after-free when the common unload_shadow(copy=False) → delete_adapter workflow is used. Two manifests share one root cause — boundary hooks hold self via bound methods and per-forward tensors are stored on self instead of as locals.

Thanks for ShadowPEFT — the design is strong, these are small lifecycle gaps. Filing as a bug for tracking; happy to open a PR after a nod.

---

Affected code

• src/peft/tuners/shadow/model.py:533-545 — _register_boundary_hooks registers self._seed_shadow_pre_hook / self._wrap_entry_pre_hook / self._unwrap_exit_hook / self._pack_shadow_cache_hook as bound methods (closures over self).
• src/peft/tuners/shadow/model.py:563-574 — _seed_shadow_pre_hook sets self._seed_shadow_state / self._shadow_past_out / self._should_pack_shadow_cache on the tuner.
• src/peft/tuners/shadow/model.py:616-627 — _pack_shadow_cache_hook only clears self._shadow_past_out when use_cache, not self._seed_shadow_state.
• src/peft/tuners/shadow/model.py:849-858 — _unload_and_optionally_merge clears handle.remove() but does not clear self._seed_shadow_state / self._shadow_past_out.
• src/peft/tuners/shadow/layers.py:276-297 — ShadowLayer.forward relies on ShadowCarrier set by the pre-hook; if the hook’s self is still alive, stale state rides the next forward.

---

Root cause

1. Per-forward tensors on self: _seed_shadow_state is the full shadow_backbone output [B, S, H] and _shadow_past_out is its KV cache. They are written in the top-of-model pre-hook and read in _wrap_entry_pre_hook, but never cleared at end of forward or on disable_adapter exit. They persist until the next forward, even when the shadow path is inactive (disable_adapter).

2. Hook holds self: register_forward_pre_hook(self._seed_shadow_pre_hook) captures self. handle.remove() deregisters the callback but the bound method still references the tuner. With unload_shadow(copy=False) (the default, documented as memory-saving) the returned DetachedShadowModel shares backbone/projection/head modules with the tuner. A subsequent delete_adapter deletes self.shadow_backbone[adapter] while the detached model still points at it — use-after-free (AttributeError or silent wrong logits).

---

Why it matters / who can trigger

• Every ShadowPEFT user who toggles disable_adapter — each with model.disable_adapter(): leaks one forward’s hidden states until next forward. On a 7B model with S=2048, H=4096, B=2 that’s ~50–200 MB per toggle that never frees until next forward (repro below).
• Documented workflow unload_shadow(copy=False) → delete_adapter — both steps are in the docstring (model.py:891 “Mutating one model then affects the other”) and in tests/test_shadow.py:619 (copy=False shares modules). Any user who detaches then cleans up hits the UAF. copy=True avoids it but costs a deep copy.
• No crash on single-adapter, no-toggle runs — so it’s easy to miss in CI (current test_shadow.py never toggles disable_adapter repeatedly or sequences unload_shadow + delete_adapter).

---

Minimal reproduction (CPU, no GPU needed)

# 1) Per-forward leak
from transformers import LlamaConfig, LlamaForCausalLM
from peft import get_peft_model, ShadowConfig
import torch, gc

config = LlamaConfig(vocab_size=32000, hidden_size=512, intermediate_size=1024,
                     num_hidden_layers=4, num_attention_heads=8, num_key_value_heads=8)
base = LlamaForCausalLM(config)
model = get_peft_model(base, ShadowConfig(task_type="CAUSAL_LM"))
ids = torch.randint(0, 32000, (2, 128))

model(ids)  # seeds self._seed_shadow_state
print(model.base_model._seed_shadow_state is not None)  # True — should be None after forward
with model.disable_adapter():
    model(ids)
print(model.base_model._seed_shadow_state is not None)  # Still True — leaked across toggle

# 2) UAF
detached = model.base_model.unload_shadow(copy=False)
print(detached.backbone is model.base_model.shadow_backbone["default"])  # True — shared
model.delete_adapter("default")
# detached(ids) now raises AttributeError or returns wrong logits (backbone deleted underneath)
try:
    detached(ids)
    print("no error — but logits are wrong (shared module deleted)")
except Exception as e:
    print("UAF:", e)

---

Proposed fix (minimal, low risk)

• In ShadowModel.forward and in _seed_shadow_pre_hook’s complementary post-hook, clear self._seed_shadow_state = self._shadow_past_out = None at end of forward (or store as locals / contextvar instead of self).
• In _unload_and_optionally_merge and delete_adapter, clear both attributes after handle.remove().
• Optionally use weakref / functools.partial for hook registration so the hook doesn’t keep the tuner alive, or document “call unload_shadow(copy=True) if you will delete_adapter afterwards.”
• No behavior change for single-forward, single-adapter runs; only affects lifecycle.

Est. diff: ~10–15 lines in src/peft/tuners/shadow/model.py, no API change.

---

Expected impact

• Fixes GPU growth on disable_adapter toggling (est. 50–200 MB/toggle on 7B, S=2048).
• Fixes UAF for the documented unload_shadow(copy=False) → delete_adapter sequence.
• Very low migration risk — only clears state that should already be transient.

---

Tests to add

• test_hook_leak_no_growth — 10× with model.disable_adapter(): model(ids) then assert model.base_model._seed_shadow_state is None and (on CUDA) torch.cuda.memory_allocated stable.
• test_unload_then_delete_isolated — detached = unload_shadow(copy=False); delete_adapter; assert detached(ids) still works or raises a clear ValueError (not AttributeError).

---

Environment

• peft@9c16ee66 (main HEAD 2026-08-31, 0 behind upstream/main), Python 3.12, torch 2.13 CPU, transformers 5.15, macOS arm64. No local modifications; read-only audit at peft@9c16ee66.
• Found via github-audit-harness 20-agent audit (Wave 2, ShadowPEFT surface tuners/shadow/* + tests/test_shadow.py), principal direct read + 4 surviving agent reports (16_shadow.md 20 findings).

Happy to open a PR with the 10-line fix + two tests after a nod — will include make style (ruff) and pytest tests/test_shadow.py -k "hook or unload or disable" pass counts.

---

cc @BenjaminBossan — please review when you have a moment, happy to hear your suggestions and ready for a PR with tests + make style as soon as you give a nod. Thanks again for the great work on PEFT!

[BenjaminBossan]

Thanks for reporting, please go ahead and create a PR with the fix.

• Optionally use weakref / functools.partial for hook registration so the hook doesn’t keep the tuner alive, or document “call unload_shadow(copy=True) if you will delete_adapter afterwards.”

Using partial seems like the best approach to me here.