huggingface / huggingface/peft
Forward hooks leak when the `_enable_peft_forward_hooks` block raises, poisoning the model
- Dominant language
- Python
- Stars
- 21.7k
- Forks
- 2.5k
- Avg merge
- 4d 12h
- Merged PRs (30d)
- 59
Description
### System Info
`peft` 0.20.0 (installed from `main`), `torch` 2.13.0, Python 3.12, CPU.
### Who can help?
_No response_
### Reproduction
`_enable_peft_forward_hooks` registers forward pre-hooks that inject `adapter_names`, yields, and then removes them — but the removal is not protected by `try`/`finally`:
```python
yield
for handle in hook_handles:
handle.remove()
```
If anything inside the `with` body raises, the removal loop never runs and the hooks stay registered on the model permanently.
```python
import torch
from torch import nn
from peft import LoraConfig, get_peft_model
class Tiny(nn.Module):
def __init__(self):
super().__init__()
self.lin = nn.Linear(4, 4)
def forward(self, x):
return self.lin(x)
def count_hooks(model):
return sum(len(m._forward_pre_hooks) for m in model.modules())
model = get_peft_model(Tiny(), LoraConfig(target_modules=["lin"], init_lora_weights=False))
model.add_adapter("other", LoraConfig(target_modules=["lin"], init_lora_weights=False))
model.eval()
# Clean exit removes the hooks, as expected.
with model.base_model._enable_peft_forward_hooks(adapter_names=["default"]):
pass
print("after clean exit:", count_hooks(model)) # 0
# An exception inside the block leaks them.
try:
with model.base_model._enable_peft_forward_hooks(adapter_names=["default"]):
raise RuntimeError("simulated OOM during generate()")
except RuntimeError:
pass
print("after exception: ", count_hooks(model)) # 1
# The model is now broken for ordinary use:
model(torch.randn(2, 4))
```
Output:
```
after clean exit: 0
after exception: 1
ValueError: Length of `adapter_names` should be the same as the number of inputs, but got 1 and 2 respectively.
```
### Expected behavior
The hooks should be removed on every exit path, so a failed call leaves the model in the state it was in beforehand.
The practical impact is larger than it looks. `adapter_names` is used for mixed-adapter batched inference, and the most common way for that call to raise is an OOM mid-`generate()` — routine when batching on a busy GPU. After one such failure the model object is silently poisoned: the stale hook keeps injecting the old `adapter_names` into every subsequent forward, so plain inference that passes no `adapter_names` at all fails with the confusing length mismatch above. Nothing points back at the earlier exception, and the only recovery is to rebuild the model.
The same unprotected pattern appears in three tuners:
- `LoraModel._enable_peft_forward_hooks`
- `RoadModel._enable_peft_forward_hooks`
- `GloraModel._enable_peft_forward_hooks`
and in `onload_layer` in `peft.tuners.tuners_utils`, where a failure during merging skips `post_forward`, leaving offloaded modules stranded on the execution device instead of being returned to CPU/disk.
I have a fix implemented and verified locally: wrap each `yield` in `try`/`finally` so cleanup always runs. Per the contributing guide I checked for overlapping issues and open PRs first and did not find any, and I am raising this before opening a PR rather than after.
Since the guide asks that an issue affecting multiple PEFT methods be fixed in one PR rather than split up, I would cover all four sites together, with a parametrized regression test over LoRA / RoAD / GLoRA added to `TestMixedAdapterBatches` in `tests/test_custom_models.py`. The test asserts the pre-hook count returns to its baseline after the block raises and that a plain forward still works; it fails on `main` for all three methods and passes with the fix.
Tests run: `pytest tests/test_custom_models.py -k "mixed_adapter_batches" --no-cov` → 35 passed, 1 failed. The one failure is `test_mixed_adapter_batches_lora_opt_timing`, which is timing-based and also fails on an unmodified checkout on this machine (2 failures there vs 1 with the change), so it looks unrelated to this bug.
Happy to open the PR if a maintainer is comfortable with the approach, or to adjust the scope if you would rather handle `onload_layer` separately.
Contributor guide
Research direction
Start with _enable_peft_forward_hooks in LoraModel, RoadModel, and GloraModel, then inspect onload_layer in peft.tuners.tuners_utils. Run tests/test_custom_models.py with -k "mixed_adapter_batches" --no-cov and review TestMixedAdapterBatches. Done means cleanup occurs after an exception, the hook count returns to baseline, and plain forward still works for all three methods.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning, testing-qa
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 66/100