inference with lora suddenly significantly slower

[a-r-r-o-w]

inference with lora suddenly significantly slower than with diffusers==0.33.1 and peft==0.14.0
Used to be 3.9it/s now 2.9it/s

Originally posted by @burgalon in #11999 (comment)

[burgalon]

actually doesn't seem related to #11999 still trying to hunt the offending commit.
This happens when upgrading from 0.34.0 to dev

[burgalon]

@a-r-r-o-w seems like the offending PR is #11180
cbc8ced is fast
0454fbb is slow

[sayakpaul]

I wonder why it would make LoRA slower :/ FWIU, it doesn't have to do anything with PEFT or the LoRA related updates in diffusers?

Also, if you could provide a short snippet on how you're benchmarking this, it would be quite helpful.

[burgalon]

sure

import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(
 '/black-forest-labs/FLUX.1-dev',
 torch_dtype=torch.bfloat16,
 local_files_only=True,
).to('cuda')

pipe.load_lora_weights('lora.safetensors', low_cpu_mem_usage=False)
pipe('test')

[sayakpaul]

Could we please keep this fully reproducible so that we can just essentially copy and paste the snippet for testing?

[burgalon]

does that work?

import torch
from diffusers import DiffusionPipeline
from huggingface_hub import hf_hub_download

pipe = DiffusionPipeline.from_pretrained(
 'black-forest-labs/FLUX.1-dev',
 torch_dtype=torch.bfloat16,
 local_files_only=True,
).to('cuda')
hf_hub_download(repo_id='burgalon/emma', filename='2910204.safetensors', local_dir='.')
pipe.load_lora_weights('./2910204.safetensors', low_cpu_mem_usage=False)
pipe('test')
# commit 0454fbb30bfbe21aa4ea29c827c396bac57dc518 SLOW 2.8it/s 9s
# commit cbc8ced20f7a67254caac36af79e666ac2379833 FAST 3.7it/s 7s

[burgalon]

@sayakpaul @a-r-r-o-w i've narrowed down the issue to the changes in src/diffusers/models/transformers/transformer_flux.py
this revert resolves the issue burgalon@a401adb

[burgalon]

Commented https://github.com/huggingface/diffusers/pull/11180/files#r2273978819

Copilot :

This diff changes the way tensors are processed and returned in the forward pass of the transformer_flux model in the HuggingFace Diffusers library. The main performance issue comes from the repeated use of torch.cat and splitting tensors, which can introduce significant overhead, especially on large tensors or in high-throughput scenarios.

Key Reasons for Slowness

1. Frequent Concatenation and Splitting
   Before: hidden_states was used as a single tensor. Forward methods returned just hidden_states.
   After:
   At the start of each block, you concatenate encoder_hidden_states and hidden_states:
   Python
   text_seq_len = encoder_hidden_states.shape[1]
   hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
   At the end of each block, you split them back:
   Python
   encoder_hidden_states, hidden_states = hidden_states[:, :text_seq_len], hidden_states[:, text_seq_len:]
   Why it’s slow:
   torch.cat allocates new memory and copies both tensors every time it’s called.
   Slicing (hidden_states[:, :text_seq_len]) also allocates new memory for the split tensors.
   If this is done in every block (as in your loop), the cost compounds quickly.
2. Increased Memory Footprint
   Every concat/split operation increases peak memory usage, which can slow down training and inference (especially on GPU with limited bandwidth and memory).
3. Potential for Duplicate Computation
   If encoder_hidden_states doesn’t change between blocks, you’re repeatedly concatenating the same tensor, which is wasteful.
   Why Was This Change Made?
   This pattern of concatenating then splitting often occurs in models that want cross-attention between different sources (e.g., text and image). It can make implementation simpler, but at the cost of speed and memory.

How to Fix or Mitigate

Avoid repeated concat/split:
Concatenate only once per forward pass if possible, not inside every block.
Pass both tensors separately and only combine when strictly necessary for attention.
Profile with torch.autograd.profiler or nvprof:
Confirm that torch.cat and slicing are taking significant time.
Summary:
This diff causes slowness because every block in your transformer repeatedly concatenates and splits large tensors, leading to high memory usage and computational overhead. Optimize by minimizing these operations and only combining tensors when essential.

[a-r-r-o-w]

Thanks for investigating and finding the faulty commit @burgalon! When I tested the change during First Block Cache addition, the speed did not seem to be impacted. However, since you're noticing a difference, I will have to test again and we might have to revert the change. I suspect supporting first block cache easily with all models might require a significant rewrite to how it's implemented

[burgalon]

@a-r-r-o-w please note that it's not the FirstBlockCache but only this part
37d88af that's restoring performance

[a-r-r-o-w]

Yes, I understand. That change was made so that FBC becomes much simpler to implement. Without that change, FBC implementation becomes a very big mess of if-else statements