nshtrainer 0.31.0__py3-none-any.whl → 0.32.1__py3-none-any.whl

nshtrainer/callbacks/__init__.py

@@ -12,6 +12,8 @@ from .checkpoint import OnExceptionCheckpoint as OnExceptionCheckpoint
 from .checkpoint import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
 )
+from .debug_flag import DebugFlagCallback as DebugFlagCallback
+from .debug_flag import DebugFlagCallbackConfig as DebugFlagCallbackConfig
 from .directory_setup import DirectorySetupCallback as DirectorySetupCallback
 from .directory_setup import DirectorySetupConfig as DirectorySetupConfig
 from .early_stopping import EarlyStopping as EarlyStopping
@@ -41,7 +43,8 @@ from .wandb_watch import WandbWatchCallback as WandbWatchCallback
 from .wandb_watch import WandbWatchConfig as WandbWatchConfig
 
 CallbackConfig = Annotated[
-    EarlyStoppingConfig
+    DebugFlagCallbackConfig
+    | EarlyStoppingConfig
     | ThroughputMonitorConfig
     | EpochTimerConfig
     | PrintTableMetricsConfig

nshtrainer/callbacks/debug_flag.py

@@ -0,0 +1,70 @@
+import logging
+from typing import TYPE_CHECKING, Literal, cast
+
+from lightning.pytorch import LightningModule, Trainer
+from lightning.pytorch.callbacks import Callback
+from typing_extensions import override
+
+from .base import CallbackConfigBase
+
+if TYPE_CHECKING:
+    from ..model.config import BaseConfig
+
+log = logging.getLogger(__name__)
+
+
+class DebugFlagCallbackConfig(CallbackConfigBase):
+    name: Literal["debug_flag"] = "debug_flag"
+
+    enabled: bool = True
+    """Whether to enable the callback."""
+
+    def __bool__(self):
+        return self.enabled
+
+    @override
+    def create_callbacks(self, root_config):
+        if not self:
+            return
+
+        yield DebugFlagCallback(self)
+
+
+class DebugFlagCallback(Callback):
+    """
+    Sets the debug flag to true in the following circumstances:
+    - fast_dev_run is enabled
+    - sanity check is running
+    """
+
+    @override
+    def __init__(self, config: DebugFlagCallbackConfig):
+        super().__init__()
+
+        self.config = config
+        del config
+
+    @override
+    def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str):
+        if not getattr(trainer, "fast_dev_run", False):
+            return
+
+        hparams = cast("BaseConfig", pl_module.hparams)
+        if not hparams.debug:
+            log.critical("Fast dev run detected, setting debug flag to True.")
+        hparams.debug = True
+
+    @override
+    def on_sanity_check_start(self, trainer: Trainer, pl_module: LightningModule):
+        hparams = cast("BaseConfig", pl_module.hparams)
+        self._debug = hparams.debug
+        if not self._debug:
+            log.critical("Enabling debug flag during sanity check routine.")
+        hparams.debug = True
+
+    @override
+    def on_sanity_check_end(self, trainer: Trainer, pl_module: LightningModule):
+        hparams = cast("BaseConfig", pl_module.hparams)
+        if not self._debug:
+            log.critical("Sanity check routine complete, disabling debug flag.")
+        hparams.debug = self._debug

nshtrainer/config.py

@@ -11,6 +11,7 @@ from nshtrainer.callbacks.checkpoint._base import BaseCheckpointCallbackConfig a
 from nshtrainer.callbacks.checkpoint.best_checkpoint import BestCheckpointCallbackConfig as BestCheckpointCallbackConfig
 from nshtrainer.callbacks.checkpoint.last_checkpoint import LastCheckpointCallbackConfig as LastCheckpointCallbackConfig
 from nshtrainer.callbacks.checkpoint.on_exception_checkpoint import OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig
+from nshtrainer.callbacks.debug_flag import DebugFlagCallbackConfig as DebugFlagCallbackConfig
 from nshtrainer.callbacks.directory_setup import DirectorySetupConfig as DirectorySetupConfig
 from nshtrainer.callbacks.early_stopping import EarlyStoppingConfig as EarlyStoppingConfig
 from nshtrainer.callbacks.ema import EMAConfig as EMAConfig

nshtrainer/loggers/wandb.py

@@ -7,8 +7,8 @@ from lightning.pytorch import Callback, LightningModule, Trainer
 from packaging import version
 from typing_extensions import override
 
-from ..callbacks import WandbWatchConfig
 from ..callbacks.base import CallbackConfigBase
+from ..callbacks.wandb_watch import WandbWatchConfig
 from ._base import BaseLoggerConfig
 
 if TYPE_CHECKING:

nshtrainer/model/base.py

@@ -7,8 +7,7 @@ from typing import IO, TYPE_CHECKING, Any, Generic, Literal, cast
 import torch
 import torch.distributed
 from lightning.fabric.utilities.types import _MAP_LOCATION_TYPE, _PATH
-from lightning.pytorch import LightningModule, Trainer
-from lightning.pytorch.callbacks import Callback
+from lightning.pytorch import LightningModule
 from lightning.pytorch.profilers import PassThroughProfiler, Profiler
 from lightning.pytorch.utilities.types import STEP_OUTPUT
 from typing_extensions import Self, TypeVar, override
@@ -16,7 +15,6 @@ from typing_extensions import Self, TypeVar, override
 from ..callbacks.rlp_sanity_checks import _RLPSanityCheckModuleMixin
 from ..util._environment_info import EnvironmentConfig
 from .config import BaseConfig
-from .mixins.callback import CallbackModuleMixin
 from .mixins.logger import LoggerLightningModuleMixin
 
 log = logging.getLogger(__name__)
@@ -24,39 +22,6 @@ log = logging.getLogger(__name__)
 THparams = TypeVar("THparams", bound=BaseConfig, infer_variance=True)
 
 
-class DebugFlagCallback(Callback):
-    """
-    Sets the debug flag to true in the following circumstances:
-    - fast_dev_run is enabled
-    - sanity check is running
-    """
-
-    @override
-    def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str):
-        if not getattr(trainer, "fast_dev_run", False):
-            return
-
-        hparams = cast(BaseConfig, pl_module.hparams)
-        if not hparams.debug:
-            log.critical("Fast dev run detected, setting debug flag to True.")
-        hparams.debug = True
-
-    @override
-    def on_sanity_check_start(self, trainer: Trainer, pl_module: LightningModule):
-        hparams = cast(BaseConfig, pl_module.hparams)
-        self._debug = hparams.debug
-        if not self._debug:
-            log.critical("Enabling debug flag during sanity check routine.")
-        hparams.debug = True
-
-    @override
-    def on_sanity_check_end(self, trainer: Trainer, pl_module: LightningModule):
-        hparams = cast(BaseConfig, pl_module.hparams)
-        if not self._debug:
-            log.critical("Sanity check routine complete, disabling debug flag.")
-        hparams.debug = self._debug
-
-
 T = TypeVar("T", infer_variance=True)
 
 ReduceOpStr = Literal[
@@ -88,7 +53,6 @@ VALID_REDUCE_OPS = (
 class LightningModuleBase(  # pyright: ignore[reportIncompatibleMethodOverride]
     _RLPSanityCheckModuleMixin,
     LoggerLightningModuleMixin,
-    CallbackModuleMixin,
     LightningModule,
     ABC,
     Generic[THparams],
@@ -288,10 +252,8 @@ class LightningModuleBase(  # pyright: ignore[reportIncompatibleMethodOverride]
         hparams.environment = EnvironmentConfig.from_current_environment(hparams, self)
         hparams = self.pre_init_update_hparams(hparams)
 
-        super().__init__(hparams)
-
+        super().__init__()
         self.save_hyperparameters(hparams)
-        self.register_callback(lambda: DebugFlagCallback())
 
     def zero_loss(self):
         """

nshtrainer/trainer/_config.py

@@ -35,6 +35,7 @@ from ..callbacks import (
     OnExceptionCheckpointCallbackConfig,
 )
 from ..callbacks.base import CallbackConfigBase
+from ..callbacks.debug_flag import DebugFlagCallbackConfig
 from ..callbacks.rlp_sanity_checks import RLPSanityChecksConfig
 from ..callbacks.shared_parameters import SharedParametersConfig
 from ..loggers import (
@@ -751,6 +752,11 @@ class TrainerConfig(C.Config):
     """If enabled, will automatically set the default root dir to [cwd/lightning_logs/<id>/]. There is basically no reason to disable this."""
     save_checkpoint_metadata: bool = True
     """If enabled, will save additional metadata whenever a checkpoint is saved."""
+    auto_set_debug_flag: DebugFlagCallbackConfig | None = DebugFlagCallbackConfig()
+    """If enabled, will automatically set the debug flag to True if:
+    - The trainer is running in fast_dev_run mode.
+    - The trainer is running a sanity check (which happens before starting the training routine).
+    """
 
     lightning_kwargs: LightningTrainerKwargs = LightningTrainerKwargs()
     """
@@ -775,4 +781,7 @@ class TrainerConfig(C.Config):
         yield self.logging
         yield self.optimizer
         yield self.hf_hub
+        yield self.shared_parameters
+        yield self.reduce_lr_on_plateau_sanity_checking
+        yield self.auto_set_debug_flag
         yield from self.callbacks

{nshtrainer-0.31.0.dist-info → nshtrainer-0.32.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.31.0
+Version: 0.32.1
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-0.31.0.dist-info → nshtrainer-0.32.1.dist-info}/RECORD

@@ -6,7 +6,7 @@ nshtrainer/_checkpoint/saver.py,sha256=MbX_WjkDtHHAf9Ms-KXDlknkjiPXVoGIe2ciO28Ad
 nshtrainer/_directory.py,sha256=RjnW6vKTeKlz2vQWT3cG0Jje5BkFXA7HpUubDhcSiq4,2993
 nshtrainer/_experimental/__init__.py,sha256=pEXPyI184UuDHvfh4p9Kg9nQZQZI41e4_HvNd4BK-yg,81
 nshtrainer/_hf_hub.py,sha256=0bkXkqhve5D1onMW-fCfuvVKlTn0i6jv_6uMNgZ7OHQ,12974
-nshtrainer/callbacks/__init__.py,sha256=NpEV8bMU12ClFN2sLKLBDXnuwIHYyZOCNxDZgjrV104,2892
+nshtrainer/callbacks/__init__.py,sha256=1SBLpMsx7BzgimO35MwQViYBcbgxlkyvTMz1JKUKK-0,3060
 nshtrainer/callbacks/_throughput_monitor_callback.py,sha256=aJo_11rc4lo0IYOd-kHmPDtzdC4ctgXyRudkRJqH4m4,23184
 nshtrainer/callbacks/actsave.py,sha256=qbnaKts4_dvjPeAaPtv7Ds12_vEWzaHUfg_--49NB9I,4041
 nshtrainer/callbacks/base.py,sha256=NpjeKmonJ1Kaz5_39XSn3LlDwvbGjk6WV8BpHSNCvI4,3508
@@ -15,6 +15,7 @@ nshtrainer/callbacks/checkpoint/_base.py,sha256=vvlwuD-20NozYVIolGGShmUdkkNYeuwN
 nshtrainer/callbacks/checkpoint/best_checkpoint.py,sha256=8BHgLAd3Tuzf5sup0guEAKF1jJiAwYsjdKBFYZw98ac,2171
 nshtrainer/callbacks/checkpoint/last_checkpoint.py,sha256=CWWv0cSwQ1VAX26N7hAyMxbNCk26Keh39oQguBEK5To,1102
 nshtrainer/callbacks/checkpoint/on_exception_checkpoint.py,sha256=ctT88EGT22_t_6tr5r7Sfo43cuve6XeroBnBYRMPOus,3372
+nshtrainer/callbacks/debug_flag.py,sha256=T7rkY9hYQ_-PsPo2XiQ4eVZ9bBsTd2knpZWctCbjxXc,2011
 nshtrainer/callbacks/directory_setup.py,sha256=c0uY0oTqLcQ3egInHO7G6BeQQgk_xvOLoHH8FR-9U0U,2629
 nshtrainer/callbacks/early_stopping.py,sha256=VWuJz0oN87b6SwBeVc32YNpeJr1wts8K45k8JJJmG9I,4617
 nshtrainer/callbacks/ema.py,sha256=8-WHmKFP3VfnzMviJaIFmVD9xHPqIPmq9NRF5xdu3c8,12131
@@ -29,7 +30,7 @@ nshtrainer/callbacks/shared_parameters.py,sha256=fqlDweFDXPV_bfcAWpRgaJIad9i5Aeh
 nshtrainer/callbacks/throughput_monitor.py,sha256=H_ocXErZxUO3dxFk8Tx_VQdpI9E_Ztvqof5WtFevLyQ,1838
 nshtrainer/callbacks/timer.py,sha256=quS79oYClDUvQxJkNWmDMe0hwRUkkREgTgqzVrnom50,4607
 nshtrainer/callbacks/wandb_watch.py,sha256=Y6SEXfIx3kDDQbI5zpP53BVq0FBLJbLd3RJsiHZk1-Y,2921
-nshtrainer/config.py,sha256=W6nAmn5Y1GVZto9vkx4v8i5XdikMSdVYDiq7kbDEWAg,5900
+nshtrainer/config.py,sha256=HJWKMFGNFHmuk92KlpYpEIhY01Ysnqr4HOWx4npGVH0,5995
 nshtrainer/data/__init__.py,sha256=7mk1tr7SWUZ7ySbsf0y0ZPszk7u4QznPhQ-7wnpH9ec,149
 nshtrainer/data/balanced_batch_sampler.py,sha256=dGBTDDtlBU6c-ZlVQOCnTW7SjTB5hczWsOWEdUWjvkA,4385
 nshtrainer/data/transform.py,sha256=6SNs3_TpNpfhcwTwvPKyEJ3opM1OT7LmMEYQNHKgRl8,2227
@@ -54,7 +55,7 @@ nshtrainer/loggers/__init__.py,sha256=C_xk0A3_qKbNdTmzK85AgjRHFD3w-jPRS2ig-iPhfE
 nshtrainer/loggers/_base.py,sha256=xiZKEK0ALJkcqf4OpVNRY0QbZsamR_WR7x7m_68YHXQ,705
 nshtrainer/loggers/csv.py,sha256=D_lYyd94bZ8jAgnRo-ARtFgVcInaD9zktxtsUD9RWCI,1052
 nshtrainer/loggers/tensorboard.py,sha256=wL2amRSdP68zbslZvBeM0ZQBnjF3hIKsz-_lBbdomaM,2216
-nshtrainer/loggers/wandb.py,sha256=8B2BMMzILRSUEiCkmp_fBpcXs69euRKViTiaV__DJZk,5128
+nshtrainer/loggers/wandb.py,sha256=C-yGX9e2FUSfbUxur7-meNUjpB3D8hIdVCOgPzGm3QM,5140
 nshtrainer/lr_scheduler/__init__.py,sha256=uEvgaFAs-4s_bAEMaildy0GT6OvgpgOEKTuzqutESHE,736
 nshtrainer/lr_scheduler/_base.py,sha256=7xOIuxQ86YHbFWG5a3gX46emQj1WN_LaY4-i0Q1TDBg,3659
 nshtrainer/lr_scheduler/linear_warmup_cosine.py,sha256=YQm84Sb4SWrofpBwa39DCslJvu2uorjbpWaGWyys1l4,5352
@@ -62,9 +63,8 @@ nshtrainer/lr_scheduler/reduce_lr_on_plateau.py,sha256=h76oTHYpMxauV_l6lviya5DW-
 nshtrainer/metrics/__init__.py,sha256=ObLIELGguIEcUpRsUkqh1ltrvZii6vglTpJGrPvoy00,50
 nshtrainer/metrics/_config.py,sha256=jgRBfDAQLFTW7AiUY7CRtdfts6CR6keeuqm0FFMWCzQ,1288
 nshtrainer/model/__init__.py,sha256=2i_VEy6u_Y1LUGKljHXWeekvhnUcanZM2QyaaBM1Bmw,261
-nshtrainer/model/base.py,sha256=hT27FtzwKQiEL0C8RcaTKYXlanfvzTxHOJpHUcWiItk,19891
+nshtrainer/model/base.py,sha256=1zVY8ybZTzVKhpp7sUC0t360Ut3YmdGxAW5PZAIBSyw,18535
 nshtrainer/model/config.py,sha256=Q4Wong6w3cp_Sq7s8iZdABKF-LZBbSCFn_TQPYkhkrI,6572
-nshtrainer/model/mixins/callback.py,sha256=lh3imlw1H3ESIG4WFA5frooSlWi6-RPUUDRFGRzEg4A,8571
 nshtrainer/model/mixins/logger.py,sha256=xOymSTofukEYZGkGojXsMEO__ZlBI5lIPZVmlotMEX8,5291
 nshtrainer/nn/__init__.py,sha256=0QPFl02a71WZQjLMGOlFNMmsYP5aa1q3eABHmnWH58Q,1427
 nshtrainer/nn/mlp.py,sha256=V0FrScpIUdg_IgIO8GMtIsGEtmHjwF14i2IWxmZrsqg,5952
@@ -80,7 +80,7 @@ nshtrainer/profiler/simple.py,sha256=MbMfsJvligd0mtGiltxJ0T8MQVDP9T9BzQZFwswl66Y
 nshtrainer/runner.py,sha256=USAjrExHkN5oVNVunsoPnLxfQrEHSaa54S3RipOe544,3605
 nshtrainer/scripts/find_packages.py,sha256=ixYivZobumyyGsf2B9oYMLyLTRcBzY_vUv-u3bNW-hs,1424
 nshtrainer/trainer/__init__.py,sha256=P2rmr8oBVTHk-HJHYPcUwWqDEArMbPR4_rPpATbWK3E,40
-nshtrainer/trainer/_config.py,sha256=2SgO1L5VBfWQ5g7Dg2dTx_vq2_Wo7dTqt2A4GlQaGo0,28673
+nshtrainer/trainer/_config.py,sha256=ZIodM5Ek1lpkWFhQ_VfmKR7q1mZFFwtjfx8FH72H8WM,29174
 nshtrainer/trainer/_runtime_callback.py,sha256=sd2cUdRJG-UCdQr9ruZvEYpNGNF1t2W2fuxwwVlQD9E,4164
 nshtrainer/trainer/checkpoint_connector.py,sha256=r0ir4xYSdf_jebM0x09qaO6nJsvsiRQDyM0fs80ppOQ,2347
 nshtrainer/trainer/signal_connector.py,sha256=2EzkVktlasl8PgWAKNLDZRUMY__gRlDy1HdinAU-tfU,10740
@@ -95,6 +95,6 @@ nshtrainer/util/seed.py,sha256=Or2wMPsnQxfnZ2xfBiyMcHFIUt3tGTNeMMyOEanCkqs,280
 nshtrainer/util/slurm.py,sha256=rofIU26z3SdL79SF45tNez6juou1cyDLz07oXEZb9Hg,1566
 nshtrainer/util/typed.py,sha256=NGuDkDzFlc1fAoaXjOFZVbmj0mRFjsQi1E_hPa7Bn5U,128
 nshtrainer/util/typing_utils.py,sha256=8ptjSSLZxlmy4FY6lzzkoGoF5fGNClo8-B_c0XHQaNU,385
-nshtrainer-0.31.0.dist-info/METADATA,sha256=99b-8IvPlMmTrjyb5EK1kKsgKj8lWhGw4gZvM5sKyzc,916
-nshtrainer-0.31.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.31.0.dist-info/RECORD,,
+nshtrainer-0.32.1.dist-info/METADATA,sha256=zGSKc6CY965hgKixgUgeAHv3VvIOsDJ4NdCeDIKTTAs,916
+nshtrainer-0.32.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.32.1.dist-info/RECORD,,

nshtrainer/model/mixins/callback.py

@@ -1,206 +0,0 @@
-import logging
-from collections.abc import Callable, Iterable, Sequence
-from typing import Any, TypeAlias, cast, final, overload
-
-from lightning.pytorch import Callback, LightningModule
-from lightning.pytorch.callbacks import LambdaCallback
-from typing_extensions import override
-
-from ...util.typing_utils import mixin_base_type
-
-log = logging.getLogger(__name__)
-
-CallbackFn: TypeAlias = Callable[[], Callback | Iterable[Callback] | None]
-
-
-class CallbackRegistrarModuleMixin:
-    @override
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-        self._nshtrainer_callbacks: list[CallbackFn] = []
-
-    @overload
-    def register_callback(
-        self, callback: Callback | Iterable[Callback] | CallbackFn | None = None, /
-    ): ...
-
-    @overload
-    def register_callback(
-        self,
-        /,
-        *,
-        setup: Callable | None = None,
-        teardown: Callable | None = None,
-        on_fit_start: Callable | None = None,
-        on_fit_end: Callable | None = None,
-        on_sanity_check_start: Callable | None = None,
-        on_sanity_check_end: Callable | None = None,
-        on_train_batch_start: Callable | None = None,
-        on_train_batch_end: Callable | None = None,
-        on_train_epoch_start: Callable | None = None,
-        on_train_epoch_end: Callable | None = None,
-        on_validation_epoch_start: Callable | None = None,
-        on_validation_epoch_end: Callable | None = None,
-        on_test_epoch_start: Callable | None = None,
-        on_test_epoch_end: Callable | None = None,
-        on_validation_batch_start: Callable | None = None,
-        on_validation_batch_end: Callable | None = None,
-        on_test_batch_start: Callable | None = None,
-        on_test_batch_end: Callable | None = None,
-        on_train_start: Callable | None = None,
-        on_train_end: Callable | None = None,
-        on_validation_start: Callable | None = None,
-        on_validation_end: Callable | None = None,
-        on_test_start: Callable | None = None,
-        on_test_end: Callable | None = None,
-        on_exception: Callable | None = None,
-        on_save_checkpoint: Callable | None = None,
-        on_load_checkpoint: Callable | None = None,
-        on_before_backward: Callable | None = None,
-        on_after_backward: Callable | None = None,
-        on_before_optimizer_step: Callable | None = None,
-        on_before_zero_grad: Callable | None = None,
-        on_predict_start: Callable | None = None,
-        on_predict_end: Callable | None = None,
-        on_predict_batch_start: Callable | None = None,
-        on_predict_batch_end: Callable | None = None,
-        on_predict_epoch_start: Callable | None = None,
-        on_predict_epoch_end: Callable | None = None,
-    ): ...
-
-    def register_callback(
-        self,
-        callback: Callback | Iterable[Callback] | CallbackFn | None = None,
-        /,
-        *,
-        setup: Callable | None = None,
-        teardown: Callable | None = None,
-        on_fit_start: Callable | None = None,
-        on_fit_end: Callable | None = None,
-        on_sanity_check_start: Callable | None = None,
-        on_sanity_check_end: Callable | None = None,
-        on_train_batch_start: Callable | None = None,
-        on_train_batch_end: Callable | None = None,
-        on_train_epoch_start: Callable | None = None,
-        on_train_epoch_end: Callable | None = None,
-        on_validation_epoch_start: Callable | None = None,
-        on_validation_epoch_end: Callable | None = None,
-        on_test_epoch_start: Callable | None = None,
-        on_test_epoch_end: Callable | None = None,
-        on_validation_batch_start: Callable | None = None,
-        on_validation_batch_end: Callable | None = None,
-        on_test_batch_start: Callable | None = None,
-        on_test_batch_end: Callable | None = None,
-        on_train_start: Callable | None = None,
-        on_train_end: Callable | None = None,
-        on_validation_start: Callable | None = None,
-        on_validation_end: Callable | None = None,
-        on_test_start: Callable | None = None,
-        on_test_end: Callable | None = None,
-        on_exception: Callable | None = None,
-        on_save_checkpoint: Callable | None = None,
-        on_load_checkpoint: Callable | None = None,
-        on_before_backward: Callable | None = None,
-        on_after_backward: Callable | None = None,
-        on_before_optimizer_step: Callable | None = None,
-        on_before_zero_grad: Callable | None = None,
-        on_predict_start: Callable | None = None,
-        on_predict_end: Callable | None = None,
-        on_predict_batch_start: Callable | None = None,
-        on_predict_batch_end: Callable | None = None,
-        on_predict_epoch_start: Callable | None = None,
-        on_predict_epoch_end: Callable | None = None,
-    ):
-        if callback is None:
-            callback = LambdaCallback(
-                setup=setup,
-                teardown=teardown,
-                on_fit_start=on_fit_start,
-                on_fit_end=on_fit_end,
-                on_sanity_check_start=on_sanity_check_start,
-                on_sanity_check_end=on_sanity_check_end,
-                on_train_batch_start=on_train_batch_start,
-                on_train_batch_end=on_train_batch_end,
-                on_train_epoch_start=on_train_epoch_start,
-                on_train_epoch_end=on_train_epoch_end,
-                on_validation_epoch_start=on_validation_epoch_start,
-                on_validation_epoch_end=on_validation_epoch_end,
-                on_test_epoch_start=on_test_epoch_start,
-                on_test_epoch_end=on_test_epoch_end,
-                on_validation_batch_start=on_validation_batch_start,
-                on_validation_batch_end=on_validation_batch_end,
-                on_test_batch_start=on_test_batch_start,
-                on_test_batch_end=on_test_batch_end,
-                on_train_start=on_train_start,
-                on_train_end=on_train_end,
-                on_validation_start=on_validation_start,
-                on_validation_end=on_validation_end,
-                on_test_start=on_test_start,
-                on_test_end=on_test_end,
-                on_exception=on_exception,
-                on_save_checkpoint=on_save_checkpoint,
-                on_load_checkpoint=on_load_checkpoint,
-                on_before_backward=on_before_backward,
-                on_after_backward=on_after_backward,
-                on_before_optimizer_step=on_before_optimizer_step,
-                on_before_zero_grad=on_before_zero_grad,
-                on_predict_start=on_predict_start,
-                on_predict_end=on_predict_end,
-                on_predict_batch_start=on_predict_batch_start,
-                on_predict_batch_end=on_predict_batch_end,
-                on_predict_epoch_start=on_predict_epoch_start,
-                on_predict_epoch_end=on_predict_epoch_end,
-            )
-
-        if not callable(callback):
-            callback_ = cast(CallbackFn, lambda: callback)
-        else:
-            callback_ = callback
-
-        self._nshtrainer_callbacks.append(callback_)
-
-
-class CallbackModuleMixin(
-    CallbackRegistrarModuleMixin,
-    mixin_base_type(LightningModule),
-):
-    def _nshtrainer_gather_all_callbacks(self):
-        modules: list[Any] = []
-        if isinstance(self, CallbackRegistrarModuleMixin):
-            modules.append(self)
-        if (
-            datamodule := getattr(self.trainer, "datamodule", None)
-        ) is not None and isinstance(datamodule, CallbackRegistrarModuleMixin):
-            modules.append(datamodule)
-        modules.extend(
-            module
-            for module in self.children()
-            if isinstance(module, CallbackRegistrarModuleMixin)
-        )
-        for module in modules:
-            yield from module._nshtrainer_callbacks
-
-    @final
-    @override
-    def configure_callbacks(self):
-        callbacks = super().configure_callbacks()
-        if not isinstance(callbacks, Sequence):
-            callbacks = [callbacks]
-
-        callbacks = list(callbacks)
-        for callback_fn in self._nshtrainer_gather_all_callbacks():
-            callback_result = callback_fn()
-            if callback_result is None:
-                continue
-
-            if not isinstance(callback_result, Iterable):
-                callback_result = [callback_result]
-
-            for callback in callback_result:
-                log.info(
-                    f"Registering {callback.__class__.__qualname__} callback {callback}"
-                )
-                callbacks.append(callback)
-
-        return callbacks