Workaround for batch size search on xpu devices (#4513)

* provide workaround for XPU batch search

* return back parameters for MaskRCNN

* fix unit test

* switch off adaprive_bs by default

Author: kprokofi

src/otx/backend/native/callbacks/batchsize_finder.py

@@ -10,6 +10,8 @@
 from lightning.pytorch.callbacks import Callback
 from lightning.pytorch.loggers.logger import DummyLogger
 
+from otx.utils.device import is_xpu_available
+
 if TYPE_CHECKING:
     from lightning import LightningModule
     from lightning.pytorch.trainer import Trainer
@@ -53,13 +55,15 @@ def _try_loop_run(trainer: Trainer) -> None:
 def _scale_batch_reset_params(trainer: Trainer, steps_per_trial: int) -> None:
     trainer.logger = DummyLogger() if trainer.logger is not None else None
     trainer.callbacks = []
+    # For XPU devices 1 epoch sometimes is not enough to catch an error
+    max_epochs = 2 if is_xpu_available() else 1
 
     loop = trainer._active_loop  # noqa: SLF001
     if loop is None:
         msg = "There is no active loop."
         raise RuntimeError(msg)
     if trainer.fit_loop.epoch_loop.max_steps == -1:  # epoch based loop
-        trainer.fit_loop.max_epochs = 1
+        trainer.fit_loop.max_epochs = max_epochs
         trainer.limit_train_batches = steps_per_trial
     else:  # iter based loop
         trainer.fit_loop.epoch_loop.max_steps = steps_per_trial

src/otx/backend/native/engine.py

@@ -162,6 +162,7 @@ def train(
         adaptive_bs: Literal["None", "Safe", "Full"] = "None",
         check_val_every_n_epoch: int | None = 1,
         num_sanity_val_steps: int | None = 0,
+        log_every_n_steps: int | None = 1,
         **kwargs,
     ) -> dict[str, Any]:
         r"""Trains the model using the provided LightningModule and OTXDataModule.
@@ -245,6 +246,7 @@ def train(
             val_check_interval=val_check_interval,
             check_val_every_n_epoch=check_val_every_n_epoch,
             num_sanity_val_steps=num_sanity_val_steps,
+            log_every_n_steps=log_every_n_steps,
             **kwargs,
         )
         fit_kwargs: dict[str, Any] = {}

src/otx/backend/native/tools/adaptive_bs/algorithm.py

@@ -269,6 +269,7 @@ def _run_trial(train_func: Callable[[int], Any], bs: int, trial_queue: mp.Queue)
             or "XPU out of memory" in str(e)
             or "UR_RESULT_ERROR_OUT_OF_DEVICE_MEMORY" in str(e)
             or "UR error" in str(e)
+            or "UR_RESULT_ERROR_UNKNOWN" in str(e)
         ):  # XPU OOM
             oom = True
         else:

src/otx/backend/native/tools/adaptive_bs/runner.py

@@ -28,7 +28,6 @@
 def adapt_batch_size(
     engine: OTXEngine,
     not_increase: bool = True,
-    callbacks: list[Callback] | Callback | None = None,
     **train_args,
 ) -> None:
     """Change the actual batch size depending on the current GPU status.
@@ -39,7 +38,6 @@ def adapt_batch_size(
     Args:
         engine (OTXEngine): engine instnace.
         not_increase (bool) : Whether adapting batch size to larger value than default value or not.
-        callbacks (list[Callback] | Callback | None, optional): callbacks used during training. Defaults to None.
     """
     if not (is_cuda_available() or is_xpu_available()):
         msg = "Adaptive batch size supports only CUDA or XPU."
@@ -55,7 +53,7 @@ def adapt_batch_size(
             _apply_new_batch_size(engine, new_batch_size)
         return
 
-    train_func = partial(_train_model, engine=engine, callbacks=callbacks, **_adjust_train_args(train_args))
+    train_func = partial(_train_model, engine=engine, **_adjust_train_args(train_args))
     bs_search_algo = BsSearchAlgo(
         train_func=train_func,
         default_bs=default_bs,
@@ -85,11 +83,12 @@ def adapt_batch_size(
 def _adjust_train_args(train_args: dict[str, Any]) -> dict[str, Any]:
     train_args.update(train_args.pop("kwargs", {}))
     train_args.pop("self", None)
-    train_args.pop("adaptive_bs")
+    train_args.pop("adaptive_bs", None)
+    train_args.pop("callbacks", None)
     return train_args
 
 
-def _train_model(bs: int, engine: OTXEngine, callbacks: list[Callback] | Callback | None = None, **train_args) -> None:
+def _train_model(bs: int, engine: OTXEngine, **train_args) -> None:
     if bs <= 0:
         msg = f"Batch size should be greater than 0, but {bs} is given."
         raise ValueError(msg)
@@ -100,7 +99,8 @@ def _train_model(bs: int, engine: OTXEngine, callbacks: list[Callback] | Callbac
     engine.datamodule.val_subset.batch_size = bs
     engine.datamodule.test_subset.batch_size = bs
     train_args["adaptive_bs"] = "None"
-    engine.train(callbacks=_register_callback(callbacks), **train_args)
+    print(f"Runnning training trial with bs = {bs} ...")
+    engine.train(callbacks=_register_callback(), **train_args)
 
 
 def _register_callback(callbacks: list[Callback] | Callback | None = None) -> list[Callback]:
@@ -114,9 +114,13 @@ def _register_callback(callbacks: list[Callback] | Callback | None = None) -> li
 
 def _apply_new_batch_size(engine: OTXEngine, new_batch_size: int) -> None:
     origin_bs = engine.datamodule.train_subset.batch_size
+    if is_xpu_available() and new_batch_size != 1:
+        new_batch_size -= 1  # for safety reasons
     if new_batch_size == origin_bs:
         return
     engine.datamodule.train_subset.batch_size = new_batch_size
     engine.datamodule.val_subset.batch_size = new_batch_size
     engine.datamodule.test_subset.batch_size = new_batch_size
-    engine.model.optimizer_callable.optimizer_kwargs["lr"] *= sqrt(new_batch_size / origin_bs)  # type: ignore[attr-defined]
+    new_lr = engine.model.optimizer_callable.optimizer_kwargs["lr"] * sqrt(new_batch_size / origin_bs)  # type: ignore[attr-defined]
+    print(f"new batch size = {new_batch_size} with learning rate = {new_lr} is set for the training and validation.")
+    engine.model.optimizer_callable.optimizer_kwargs["lr"] = new_lr  # type: ignore[attr-defined]

src/otx/recipe/instance_segmentation/maskrcnn_r50_tv.yaml

@@ -57,7 +57,7 @@ overrides:
   data:
     train_subset:
       batch_size: 4
-      num_workers: 8
+      num_workers: 4
 
     val_subset:
       num_workers: 4

tests/unit/backend/native/tools/adaptive_bs/test_adaptive_bs_api.py

@@ -17,6 +17,7 @@
     _train_model,
     adapt_batch_size,
 )
+from otx.utils.device import is_xpu_available
 
 
 @pytest.fixture()
@@ -263,7 +264,7 @@ def test_on_fit_start(self, mock_trainer, mock_active_loop):
         # check steps_per_trial is set well
         assert mock_trainer.limit_val_batches == steps_per_trial
         assert mock_trainer.fit_loop.epoch_loop.max_steps == -1
-        assert mock_trainer.fit_loop.max_epochs == 1
+        assert mock_trainer.fit_loop.max_epochs == 1 if not is_xpu_available() else 2
         assert mock_trainer.limit_train_batches == steps_per_trial
         # check active_loop is run
         assert mock_active_loop.restarting is False
@@ -281,7 +282,7 @@ def test_on_fit_start_no_val(self, mock_trainer, mock_active_loop):
         # check steps_per_trial is set well
         assert mock_trainer.limit_val_batches == 0
         assert mock_trainer.fit_loop.epoch_loop.max_steps == -1
-        assert mock_trainer.fit_loop.max_epochs == 1
+        assert mock_trainer.fit_loop.max_epochs == 1 if not is_xpu_available() else 2
         assert mock_trainer.limit_train_batches == steps_per_trial
         # check active_loop is run
         assert mock_active_loop.restarting is False

tests/unit/tools/test_converter.py

@@ -114,3 +114,5 @@ def test_instantiate(self, tmp_path):
         if "logger" in train_kwargs and train_kwargs["logger"] is not None:
             assert len(train_kwargs["logger"]) == len(config["logger"])
         assert train_kwargs["max_epochs"] == 100
+        assert "adaptive_bs" in train_kwargs
+        assert train_kwargs["adaptive_bs"] == "Safe"