Add auto_device_count and device name support (#13423)

Co-authored-by: Kaushik B <[email protected]>
Co-authored-by: Rohit Gupta <[email protected]>
Co-authored-by: Justus Schock <[email protected]>
Co-authored-by: Jirka <[email protected]>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Akihiro Nitta <[email protected]>
Co-authored-by: awaelchli <[email protected]>
Co-authored-by: Carlos Mocholí <[email protected]>
Co-authored-by: Adrian Wälchli <[email protected]>
Co-authored-by: ananthsub <[email protected]>
Co-authored-by: mansy <[email protected]>
Co-authored-by: manskx <[email protected]>
Co-authored-by: Jirka Borovec <[email protected]>
Co-authored-by: Mansy <[email protected]>
Co-authored-by: otaj <[email protected]>
Co-authored-by: Sean Naren <[email protected]>
Co-authored-by: Keiichi Kuroyanagi <[email protected]>
Co-authored-by: Martino Sorbaro <[email protected]>
Co-authored-by: Wang Ran (汪然) <[email protected]>
Co-authored-by: Rhys Goodall <[email protected]>
Co-authored-by: Siyuan Li <[email protected]>
Co-authored-by: Ekagra Ranjan <[email protected]>
Co-authored-by: S. Kumano <[email protected]>
Co-authored-by: otaj <[email protected]>
Co-authored-by: Gautier Dagan <[email protected]>
Co-authored-by: Sherin Thomas <[email protected]>
Co-authored-by: Cyprien Ricque <[email protected]>
Co-authored-by: Masahiro Wada <[email protected]>
Co-authored-by: nitinramvelraj <[email protected]>
Co-authored-by: donlapark <[email protected]>
Co-authored-by: Justin Goheen <[email protected]>
Co-authored-by: Shantam Gilra <[email protected]>
Co-authored-by: Bibhabasu Mohapatra <[email protected]>
Co-authored-by: Jimmy Yao <[email protected]>
Co-authored-by: Nikhil Shenoy <[email protected]>
Co-authored-by: Sanjay Aradhyamath <[email protected]>

Author: 35 people

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