[SW-51190] Enable HPU accelerator on pytorch lightning 1.4.0

Change-Id: Ia8128a48ae0239cfc371c7cd983339219d29b0e6

Author: kalyanjkk

pl_examples/hpu_examples/mnist-test.py

@@ -0,0 +1,44 @@
+import os
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.utils.data import DataLoader, random_split
+from torchvision.datasets import MNIST
+from torchvision import transforms
+import pytorch_lightning as pl
+from pytorch_lightning.metrics.functional import accuracy
+import sys
+
+from habana_frameworks.torch.utils.library_loader import load_habana_module
+load_habana_module()
+
+class MNISTModel(pl.LightningModule):
+
+    def __init__(self):
+        super(MNISTModel, self).__init__()
+        self.l1 = torch.nn.Linear(28 * 28, 10)
+
+    def forward(self, x):
+        return torch.relu(self.l1(x.view(x.size(0), -1)))
+
+    def training_step(self, batch, batch_nb):
+        x, y = batch
+        loss = F.cross_entropy(self(x), y)
+        return loss
+
+    def configure_optimizers(self):
+        return torch.optim.Adam(self.parameters(), lr=0.02)
+
+# Init our model
+mnist_model = MNISTModel()
+
+# Init DataLoader from MNIST Dataset
+train_ds = MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor())
+train_loader = DataLoader(train_ds, batch_size=32)
+
+# Initialize a trainer
+trainer = pl.Trainer(hpus=1, max_epochs=3, progress_bar_refresh_rate=20)
+
+# Train the model ⚡
+trainer.fit(mnist_model, train_loader)

pytorch_lightning/accelerators/__init__.py

@@ -15,3 +15,4 @@
 from pytorch_lightning.accelerators.gpu import GPUAccelerator  # noqa F401
 from pytorch_lightning.accelerators.ipu import IPUAccelerator  # noqa F401
 from pytorch_lightning.accelerators.tpu import TPUAccelerator  # noqa F401
+from pytorch_lightning.accelerators.hpu import HPUAccelerator  # noqa F401

pytorch_lightning/accelerators/accelerator.py

@@ -44,6 +44,7 @@ class Accelerator:
     - CPU
     - GPU
     - TPU
+    - HPU
 
     Each Accelerator gets two plugins upon initialization:
     One to handle differences from the training routine and one to handle different precisions.

pytorch_lightning/accelerators/hpu.py

@@ -0,0 +1,50 @@
+# Copyright The PyTorch Lightning team.
+#
+# 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 logging
+import os
+from typing import Any
+
+import torch
+
+import pytorch_lightning as pl
+from pytorch_lightning.accelerators.accelerator import Accelerator
+from pytorch_lightning.plugins import DataParallelPlugin
+from pytorch_lightning.plugins.training_type.single_hpu import SingleHPUPlugin
+from pytorch_lightning.utilities.exceptions import MisconfigurationException
+
+_log = logging.getLogger(__name__)
+
+
+
+class HPUAccelerator(Accelerator):
+    """ Accelerator for HPU devices. """
+
+    def setup(self, trainer: 'pl.Trainer', model: 'pl.LightningModule') -> None:
+        """
+        Raises:
+            MisconfigurationException:
+                If the selected device is not HPU.
+        """
+        if "hpu" not in str(self.root_device):
+            raise MisconfigurationException(f"Device should be HPU, got {self.root_device} instead")
+        return super().setup(trainer, model)
+
+    def to_device(self, batch: Any) -> Any:
+        # no need to transfer batch to device in DP mode
+        # TODO: Add support to allow batch transfer to device in Lightning for DP mode.
+        #if isinstance(self.training_type_plugin, SingleHPUPlugin):
+        if not isinstance(self.training_type_plugin, DataParallelPlugin):
+            batch = super().to_device(batch)
+
+        return batch

pytorch_lightning/core/lightning.py

@@ -242,6 +242,14 @@ def on_gpu(self):
         """
         return self.device.type == "cuda"
 
+    @property
+    def on_hpu(self):
+        """
+        True if your model is currently running on HPUs.
+        Useful to set flags around the LightningModule for different CPU vs GPU vs HPU behavior.
+        """
+        return self.device.type == "hpu"
+
     @property
     def automatic_optimization(self) -> bool:
         """
@@ -1525,6 +1533,7 @@ def optimizer_step(
         optimizer_idx: int = None,
         optimizer_closure: Optional[Callable] = None,
         on_tpu: bool = None,
+        on_hpu: bool = None,
         using_native_amp: bool = None,
         using_lbfgs: bool = None,
     ) -> None:

pytorch_lightning/plugins/__init__.py

@@ -27,6 +27,7 @@
 from pytorch_lightning.plugins.training_type.sharded_spawn import DDPSpawnShardedPlugin  # noqa: F401
 from pytorch_lightning.plugins.training_type.single_device import SingleDevicePlugin  # noqa: F401
 from pytorch_lightning.plugins.training_type.single_tpu import SingleTPUPlugin  # noqa: F401
+from pytorch_lightning.plugins.training_type.single_hpu import SingleHPUPlugin  # noqa: F401
 from pytorch_lightning.plugins.training_type.tpu_spawn import TPUSpawnPlugin  # noqa: F401
 from pytorch_lightning.plugins.training_type.training_type_plugin import TrainingTypePlugin  # noqa: F401
 
@@ -49,6 +50,7 @@
     "FullyShardedNativeMixedPrecisionPlugin",
     "SingleDevicePlugin",
     "SingleTPUPlugin",
+    "SingleHPUPlugin",
     "TPUHalfPrecisionPlugin",
     "TPUSpawnPlugin",
     "TrainingTypePlugin",

pytorch_lightning/plugins/training_type/__init__.py

@@ -10,5 +10,6 @@
 from pytorch_lightning.plugins.training_type.sharded_spawn import DDPSpawnShardedPlugin  # noqa: F401
 from pytorch_lightning.plugins.training_type.single_device import SingleDevicePlugin  # noqa: F401
 from pytorch_lightning.plugins.training_type.single_tpu import SingleTPUPlugin  # noqa: F401
+from pytorch_lightning.plugins.training_type.single_hpu import SingleHPUPlugin  # noqa: F401
 from pytorch_lightning.plugins.training_type.tpu_spawn import TPUSpawnPlugin  # noqa: F401
 from pytorch_lightning.plugins.training_type.training_type_plugin import TrainingTypePlugin  # noqa: F401

pytorch_lightning/plugins/training_type/ddp.py

@@ -308,7 +308,7 @@ def configure_ddp(self):
         self._register_ddp_hooks()
 
     def determine_ddp_device_ids(self):
-        if self.root_device.type == "cpu":
+        if self.root_device.type == "cpu" or self.root_device.type == "hpu" :
             return None
         return [self.root_device.index]
 
@@ -317,6 +317,11 @@ def init_ddp_connection(self, global_rank: Optional[int] = None, world_size: Opt
         world_size = world_size if world_size is not None else self.cluster_environment.world_size()
         os.environ["MASTER_ADDR"] = self.cluster_environment.master_address()
         os.environ["MASTER_PORT"] = str(self.cluster_environment.master_port())
+
+        if self.torch_distributed_backend == 'hcl':
+            os.environ["ID"] = str(self.local_rank)
+            import habana_torch_hcl
+
         if torch.distributed.is_available() and not torch.distributed.is_initialized():
             log.info(f"initializing ddp: GLOBAL_RANK: {global_rank}, MEMBER: {global_rank + 1}/{world_size}")
             torch.distributed.init_process_group(

pytorch_lightning/plugins/training_type/single_device.py

@@ -33,6 +33,10 @@ def __init__(self, device: torch.device):
     def on_tpu(self) -> bool:
         return self.root_device.type == "xla" and _XLA_AVAILABLE
 
+    @property
+    def on_hpu(self) -> bool:
+        return self.device.type == "hpu"
+
     @property
     def on_gpu(self) -> bool:
         return self.root_device.type == "cuda" and torch.cuda.is_available()

pytorch_lightning/plugins/training_type/single_hpu.py

@@ -0,0 +1,61 @@
+# Copyright The PyTorch Lightning team.
+#
+# 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 torch
+
+from pytorch_lightning.plugins.training_type.single_device import SingleDevicePlugin
+from pytorch_lightning.utilities import _HPU_AVAILABLE
+from pytorch_lightning.utilities.apply_func import move_data_to_device
+
+
+class SingleHPUPlugin(SingleDevicePlugin):
+
+    def __init__(self, device: int):
+
+        device = torch.device("hpu")
+        super().__init__(device)
+
+        self.hpu_local_core_rank = 0
+        self.hpu_global_core_rank = 0
+
+    @property
+    def on_hpu(self) -> bool:
+        return True
+
+    def connect(self, model: torch.nn.Module) -> torch.nn.Module:
+        self._model = model
+        self.model_to_device()
+        return self._model
+
+    @property
+    def is_distributed(self) -> bool:
+        return False
+
+    def model_to_device(self) -> None:
+        self._model.to(self.root_device)
+
+    def pre_dispatch(self) -> None:
+        if isinstance(self.device, int):
+            self.device = torch.device(self.device)
+
+        self.hpu_local_core_rank = 0
+        self.hpu_global_core_rank = 0
+
+    def on_save(self, checkpoint: dict) -> dict:
+        """
+        Need to check how this part works
+        Move XLA tensors to CPU before saving
+        Recommended on XLA Guide:
+        https://github.com/pytorch/xla/blob/master/API_GUIDE.md#saving-and-loading-xla-tensors
+        """
+        return move_data_to_device(checkpoint, torch.device("cpu"))