Add auto_device_count and device name support

src/pytorch_lightning/CHANGELOG.md

@@ -147,6 +147,12 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Updated `val_check_interval`(int) to consider total train batches processed instead of `_batches_that_stepped` for validation check during training ([#12832](https://github.com/Lightning-AI/lightning/pull/12832)
 
 
+- Updated Habana Accelerator's `auto_device_count`, `is_available` & `get_device_name` methods based on the latest torch habana package ([#13423](https://github.com/PyTorchLightning/pytorch-lightning/pull/13423))
+
+
+-
+
+
 ### Deprecated
 
 - Deprecated `pytorch_lightning.loggers.base.LightningLoggerBase` in favor of `pytorch_lightning.loggers.logger.Logger`, and deprecated `pytorch_lightning.loggers.base` in favor of `pytorch_lightning.loggers.logger` ([#120148](https://github.com/PyTorchLightning/pytorch-lightning/pull/12014))

src/pytorch_lightning/accelerators/hpu.py

@@ -21,6 +21,9 @@
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.rank_zero import rank_zero_debug
 
+if _HPU_AVAILABLE:
+    import habana_frameworks.torch.hpu as torch_hpu
+
 
 class HPUAccelerator(Accelerator):
     """Accelerator for HPU devices."""
@@ -52,13 +55,28 @@ def get_parallel_devices(devices: int) -> List[torch.device]:
 
     @staticmethod
     def auto_device_count() -> int:
-        """Get the devices when set to auto."""
-        # TODO(@kaushikb11): Update this when api is exposed by the Habana team
-        return 8
+        """Returns the number of HPU devices when the devices is set to auto."""
+        try:
+            return torch_hpu.device_count()
+        except (AttributeError, NameError):
+            rank_zero_debug("HPU `auto_device_count` failed, returning default count of 8.")
+            return 8
 
     @staticmethod
     def is_available() -> bool:
-        return _HPU_AVAILABLE
+        """Returns a bool indicating if HPU is currently available."""
+        try:
+            return torch_hpu.is_available()
+        except (AttributeError, NameError):
+            return False
+
+    @staticmethod
+    def get_device_name() -> str:
+        """Returns the name of the HPU device."""
+        try:
+            return torch_hpu.get_device_name()
+        except (AttributeError, NameError):
+            return ""
 
     @classmethod
     def register_accelerators(cls, accelerator_registry: Dict) -> None:

src/pytorch_lightning/strategies/hpu_parallel.py

@@ -32,7 +32,7 @@
 
 if _HPU_AVAILABLE:
     import habana_frameworks.torch.core as htcore
-    import habana_frameworks.torch.core.hccl  # noqa: F401
+    import habana_frameworks.torch.distributed.hccl  # noqa: F401
 
 log = logging.getLogger(__name__)
 

src/pytorch_lightning/strategies/single_hpu.py

@@ -24,7 +24,6 @@
 
 if _HPU_AVAILABLE:
     import habana_frameworks.torch.core as htcore
-    import habana_frameworks.torch.core.hccl  # noqa: F401
 
 
 class SingleHPUStrategy(SingleDeviceStrategy):

tests/tests_pytorch/accelerators/test_hpu.py

@@ -40,6 +40,11 @@ def test_availability():
     assert HPUAccelerator.is_available()
 
 
+@RunIf(hpu=True)
+def test_device_name():
+    assert HPUAccelerator.get_device_name() == "GAUDI"
+
+
 @pytest.mark.skipif(_HPU_AVAILABLE, reason="test requires non-HPU machine")
 def test_fail_if_no_hpus():
     with pytest.raises(MisconfigurationException, match="HPUAccelerator can not run on your system"):
@@ -239,6 +244,7 @@ def test_inference_only(tmpdir, hpus):
     trainer.predict(model)
 
 
+@RunIf(hpu=True)
 def test_hpu_auto_device_count():
     assert HPUAccelerator.auto_device_count() == 8