Skip to content

[core] DiffusionPipeline.device deduction for split-device pipelines - #14383

Open
JingyaHuang wants to merge 2 commits into
huggingface:mainfrom
JingyaHuang:impv-execution-device-split-placement
Open

[core] DiffusionPipeline.device deduction for split-device pipelines#14383
JingyaHuang wants to merge 2 commits into
huggingface:mainfrom
JingyaHuang:impv-execution-device-split-placement

Conversation

@JingyaHuang

Copy link
Copy Markdown
Contributor

Description

What does this PR do?

Follow-up to the discussion in #13718 (cc @sayakpaul @DN6), split out as requested.

  • Issue

DiffusionPipeline.device returns the device of the first module component in alphabetical order of the __init__ signature (eg. text_encoder for Flux, Flux.2, Qwen-Image). When a pipeline is deliberately split across devices, DiffusionPipeline.device reports the device where locates the text encoder, for example:

# text_encoder stays on CPU
pipe.unet.to("cuda") 
pipe.vae.to("cuda")

pipe.device  # -> cpu, even though everything that computes is on cuda

This PR prefers any non-CPU, non-meta component, falling back to the previous first-component behaviour when there is nothing else to prefer. Backend-agnostic CUDA, XPU, and other accelerators, including TP-sharded models whose DTensor
parameters report their local shard device.

@JingyaHuang
JingyaHuang requested a review from DN6 August 4, 2026 14:54
@github-actions github-actions Bot added size/S PR with diff < 50 LOC tests pipelines and removed size/S PR with diff < 50 LOC labels Aug 4, 2026
@JingyaHuang
JingyaHuang requested a review from sayakpaul August 4, 2026 14:54
@HuggingFaceDocBuilderDev

Copy link
Copy Markdown

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Comment on lines +1975 to +1977
# With every component on CPU there is no accelerator to prefer.
sd.to("cpu")
assert sd.device.type == "cpu"

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

But will we be able to run the pipeline with this configuration? For the test, I would also try to mimic the scenario in which this will be used in a minimal manner.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants