Restructure package

.azure/hpu-tests-fabric.yml

@@ -86,7 +86,7 @@ jobs:
 
     - task: PublishTestResults@2
       inputs:
-        testResultsFiles: 'tests/hpu*_test-results.xml'
+        testResultsFiles: 'tests/fabric_hpu*_test-results.xml'
         testRunTitle: '$(Build.DefinitionName) - Python $(python.version)'
       condition: succeededOrFailed()
       displayName: 'Publish test results'

.azure/hpu-tests-pl.yml

@@ -73,19 +73,19 @@ jobs:
       displayName: 'Check the driver status'
 
     - bash: |
-        python -m pytest -sv test_accelerator.py --forked --junitxml=hpu1_test-results.xml
+        python -m pytest -sv tests_pytorch/test_accelerator.py --forked --junitxml=hpu1_test-results.xml
       workingDirectory: tests/
       displayName: 'Single card HPU test'
 
     - bash: |
-        python -m pytest -sv test_accelerator.py --forked --hpus 8 --junitxml=hpu8_test-results.xml
+        python -m pytest -sv tests_pytorch/test_accelerator.py --forked --hpus 8 --junitxml=hpu8_test-results.xml
       workingDirectory: tests/
       displayName: 'Multi card(8) HPU test'
 
     - bash: |
-        python -m pytest -sv plugins/test_precision.py --hmp-bf16 \
-          'plugins/ops_bf16.txt' --hmp-fp32 \
-          'plugins/ops_fp32.txt' --forked \
+        python -m pytest -sv tests_pytorch/test_precision.py --hmp-bf16 \
+          'tests_pytorch/ops_bf16.txt' --hmp-fp32 \
+          'tests_pytorch/ops_fp32.txt' --forked \
           --junitxml=hpu1_precision_test-results.xml
       workingDirectory: tests/
       displayName: 'HPU precision test'
@@ -102,7 +102,7 @@ jobs:
 
     - task: PublishTestResults@2
       inputs:
-        testResultsFiles: 'tests/hpu*_test-results.xml'
+        testResultsFiles: 'tests/pl_hpu*_test-results.xml'
         testRunTitle: '$(Build.DefinitionName) - Python $(python.version)'
       condition: succeededOrFailed()
       displayName: 'Publish test results'

CHANGELOG.md

@@ -13,6 +13,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 -
 ### Changed
 
+- Changed code hierarchy in compliance with base lightning code for pytorch ([#12](https://github.com/Lightning-AI/lightning-Habana/pull/12))
+-
 ### Fixed
 
 ### Removed

_requirements/test.txt

@@ -1,5 +1,4 @@
 coverage>=5.0
-codecov>=2.1
 pytest>=6.0
 pytest-cov
 pytest-forked==1.6.0

examples/mnist_sample.py → examples/pytorch/mnist_sample.py

@@ -11,14 +11,22 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
 import torch
 from jsonargparse import lazy_instance
-from pytorch_lightning import LightningModule
-from pytorch_lightning.cli import LightningCLI
-from pytorch_lightning.demos.mnist_datamodule import MNISTDataModule
+from lightning_utilities import module_available
 from torch.nn import functional as F  # noqa: N812
 
-from lightning_habana.plugins.precision import HPUPrecisionPlugin
+if module_available("lightning"):
+    from lightning.pytorch import LightningModule
+    from lightning.pytorch.cli import LightningCLI
+    from lightning.pytorch.demos.mnist_datamodule import MNISTDataModule
+elif module_available("pytorch_lightning"):
+    from pytorch_lightning import LightningModule
+    from pytorch_lightning.cli import LightningCLI
+    from pytorch_lightning.demos.mnist_datamodule import MNISTDataModule
+
+from lightning_habana.pytorch.plugins.precision import HPUPrecisionPlugin
 
 
 class LitClassifier(LightningModule):
@@ -61,7 +69,7 @@ def configure_optimizers(self):
             "accelerator": "hpu",
             "devices": 1,
             "max_epochs": 1,
-            "plugins": lazy_instance(HPUPrecisionPlugin, precision="16-mixed"),
+            "plugins": lazy_instance(HPUPrecisionPlugin, precision="bf16-mixed"),
         },
         run=False,
         save_config_kwargs={"overwrite": True},

src/lightning_habana/__about__.py

@@ -1,10 +1,10 @@
-__version__ = "0.1.0rc1"
+__version__ = "0.2.0dev"
 __author__ = "Lightning-AI et al."
 __author_email__ = "name@lightning.ai"
 __license__ = "Apache-2.0"
-__copyright__ = f"Copyright (c) 2020-2022, {__author__}."
+__copyright__ = f"Copyright (c) 2020-2023, {__author__}."
 __homepage__ = "https://github.com/Lightning-AI/lightning-habana"
-__docs__ = "PyTorch Lightning Sample project."
+__docs__ = "Lightning suport for Intel Habana accelerators"
 
 __all__ = [
     "__author__",

src/lightning_habana/__init__.py

@@ -1,31 +1,16 @@
-"""Root package info."""
+# Copyright (c) 2023 Habana Labs, Ltd. an Intel Company
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 
-import os
-
-_PACKAGE_ROOT = os.path.dirname(__file__)
-_PROJECT_ROOT = os.path.dirname(_PACKAGE_ROOT)
-
-from lightning_utilities.core.imports import package_available  # noqa: E402
-
-_HABANA_FRAMEWORK_AVAILABLE = package_available("habana_frameworks")
-if _HABANA_FRAMEWORK_AVAILABLE:
-    from habana_frameworks.torch.utils.library_loader import is_habana_available
-
-    _HPU_AVAILABLE: bool = is_habana_available()
-else:
-    _HPU_AVAILABLE = False
-
-from lightning_habana.__about__ import *  # noqa: E402, F401, F403
-from lightning_habana.accelerator import HPUAccelerator  # noqa: E402
-from lightning_habana.plugins.io_plugin import HPUCheckpointIO  # noqa: E402
-from lightning_habana.plugins.precision import HPUPrecisionPlugin  # noqa: E402
-from lightning_habana.strategies.parallel import HPUParallelStrategy  # noqa: E402
-from lightning_habana.strategies.single import SingleHPUStrategy  # noqa: E402
-
-__all__ = [
-    "HPUAccelerator",
-    "HPUParallelStrategy",
-    "SingleHPUStrategy",
-    "HPUPrecisionPlugin",
-    "HPUCheckpointIO",
-]
+from lightning_habana.__about__ import *  # noqa: F401, F403
+from lightning_habana.utils import *  # noqa: F401, F403

src/lightning_habana/pytorch/__init__.py

@@ -0,0 +1,28 @@
+# Copyright (c) 2023 Habana Labs, Ltd. an Intel Company
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from lightning_habana.__about__ import *  # noqa: F401, F403
+from lightning_habana.pytorch.accelerator import HPUAccelerator
+from lightning_habana.pytorch.plugins.io_plugin import HPUCheckpointIO
+from lightning_habana.pytorch.plugins.precision import HPUPrecisionPlugin
+from lightning_habana.pytorch.strategies.parallel import HPUParallelStrategy
+from lightning_habana.pytorch.strategies.single import SingleHPUStrategy
+
+__all__ = [
+    "HPUAccelerator",
+    "HPUParallelStrategy",
+    "SingleHPUStrategy",
+    "HPUPrecisionPlugin",
+    "HPUCheckpointIO",
+]

src/lightning_habana/pytorch/accelerator/__init__.py

@@ -0,0 +1,17 @@
+# Copyright (c) 2023 Habana Labs, Ltd. an Intel Company
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from lightning_habana.pytorch.accelerator.hpu import HPUAccelerator
+
+__all__ = ["HPUAccelerator"]

src/lightning_habana/accelerator.py → src/lightning_habana/pytorch/accelerator/hpu.py

@@ -30,7 +30,7 @@
 else:
     raise ModuleNotFoundError("You are missing `lightning` or `pytorch-lightning` package, please install it.")
 
-from lightning_habana import _HPU_AVAILABLE
+from lightning_habana.utils.imports import _HPU_AVAILABLE
 
 if _HPU_AVAILABLE:
     import habana_frameworks.torch.hpu as torch_hpu

src/lightning_habana/plugins/__init__.py → src/lightning_habana/pytorch/plugins/__init__.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from lightning_habana.plugins.io_plugin import HPUCheckpointIO
-from lightning_habana.plugins.precision import HPUPrecisionPlugin
+from lightning_habana.pytorch.plugins.io_plugin import HPUCheckpointIO
+from lightning_habana.pytorch.plugins.precision import HPUPrecisionPlugin
 
 __all__ = ["HPUPrecisionPlugin", "HPUCheckpointIO"]

src/lightning_habana/plugins/precision.py → src/lightning_habana/pytorch/plugins/precision.py

@@ -23,12 +23,12 @@
 else:
     raise ModuleNotFoundError("You are missing `lightning` or `pytorch-lightning` package, please install it.")
 
-from lightning_habana import _HPU_AVAILABLE
+from lightning_habana.utils.imports import _HPU_AVAILABLE
 
 if _HPU_AVAILABLE:
     from habana_frameworks.torch.hpex import hmp
 
-_PRECISION_INPUT = Literal["32-true", "16-mixed", "bf16-mixed"]
+_PRECISION_INPUT = Literal["32", "32-true", "bf16", "bf16-mixed"]
 
 
 class HPUPrecisionPlugin(PrecisionPlugin):

src/lightning_habana/strategies/__init__.py → src/lightning_habana/pytorch/strategies/__init__.py

@@ -11,7 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from lightning_habana.strategies.parallel import HPUParallelStrategy
-from lightning_habana.strategies.single import SingleHPUStrategy
+
+from lightning_habana.pytorch.strategies.parallel import HPUParallelStrategy
+from lightning_habana.pytorch.strategies.single import SingleHPUStrategy
 
 __all__ = ["HPUParallelStrategy", "SingleHPUStrategy"]

src/lightning_habana/strategies/parallel.py → src/lightning_habana/pytorch/strategies/parallel.py

@@ -17,13 +17,13 @@
 
 import torch.distributed
 from lightning_utilities import module_available
-from torch.distributed import broadcast_object_list
 
 if module_available("lightning"):
     from lightning.fabric.plugins import CheckpointIO, ClusterEnvironment
     from lightning.fabric.utilities.distributed import group as _group
     from lightning.pytorch import LightningModule
     from lightning.pytorch.accelerators import Accelerator
+    from lightning.pytorch.plugins.io.hpu_plugin import HPUCheckpointIO
     from lightning.pytorch.plugins.io.wrapper import _WrappingCheckpointIO
     from lightning.pytorch.plugins.precision import PrecisionPlugin
     from lightning.pytorch.strategies.ddp import DDPStrategy
@@ -41,7 +41,7 @@
 from torch.nn import Module
 from torch.optim.optimizer import Optimizer
 
-from lightning_habana import _HPU_AVAILABLE
+from lightning_habana.utils.imports import _HPU_AVAILABLE
 
 if _HPU_AVAILABLE:
     import habana_frameworks.torch.core as htcore
@@ -94,7 +94,7 @@ def checkpoint_io(self) -> CheckpointIO:
 
     @checkpoint_io.setter
     def checkpoint_io(self, io: Optional[CheckpointIO]) -> None:
-        self._checkpoint_io = io  # type: ignore[assignment]
+        self._checkpoint_io = io
 
     def setup_environment(self) -> None:
         os.environ["ID"] = str(self.local_rank)
@@ -111,7 +111,7 @@ def broadcast(self, obj: object, src: int = 0) -> object:
         if self.global_rank != src:
             obj = [None]
 
-        broadcast_object_list(obj, src, group=_group.WORLD)
+        _hpu_broadcast_object_list(obj, src, group=_group.WORLD)
         return obj[0]
 
     def on_after_backward(self) -> None:
@@ -143,3 +143,80 @@ def teardown(self) -> None:
         # Was set to local rank
         os.environ.pop("ID", None)
         os.environ.pop("HCCL_DISTRIBUTED_BACKEND", None)
+
+
+# The code underneath is taken from PyTorch `torch/distributed/distributed_c10d.py`
+# the distributed backend and tensor type updates for habana backend is done here before broadcast
+def _hpu_broadcast_object_list(object_list, src=0, group=None, device=None):  # type: ignore
+    from torch.distributed import Backend, _rank_not_in_group, broadcast, get_backend, get_rank
+    from torch.distributed.distributed_c10d import _object_to_tensor, _tensor_to_object
+
+    if _rank_not_in_group(group):
+        return
+
+    my_rank = get_rank()
+    # Serialize object_list elements to tensors on src rank.
+    if my_rank == src:
+        tensor_list, size_list = zip(*[_object_to_tensor(obj, device) for obj in object_list])
+        object_sizes_tensor = torch.cat(size_list)
+    else:
+        object_sizes_tensor = torch.empty(len(object_list), dtype=torch.long)
+
+    # Current device selection.
+    # To preserve backwards compatibility, ``device`` is default to ``None``
+    # in which case we run current logic of device selection, i.e.
+    # ``current_device`` is CUDA if backend is NCCL otherwise CPU device. In the
+    # case it is not ``None`` we move the size and object tensors to be
+    # broadcasted to this device.
+    group_backend = get_backend(group)
+    is_nccl_backend = group_backend == Backend.NCCL
+    is_hpu_backend = os.environ.get("HCCL_DISTRIBUTED_BACKEND") == "1"
+    if device is not None:
+        if is_nccl_backend and device.type != "cuda":
+            raise ValueError("device type must be cuda for nccl backend")
+        current_device = device
+    else:
+        current_device = torch.device("cpu")
+        if is_nccl_backend:
+            # See note about using torch.cuda.current_device() here in
+            # docstring. We cannot simply use my_rank since rank == device is
+            # not necessarily true.
+            current_device = torch.device("cuda", torch.cuda.current_device())
+    if is_nccl_backend:
+        object_sizes_tensor = object_sizes_tensor.to(current_device)
+
+    elif is_hpu_backend:
+        current_device = torch.device("hpu")
+        # Workaround: HPU doesn't not support long tensors for collectives
+        if (object_sizes_tensor.type() == "torch.LongTensor") or (object_sizes_tensor.type() == "torch.hpu.LongTensor"):
+            object_sizes_tensor = object_sizes_tensor.int()
+        else:
+            print("unhandled hpu object_sizes_tensor type :: ", object_sizes_tensor.type())
+        object_sizes_tensor = object_sizes_tensor.to(current_device)
+
+    # Broadcast object sizes
+    broadcast(object_sizes_tensor, src=src, group=group)
+
+    # Concatenate and broadcast serialized object tensors
+    if my_rank == src:
+        object_tensor = torch.cat(tensor_list)
+    else:
+        object_tensor = torch.empty(
+            torch.sum(object_sizes_tensor).int().item(),
+            dtype=torch.uint8,
+        )
+
+    if is_nccl_backend or is_hpu_backend:
+        object_tensor = object_tensor.to(current_device)
+
+    broadcast(object_tensor, src=src, group=group)
+    # Deserialize objects using their stored sizes.
+    offset = 0
+    if my_rank != src:
+        for i, obj_size in enumerate(object_sizes_tensor):
+            obj_view = object_tensor[offset : offset + obj_size]
+            obj_view = obj_view.type(torch.uint8)
+            if obj_view.device != torch.device("cpu"):
+                obj_view = obj_view.cpu()
+            offset += obj_size
+            object_list[i] = _tensor_to_object(obj_view, obj_size)

src/lightning_habana/strategies/single.py → src/lightning_habana/pytorch/strategies/single.py

@@ -21,6 +21,7 @@
     from lightning.fabric.utilities.types import _DEVICE
     from lightning.pytorch import LightningModule, Trainer
     from lightning.pytorch.accelerators import Accelerator
+    from lightning.pytorch.plugins.io.hpu_plugin import HPUCheckpointIO
     from lightning.pytorch.plugins.io.wrapper import _WrappingCheckpointIO
     from lightning.pytorch.plugins.precision import PrecisionPlugin
     from lightning.pytorch.strategies.single_device import SingleDeviceStrategy
@@ -39,7 +40,7 @@
 from torch.nn import Module
 from torch.optim.optimizer import Optimizer
 
-from lightning_habana import _HPU_AVAILABLE
+from lightning_habana.utils.imports import _HPU_AVAILABLE
 
 if _HPU_AVAILABLE:
     import habana_frameworks.torch.core as htcore
@@ -75,7 +76,7 @@ def checkpoint_io(self) -> CheckpointIO:
 
     @checkpoint_io.setter
     def checkpoint_io(self, io: Optional[CheckpointIO]) -> None:
-        self._checkpoint_io = io  # type: ignore[assignment]
+        self._checkpoint_io = io
 
     @property
     def is_distributed(self) -> bool: