Skip to content

Add generic support for Intel Gaudi accelerator (hpu device) #11328

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
Apr 30, 2025
Merged

Add generic support for Intel Gaudi accelerator (hpu device) #11328

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
2c6a028
Add generic support for Intel Gaudi accelerator (hpu device)
dsocek Mar 12, 2025
fd2e7bf
Merge branch 'main' into hpu_support
sayakpaul Apr 16, 2025
ae3c497
Add loggers for generic HPU support
dsocek Apr 16, 2025
baf75db
Refactor hpu support with is_hpu_available() logic
dsocek Apr 17, 2025
04f2f55
Fix style for hpu support update
dsocek Apr 17, 2025
15725c4
Decouple soft HPU check from hard device validation to support HPU mi…
dsocek Apr 24, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions src/diffusers/pipelines/pipeline_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@
_is_valid_type,
is_accelerate_available,
is_accelerate_version,
is_hpu_available,
is_torch_npu_available,
is_torch_version,
is_transformers_version,
Expand Down Expand Up @@ -445,6 +446,11 @@ def module_is_offloaded(module):
f"It seems like you have activated model offloading by calling `enable_model_cpu_offload`, but are now manually moving the pipeline to GPU. It is strongly recommended against doing so as memory gains from offloading are likely to be lost. Offloading automatically takes care of moving the individual components {', '.join(self.components.keys())} to GPU when needed. To make sure offloading works as expected, you should consider moving the pipeline back to CPU: `pipeline.to('cpu')` or removing the move altogether if you use offloading."
)

# Enable generic support for Intel Gaudi accelerator using GPU/HPU migration
if kwargs.pop("hpu_migration", True) and is_hpu_available():
Copy link
Collaborator

Choose a reason for hiding this comment

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

Shouldn't we keep the device_type check that was here earlier? e.g. If HPU is available on the machine, and we set pipe.to(torch.float16) this path would still run and set the device silently right?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

@DN6 with the new is_hpu_available() we dont need explicit device check. Device will silently be set to hpu within is_hpu_available() when all checks for HPU env pass

Copy link
Collaborator

Choose a reason for hiding this comment

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

ohh this would not be expected by our users and not aligned with our design philosophy: they would need to explicitly set device_type if they want to use the non-default one

Copy link
Contributor Author

Choose a reason for hiding this comment

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

If habana_frameworks is hooked into torch, then HPU would be the default device

Copy link
Collaborator

Choose a reason for hiding this comment

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

@dsocek So if I understand correctly, if an HPU is available and the pipeline needs to be run on CPU it wont unless it's explicitly moved?

Assuming the following snippet runs on an HPU machine

e.g

# automatically run on HPU if it is available? 
pipe = DiffusionPipeline.from_pretrained("..")
pipe(**args)

To run on CPU you would have to explicitly set pipe.to(cpu)?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

@DN6 @yiyixuxu
Apologies for the earlier confusion - my previous comment was incorrect. 🙇

I was referring to the if statement in the to() function of this PR where I originally used if device == "hpu", and after refactoring now use is_hpu_available(). We don't explicitly have "hpu" check anymore in the if statement.
This new check will not change any user expected behavior:

Current PR's behavior is as follows:

  • pipe.to() not called: then device used will be (default) CPU
  • pipe.to("cpu"): then device used will be CPU
  • pipe.to("hpu"): then device used will be HPU
def is_hpu_available():
    if (
        importlib.util.find_spec("habana_frameworks") is None
        or importlib.util.find_spec("habana_frameworks.torch") is None
    ):
        return False

    os.environ["PT_HPU_GPU_MIGRATION"] = "1"
    logger.debug("Environment variable set: PT_HPU_GPU_MIGRATION=1")

    import habana_frameworks.torch  # noqa: F401
    import torch

    return hasattr(torch, "hpu") and torch.hpu.is_available()

Here we 1st check is if "habana_frameworks" or "habana_frameworks.torch" are in the environment.

If they are, then we set hpu_migration RT var to true, and import habana_frameworks.torch.
We must define this run-time var before importing habana_frameworks.torch.

All of this will still keep device on CPU unless user explicitly set it to HPU

Finally we also do hard checks hasattr(torch, "hpu") and torch.hpu.is_available()

Let me know if further adjustments are needed :)

os.environ["PT_HPU_MAX_COMPOUND_OP_SIZE"] = "1"
logger.debug("Environment variable set: PT_HPU_MAX_COMPOUND_OP_SIZE=1")

module_names, _ = self._get_signature_keys(self)
modules = [getattr(self, n, None) for n in module_names]
modules = [m for m in modules if isinstance(m, torch.nn.Module)]
Expand Down
1 change: 1 addition & 0 deletions src/diffusers/utils/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,7 @@
is_gguf_version,
is_google_colab,
is_hf_hub_version,
is_hpu_available,
is_inflect_available,
is_invisible_watermark_available,
is_k_diffusion_available,
Expand Down
16 changes: 16 additions & 0 deletions src/diffusers/utils/import_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -336,6 +336,22 @@ def is_timm_available():
return _timm_available


def is_hpu_available():
if (
importlib.util.find_spec("habana_frameworks") is None
or importlib.util.find_spec("habana_frameworks.torch") is None
):
return False

os.environ["PT_HPU_GPU_MIGRATION"] = "1"
logger.debug("Environment variable set: PT_HPU_GPU_MIGRATION=1")

import habana_frameworks.torch # noqa: F401
import torch

return hasattr(torch, "hpu") and torch.hpu.is_available()
Copy link
Collaborator

Choose a reason for hiding this comment

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

Suggested change
os.environ["PT_HPU_GPU_MIGRATION"] = "1"
logger.debug("Environment variable set: PT_HPU_GPU_MIGRATION=1")
import habana_frameworks.torch # noqa: F401
import torch
return hasattr(torch, "hpu") and torch.hpu.is_available()

ohh, the check should just return True or False, indicate if hpu is available or not

Copy link
Contributor Author

Choose a reason for hiding this comment

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

@yiyixuxu Yes this does return True or False. We need to set migration to True for HPU before import habana_frameworks.torch so this is why I added it as well



# docstyle-ignore
FLAX_IMPORT_ERROR = """
{0} requires the FLAX library but it was not found in your environment. Checkout the instructions on the
Expand Down