Repository navigation
Loading multiple LoRAs to 1 pipeline in parallel, 1 LoRA to 2-pipelines on 2-GPUs #11914
Description
Hi everyone,
I have the following scenario.
I have a machine with 2-GPUs and a running service that keep has two pipelines loaded to their corresponding devices. Also I have a list of LoRAs (say 10). On each request I split the batch into 2 parts (request also has the corresponding information about LoRA), load LoRAs and run the forward pass.
The problem I encounter is that whatever parallelization method I have tried (threading, multi-processing), the maximum I have achieved is pre-loading LoRAs on the cpu and then, moving them to GPU and only after that load_lora_weights from the state_dict.
Even if I attempt to achieve parallelization in by calling the chunk where I load in parallel in threads, the pipe starts to produce either a complete noise or a black image.
Where I would appreciate a lot the help is:
- To get an advice of elegantly loading multiple LoRAs at once into one pipe (all examples in the documentation indicate that one needs to do it 1 by 1)
- If I have 2 pipes on 2 different devices, how to parallelize the process of loading 1 LoRA to pipes on their corresponding devices.
def apply_multiple_loras_from_cache(pipes, adapter_names, lora_cache, lora_names, lora_strengths, devices):
for device_index, pipe in enumerate(pipes):
logger.info(f"Starting setup for device {devices[device_index]}")
# Step 1: Unload LoRAs
start = time.time()
pipe.unload_lora_weights(reset_to_overwritten_params=False)
logger.info(f"[Device {device_index}] Unload time: {time.time() - start:.3f}s")
# Step 2: Parallelize CPU → GPU state_dict move
def move_to_device(name):
return name, {
k: v.to(devices[device_index], non_blocking=True).to(pipe.dtype)
for k, v in lora_cache[name]['state_dict'].items()
}
start = time.time()
with ThreadPoolExecutor() as executor:
future_to_name = {executor.submit(move_to_device, name): name for name in adapter_names}
results = [future.result() for future in as_completed(future_to_name)]
logger.info(f"[Device {device_index}] State dict move + dtype conversion time: {time.time() - start:.3f}s")
# Step 3: Load adapters
start = time.time()
for adapter_name, state_dict in results:
pipe.load_lora_weights(
pretrained_model_name_or_path_or_dict=state_dict,
adapter_name=adapter_name
)
logger.info(f"[Device {device_index}] Load adapter weights time: {time.time() - start:.3f}s")
# Step 4: Set adapter weights
start = time.time()
pipe.set_adapters(lora_names, adapter_weights=lora_strengths)
logger.info(f"[Device {device_index}] Set adapter weights time: {time.time() - start:.3f}s")
torch.cuda.empty_cache()
logger.info("All LoRAs applied and GPU cache cleared.")
Activity
asomoza commentedon Jul 12, 2025
Hi, do you have a specific reason to maintain the loras in the CPU? are you VRAM constrained?
Without knowing more context, the best answer here is to just load all the loras in both pipelines and set the scale to 0, then use set_adapters to activate them on inference when needed.
But I suspect that you're getting the black images or noise because there's some error in your code, it should work, it's not clear to me how are you splitting the batches or if you're making sure the pipelines aren't sharing anything between them.
vahe-toffee commentedon Jul 13, 2025
that is correct, I am VRAM constrained.
The example above works, by the way - not sure if I wass clear about in the my original message.
As soon as I change it to something like:
and call this function in another threadPool (e.g. to parallelize this over the pipelines) - the black images appear
asomoza commentedon Jul 14, 2025
My mistake, usually the issues are posted with the code that doesn't work and not with the one that actually works so that's why I though that code had the issue.
I won't have time to test your code but if your code works with those pipelines without loading loras, meaning that the lora loading is what breaks it, the best choice (without looking at your whole codebase) is to
fuseandunfusethe loras, I think this is the only way it also works with a distributed environment, like separating the transformer model in multiple GPUs.pinging @sayakpaul and @a-r-r-o-w in case they have more insights because I still haven't tested running multiple pipelines in multiple GPUs at the same time yet.
vahe-toffee commentedon Jul 15, 2025
The lora loading is not the problem per-se. When I load LoRAs to each pipeline sequentially (first code snippet) - everything works.
When I try to load them in parallel to two pipelines - e.g. invoking second code snippet in a threadpool - everything breaks.
a-r-r-o-w commentedon Jul 15, 2025
I have only looked through the code snippets roughly and don't have a full context yet. I'm only responding because some things look fundamentally wrong to me. I'm occupied with some other things at the moment, but will try to take a better look over the weekend to understand what is expected.
Regarding the black images, not quite sure what's happening. You will have to look at the profiles to identify the problem.
The right programming model for multi-GPU setup with pytorch is to make use of
torchrun(or similar distributed environment handler). What you're effectively doing here is data parallelism. The request handler code will be same on each GPU, but you can use some kind of schedule to tell which GPU continues execution with which request.If you'd like to see speedup in device transfer times, you could look into pinning the state dict tensors on CPU during the first load from disk, and combining that with non-blocking. Note that pinning assumes you have enough CPU RAM to handle all the loras you want to keep in memory.