Broken group offloading using block_level

[vladmandic]

Describe the bug

pr #12283 introduced group offloading on a pipeline
but its currently broken when using offload_type="block_level" as it does not handle vae module correctly
and result is failure with input vs weight device mismatch
using offload_type="leaf_level" it seems to work

Reproduction

import torch
import diffusers

model = '/mnt/models/stable-diffusion/mine/tempest-by-vlad-0.1.safetensors'
cache_dir = '/mnt/models/Diffusers'
pipe = diffusers.StableDiffusionXLPipeline.from_single_file(
    model,
    torch_dtype=torch.bfloat16,
    cache_dir=cache_dir,
)
pipe.enable_group_offload(
    onload_device=torch.device('cuda:0'),
    offload_device=torch.device('cpu'),
    offload_type="block_level",
    num_blocks_per_group=1,
    use_stream=True
)
image = pipe(
    prompt='A beautiful painting of a futuristic cityscape at sunset',
    width=1024,
    height=1024,
    num_inference_steps=10,
).images[0]
image.save('/tmp/output.png')

Logs

Traceback (most recent call last):
...
File: diffusers/models/autoencoders/autoencoder_kl.py", line 294, in _decode
> image = self.vae.decode(latents, return_dict=False)[0]
...
RuntimeError: Input type (CUDABFloat16Type) and weight type (CPUBFloat16Type) should be the same

System Info

python==3.12, diffusers==5e181eddfe7e44c1444a2511b0d8e21d177850a0

Who can help?

@sayakpaul @yiyixuxu @DN6 @asomoza @a-r-r-o-w

[sayakpaul]

Does it work when you apply it on the models individually?

[vladmandic]

i haven't had luck running in either way...
also, when using any type of quantization, it gets even worse since vae performs additional operation just before main decode which on quant weights which are also not moved back to gpu.

[sayakpaul]

Thanks for confirming (and also for updating the PR description). I can reproduce this. Till we fix, some comments below.

also, when using any type of quantization, it gets even worse since vae performs additional operation just before main decode which on quant weights which are also not moved back to gpu.

We have various offloading tests with quantization that should confirm offloading, working with quantization (at least with leaf):

diffusers/tests/quantization/bnb/test_4bit.py

Line 876 in f5c113e

class Bnb4BitCompileTests(QuantCompileTests, unittest.TestCase):

diffusers/tests/quantization/test_torch_compile_utils.py

Line 102 in f5c113e

def test_torch_compile_with_group_offload_leaf(self, use_stream=False):

So, maybe you could try that? The thing with many of these quantization methods is that the tensors need to be present on the GPU when performing matmuls. So, you might have noticed an indicative error when using num_blocks_per_group=1 with block_level group offloading.

Also, if just the VAE is the problem, we could consider excluding it from the offloading sequence by using:

pipe.enable_group_offload(
    onload_device=torch.device('cuda:0'),
    offload_device=torch.device('cpu'),
    offload_type="block_level",
    num_blocks_per_group=3,
    use_stream=True,
    exclude_modules=["vae"],
)

In general, did you notice any speed degradations when using leaf?

[vladmandic]

In general, did you notice any speed degradations when using leaf?

speed degradation compared to what?
group offloading with leaf is faster than sequential, but still quite a bit slower than pretty much anything else.
i did try different settings, streams, etc.)

[sayakpaul]

Than using block level offloading.

Do you have a setting where both block and leaf level group offloading work? That would be amazing

[vladmandic]

• none: its=5.36
• balanced: its=4.64
• model: its=3.70
• sequential: its=0.64
• group-leaf: its=0.61
• group-block: its=0.74

settings: exclude_modules=vae, num_blocks_per_group=1, use_stream=False

sequental, group-leaf, group-block are all very slow and within margin of error of each other

btw, balanced offloading is what we developed for sdnext, what makes it different is that it

• offloads based on user defined thresholds
  you can say start offloading if vram use is above 80% and stop offloading when it drops below 20%
• offloads largest modules first
• is aware of module pairs
  for example, unet and controlnet should be on the same device to avoid going on/off device for each step

[sayakpaul]

Not to digress but some further comments on speed:

sequental, group-leaf, group-block are all very slow and within margin of error of each other

Strange. Although not entirely similar, I used group offloading on just the denoiser and text encoder 2 of Flux, and group offloading performed reasonably well. Much better than sequential offloading and, of course, still worse than model-level offloading as one would expect: