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

Changes from 3 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
14 changes: 14 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 @@ -445,6 +445,20 @@ 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 device_type == "hpu" and kwargs.pop("hpu_migration", True):
os.environ["PT_HPU_GPU_MIGRATION"] = "1"
logger.debug('Environment variable set: PT_HPU_GPU_MIGRATION=1')

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

try:
import habana_frameworks.torch.core as htcore # noqa: F401
Copy link
Collaborator

@DN6 DN6 Apr 16, 2025

Choose a reason for hiding this comment

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

Can this package be installed via PyPI? I wasn't able to find any instructions here?
https://docs.habana.ai/en/latest/PyTorch/PyTorch_Model_Porting/GPU_Migration_Toolkit/GPU_Migration_Toolkit.html#enabling-the-gpu-migration-toolkit

Could we add an import check like is_habana_available? Similar to how we have it here

diffusers/src/diffusers/utils/import_utils.py

Line 165 in d63e6fc

_torch_xla_available, _torch_xla_version = _is_package_available("torch_xla")

Copy link
Contributor Author

@dsocek dsocek Apr 16, 2025

Choose a reason for hiding this comment

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

@DN6 Thanks for review!

The habana_frameworks.torch package is Gaudi PyTorch bridge, part of complex SW stack for Intel Gaudi accelerators. You can see installation instructions (non-trivial) here and more specifically here. However, for all practical purposes users will be working under official release docker which will have this already installed. @regisss, can you can pitch in your view on this?

Thanks for good suggestion about using importlib-based check is_habana_available(), I will add this next.

Copy link
Contributor Author

@dsocek dsocek Apr 16, 2025

Choose a reason for hiding this comment

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

@DN6 Unfortunately, we can't use importlib here. The _is_package_available falsely returns that "habana_frameworks.torch.core" is not available because it is not able to get version via importlib_metadata, even importlib.util.find_spec("habana_frameworks.torch.core") finds the package OK.

import importlib.util
import sys

# The package importlib_metadata is in a different place, depending on the python version.
if sys.version_info < (3, 8):
    import importlib_metadata
else:
    import importlib.metadata as importlib_metadata

def _is_package_available(pkg_name: str):
    pkg_exists = importlib.util.find_spec(pkg_name) is not None
    pkg_version = "N/A"

    if pkg_exists:
        try:
            pkg_version = importlib_metadata.version(pkg_name)
            print(f"Successfully imported {pkg_name} version {pkg_version}")
        except (ImportError, importlib_metadata.PackageNotFoundError):
            pkg_exists = False

    return pkg_exists, pkg_version

print(importlib.util.find_spec("habana_frameworks.torch.core") is not None)
print(_is_package_available("habana_frameworks.torch.core"))

Output:
True
(False, 'N/A')

Any other suggestion for handling this import?

Copy link
Contributor

@regisss regisss Apr 17, 2025

Choose a reason for hiding this comment

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

@dsocek Here is how we do it in Accelerate: https://github.com/huggingface/accelerate/blob/34c1779828b3d0769992e6492e6de93d869f71b5/src/accelerate/utils/imports.py#L435

habana_frameworks.torch.core should be there anyway if habana_frameworks is there right?

Copy link
Contributor Author

@dsocek dsocek Apr 17, 2025

Choose a reason for hiding this comment

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

@regisss great example how its done in accelerate!

I did a quick test, seems at the minimum we need to import habana_frameworks.torch (importing only habana_frameworks would cause an error in migration ModuleNotFoundError: No module named 'torch.hpu').

I will copy most of what we do in accelerate and update this PR.

Copy link
Collaborator

@DN6 DN6 Apr 17, 2025

Choose a reason for hiding this comment

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

@dsocek Also fine to implement an is_hpu_available check the way it's done with accelerate
https://github.com/huggingface/accelerate/blob/34c1779828b3d0769992e6492e6de93d869f71b5/src/accelerate/utils/imports.py#L435

logger.debug("Successfully imported habana_frameworks.torch.core")
except ImportError as e:
logger.warning("Could not import habana_frameworks.torch.core: %s", e)

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