Raise an error if batch transfer hooks are overridden with IPUAccelerator (#13961)

Co-authored-by: Adrian Wälchli <[email protected]>

Author: rohitgr7

docs/source-pytorch/accelerators/gpu_intermediate.rst

@@ -47,8 +47,9 @@ after which the root node will aggregate the results.
     :doc:`Manual Optimization <../model/manual_optimization>` with DP. Use DDP which is more stable and at least 3x faster.
 
 .. warning:: DP only supports scattering and gathering primitive collections of tensors like lists, dicts, etc.
-    Therefore the :meth:`~pytorch_lightning.core.hooks.ModelHooks.transfer_batch_to_device` hook does not apply in
-    this mode and if you have overridden it, it will not be called.
+    Therefore the hooks :meth:`~pytorch_lightning.core.hooks.ModelHooks.on_before_batch_transfer`,
+    :meth:`~pytorch_lightning.core.hooks.ModelHooks.transfer_batch_to_device` and :meth:`~pytorch_lightning.core.hooks.ModelHooks.on_after_batch_transfer`
+    do not apply in this mode and if you have overridden any of them, an exception will be raised.
 
 .. testcode::
     :skipif: torch.cuda.device_count() < 2

docs/source-pytorch/accelerators/ipu_basic.rst

@@ -67,3 +67,6 @@ Please see the `MNIST example <https://github.com/Lightning-AI/lightning/blob/ma
 * Since the step functions are traced, branching logic or any form of primitive values are traced into constants. Be mindful as this could lead to errors in your custom code.
 * Clipping gradients is not supported.
 * It is not possible to use :class:`torch.utils.data.BatchSampler` in your dataloaders if you are using multiple IPUs.
+* IPUs handles the data transfer to the device on the host, hence the hooks :meth:`~pytorch_lightning.core.hooks.ModelHooks.on_before_batch_transfer`,
+  :meth:`~pytorch_lightning.core.hooks.ModelHooks.transfer_batch_to_device` and :meth:`~pytorch_lightning.core.hooks.ModelHooks.on_after_batch_transfer`
+  do not apply here and if you have overridden any of them, an exception will be raised.

src/pytorch_lightning/CHANGELOG.md

@@ -19,7 +19,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - The `Trainer.{fit,validate,test,predict,tune}` methods now raise a useful error message if the input is not a `LightningModule` ([#13892](https://github.com/Lightning-AI/lightning/pull/13892))
 
 
--
+- Raised a `MisconfigurationException` if batch transfer hooks are overriden with `IPUAccelerator` ([13961](https://github.com/Lightning-AI/lightning/pull/13961))
 
 
 ### Deprecated

src/pytorch_lightning/core/hooks.py

@@ -665,6 +665,9 @@ def transfer_batch_to_device(self, batch, device, dataloader_idx):
             MisconfigurationException:
                 If using data-parallel, ``Trainer(strategy='dp')``.
 
+            MisconfigurationException:
+                If using IPUs, ``Trainer(accelerator='ipu')``.
+
         See Also:
             - :meth:`move_data_to_device`
             - :meth:`apply_to_collection`
@@ -700,6 +703,9 @@ def on_before_batch_transfer(self, batch, dataloader_idx):
             MisconfigurationException:
                 If using data-parallel, ``Trainer(strategy='dp')``.
 
+            MisconfigurationException:
+                If using IPUs, ``Trainer(accelerator='ipu')``.
+
         See Also:
             - :meth:`on_after_batch_transfer`
             - :meth:`transfer_batch_to_device`
@@ -735,6 +741,9 @@ def on_after_batch_transfer(self, batch, dataloader_idx):
             MisconfigurationException:
                 If using data-parallel, ``Trainer(strategy='dp')``.
 
+            MisconfigurationException:
+                If using IPUs, ``Trainer(accelerator='ipu')``.
+
         See Also:
             - :meth:`on_before_batch_transfer`
             - :meth:`transfer_batch_to_device`

src/pytorch_lightning/trainer/configuration_validator.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import pytorch_lightning as pl
+from pytorch_lightning.accelerators.ipu import IPUAccelerator
 from pytorch_lightning.plugins.precision.precision_plugin import PrecisionPlugin
 from pytorch_lightning.strategies import DataParallelStrategy
 from pytorch_lightning.trainer.states import TrainerFn
@@ -45,7 +46,7 @@ def verify_loop_configurations(trainer: "pl.Trainer") -> None:
     elif trainer.state.fn == TrainerFn.PREDICTING:
         __verify_eval_loop_configuration(trainer, model, "predict")
 
-    __verify_dp_batch_transfer_support(trainer, model)
+    __verify_batch_transfer_support(trainer, model)
     _check_deprecated_callback_hooks(trainer)
     # TODO: Delete _check_on_hpc_hooks in v1.8
     _check_on_hpc_hooks(model)
@@ -148,17 +149,22 @@ def __verify_eval_loop_configuration(trainer: "pl.Trainer", model: "pl.Lightning
             raise MisconfigurationException(f"No `{step_name}()` method defined to run `Trainer.{trainer_method}`.")
 
 
-def __verify_dp_batch_transfer_support(trainer: "pl.Trainer", model: "pl.LightningModule") -> None:
+def __verify_batch_transfer_support(trainer: "pl.Trainer", model: "pl.LightningModule") -> None:
     """Raise Misconfiguration exception since these hooks are not supported in DP mode."""
-    # TODO: Remove this blocker once batch transfer to device is integrated in Lightning for DP mode.
     batch_transfer_hooks = ("on_before_batch_transfer", "transfer_batch_to_device", "on_after_batch_transfer")
     datahook_selector = trainer._data_connector._datahook_selector
     for hook in batch_transfer_hooks:
+        # TODO: Remove this blocker once batch transfer to device is integrated in Lightning for DP mode.
         if isinstance(trainer.strategy, DataParallelStrategy) and (
             is_overridden(hook, datahook_selector.model) or is_overridden(hook, datahook_selector.datamodule)
         ):
             raise MisconfigurationException(f"Overriding `{hook}` is not supported in DP mode.")
 
+        if isinstance(trainer.accelerator, IPUAccelerator) and (
+            is_overridden(hook, datahook_selector.model) or is_overridden(hook, datahook_selector.datamodule)
+        ):
+            raise MisconfigurationException(f"Overriding `{hook}` is not supported with IPUs.")
+
 
 def __verify_manual_optimization_support(trainer: "pl.Trainer", model: "pl.LightningModule") -> None:
     if model.automatic_optimization:

tests/tests_pytorch/strategies/test_dp.py

@@ -11,20 +11,16 @@
 # 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 unittest import mock
 
-import pytest
 import torch
 import torch.nn.functional as F
 from torch.utils.data import DataLoader
 
 import pytorch_lightning as pl
 import tests_pytorch.helpers.pipelines as tpipes
 import tests_pytorch.helpers.utils as tutils
-from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks import EarlyStopping
 from pytorch_lightning.demos.boring_classes import BoringModel, RandomDataset
-from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from tests_pytorch.helpers.datamodules import ClassifDataModule
 from tests_pytorch.helpers.runif import RunIf
 from tests_pytorch.helpers.simple_models import ClassificationModel
@@ -154,47 +150,6 @@ def _assert_extra_outputs(self, outputs):
         assert out.dtype is torch.float
 
 
-@mock.patch("pytorch_lightning.utilities.device_parser.num_cuda_devices", return_value=2)
-@mock.patch("pytorch_lightning.utilities.device_parser.is_cuda_available", return_value=True)
-def test_dp_raise_exception_with_batch_transfer_hooks(mock_is_available, mock_device_count, tmpdir):
-    """Test that an exception is raised when overriding batch_transfer_hooks in DP model."""
-
-    class CustomModel(BoringModel):
-        def transfer_batch_to_device(self, batch, device, dataloader_idx):
-            batch = batch.to(device)
-            return batch
-
-    trainer_options = dict(default_root_dir=tmpdir, max_steps=7, accelerator="gpu", devices=[0, 1], strategy="dp")
-
-    trainer = Trainer(**trainer_options)
-    model = CustomModel()
-
-    with pytest.raises(MisconfigurationException, match=r"Overriding `transfer_batch_to_device` is not .* in DP"):
-        trainer.fit(model)
-
-    class CustomModel(BoringModel):
-        def on_before_batch_transfer(self, batch, dataloader_idx):
-            batch += 1
-            return batch
-
-    trainer = Trainer(**trainer_options)
-    model = CustomModel()
-
-    with pytest.raises(MisconfigurationException, match=r"Overriding `on_before_batch_transfer` is not .* in DP"):
-        trainer.fit(model)
-
-    class CustomModel(BoringModel):
-        def on_after_batch_transfer(self, batch, dataloader_idx):
-            batch += 1
-            return batch
-
-    trainer = Trainer(**trainer_options)
-    model = CustomModel()
-
-    with pytest.raises(MisconfigurationException, match=r"Overriding `on_after_batch_transfer` is not .* in DP"):
-        trainer.fit(model)
-
-
 @RunIf(min_cuda_gpus=2)
 def test_dp_training_step_dict(tmpdir):
     """This test verifies that dp properly reduces dictionaries."""

tests/tests_pytorch/trainer/test_config_validator.py

@@ -14,9 +14,11 @@
 import pytest
 import torch
 
+import pytorch_lightning as pl
 from pytorch_lightning import LightningDataModule, LightningModule, Trainer
 from pytorch_lightning.callbacks.callback import Callback
 from pytorch_lightning.demos.boring_classes import BoringDataModule, BoringModel, RandomDataset
+from pytorch_lightning.utilities import device_parser
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.warnings import PossibleUserWarning
 
@@ -192,3 +194,29 @@ def setup(self, pl_module, trainer):
 
     with pytest.raises(MisconfigurationException, match="does not have a `stage` argument"):
         trainer.fit(model)
+
+
+@pytest.mark.parametrize("trainer_kwargs", [{"accelerator": "ipu"}, {"accelerator": "gpu", "strategy": "dp"}])
+@pytest.mark.parametrize("hook", ["on_before_batch_transfer", "transfer_batch_to_device", "on_after_batch_transfer"])
+def test_raise_exception_with_batch_transfer_hooks(monkeypatch, hook, trainer_kwargs, tmpdir):
+    """Test that an exception is raised when overriding batch_transfer_hooks."""
+    if trainer_kwargs.get("accelerator") == "gpu":
+        match_pattern = rf"Overriding `{hook}` is not .* in DP mode."
+        monkeypatch.setattr(device_parser, "is_cuda_available", lambda: True)
+        monkeypatch.setattr(device_parser, "num_cuda_devices", lambda: 2)
+    elif trainer_kwargs.get("accelerator") == "ipu":
+        match_pattern = rf"Overriding `{hook}` is not .* with IPUs"
+        monkeypatch.setattr(pl.accelerators.ipu.IPUAccelerator, "is_available", lambda _: True)
+        monkeypatch.setattr(pl.strategies.ipu, "_IPU_AVAILABLE", lambda: True)
+
+    def custom_method(self, batch, *_, **__):
+        batch = batch + 1
+        return batch
+
+    trainer = Trainer(default_root_dir=tmpdir, **trainer_kwargs)
+
+    model = BoringModel()
+    setattr(model, hook, custom_method)
+
+    with pytest.raises(MisconfigurationException, match=match_pattern):
+        trainer.fit(model)