Skip to content

Hooks conflicts: Context Parallelism and CPU Offload #12533

Open
Open
Hooks conflicts: Context Parallelism and CPU Offload#12533
Assignees
yiyixuxu
Labels
bugSomething isn't workingcontext-parallelroadmapAdd to current release roadmap

Description

@DefTruth
DefTruth
opened on Oct 23, 2025

Describe the bug

Enable cpu offload before enabling parallelism will raise shape error after first pipe call. It seems a bug of diffusers that cpu offload is not fully compatible with context parallelism, visa versa.

  • cpu offload before context parallelism (not work)
pipe.enable_model_cpu_offload(device=device)

# pipe.transformer.set_attention_backend("flash")
pipe.transformer.set_attention_backend("_native_cudnn")
pipe.transformer.enable_parallelism(
    config=ContextParallelConfig(ulysses_degree=dist.get_world_size())
)
  • cpu offload after context parallelism (work)
# pipe.transformer.set_attention_backend("flash")
pipe.transformer.set_attention_backend("_native_cudnn")
pipe.transformer.enable_parallelism(
    config=ContextParallelConfig(ulysses_degree=dist.get_world_size())
)
pipe.enable_model_cpu_offload(device=device)

Reproduction

import os
import time
import torch
import torch.distributed as dist
from diffusers import (
    QwenImagePipeline,
    QwenImageTransformer2DModel,
    ContextParallelConfig,
)


def maybe_init_distributed():
    if not dist.is_initialized():
        dist.init_process_group("nccl")
    rank = dist.get_rank()
    device = torch.device("cuda", rank % torch.cuda.device_count())
    torch.cuda.set_device(device)
    return rank, device


def maybe_destroy_distributed():
    if dist.is_initialized():
        dist.destroy_process_group()


rank, device = maybe_init_distributed()

pipe = QwenImagePipeline.from_pretrained(
    os.environ.get(
        "QWEN_IMAGE_DIR",
        "Qwen/Qwen-Image",
    ),
    torch_dtype=torch.bfloat16,
)

# NOTE: Enable cpu offload before enabling parallelism will
# raise shape error after first pipe call, so we enable it after.
# It seems a bug of diffusers that cpu offload is not fully
# compatible with context parallelism, visa versa.
pipe.enable_model_cpu_offload(device=device)

assert isinstance(pipe.transformer, QwenImageTransformer2DModel)
# pipe.transformer.set_attention_backend("flash")
pipe.transformer.set_attention_backend("_native_cudnn")
pipe.transformer.enable_parallelism(
    config=ContextParallelConfig(ulysses_degree=dist.get_world_size())
)

# NOTE: Enable cpu offload after enabling parallelism
# pipe.enable_model_cpu_offload(device=device)

# assert isinstance(pipe.vae, AutoencoderKLQwenImage)
# pipe.vae.enable_tiling()

positive_magic = {
    "en": ", Ultra HD, 4K, cinematic composition.",  # for english prompt
    "zh": ", 超清，4K，电影级构图.",  # for chinese prompt
}

# Generate image
prompt = """A coffee shop entrance features a chalkboard sign reading "Qwen Coffee 😊 $2 per cup," with a neon light beside it displaying "通义千问". Next to it hangs a poster showing a beautiful Chinese woman, and beneath the poster is written "π≈3.1415926-53589793-23846264-33832795-02384197". Ultra HD, 4K, cinematic composition"""

# using an empty string if you do not have specific concept to remove
negative_prompt = " "

pipe.set_progress_bar_config(disable=rank != 0)


def run_pipe():
    # do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
    image = pipe(
        prompt=prompt + positive_magic["en"],
        negative_prompt=negative_prompt,
        width=1024,
        height=1024,
        num_inference_steps=50,
        true_cfg_scale=4.0,
        generator=torch.Generator(device="cpu").manual_seed(42),
    ).images[0]

    return image


# warmup
_ = run_pipe() # always work

start = time.time()
image = run_pipe() # raise error here if cpu offload is enabled before parallelism
end = time.time()


if rank == 0:
    time_cost = end - start
    save_path = f"qwen-image.cp{dist.get_world_size()}.png"
    print(f"Time cost: {time_cost:.2f}s")
    print(f"Saving image to {save_path}")
    image.save(save_path)

maybe_destroy_distributed()

Logs

Error:


[rank0]: Traceback (most recent call last):
[rank0]:   File "/workspace/dev/vipshop/cache-dit/examples/parallelism/run_qwen_image_cp_naive.py", line 71, in <module>
[rank0]:     start = time.time()
[rank0]:             ^^^^^^^^^^
[rank0]:   File "/workspace/dev/vipshop/cache-dit/examples/parallelism/run_qwen_image_cp_naive.py", line 54, in run_pipe
[rank0]:     # do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
[rank0]:             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/utils/_contextlib.py", line 120, in decorate_context
[rank0]:     return func(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/workspace/dev/vipshop/diffusers/src/diffusers/pipelines/qwenimage/pipeline_qwenimage.py", line 691, in __call__
[rank0]:     noise_pred = self.transformer(
[rank0]:                  ^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1786, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/accelerate/hooks.py", line 175, in new_forward
[rank0]:     output = module._old_forward(*args, **kwargs)
[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/workspace/dev/vipshop/diffusers/src/diffusers/models/transformers/transformer_qwenimage.py", line 647, in forward
[rank0]:     encoder_hidden_states, hidden_states = block(
[rank0]:                                            ^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1786, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/workspace/dev/vipshop/diffusers/src/diffusers/models/transformers/transformer_qwenimage.py", line 443, in forward
[rank0]:     attn_output = self.attn(
[rank0]:                   ^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1786, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/workspace/dev/vipshop/diffusers/src/diffusers/models/attention_processor.py", line 605, in forward
[rank0]:     return self.processor(
[rank0]:            ^^^^^^^^^^^^^^^
[rank0]:   File "/workspace/dev/vipshop/diffusers/src/diffusers/models/transformers/transformer_qwenimage.py", line 322, in __call__
[rank0]:     img_query = apply_rotary_emb_qwen(img_query, img_freqs, use_real=False)
[rank0]:                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/workspace/dev/vipshop/diffusers/src/diffusers/models/transformers/transformer_qwenimage.py", line 139, in apply_rotary_emb_qwen
[rank0]:     x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3)
[rank0]:                                ~~~~~~~~~~^~~~~~~~~~~
[rank0]: RuntimeError: The size of tensor a (4096) must match the size of tensor b (2048) at non-singleton dimension 1

System Info

diffusers 0.36.dev0 (latest main branch), pytorch 2.9.0

Who can help?

@yiyixuxu @DN6

Activity

  1. DefTruth
    added
    bugSomething isn't working
    on Oct 23, 2025

  2. xidiancpy commented on Oct 23, 2025

    @xidiancpy
    xidiancpy
    on Oct 23, 2025

    follow it

  3. yiyixuxu commented on Nov 4, 2025

    @yiyixuxu
    yiyixuxu
    on Nov 4, 2025
    Collaborator

    this is because Accelerate stores a reference to module.forward before wrapping it:

    old_forward = module.forward
module._old_forward = old_forward  # Stores reference
module.forward = new_forward        # Wraps with Accelerate hook

    If it is applied after the CP hooks, the forward here already has CP hooks applied.

    At the end of each pipeline run, enable_model_cpu_offload will reset the Accelerate hooks by setting the forward method back to _old_forward:

    module.forward = module._old_forward

    https://github.com/huggingface/accelerate/blob/a12beee389f6bd37cfae0aba233db03f375f7f80/src/accelerate/hooks.py#L188

    So depending on which one is called first: if CP hooks are called before the Accelerate hooks, the _old_forward here already has CP applied. Otherwise, if CP is called after the Accelerate hooks, the _old_forward here is the model's original forward and we basically remove the CP hooks on the second run.

    I think the easiest way here is to always remove and re-apply after each pipeline run. We can send a warning or something, but let me know if you have better ideas in mind.

  4. yiyixuxu
    added
    context-parallel
    on Nov 4, 2025
  5. yiyixuxu
    self-assigned this
    on Nov 4, 2025

  6. DefTruth commented on Nov 4, 2025

    @DefTruth
    DefTruth
    on Nov 4, 2025
    ContributorAuthor

    get

  7. github-actions commented on Jan 9, 2026

    @github-actions
    github-actionsbot
    on Jan 9, 2026 – with GitHub Actions
    Contributor

    This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

    Please note that issues that do not follow the contributing guidelines are likely to be ignored.

  8. github-actions
    added
    staleIssues that haven't received updates
    on Jan 9, 2026
  9. yiyixuxu
    removed
    staleIssues that haven't received updates
    on Jan 10, 2026

  10. github-actions commented on Feb 3, 2026

    @github-actions
    github-actionsbot
    on Feb 3, 2026 – with GitHub Actions
    Contributor

    This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

    Please note that issues that do not follow the contributing guidelines are likely to be ignored.

  11. github-actions
    added
    staleIssues that haven't received updates
    on Feb 3, 2026
  12. DN6
    added
    roadmapAdd to current release roadmap
    and removed
    staleIssues that haven't received updates
    on Mar 10, 2026
  13. github-project-automation
    added this to Diffusers Roadmapon Mar 10, 2026
  14. DN6
    moved this to Todo in Diffusers Roadmapon Mar 10, 2026
  15. DN6
    moved this from Todo to Future Release in Diffusers Roadmapon Aug 20, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

bugSomething isn't workingcontext-parallelroadmapAdd to current release roadmap

Type

No type

Projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions