huggingface / huggingface/diffusers

[Bug] low_cpu_mem_usage=True is not thread-safe and can leak the global nn.Module.register_parameter patch

Open
#14,347 1 comment 0 reactions 0 assignees View on GitHub
bug
Dominant language
Python
Stars
34.5k
Forks
7.3k
Avg merge
3d 3h
Merged PRs (30d)
91

Description

### Describe the bug

# [Bug] `low_cpu_mem_usage=True` is not thread-safe and can leak the global `nn.Module.register_parameter` patch

When `low_cpu_mem_usage=True`, adapter injection enters PEFT's
`init_empty_weights()` context for each target module
([`tuners_utils.py`](https://github.com/huggingface/peft/blob/9f1fe21d8131a24634d6d23c13efa2aae72b6cca/src/peft/tuners/tuners_utils.py#L915-L919)).
The context replaces `torch.nn.Module.register_parameter` process-wide and
restores the value captured on entry
([`integrations.py`](https://github.com/huggingface/peft/blob/9f1fe21d8131a24634d6d23c13efa2aae72b6cca/src/peft/utils/integrations.py#L209-L257)).

## Actual behavior

- An active `init_empty_weights()` context can lose its patch when another
thread exits.
- After both contexts exit, `nn.Module.register_parameter` can remain patched.
- Unrelated modules subsequently created in the process receive `meta`
parameters.
- The process remains corrupted until the function is repaired or the process
is restarted.

There is a second process-global value in the same implementation,
`_init_on_device._skip`, which may also allow cross-thread interference.

## Downstream impact in Diffusers

Diffusers enables the PEFT low-memory LoRA path by default for compatible
versions
([default selection](https://github.com/huggingface/diffusers/blob/7685bffe89041496c2c0ae07ea933df1c80d1f43/src/diffusers/loaders/lora_pipeline.py#L67-L75),
[`load_lora_weights()` use](https://github.com/huggingface/diffusers/blob/7685bffe89041496c2c0ae07ea933df1c80d1f43/src/diffusers/loaders/lora_pipeline.py#L205-L230)).
With PEFT 0.20.0 and Diffusers 0.39.0,
`_LOW_CPU_MEM_USAGE_DEFAULT_LORA` evaluates to `True`.

In an end-to-end SDXL test with two independent pipelines in one process:

- Sequential LoRA loads generated valid images.
- The first pair of concurrent `load_lora_weights()` calls caused both
pipelines' final latent tensors to become entirely NaN (`finite_fraction=0`).
- Both images were black, and later sequential requests remained black.
- Loading the adapters sequentially and then running inference concurrently did
not fail.
- Serializing only adapter load/set/unload operations did not fail.
- Passing `low_cpu_mem_usage=False` did not fail: a full-resolution control run
at 896x1152 and 28 steps completed 10 generations, including three concurrent
pairs, with 0 black images, 0 errors, and finite latents throughout.

The end-to-end test used Diffusers 0.32.2, PEFT 0.20.0, PyTorch 2.13.0+cu130,
and an NVIDIA H100 NVL. The deterministic CPU reproducer above confirms the
same root defect on the latest Diffusers 0.39.0 / PEFT 0.20.0 stack, and the
relevant implementation is unchanged on both current `main` branches.

## Expected behavior

Concurrent adapter injection on independent models should not corrupt global
PyTorch behavior. In particular:

- the patch must remain effective until its owning context exits;
- after all contexts exit, `nn.Module.register_parameter` must be restored to
its original value; and
- unrelated modules created afterward must use their requested device rather
than silently receiving `meta` parameters.

If concurrent low-memory adapter injection cannot be supported, the operation
should be internally serialized or explicitly rejected/documented rather than
silently corrupting the process.

## Workaround

For Diffusers, explicitly disabling the low-memory adapter initialization path
avoids this context manager:

```python
pipe.load_lora_weights(
lora_path,
adapter_name=adapter_name,
low_cpu_mem_usage=False,
)
```

### Reproduction

Install the latest PEFT release:

```shell
python -m pip install "peft==0.20.0"
```

The following model-free reproducer forces two contexts to overlap and exit in
the opposite order:

```python
import platform
import threading

import peft
import torch
from peft.utils.integrations import init_empty_weights

original = torch.nn.Module.register_parameter
a_entered = threading.Event()
b_entered = threading.Event()
a_exited = threading.Event()
observed = {}

def worker_a():
with init_empty_weights():
a_entered.set()
assert b_entered.wait(5)
a_exited.set()

def worker_b():
assert a_entered.wait(5)
with init_empty_weights():
b_entered.set()
assert a_exited.wait(5)
observed["active_context_lost_patch"] = (
torch.nn.Module.register_parameter is original
)

a = threading.Thread(target=worker_a)
b = threading.Thread(target=worker_b)
a.start()
b.start()
a.join()
b.join()

try:
observed["patch_leaked_after_both_contexts"] = (
torch.nn.Module.register_parameter is not original
)
observed["new_parameter_device"] = torch.nn.Linear(2, 2).weight.device.type

print(f"Python: {platform.python_version()}")
print(f"PyTorch: {torch.__version__}")
print(f"PEFT: {peft.__version__}")
for key, value in observed.items():
print(f"{key}: {value}")
finally:
# Repair the process so the reproducer exits cleanly.
torch.nn.Module.register_parameter = original

assert observed == {
"active_context_lost_patch": True,
"patch_leaked_after_both_contexts": True,
"new_parameter_device": "meta",
}
```

Observed output:

```text
Python: 3.12.3
PyTorch: 2.13.0+cpu
PEFT: 0.20.0
active_context_lost_patch: True
patch_leaked_after_both_contexts: True
new_parameter_device: meta
```

The failure follows directly from the interleaving:

1. Thread A captures the original function and installs wrapper A.
2. Thread B captures wrapper A and installs wrapper B.
3. Thread A exits first and restores the original function, even though B's
context is still active.
4. Thread B exits and restores wrapper A, leaving PEFT's meta-device wrapper
installed after every context has finished.

The case uses two
independent model/pipeline instances, but they collide through this process-wide
function replacement.

### Logs

```shell

```

### System Info

## System Info

- PEFT: `0.20.0` (latest PyPI release as of 2026-07-30)
- Diffusers: `0.39.0` (latest PyPI release as of 2026-07-30)
- Transformers: `5.14.1`
- Accelerate: `1.14.0`
- PyTorch: `2.13.0+cpu`
- Python: `3.12.3`
- Platform: Linux
- Minimal reproduction requires no GPU or model download.
- Current upstream source was also inspected:
- PEFT `main`: [`9f1fe21d`](https://github.com/huggingface/peft/commit/9f1fe21d8131a24634d6d23c13efa2aae72b6cca)
- Diffusers `main`: [`7685bffe`](https://github.com/huggingface/diffusers/commit/7685bffe89041496c2c0ae07ea933df1c80d1f43)

Both current `main` branches still contain the behavior described below.

### Who can help?

_No response_

Contributor guide

Open the contributing guide

Research direction

Start with the linked PEFT implementations in tuners_utils.py and integrations.py, then run the model-free threading reproducer against init_empty_weights(). Verify the context patch remains active until its owner exits, is restored after overlapping contexts finish, and does not make a later torch.nn.Linear use meta parameters; add focused regression coverage for those cases.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
38/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.