CPU Offloading Inefficiency: UNet/Transformer Repeatedly Moved Between CPU/GPU During Loops

[JoeGaffney]

Describe the bug

CPU Offloading Inefficiency: UNet/Transformer Repeatedly Moved Between CPU/GPU During Denoising Loop

Issue

The current CPU offloading mechanism (enable_model_cpu_offload()) in Diffusers is inefficient for models with denoising loops. It repeatedly moves the transformer/UNet component between CPU and GPU for each step, creating unnecessary data transfer overhead and potentially slowing down inference or causing uneeded back and forth transfer.

Observed Behavior

After enabling CPU offloading with pipe.enable_model_cpu_offload(), the transformer/UNet is:

1. Loaded to GPU for a single denoising step
2. Immediately moved back to CPU after use
3. Loaded to GPU again for the next step
4. Repeat for every step in the denoising process

This pattern appears in both image models (SD3, FLUX) and video models (WAN, LTX), but the performance impact is more severe for video models due to their larger transformer sizes and 3D tensors.

Log Evidence

✅ Moved FluxTransformer2DModel to cuda:0
🔄 Moved FluxTransformer2DModel back to CPU
✅ Moved FluxTransformer2DModel to cuda:0
🔄 Moved FluxTransformer2DModel back to CPU
[...repeats for each denoising step...]

Performance Impact

• Each GPU↔CPU transfer adds significant latency (milliseconds to seconds depending on model size)
• Maybe data is still in the current cuda cache so some performance is mitigated?

Proposed Solutions

1. Add an exclude_from_offload parameter to enable_model_cpu_offload():

pipe.enable_model_cpu_offload(exclude_from_offload=["transformer"])

2. Make the offloading mechanism smarter about components used in loops or used sequentially

3. Add a keep_on_device_during_generation parameter to control this behavior

The issue is evident in the core loop of diffusion pipelines (e.g., in WanImageToVideoPipeline.call) where the transformer is called repeatedly for each timestep but gets moved back to CPU after each use.

Reproduction

Monkey patch accelerate.hooks and you still see this pattern in most model pipelines - or at least pretty common ones I've tried.

import accelerate.hooks

# Save the original methods
original_pre_forward = accelerate.hooks.CpuOffload.pre_forward
original_post_forward = accelerate.hooks.CpuOffload.post_forward


# Create logging wrappers
def logged_pre_forward(self, module, *args, **kwargs):
    module_name = module.__class__.__name__
    result = original_pre_forward(self, module, *args, **kwargs)
    print(f"✅ Moved {module_name} to {self.execution_device}")
    return result


def logged_post_forward(self, module, output):
    module_name = module.__class__.__name__
    result = original_post_forward(self, module, output)
    print(f"🔄 Moved {module_name} back to CPU")
    return result


# Apply the monkey patch
accelerate.hooks.CpuOffload.pre_forward = logged_pre_forward
accelerate.hooks.CpuOffload.post_forward = logged_post_forward

Apply on any diffusers pipe that supports offloading.

pipe.enable_model_cpu_offload()

Logs

-------------------------------------------------------------------------------------------------- live log call --------------------------------------------------------------------------------------------------INFO     common.logger:utils.py:85 Image loaded from Base64 bytes, size: (1897, 788)
INFO     common.logger:utils.py:74 Image Resized from: (1897, 788) to 1897x788
WARNING  common.logger:pipeline_helpers.py:100 Evicting all 2 models from cache
INFO     common.logger:pipeline_helpers.py:228 Loading quantized model from /WORKSPACE/quantized/black-forest-labs/FLUX.1-dev/8bit/transformer
Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:01<00:00,  1.48it/s]INFO     common.logger:utils.py:151 Calling get_quantized_model with args: (), kwargs: {'model_id': 'black-forest-labs/FLUX.1-dev', 'subfolder': 'transformer', 'model_class': <class 'diffusers.models.transformers.transformer_flux.FluxTransformer2DModel'>, 'target_precision': 8, 'torch_dtype': torch.bfloat16}, took: 1.44s
INFO     common.logger:pipeline_helpers.py:228 Loading quantized model from /WORKSPACE/quantized/black-forest-labs/FLUX.1-schnell/8bit/text_encoder_2
Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:00<00:00, 12.12it/s]Loading pipeline components...: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 12.03it/s]WARNING  common.logger:pipeline_helpers.py:77 Cache miss for ('get_pipeline', PipelineConfig(model_id='black-forest-labs/FLUX.1-dev', model_family='flux', ip_adapter_models=(), ip_adapter_subfolders=(), ip_adapter_weights=(), ip_adapter_image_encoder_model='', ip_adapter_image_encoder_subfolder='')) - took: 2.93s - Cache size: 1/2
INFO     common.logger:flux.py:128 Image to image call {'width': 1897, 'height': 788, 'prompt': 'Change to night time and add rain and lighting', 'negative_prompt': 'worst quality, inconsistent motion, blurry, jittery, distorted', 'image': <PIL.Image.Image image mode=RGB size=1897x788 at 0x7FB9D0F31650>, 'num_inference_steps': 10, 'generator': <torch._C.Generator object at 0x7fbc5ead1410>, 'strength': 0.5, 'guidance_scale': 3.5}
`height` and `width` have to be divisible by 16 but are 788 and 1897. Dimensions will be resized accordingly
✅ Moved CLIPTextModel to cuda:0
🔄 Moved CLIPTextModel back to CPU
✅ Moved T5EncoderModel to cuda:0
🔄 Moved T5EncoderModel back to CPU
✅ Moved AutoencoderKL to cuda:0
  0%|                                                                                                                                                                                        | 0/5 [00:00<?, ?it/s]
✅ Moved FluxTransformer2DModel to cuda:0
🔄 Moved FluxTransformer2DModel back to CPU
 20%|███████████████████████████████████▏                                                                                                                                            | 1/5 [00:23<01:32, 23.13s/it]
✅ Moved FluxTransformer2DModel to cuda:0
🔄 Moved FluxTransformer2DModel back to CPU
 40%|██████████████████████████████████████████████████████████████████████▍                                                                                                         | 2/5 [00:25<00:32, 10.99s/it]
✅ Moved FluxTransformer2DModel to cuda:0
🔄 Moved FluxTransformer2DModel back to CPU
 60%|█████████████████████████████████████████████████████████████████████████████████████████████████████████▌                                                                      | 3/5 [00:28<00:14,  7.11s/it]
✅ Moved FluxTransformer2DModel to cuda:0
🔄 Moved FluxTransformer2DModel back to CPU
 80%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▊                                   | 4/5 [00:30<00:05,  5.30s/it]
✅ Moved FluxTransformer2DModel to cuda:0
🔄 Moved FluxTransformer2DModel back to CPU
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:33<00:00,  6.63s/it]
✅ Moved AutoencoderKL to cuda:0
WARNING  common.logger:memory.py:58 GPU Memory Usage: 4.91GB / 25.77GB,  Reserved: 4.29GB, Allocated: 0.51GB, Usage: 19.05%
PASSED
tests/images/test_image_to_image.py::test_image_to_image[flux-kontext-1-image_to_image] Base64: /9j/4AAQSkZJRgABAQEAkACQAAD/4QLcRXhpZgAATU0AKgAAAAgABAE7AAIAAAAGAAABSodpAAQAAAABAAABUJydAAEAAAAMAAAC.

System Info

• 🤗 Diffusers version: 0.32.2
• Platform: Windows-10-10.0.19045-SP0
• Running on Google Colab?: No
• Python version: 3.12.7
• PyTorch version (GPU?): 2.5.1+cu124 (True)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Huggingface_hub version: 0.30.1
• Transformers version: 4.51.0
• Accelerate version: 1.6.0
• PEFT version: not installed
• Bitsandbytes version: 0.45.5
• Safetensors version: 0.5.3
• xFormers version: not installed
• Accelerator: NVIDIA GeForce RTX 3090, 24576 MiB
• Using GPU in script?: Docker mostly
• Using distributed or parallel set-up in script?: No

Who can help?

No response

[asomoza]

Hi, doing a quick test, even if your code reports that the transformer model gets moved back and forth from CPU and the GPU, that's not really what's happening.

To check this, you can just manually do the moving yourself in the denoise loop, move it to GPU at the beginning and then back to the CPU at the end, you will see that the resulting time is a lot more than with just enable_model_cpu_offload()

Since you mentioned SD3 I used that model because it's the smaller one, these are my tests with a local 3090, timing just the denoising loop since its the only relevant part:

GPU only	model_cpu_offload()	Manual offload
12.18s	14.65s	39.17s

Nevertheless, it actually adds a couple of seconds to the inference, so there's something happening here, maybe @SunMarc can help.

[JoeGaffney]

Hey @asomoza

Looking at the code more closely it seems the offloading is actually done in the pre_forward as well and the CpuOffLoad does not override the base hooks post_forward.

Will run a few more tests.

[JoeGaffney]

Hey, I ran my SD3 test timing the pre_forward with enable_model_cpu_offload()

def time_info_decorator(func):
    def wrapper(*args, **kwargs):

        start = time.time()
        result = func(*args, **kwargs)
        end = time.time()

        elapsed = end - start
        if elapsed > 0.1:  # Only log if execution took more than 0.1 second
            args_str = str(args)[:100] + ("..." if len(str(args)) > 100 else "")
            kwargs_str = str(kwargs)[:100] + ("..." if len(str(kwargs)) > 100 else "")

            info = f"Calling {func.__name__} with args: {args_str}, kwargs: {kwargs_str}"
            logger.info(f"{info}, took: {elapsed:.2f}s")

        return result

    return wrapper

original_pre_forward = accelerate.hooks.CpuOffload.pre_forward

@time_info_decorator
def logged_pre_forward(self, module, *args, **kwargs):
    module_name = module.__class__.__name__
    result = original_pre_forward(self, module, *args, **kwargs)
    print(f"✅ Moved {module_name} to {self.execution_device}")
    if self.prev_module_hook is not None:
        print(f"🔄 Offloading previous module {self.prev_module_hook.model.__class__.__name__} to CPU")
    return result

Maybe there is still some overhead here calling each time. It seems like there is as on the initial call its longer then subsequent there is like 0.* s timing. But I could be profiling wrong with the patching. But seems it is only trying to move the previous model to CPU which is good.

venv\Lib\site-packages\accelerate\hooks.py

    def pre_forward(self, module, *args, **kwargs):
        if self.prev_module_hook is not None:
            self.prev_module_hook.offload()
            clear_device_cache()
        module.to(self.execution_device)
        return send_to_device(args, self.execution_device), send_to_device(kwargs, self.execution_device)

Sorry logs are bit long.

tests/images/test_text_to_image.py::test_text_to_image[sd-3-text_to_image]
--------------------------------------------------------------------------------------------------------- live log call ----------------------------------------------------------------------------------------------------------INFO common.logger:pipeline_helpers.py:231 Loading quantized model from /WORKSPACE/quantized/stabilityai/stable-diffusion-3.5-large/8bit/transformer
INFO common.logger:utils.py:154 Calling get_quantized_model with args: (), kwargs: {'model_id': 'stabilityai/stable-diffusion-3.5-large', 'subfolder': 'transformer', 'model_class': <c..., took: 1.17s
INFO common.logger:pipeline_helpers.py:231 Loading quantized model from /WORKSPACE/quantized/black-forest-labs/FLUX.1-schnell/8bit/text_encoder_2
Loading checkpoint shards: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:00<00:00, 11.39it/s]INFO common.logger:utils.py:154 Calling get_quantized_model with args: (), kwargs: {'model_id': 'black-forest-labs/FLUX.1-schnell', 'subfolder': 'text_encoder_2', 'model_class': <clas..., took: 0.33s
Loading pipeline components...: 44%|██████████████████████████████████████████████████████████████████████▋ | 4/9 [00:01<00:01, 4.08it/s]You set add_prefix_space. The tokenizer needs to be converted from the slow tokenizers
Loading pipeline components...: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 9/9 [00:01<00:00, 5.19it/s]WARNING common.logger:pipeline_helpers.py:80 Cache miss for ('get_pipeline', PipelineConfig(model_id='stabilityai/stable-diffusion-3.5-large', model_family='sd3', ip_adapter_models=(), ip_adapter_subfolders=(), ip_adapter_weights=(), ip_adapter_image_encoder_model='', ip_adapter_image_encoder_subfolder='')) - took: 3.94s - Cache size: 1/2

INFO common.logger:sd3.py:79 Text to image call {'width': 960, 'height': 528, 'prompt': "A serene scene of a woman lying on lush green grass in a sunlit meadow. She has long flowing hair spread out around her, eyes closed, with a peaceful expression on her face. She's wearing a light summer dress that gently ripples in the breeze. Around her, wildflowers bloom in soft pastel colors, and sunlight filters through the leaves of nearby trees, casting dappled shadows. The mood is calm, dreamy, and connected to nature.", 'negative_prompt': 'worst quality, inconsistent motion, blurry, jittery, distorted', 'num_inference_steps': 15, 'generator': <torch._C.Generator object at 0x7fa677d1b9f0>, 'guidance_scale': 3.5}
Token indices sequence length is longer than the specified maximum sequence length for this model (92 > 77). Running this sequence through the model will result in indexing errors
The following part of your input was truncated because CLIP can only handle sequences up to 77 tokens: ['led shadows. the mood is calm, dreamy, and connected to nature.']

✅ Moved CLIPTextModelWithProjection to cuda:0
Token indices sequence length is longer than the specified maximum sequence length for this model (92 > 77). Running this sequence through the model will result in indexing errors
The following part of your input was truncated because CLIP can only handle sequences up to 77 tokens: ['led shadows. the mood is calm, dreamy, and connected to nature.']
✅ Moved CLIPTextModelWithProjection to cuda:0
🔄 Offloading previous module CLIPTextModelWithProjection to CPU
INFO common.logger:utils.py:154 Calling logged_pre_forward with args: (<accelerate.hooks.CpuOffload object at 0x7fa3e3894a50>, CLIPTextModelWithProjection(
(text_model)..., kwargs: {'output_hidden_states': True}, took: 0.53s
✅ Moved T5EncoderModel to cuda:0
🔄 Offloading previous module CLIPTextModelWithProjection to CPU
INFO common.logger:utils.py:154 Calling logged_pre_forward with args: (<accelerate.hooks.CpuOffload object at 0x7fa3e37ca1d0>, T5EncoderModel(
(shared): Embedding(32128..., kwargs: {}, took: 2.38s
✅ Moved CLIPTextModelWithProjection to cuda:0
✅ Moved CLIPTextModelWithProjection to cuda:0
🔄 Offloading previous module CLIPTextModelWithProjection to CPU
INFO common.logger:utils.py:154 Calling logged_pre_forward with args: (<accelerate.hooks.CpuOffload object at 0x7fa3e3894a50>, CLIPTextModelWithProjection(
(text_model)..., kwargs: {'output_hidden_states': True}, took: 0.45s
✅ Moved T5EncoderModel to cuda:0
🔄 Offloading previous module CLIPTextModelWithProjection to CPU
INFO common.logger:utils.py:154 Calling logged_pre_forward with args: (<accelerate.hooks.CpuOffload object at 0x7fa3e37ca1d0>, T5EncoderModel(
(shared): Embedding(32128..., kwargs: {}, took: 0.55s
0%| | 0/15 [00:00<?, ?it/s]✅ Moved SD3Transformer2DModel to cuda:0
🔄 Offloading previous module T5EncoderModel to CPU
INFO common.logger:utils.py:154 Calling logged_pre_forward with args: (<accelerate.hooks.CpuOffload object at 0x7fa3e18b92d0>, SD3Transformer2DModel(
(pos_embed): Patch..., kwargs: {'hidden_states': tensor([[[[-8.0859e-01, -1.5312e+00, 4.0625e-01, ..., 1.3828e+00,
-4..., took: 5.74s
7%|████████████▋ | 1/15 [00:06<01:35, 6.84s/it]✅ Moved SD3Transformer2DModel to cuda:0
🔄 Offloading previous module T5EncoderModel to CPU
INFO common.logger:utils.py:154 Calling logged_pre_forward with args: (<accelerate.hooks.CpuOffload object at 0x7fa3e18b92d0>, SD3Transformer2DModel(
(pos_embed): Patch..., kwargs: {'hidden_states': tensor([[[[-7.8516e-01, -1.4922e+00, 3.9453e-01, ..., 1.3438e+00,
-4..., took: 0.17s
13%|█████████████████████████▎ | 2/15 [00:07<00:41, 3.18s/it]✅ Moved SD3Transformer2DModel to cuda:0
🔄 Offloading previous module T5EncoderModel to CPU
INFO common.logger:utils.py:154 Calling logged_pre_forward with args: (<accelerate.hooks.CpuOffload object at 0x7fa3e18b92d0>, SD3Transformer2DModel(
(pos_embed): Patch..., kwargs: {'hidden_states': tensor([[[[-7.6172e-01, -1.4453e+00, 3.8672e-01, ..., 1.3125e+00,
-4..., took: 0.58s
20%|██████████████████████████████████████ | 3/15 [00:08<00:26, 2.21s/it]✅ Moved SD3Transformer2DModel to cuda:0
🔄 Offloading previous module T5EncoderModel to CPU
INFO common.logger:utils.py:154 Calling logged_pre_forward with args: (<accelerate.hooks.CpuOffload object at 0x7fa3e18b92d0>, SD3Transformer2DModel(
(pos_embed): Patch..., kwargs: {'hidden_states': tensor([[[[-0.7422, -1.4062, 0.3711, ..., 1.2812, -0.4160, 0.9922],
..., took: 0.56s
27%|██████████████████████████████████████████████████▋ | 4/15 [00:09<00:19, 1.73s/it]✅ Moved SD3Transformer2DModel to cuda:0
🔄 Offloading previous module T5EncoderModel to CPU
INFO common.logger:utils.py:154 Calling logged_pre_forward with args: (<accelerate.hooks.CpuOffload object at 0x7fa3e18b92d0>, SD3Transformer2DModel(
(pos_embed): Patch..., kwargs: {'hidden_states': tensor([[[[-7.1875e-01, -1.3594e+00, 3.5547e-01, ..., 1.2422e+00,
-3..., took: 0.56s
33%|███████████████████████████████████████████████████████████████▎ | 5/15 [00:10<00:14, 1.48s/it]✅ Moved SD3Transformer2DModel to cuda:0
🔄 Offloading previous module T5EncoderModel to CPU
INFO common.logger:utils.py:154 Calling logged_pre_forward with args: (<accelerate.hooks.CpuOffload object at 0x7fa3e18b92d0>, SD3Transformer2DModel(
(pos_embed): Patch..., kwargs: {'hidden_states': tensor([[[[-6.9922e-01, -1.3125e+00, 3.3203e-01, ..., 1.2031e+00,
-3..., took: 0.58s
40%|████████████████████████████████████████████████████████████████████████████ | 6/15 [00:11<00:12, 1.33s/it]✅ Moved SD3Transformer2DModel to cuda:0
🔄 Offloading previous module T5EncoderModel to CPU
INFO common.logger:utils.py:154 Calling logged_pre_forward with args: (<accelerate.hooks.CpuOffload object at 0x7fa3e18b92d0>, SD3Transformer2DModel(
(pos_embed): Patch..., kwargs: {'hidden_states': tensor([[[[-0.6719, -1.2578, 0.3066, ..., 1.1562, -0.3359, 0.8906],
..., took: 0.57s
47%|████████████████████████████████████████████████████████████████████████████████████████▋ | 7/15 [00:12<00:09, 1.23s/it]✅ Moved SD3Transformer2DModel to cuda:0
🔄 Offloading previous module T5EncoderModel to CPU
INFO common.logger:utils.py:154 Calling logged_pre_forward with args: (<accelerate.hooks.CpuOffload object at 0x7fa3e18b92d0>, SD3Transformer2DModel(
(pos_embed): Patch..., kwargs: {'hidden_states': tensor([[[[-0.6328, -1.1953, 0.2832, ..., 1.1016, -0.3008, 0.8516],
..., took: 0.57s
53%|█████████████████████████████████████████████████████████████████████████████████████████████████████▎ | 8/15 [00:13<00:08, 1.17s/it]✅ Moved SD3Transformer2DModel to cuda:0
🔄 Offloading previous module T5EncoderModel to CPU
INFO common.logger:utils.py:154 Calling logged_pre_forward with args: (<accelerate.hooks.CpuOffload object at 0x7fa3e18b92d0>, SD3Transformer2DModel(
(pos_embed): Patch..., kwargs: {'hidden_states': tensor([[[[-0.5977, -1.0938, 0.2852, ..., 1.0469, -0.2520, 0.8047],
..., took: 0.58s
60%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████ | 9/15 [00:14<00:06, 1.13s/it]✅ Moved SD3Transformer2DModel to cuda:0
🔄 Offloading previous module T5EncoderModel to CPU
INFO common.logger:utils.py:154 Calling logged_pre_forward with args: (<accelerate.hooks.CpuOffload object at 0x7fa3e18b92d0>, SD3Transformer2DModel(
(pos_embed): Patch..., kwargs: {'hidden_states': tensor([[[[-0.5469, -0.9766, 0.2949, ..., 0.9844, -0.1934, 0.7578],
..., took: 0.57s
67%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████ | 10/15 [00:15<00:05, 1.10s/it]✅ Moved SD3Transformer2DModel to cuda:0
🔄 Offloading previous module T5EncoderModel to CPU
INFO common.logger:utils.py:154 Calling logged_pre_forward with args: (<accelerate.hooks.CpuOffload object at 0x7fa3e18b92d0>, SD3Transformer2DModel(
(pos_embed): Patch..., kwargs: {'hidden_states': tensor([[[[-0.4785, -0.8516, 0.2832, ..., 0.9141, -0.1128, 0.7148],
..., took: 0.58s
73%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▌ | 11/15 [00:16<00:04, 1.08s/it]✅ Moved SD3Transformer2DModel to cuda:0
🔄 Offloading previous module T5EncoderModel to CPU
INFO common.logger:utils.py:154 Calling logged_pre_forward with args: (<accelerate.hooks.CpuOffload object at 0x7fa3e18b92d0>, SD3Transformer2DModel(
(pos_embed): Patch..., kwargs: {'hidden_states': tensor([[[[-0.3945, -0.7148, 0.2559, ..., 0.8242, -0.0098, 0.6875],
..., took: 0.56s
80%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▏ | 12/15 [00:17<00:03, 1.06s/it]✅ Moved SD3Transformer2DModel to cuda:0
🔄 Offloading previous module T5EncoderModel to CPU
INFO common.logger:utils.py:154 Calling logged_pre_forward with args: (<accelerate.hooks.CpuOffload object at 0x7fa3e18b92d0>, SD3Transformer2DModel(
(pos_embed): Patch..., kwargs: {'hidden_states': tensor([[[[-0.2930, -0.5742, 0.2295, ..., 0.7344, 0.1108, 0.6719],
..., took: 0.57s
87%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▊ | 13/15 [00:18<00:02, 1.05s/it]✅ Moved SD3Transformer2DModel to cuda:0
🔄 Offloading previous module T5EncoderModel to CPU
INFO common.logger:utils.py:154 Calling logged_pre_forward with args: (<accelerate.hooks.CpuOffload object at 0x7fa3e18b92d0>, SD3Transformer2DModel(
(pos_embed): Patch..., kwargs: {'hidden_states': tensor([[[[-1.5625e-01, -4.5117e-01, 1.8555e-01, ..., 6.3672e-01,
2..., took: 0.56s
93%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▍ | 14/15 [00:19<00:01, 1.04s/it]✅ Moved SD3Transformer2DModel to cuda:0
🔄 Offloading previous module T5EncoderModel to CPU
INFO common.logger:utils.py:154 Calling logged_pre_forward with args: (<accelerate.hooks.CpuOffload object at 0x7fa3e18b92d0>, SD3Transformer2DModel(
(pos_embed): Patch..., kwargs: {'hidden_states': tensor([[[[-0.0088, -0.3867, 0.1787, ..., 0.5625, 0.5156, 0.7695],
..., took: 0.56s
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 15/15 [00:20<00:00, 1.39s/it]✅ Moved AutoencoderKL to cuda:0
🔄 Offloading previous module SD3Transformer2DModel to CPU
INFO common.logger:utils.py:154 Calling logged_pre_forward with args: (<accelerate.hooks.CpuOffload object at 0x7fa3e37e4910>, AutoencoderKL(
(encoder): Encoder(
(c..., kwargs: {}, took: 4.29s
WARNING common.logger:memory.py:58 GPU Memory Usage: 3.61GB / 25.77GB, Reserved: 2.39GB, Allocated: 0.18GB, Usage: 14.01%
PASSED

================================================================================================= 1 passed, 7 warnings in 43.34s ================================================================================================

[JoeGaffney]

With out offloading and putting the pipe.to("cuda") I do get much better performance which is expected. But I'm wondering if there is some more efficiency to be gained from the pre_forward calls during repeat steps.

docker-compose exec gpu-workers pytest tests/images/test_text_to_image.py -vs
====================================================================================================== test session starts =======================================================================================================platform linux -- Python 3.11.12, pytest-8.4.1, pluggy-1.5.0 -- /opt/conda/bin/python3.11
cachedir: .pytest_cache
hypothesis profile 'default' -> database=DirectoryBasedExampleDatabase(PosixPath('/app/workers/.hypothesis/examples'))
rootdir: /app/workers
configfile: pytest.ini
plugins: hypothesis-6.131.7, hydra-core-1.3.2, anyio-4.9.0
collecting 0 items
------------------------------------------------------------------------------------------------------ live log collection -------------------------------------------------------------------------------------------------------INFO qwen_vl_utils.vision_process:vision_process.py:42 set VIDEO_TOTAL_PIXELS: 90316800
collected 1 item

tests/images/test_text_to_image.py::test_text_to_image[sd-3-text_to_image]
--------------------------------------------------------------------------------------------------------- live log call ----------------------------------------------------------------------------------------------------------INFO common.logger:pipeline_helpers.py:232 Loading quantized model from /WORKSPACE/quantized/stabilityai/stable-diffusion-3.5-large/8bit/transformer
INFO common.logger:utils.py:154 Calling get_quantized_model with args: (), kwargs: {'model_id': 'stabilityai/stable-diffusion-3.5-large', 'subfolder': 'transformer', 'model_class': <c..., took: 0.96s
INFO common.logger:pipeline_helpers.py:232 Loading quantized model from /WORKSPACE/quantized/black-forest-labs/FLUX.1-schnell/8bit/text_encoder_2
Loading checkpoint shards: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:00<00:00, 11.99it/s]INFO common.logger:utils.py:154 Calling get_quantized_model with args: (), kwargs: {'model_id': 'black-forest-labs/FLUX.1-schnell', 'subfolder': 'text_encoder_2', 'model_class': <clas..., took: 0.29s
Loading pipeline components...: 44%|██████████████████████████████████████████████████████████████████████▋ | 4/9 [00:00<00:01, 3.87it/s]You set add_prefix_space. The tokenizer needs to be converted from the slow tokenizers
Loading pipeline components...: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 9/9 [00:01<00:00, 6.88it/s]WARNING common.logger:pipeline_helpers.py:80 Cache miss for ('get_pipeline', PipelineConfig(model_id='stabilityai/stable-diffusion-3.5-large', model_family='sd3', ip_adapter_models=(), ip_adapter_subfolders=(), ip_adapter_weights=(), ip_adapter_image_encoder_model='', ip_adapter_image_encoder_subfolder='')) - took: 8.58s - Cache size: 1/2

INFO common.logger:sd3.py:79 Text to image call {'width': 960, 'height': 528, 'prompt': "A serene scene of a woman lying on lush green grass in a sunlit meadow. She has long flowing hair spread out around her, eyes closed, with a peaceful expression on her face. She's wearing a light summer dress that gently ripples in the breeze. Around her, wildflowers bloom in soft pastel colors, and sunlight filters through the leaves of nearby trees, casting dappled shadows. The mood is calm, dreamy, and connected to nature.", 'negative_prompt': 'worst quality, inconsistent motion, blurry, jittery, distorted', 'num_inference_steps': 15, 'generator': <torch._C.Generator object at 0x7fa44d97be90>, 'guidance_scale': 3.5}
Token indices sequence length is longer than the specified maximum sequence length for this model (92 > 77). Running this sequence through the model will result in indexing errors
The following part of your input was truncated because CLIP can only handle sequences up to 77 tokens: ['led shadows. the mood is calm, dreamy, and connected to nature.']
Token indices sequence length is longer than the specified maximum sequence length for this model (92 > 77). Running this sequence through the model will result in indexing errors
The following part of your input was truncated because CLIP can only handle sequences up to 77 tokens: ['led shadows. the mood is calm, dreamy, and connected to nature.']

100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 15/15 [00:11<00:00, 1.36it/s]WARNING common.logger:memory.py:58 GPU Memory Usage: 18.97GB / 25.77GB, Reserved: 19.09GB, Allocated: 16.02GB, Usage: 73.60%
WARNING common.logger:memory.py:75 Clean GPU Memory Usage: 16.73GB / 25.77GB, Reserved: 16.69GB, Allocated: 16.02GB, Usage: 64.92%
PASSED

================================================================================================= 1 passed, 7 warnings in 30.25s =================================================================================================

[asomoza]

With out offloading and putting the pipe.to("cuda") I do get much better performance

you should never use cpu offloading if you have the VRAM to run inference, is expected for it to be slower when you switch between models, e.g.: between the text encoders, the transformer and the vae, but it shouldn't be the case with the denoise loop which doesn't need to switch models, in my test it makes the inference aprox. 20% slower and this could mean minutes when doing inference with a video model like WAN, so I think there's room for efficiency.

Still, I don't know if this is expected or if it's something that we can optimize in diffusers or accelerate, I already pinged @SunMarc but adding @sayakpaul for more insights.

[JoeGaffney]

Hey @asomoza thanks for the follow up and linking people.

Yes in an ideal world keeping everything in vram would be perfect. But with mixed spec machines and also in some cases bolting on other things adapters, control nets etc. Being able to offload as stages is very valuable as then you only need to worry about fitting the largest models that need to exist at the same time in vram with some buffer.

In theory loading all the models into VRAM at once then putting them back to RAM in one go should not but be that much more than doing so in blocks of when they are used if they are used sequentially and there isn't a need to reload them.

I'm testing out this which should only trigger things once when needed. But haven't fully profiled yet

import time
from accelerate.hooks import CpuOffload, clear_device_cache, send_to_device

_original_pre_forward = CpuOffload.pre_forward


def patched_pre_forward(self, module, *args, **kwargs):
    print(f"🔄 Pre-forward hook for {module.__class__.__name__}")
    target_device = self.execution_device
    current_device = next(module.parameters()).device

    # Handle previous module offload
    if self.prev_module_hook is not None:
        prev_module = self.prev_module_hook.model  
        prev_device = next(prev_module.parameters()).device
        if prev_device != torch.device("cpu"):
            start = time.perf_counter()
            self.prev_module_hook.offload()
            clear_device_cache()
            print(f"⬅️ Offloaded {prev_module.__class__.__name__} to CPU ({time.perf_counter() - start:.2f}s)")

    if current_device == target_device:
        return args, kwargs

    # Move current module to target device only if needed
    start = time.perf_counter()
    module.to(target_device)
    result_args, result_kwargs = send_to_device(args, target_device), send_to_device(kwargs, target_device)
    print(f"➡️ Loaded {module.__class__.__name__} to {target_device} ({time.perf_counter() - start:.2f}s)")

    return result_args, result_kwargs


# Apply patch
CpuOffload.pre_forward = patched_pre_forward

[sayakpaul]

Thanks for the detailed investigations so far.

but it shouldn't be the case with the denoise loop which doesn't need to switch models, in my test it makes the inference aprox. 20% slower and this could mean minutes when doing inference with a video model like WAN, so I think there's room for efficiency.

For enable_model_cpu_offload(), I think we should definitely not offload the denoiser to CPU and load it back to GPU again before the entire denoising loop is finished. It would be meaningful to keep it on the GPU for the entirety of the denoising process (IFF the VRAM allows for it). I have confirmed that this is the case and as expected:

from diffusers import DiffusionPipeline
import torch 

pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev", 
    torch_dtype=torch.bfloat16
)
pipeline.enable_model_cpu_offload()

pipeline.set_progress_bar_config(disable=True)
pipe_kwargs = {
    "prompt": "A cat holding a sign that says hello world",
    "height": 1024,
    "width": 1024,
    "guidance_scale": 3.5,
    "num_inference_steps": 5, # small number for testing
    "max_sequence_length": 512,
}
image = pipeline(**pipe_kwargs).images[0]

I made the following changes to literally print the device of the transformer:

diff --git a/src/diffusers/pipelines/flux/pipeline_flux.py b/src/diffusers/pipelines/flux/pipeline_flux.py
index 4c83ae740..516d44be1 100644
--- a/src/diffusers/pipelines/flux/pipeline_flux.py
+++ b/src/diffusers/pipelines/flux/pipeline_flux.py
@@ -900,6 +900,7 @@ class FluxPipeline(
         # 6. Denoising loop
         # We set the index here to remove DtoH sync, helpful especially during compilation.
         # Check out more details here: https://github.com/huggingface/diffusers/pull/11696
+        print(f"{self.transformer.device=} before for loop.")
         self.scheduler.set_begin_index(0)
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):
@@ -912,6 +913,7 @@ class FluxPipeline(
                 # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
                 timestep = t.expand(latents.shape[0]).to(latents.dtype)
 
+                print(f"{self.transformer.device=} before invocation.")
                 noise_pred = self.transformer(
                     hidden_states=latents,
                     timestep=timestep / 1000,
@@ -923,6 +925,7 @@ class FluxPipeline(
                     joint_attention_kwargs=self.joint_attention_kwargs,
                     return_dict=False,
                 )[0]
+                print(f"{self.transformer.device=} after invocation.")
 
                 if do_true_cfg:
                     if negative_image_embeds is not None:

I got:

self.transformer.device=device(type='cpu') before for loop.
self.transformer.device=device(type='cpu') before invocation.
self.transformer.device=device(type='cuda', index=0) after invocation.
self.transformer.device=device(type='cuda', index=0) before invocation.
self.transformer.device=device(type='cuda', index=0) after invocation.
self.transformer.device=device(type='cuda', index=0) before invocation.
self.transformer.device=device(type='cuda', index=0) after invocation.
self.transformer.device=device(type='cuda', index=0) before invocation.
self.transformer.device=device(type='cuda', index=0) after invocation.
self.transformer.device=device(type='cuda', index=0) before invocation.
self.transformer.device=device(type='cuda', index=0) after invocation.

We can clearly see that the transformer is kept on the GPU for the entirety of the denoising loop. What's wrong here?