From 0ba5ce28a5c793b8108c32ea63802cc6c5369f6f Mon Sep 17 00:00:00 2001 From: lorenzozanee Date: Sat, 8 Aug 2026 04:07:48 +0800 Subject: [PATCH 1/2] Serialize low_cpu_mem_usage LoRA injection to prevent register_parameter patch leak PEFT's init_empty_weights() context monkey-patches torch.nn.Module.register_parameter process-wide and restores the value captured at entry on exit, which is not thread-safe. Concurrent adapter injection with low_cpu_mem_usage=True can interleave these capture/restore operations and leak the patch, leaving newly created modules on the meta device. Wrap the low-memory injection path in PeftAdapterMixin.load_lora_adapter and _load_lora_into_text_encoder in a process-wide threading.Lock, and add a regression test that asserts the global patch is not leaked under concurrent injection. Closes #14347 --- src/diffusers/loaders/lora_base.py | 24 ++- src/diffusers/loaders/peft.py | 19 ++- .../test_lora_low_cpu_mem_thread_safety.py | 140 ++++++++++++++++++ 3 files changed, 174 insertions(+), 9 deletions(-) create mode 100644 tests/lora/test_lora_low_cpu_mem_thread_safety.py diff --git a/src/diffusers/loaders/lora_base.py b/src/diffusers/loaders/lora_base.py index d4c88d35924f..eebb8e48b1e8 100644 --- a/src/diffusers/loaders/lora_base.py +++ b/src/diffusers/loaders/lora_base.py @@ -332,6 +332,9 @@ def _load_lora_into_text_encoder( ): from ..hooks.group_offloading import _maybe_remove_and_reapply_group_offloading + # Imported lazily to avoid a circular import (peft.py imports from this module). + from .peft import _LOW_CPU_MEM_USAGE_INJECTION_LOCK + if not USE_PEFT_BACKEND: raise ValueError("PEFT backend is required for this method.") @@ -398,12 +401,21 @@ def _load_lora_into_text_encoder( ) # inject LoRA layers and load the state dict # in transformers we automatically check whether the adapter name is already in use or not - text_encoder.load_adapter( - adapter_name=adapter_name, - adapter_state_dict=state_dict, - peft_config=lora_config, - **peft_kwargs, - ) + if low_cpu_mem_usage: + with _LOW_CPU_MEM_USAGE_INJECTION_LOCK: + text_encoder.load_adapter( + adapter_name=adapter_name, + adapter_state_dict=state_dict, + peft_config=lora_config, + **peft_kwargs, + ) + else: + text_encoder.load_adapter( + adapter_name=adapter_name, + adapter_state_dict=state_dict, + peft_config=lora_config, + **peft_kwargs, + ) # scale LoRA layers with `lora_scale` scale_lora_layers(text_encoder, weight=lora_scale) diff --git a/src/diffusers/loaders/peft.py b/src/diffusers/loaders/peft.py index daa078bc25d5..a7eeb8d4d6cd 100644 --- a/src/diffusers/loaders/peft.py +++ b/src/diffusers/loaders/peft.py @@ -15,6 +15,7 @@ import inspect import json import os +import threading from collections import defaultdict from functools import partial from pathlib import Path @@ -45,6 +46,12 @@ logger = logging.get_logger(__name__) +# PEFT's `init_empty_weights()` context monkey-patches `torch.nn.Module.register_parameter` +# process-wide on entry and restores the value captured at entry on exit, which is not +# thread-safe. Concurrent adapter injection with `low_cpu_mem_usage=True` can interleave these +# capture/restore operations and leak the patch, so the low-memory injection path is serialized. +_LOW_CPU_MEM_USAGE_INJECTION_LOCK = threading.Lock() + _SET_ADAPTER_SCALE_FN_MAPPING = defaultdict( lambda: (lambda model_cls, weights: weights), { @@ -324,9 +331,15 @@ def map_state_dict_for_hotswap(sd): # it to None incompatible_keys = None else: - inject_adapter_in_model( - lora_config, self, adapter_name=adapter_name, state_dict=state_dict, **peft_kwargs - ) + if low_cpu_mem_usage: + with _LOW_CPU_MEM_USAGE_INJECTION_LOCK: + inject_adapter_in_model( + lora_config, self, adapter_name=adapter_name, state_dict=state_dict, **peft_kwargs + ) + else: + inject_adapter_in_model( + lora_config, self, adapter_name=adapter_name, state_dict=state_dict, **peft_kwargs + ) incompatible_keys = set_peft_model_state_dict(self, state_dict, adapter_name, **peft_kwargs) if self._prepare_lora_hotswap_kwargs is not None: diff --git a/tests/lora/test_lora_low_cpu_mem_thread_safety.py b/tests/lora/test_lora_low_cpu_mem_thread_safety.py new file mode 100644 index 000000000000..a67c44c598f6 --- /dev/null +++ b/tests/lora/test_lora_low_cpu_mem_thread_safety.py @@ -0,0 +1,140 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import gc +import threading +import unittest + +import torch + +from diffusers import UNet2DConditionModel + +from ..testing_utils import require_peft_backend + + +def _build_lora_state_dict(unet, rank=4): + """Build a minimal LoRA state dict for the attention linear layers of the given UNet.""" + state_dict = {} + for name, module in unet.named_modules(): + if isinstance(module, torch.nn.Linear) and ( + name.endswith("to_q") or name.endswith("to_k") or name.endswith("to_v") or name.endswith("to_out.0") + ): + state_dict[f"{name}.lora_A.weight"] = torch.randn(rank, module.in_features) + state_dict[f"{name}.lora_B.weight"] = torch.randn(module.out_features, rank) + return state_dict + + +@require_peft_backend +class LowCpuMemUsageThreadSafetyTests(unittest.TestCase): + unet_kwargs = { + "block_out_channels": (32, 64), + "layers_per_block": 2, + "sample_size": 32, + "in_channels": 4, + "out_channels": 4, + "down_block_types": ("DownBlock2D", "CrossAttnDownBlock2D"), + "up_block_types": ("CrossAttnUpBlock2D", "UpBlock2D"), + "cross_attention_dim": 32, + } + + num_threads = 4 + num_rounds = 10 + + @classmethod + def setUpClass(cls): + super().setUpClass() + torch.manual_seed(0) + probe_unet = UNet2DConditionModel(**cls.unet_kwargs) + cls.lora_state_dict = _build_lora_state_dict(probe_unet) + del probe_unet + gc.collect() + + def _run_concurrent_injections(self, low_cpu_mem_usage): + original = torch.nn.Module.register_parameter + errors = [] + try: + barrier = threading.Barrier(self.num_threads) + + def worker(thread_id): + for _ in range(self.num_rounds): + try: + barrier.wait() + except threading.BrokenBarrierError: + return + try: + unet = UNet2DConditionModel(**self.unet_kwargs) + unet.load_lora_adapter( + self.lora_state_dict, + adapter_name=f"adapter-{thread_id}", + low_cpu_mem_usage=low_cpu_mem_usage, + prefix=None, + ) + except Exception as e: + errors.append((thread_id, type(e).__name__, str(e))) + return + + threads = [threading.Thread(target=worker, args=(i,)) for i in range(self.num_threads)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + return { + "register_parameter_patched": torch.nn.Module.register_parameter is not original, + "fresh_linear_device": torch.nn.Linear(2, 2).weight.device.type, + "thread_errors": errors, + } + finally: + torch.nn.Module.register_parameter = original + + def test_concurrent_low_cpu_mem_usage_injection_does_not_leak_register_parameter(self): + """Concurrent `low_cpu_mem_usage=True` LoRA injection must not leak the global patch.""" + observed = self._run_concurrent_injections(low_cpu_mem_usage=True) + self.assertFalse( + observed["register_parameter_patched"], + f"Concurrent `low_cpu_mem_usage=True` injection leaked the global " + f"`torch.nn.Module.register_parameter` monkey patch. Observed: {observed}", + ) + self.assertNotEqual( + observed["fresh_linear_device"], + "meta", + f"Concurrent `low_cpu_mem_usage=True` injection left newly created modules on the meta " + f"device. Observed: {observed}", + ) + self.assertEqual( + observed["thread_errors"], + [], + f"Worker threads raised during concurrent `low_cpu_mem_usage=True` injection: {observed['thread_errors']}", + ) + + def test_concurrent_injection_without_low_cpu_mem_usage_is_safe(self): + """Sanity control: concurrent `low_cpu_mem_usage=False` injection must be safe.""" + observed = self._run_concurrent_injections(low_cpu_mem_usage=False) + self.assertFalse( + observed["register_parameter_patched"], + f"Concurrent `low_cpu_mem_usage=False` injection leaked the global " + f"`torch.nn.Module.register_parameter` monkey patch. Observed: {observed}", + ) + self.assertNotEqual( + observed["fresh_linear_device"], + "meta", + f"Concurrent `low_cpu_mem_usage=False` injection left newly created modules on the meta " + f"device. Observed: {observed}", + ) + self.assertEqual( + observed["thread_errors"], + [], + f"Worker threads raised during concurrent `low_cpu_mem_usage=False` injection: " + f"{observed['thread_errors']}", + ) From 8e8dd378660175e51916469c891932ab1e1806cb Mon Sep 17 00:00:00 2001 From: lorenzozanee Date: Sat, 8 Aug 2026 16:49:39 +0800 Subject: [PATCH 2/2] [docs] Document thread-safe low_cpu_mem_usage LoRA loading --- src/diffusers/loaders/peft.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/diffusers/loaders/peft.py b/src/diffusers/loaders/peft.py index a7eeb8d4d6cd..3caf4ac1e4af 100644 --- a/src/diffusers/loaders/peft.py +++ b/src/diffusers/loaders/peft.py @@ -129,7 +129,8 @@ def load_lora_adapter( link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning). low_cpu_mem_usage (`bool`, *optional*): Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. + weights. When enabled, the low-memory loading path is serialized with a process-wide lock, making it + safe to call concurrently across threads. hotswap : (`bool`, *optional*) Defaults to `False`. Whether to substitute an existing (LoRA) adapter with the newly loaded adapter in-place. This means that, instead of loading an additional adapter, this will take the existing