Skip to content

Consolidate torch device backend dispatch - #14792

Open
DN6 wants to merge 2 commits into
mainfrom
device_dtype_routing
Open

DN6 wants to merge 2 commits into
mainfrom
device_dtype_routing

Conversation

@DN6

@DN6 DN6 commented Sep 16, 2026

Copy link
Copy Markdown
Collaborator

What does this PR do?

Support for torch device backends keeps growing, and each new backend adds a branch at every site that touches the device module. e.g. Adding NPU streams to group offloading (#14785) looks like this:

if torch.cuda.is_available():
    stream = torch.cuda.Stream()
elif hasattr(torch, "xpu") and torch.xpu.is_available():
    stream = torch.Stream()
elif hasattr(torch, "npu") and torch.npu.is_available():
    stream = torch.npu.Stream()

This PR introduces a TorchDeviceBackend object that automatically routes to the appropriate backend module for a given device type. e.g.

backend = TorchDeviceBackend(device)   # torch.cuda for "cuda", torch.xpu for "xpu", torch.npu for "npu", ...

# These methods will now work regardless of the device. We can avoid the need for multiple if/else checks
backend.Stream()                       
backend.empty_cache()                  
backend.max_memory_allocated()

Fixes # (issue)

Before submitting

  • Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
    • Did you read the Coding with AI agents guide?
    • Did you run the self-review skill on the diff?
    • Did you share the final self-review notes in the PR description or a comment?
  • Did you read the contributor guideline?
  • Did you read our philosophy doc? (important for complex PRs)
  • Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
  • Did you make sure to update the documentation with your changes? Here are the
    documentation guidelines, and
    here are tips on formatting docstrings.
  • Did you write any new necessary tests?
  • Are you the author (or part of the team) of the model/pipeline (only applicable for model/pipeline related PRs)?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@github-actions github-actions Bot added the size/L PR with diff > 200 LOC label Sep 16, 2026
image = latents
else:
latents = latents.to(self.vae.dtype)
torch_accelerator_module = getattr(torch, get_device(), torch.cuda)

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

Deadcode since min supported torch version is 2.6

@DN6
DN6 requested review from sayakpaul and yiyixuxu September 16, 2026 12:35
@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.

@sayakpaul sayakpaul left a comment

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.

Clean

raise ValueError("Using streams for data transfer requires a CUDA device, or an Intel XPU device.")
if use_stream and onload_device.type == "cpu":
raise ValueError("Using streams for data transfer requires an accelerator onload device, got `cpu`.")
stream = TorchDeviceBackend(onload_device).Stream() if use_stream else None

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.

What if the underlying device doesn't support streams?

Comment on lines -1056 to -1061
torch_accelerator_module = getattr(torch, get_device(), torch.cuda)
oom_error = (
torch.OutOfMemoryError
if is_torch_version(">=", "2.5.0")
else torch_accelerator_module.OutOfMemoryError
)

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.

Safe because we pin on >=2.6.

Comment on lines +224 to +228
def empty_cache(self) -> None:
# Backends without a caching allocator (cpu, neuron) have nothing to clear.
empty_cache = getattr(self.module, "empty_cache", None)
if empty_cache is not None:
empty_cache()

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.

Should we also warn in case empty_cache is None?

"""Tests for :class:`diffusers.utils.torch_utils.TorchDeviceBackend` on the CPU backend."""

def test_module_resolves_from_str_and_torch_device(self):
import torch

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.

I think we can keep all the imports at the top. They seem harmless.

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

Labels

Projects

Status: In Progress

Development

Successfully merging this pull request may close these issues.

3 participants