nshtrainer 1.0.0b32__py3-none-any.whl → 1.0.0b36__py3-none-any.whl

nshtrainer/callbacks/norm_logging.py

@@ -7,13 +7,47 @@ import torch
 import torch.nn as nn
 from lightning.pytorch import Callback, LightningModule, Trainer
 from torch.optim import Optimizer
-from typing_extensions import override
+from typing_extensions import final, override
 
-from .base import CallbackConfigBase
+from .base import CallbackConfigBase, callback_registry
 
 log = logging.getLogger(__name__)
 
 
+@final
+@callback_registry.register
+class NormLoggingCallbackConfig(CallbackConfigBase):
+    name: Literal["norm_logging"] = "norm_logging"
+
+    log_grad_norm: bool | str | float = False
+    """If enabled, will log the gradient norm (averaged across all model parameters) to the logger."""
+    log_grad_norm_per_param: bool | str | float = False
+    """If enabled, will log the gradient norm for each model parameter to the logger."""
+
+    log_param_norm: bool | str | float = False
+    """If enabled, will log the parameter norm (averaged across all model parameters) to the logger."""
+    log_param_norm_per_param: bool | str | float = False
+    """If enabled, will log the parameter norm for each model parameter to the logger."""
+
+    def __bool__(self):
+        return any(
+            v
+            for v in (
+                self.log_grad_norm,
+                self.log_grad_norm_per_param,
+                self.log_param_norm,
+                self.log_param_norm_per_param,
+            )
+        )
+
+    @override
+    def create_callbacks(self, trainer_config):
+        if not self:
+            return
+
+        yield NormLoggingCallback(self)
+
+
 def grad_norm(
     module: nn.Module,
     norm_type: float | int | str,
@@ -155,35 +189,3 @@ class NormLoggingCallback(Callback):
                 self._perform_norm_logging(
                     pl_module, optimizer, prefix=f"train/optimizer_{i}/"
                 )
-
-
-class NormLoggingCallbackConfig(CallbackConfigBase):
-    name: Literal["norm_logging"] = "norm_logging"
-
-    log_grad_norm: bool | str | float = False
-    """If enabled, will log the gradient norm (averaged across all model parameters) to the logger."""
-    log_grad_norm_per_param: bool | str | float = False
-    """If enabled, will log the gradient norm for each model parameter to the logger."""
-
-    log_param_norm: bool | str | float = False
-    """If enabled, will log the parameter norm (averaged across all model parameters) to the logger."""
-    log_param_norm_per_param: bool | str | float = False
-    """If enabled, will log the parameter norm for each model parameter to the logger."""
-
-    def __bool__(self):
-        return any(
-            v
-            for v in (
-                self.log_grad_norm,
-                self.log_grad_norm_per_param,
-                self.log_param_norm,
-                self.log_param_norm_per_param,
-            )
-        )
-
-    @override
-    def create_callbacks(self, trainer_config):
-        if not self:
-            return
-
-        yield NormLoggingCallback(self)

nshtrainer/callbacks/print_table.py

@@ -9,13 +9,31 @@ from typing import Literal
 import torch
 from lightning.pytorch import LightningModule, Trainer
 from lightning.pytorch.callbacks import Callback
-from typing_extensions import override
+from typing_extensions import final, override
 
-from .base import CallbackConfigBase
+from .base import CallbackConfigBase, callback_registry
 
 log = logging.getLogger(__name__)
 
 
+@final
+@callback_registry.register
+class PrintTableMetricsCallbackConfig(CallbackConfigBase):
+    """Configuration class for PrintTableMetricsCallback."""
+
+    name: Literal["print_table_metrics"] = "print_table_metrics"
+
+    enabled: bool = True
+    """Whether to enable the callback or not."""
+
+    metric_patterns: list[str] | None = None
+    """List of patterns to filter the metrics to be displayed. If None, all metrics are displayed."""
+
+    @override
+    def create_callbacks(self, trainer_config):
+        yield PrintTableMetricsCallback(metric_patterns=self.metric_patterns)
+
+
 class PrintTableMetricsCallback(Callback):
     """Prints a table with the metrics in columns on every epoch end."""
 
@@ -74,19 +92,3 @@ class PrintTableMetricsCallback(Callback):
             table.add_row(*values)
 
         return table
-
-
-class PrintTableMetricsCallbackConfig(CallbackConfigBase):
-    """Configuration class for PrintTableMetricsCallback."""
-
-    name: Literal["print_table_metrics"] = "print_table_metrics"
-
-    enabled: bool = True
-    """Whether to enable the callback or not."""
-
-    metric_patterns: list[str] | None = None
-    """List of patterns to filter the metrics to be displayed. If None, all metrics are displayed."""
-
-    @override
-    def create_callbacks(self, trainer_config):
-        yield PrintTableMetricsCallback(metric_patterns=self.metric_patterns)

nshtrainer/callbacks/rlp_sanity_checks.py

@@ -11,13 +11,15 @@ from lightning.pytorch.utilities.types import (
     LRSchedulerConfigType,
     LRSchedulerTypeUnion,
 )
-from typing_extensions import Protocol, override, runtime_checkable
+from typing_extensions import Protocol, final, override, runtime_checkable
 
-from .base import CallbackConfigBase
+from .base import CallbackConfigBase, callback_registry
 
 log = logging.getLogger(__name__)
 
 
+@final
+@callback_registry.register
 class RLPSanityChecksCallbackConfig(CallbackConfigBase):
     """
     If enabled, will do some sanity checks if the `ReduceLROnPlateau` scheduler is used:

nshtrainer/callbacks/shared_parameters.py

@@ -7,18 +7,15 @@ from typing import Literal, Protocol, runtime_checkable
 import torch.nn as nn
 from lightning.pytorch import LightningModule, Trainer
 from lightning.pytorch.callbacks import Callback
-from typing_extensions import TypeAliasType, override
+from typing_extensions import TypeAliasType, final, override
 
-from .base import CallbackConfigBase
+from .base import CallbackConfigBase, callback_registry
 
 log = logging.getLogger(__name__)
 
 
-def _parameters_to_names(parameters: Iterable[nn.Parameter], model: nn.Module):
-    mapping = {id(p): n for n, p in model.named_parameters()}
-    return [mapping[id(p)] for p in parameters]
-
-
+@final
+@callback_registry.register
 class SharedParametersCallbackConfig(CallbackConfigBase):
     """A callback that allows scaling the gradients of shared parameters that
     are registered in the ``self.shared_parameters`` list of the root module.
@@ -34,6 +31,11 @@ class SharedParametersCallbackConfig(CallbackConfigBase):
         yield SharedParametersCallback(self)
 
 
+def _parameters_to_names(parameters: Iterable[nn.Parameter], model: nn.Module):
+    mapping = {id(p): n for n, p in model.named_parameters()}
+    return [mapping[id(p)] for p in parameters]
+
+
 SharedParametersList = TypeAliasType(
     "SharedParametersList", list[tuple[nn.Parameter, int | float]]
 )

nshtrainer/callbacks/timer.py

@@ -7,13 +7,23 @@ from typing import Any, Literal
 from lightning.pytorch import LightningModule, Trainer
 from lightning.pytorch.callbacks import Callback
 from lightning.pytorch.utilities.types import STEP_OUTPUT
-from typing_extensions import override
+from typing_extensions import final, override
 
-from .base import CallbackConfigBase
+from .base import CallbackConfigBase, callback_registry
 
 log = logging.getLogger(__name__)
 
 
+@final
+@callback_registry.register
+class EpochTimerCallbackConfig(CallbackConfigBase):
+    name: Literal["epoch_timer"] = "epoch_timer"
+
+    @override
+    def create_callbacks(self, trainer_config):
+        yield EpochTimerCallback()
+
+
 class EpochTimerCallback(Callback):
     def __init__(self):
         super().__init__()
@@ -149,11 +159,3 @@ class EpochTimerCallback(Callback):
     def load_state_dict(self, state_dict: dict[str, Any]) -> None:
         self._elapsed_time = state_dict["elapsed_time"]
         self._total_batches = state_dict["total_batches"]
-
-
-class EpochTimerCallbackConfig(CallbackConfigBase):
-    name: Literal["epoch_timer"] = "epoch_timer"
-
-    @override
-    def create_callbacks(self, trainer_config):
-        yield EpochTimerCallback()

nshtrainer/callbacks/wandb_upload_code.py

@@ -9,13 +9,15 @@ from lightning.pytorch import LightningModule, Trainer
 from lightning.pytorch.callbacks.callback import Callback
 from lightning.pytorch.loggers import WandbLogger
 from nshrunner._env import SNAPSHOT_DIR
-from typing_extensions import override
+from typing_extensions import final, override
 
-from .base import CallbackConfigBase
+from .base import CallbackConfigBase, callback_registry
 
 log = logging.getLogger(__name__)
 
 
+@final
+@callback_registry.register
 class WandbUploadCodeCallbackConfig(CallbackConfigBase):
     name: Literal["wandb_upload_code"] = "wandb_upload_code"
 

nshtrainer/callbacks/wandb_watch.py

@@ -7,13 +7,15 @@ import torch.nn as nn
 from lightning.pytorch import LightningModule, Trainer
 from lightning.pytorch.callbacks.callback import Callback
 from lightning.pytorch.loggers import WandbLogger
-from typing_extensions import override
+from typing_extensions import final, override
 
-from .base import CallbackConfigBase
+from .base import CallbackConfigBase, callback_registry
 
 log = logging.getLogger(__name__)
 
 
+@final
+@callback_registry.register
 class WandbWatchCallbackConfig(CallbackConfigBase):
     name: Literal["wandb_watch"] = "wandb_watch"
 

nshtrainer/configs/__init__.py

@@ -5,6 +5,7 @@ __codegen__ = True
 from nshtrainer import MetricConfig as MetricConfig
 from nshtrainer import TrainerConfig as TrainerConfig
 from nshtrainer import accelerator_registry as accelerator_registry
+from nshtrainer import callback_registry as callback_registry
 from nshtrainer import plugin_registry as plugin_registry
 from nshtrainer._checkpoint.metadata import CheckpointMetadata as CheckpointMetadata
 from nshtrainer._directory import DirectoryConfig as DirectoryConfig
@@ -13,6 +14,7 @@ from nshtrainer._hf_hub import (
     HuggingFaceHubAutoCreateConfig as HuggingFaceHubAutoCreateConfig,
 )
 from nshtrainer._hf_hub import HuggingFaceHubConfig as HuggingFaceHubConfig
+from nshtrainer.callbacks import ActSaveConfig as ActSaveConfig
 from nshtrainer.callbacks import (
     BestCheckpointCallbackConfig as BestCheckpointCallbackConfig,
 )
@@ -35,6 +37,7 @@ from nshtrainer.callbacks import (
 from nshtrainer.callbacks import (
     LastCheckpointCallbackConfig as LastCheckpointCallbackConfig,
 )
+from nshtrainer.callbacks import LearningRateMonitorConfig as LearningRateMonitorConfig
 from nshtrainer.callbacks import LogEpochCallbackConfig as LogEpochCallbackConfig
 from nshtrainer.callbacks import NormLoggingCallbackConfig as NormLoggingCallbackConfig
 from nshtrainer.callbacks import (
@@ -49,14 +52,10 @@ from nshtrainer.callbacks import (
 from nshtrainer.callbacks import (
     SharedParametersCallbackConfig as SharedParametersCallbackConfig,
 )
-from nshtrainer.callbacks import (
-    TimeCheckpointCallbackConfig as TimeCheckpointCallbackConfig,
-)
 from nshtrainer.callbacks import (
     WandbUploadCodeCallbackConfig as WandbUploadCodeCallbackConfig,
 )
 from nshtrainer.callbacks import WandbWatchCallbackConfig as WandbWatchCallbackConfig
-from nshtrainer.callbacks.actsave import ActSaveConfig as ActSaveConfig
 from nshtrainer.callbacks.checkpoint._base import (
     BaseCheckpointCallbackConfig as BaseCheckpointCallbackConfig,
 )
@@ -106,9 +105,6 @@ from nshtrainer.trainer._config import (
 from nshtrainer.trainer._config import CheckpointSavingConfig as CheckpointSavingConfig
 from nshtrainer.trainer._config import EnvironmentConfig as EnvironmentConfig
 from nshtrainer.trainer._config import GradientClippingConfig as GradientClippingConfig
-from nshtrainer.trainer._config import (
-    LearningRateMonitorConfig as LearningRateMonitorConfig,
-)
 from nshtrainer.trainer._config import SanityCheckingConfig as SanityCheckingConfig
 from nshtrainer.trainer._config import StrategyConfig as StrategyConfig
 from nshtrainer.trainer.accelerator import CPUAcceleratorConfig as CPUAcceleratorConfig
@@ -320,7 +316,6 @@ __all__ = [
     "SwishNonlinearityConfig",
     "TanhNonlinearityConfig",
     "TensorboardLoggerConfig",
-    "TimeCheckpointCallbackConfig",
     "TorchCheckpointIOPlugin",
     "TorchElasticEnvironmentPlugin",
     "TorchSyncBatchNormPlugin",
@@ -337,6 +332,7 @@ __all__ = [
     "_directory",
     "_hf_hub",
     "accelerator_registry",
+    "callback_registry",
     "callbacks",
     "loggers",
     "lr_scheduler",

nshtrainer/configs/_hf_hub/__init__.py

@@ -7,9 +7,11 @@ from nshtrainer._hf_hub import (
     HuggingFaceHubAutoCreateConfig as HuggingFaceHubAutoCreateConfig,
 )
 from nshtrainer._hf_hub import HuggingFaceHubConfig as HuggingFaceHubConfig
+from nshtrainer._hf_hub import callback_registry as callback_registry
 
 __all__ = [
     "CallbackConfigBase",
     "HuggingFaceHubAutoCreateConfig",
     "HuggingFaceHubConfig",
+    "callback_registry",
 ]

nshtrainer/configs/callbacks/__init__.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 __codegen__ = True
 
+from nshtrainer.callbacks import ActSaveConfig as ActSaveConfig
 from nshtrainer.callbacks import (
     BestCheckpointCallbackConfig as BestCheckpointCallbackConfig,
 )
@@ -25,6 +26,7 @@ from nshtrainer.callbacks import (
 from nshtrainer.callbacks import (
     LastCheckpointCallbackConfig as LastCheckpointCallbackConfig,
 )
+from nshtrainer.callbacks import LearningRateMonitorConfig as LearningRateMonitorConfig
 from nshtrainer.callbacks import LogEpochCallbackConfig as LogEpochCallbackConfig
 from nshtrainer.callbacks import NormLoggingCallbackConfig as NormLoggingCallbackConfig
 from nshtrainer.callbacks import (
@@ -39,14 +41,11 @@ from nshtrainer.callbacks import (
 from nshtrainer.callbacks import (
     SharedParametersCallbackConfig as SharedParametersCallbackConfig,
 )
-from nshtrainer.callbacks import (
-    TimeCheckpointCallbackConfig as TimeCheckpointCallbackConfig,
-)
 from nshtrainer.callbacks import (
     WandbUploadCodeCallbackConfig as WandbUploadCodeCallbackConfig,
 )
 from nshtrainer.callbacks import WandbWatchCallbackConfig as WandbWatchCallbackConfig
-from nshtrainer.callbacks.actsave import ActSaveConfig as ActSaveConfig
+from nshtrainer.callbacks import callback_registry as callback_registry
 from nshtrainer.callbacks.checkpoint._base import (
     BaseCheckpointCallbackConfig as BaseCheckpointCallbackConfig,
 )
@@ -54,9 +53,6 @@ from nshtrainer.callbacks.checkpoint._base import (
     CheckpointMetadata as CheckpointMetadata,
 )
 from nshtrainer.callbacks.early_stopping import MetricConfig as MetricConfig
-from nshtrainer.callbacks.lr_monitor import (
-    LearningRateMonitorConfig as LearningRateMonitorConfig,
-)
 
 from . import actsave as actsave
 from . import base as base
@@ -100,11 +96,11 @@ __all__ = [
     "PrintTableMetricsCallbackConfig",
     "RLPSanityChecksCallbackConfig",
     "SharedParametersCallbackConfig",
-    "TimeCheckpointCallbackConfig",
     "WandbUploadCodeCallbackConfig",
     "WandbWatchCallbackConfig",
     "actsave",
     "base",
+    "callback_registry",
     "checkpoint",
     "debug_flag",
     "directory_setup",

nshtrainer/configs/callbacks/actsave/__init__.py

@@ -4,8 +4,10 @@ __codegen__ = True
 
 from nshtrainer.callbacks.actsave import ActSaveConfig as ActSaveConfig
 from nshtrainer.callbacks.actsave import CallbackConfigBase as CallbackConfigBase
+from nshtrainer.callbacks.actsave import callback_registry as callback_registry
 
 __all__ = [
     "ActSaveConfig",
     "CallbackConfigBase",
+    "callback_registry",
 ]

nshtrainer/configs/callbacks/base/__init__.py

@@ -3,7 +3,9 @@ from __future__ import annotations
 __codegen__ = True
 
 from nshtrainer.callbacks.base import CallbackConfigBase as CallbackConfigBase
+from nshtrainer.callbacks.base import callback_registry as callback_registry
 
 __all__ = [
     "CallbackConfigBase",
+    "callback_registry",
 ]

nshtrainer/configs/callbacks/checkpoint/__init__.py

@@ -11,9 +11,6 @@ from nshtrainer.callbacks.checkpoint import (
 from nshtrainer.callbacks.checkpoint import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
 )
-from nshtrainer.callbacks.checkpoint import (
-    TimeCheckpointCallbackConfig as TimeCheckpointCallbackConfig,
-)
 from nshtrainer.callbacks.checkpoint._base import (
     BaseCheckpointCallbackConfig as BaseCheckpointCallbackConfig,
 )
@@ -24,12 +21,14 @@ from nshtrainer.callbacks.checkpoint._base import (
     CheckpointMetadata as CheckpointMetadata,
 )
 from nshtrainer.callbacks.checkpoint.best_checkpoint import MetricConfig as MetricConfig
+from nshtrainer.callbacks.checkpoint.best_checkpoint import (
+    callback_registry as callback_registry,
+)
 
 from . import _base as _base
 from . import best_checkpoint as best_checkpoint
 from . import last_checkpoint as last_checkpoint
 from . import on_exception_checkpoint as on_exception_checkpoint
-from . import time_checkpoint as time_checkpoint
 
 __all__ = [
     "BaseCheckpointCallbackConfig",
@@ -39,10 +38,9 @@ __all__ = [
     "LastCheckpointCallbackConfig",
     "MetricConfig",
     "OnExceptionCheckpointCallbackConfig",
-    "TimeCheckpointCallbackConfig",
     "_base",
     "best_checkpoint",
+    "callback_registry",
     "last_checkpoint",
     "on_exception_checkpoint",
-    "time_checkpoint",
 ]

nshtrainer/configs/callbacks/checkpoint/best_checkpoint/__init__.py

@@ -12,10 +12,14 @@ from nshtrainer.callbacks.checkpoint.best_checkpoint import (
     CheckpointMetadata as CheckpointMetadata,
 )
 from nshtrainer.callbacks.checkpoint.best_checkpoint import MetricConfig as MetricConfig
+from nshtrainer.callbacks.checkpoint.best_checkpoint import (
+    callback_registry as callback_registry,
+)
 
 __all__ = [
     "BaseCheckpointCallbackConfig",
     "BestCheckpointCallbackConfig",
     "CheckpointMetadata",
     "MetricConfig",
+    "callback_registry",
 ]

nshtrainer/configs/callbacks/checkpoint/last_checkpoint/__init__.py

@@ -11,9 +11,13 @@ from nshtrainer.callbacks.checkpoint.last_checkpoint import (
 from nshtrainer.callbacks.checkpoint.last_checkpoint import (
     LastCheckpointCallbackConfig as LastCheckpointCallbackConfig,
 )
+from nshtrainer.callbacks.checkpoint.last_checkpoint import (
+    callback_registry as callback_registry,
+)
 
 __all__ = [
     "BaseCheckpointCallbackConfig",
     "CheckpointMetadata",
     "LastCheckpointCallbackConfig",
+    "callback_registry",
 ]

nshtrainer/configs/callbacks/checkpoint/on_exception_checkpoint/__init__.py

@@ -8,8 +8,12 @@ from nshtrainer.callbacks.checkpoint.on_exception_checkpoint import (
 from nshtrainer.callbacks.checkpoint.on_exception_checkpoint import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
 )
+from nshtrainer.callbacks.checkpoint.on_exception_checkpoint import (
+    callback_registry as callback_registry,
+)
 
 __all__ = [
     "CallbackConfigBase",
     "OnExceptionCheckpointCallbackConfig",
+    "callback_registry",
 ]

nshtrainer/configs/callbacks/debug_flag/__init__.py

@@ -6,8 +6,10 @@ from nshtrainer.callbacks.debug_flag import CallbackConfigBase as CallbackConfig
 from nshtrainer.callbacks.debug_flag import (
     DebugFlagCallbackConfig as DebugFlagCallbackConfig,
 )
+from nshtrainer.callbacks.debug_flag import callback_registry as callback_registry
 
 __all__ = [
     "CallbackConfigBase",
     "DebugFlagCallbackConfig",
+    "callback_registry",
 ]

nshtrainer/configs/callbacks/directory_setup/__init__.py

@@ -8,8 +8,10 @@ from nshtrainer.callbacks.directory_setup import (
 from nshtrainer.callbacks.directory_setup import (
     DirectorySetupCallbackConfig as DirectorySetupCallbackConfig,
 )
+from nshtrainer.callbacks.directory_setup import callback_registry as callback_registry
 
 __all__ = [
     "CallbackConfigBase",
     "DirectorySetupCallbackConfig",
+    "callback_registry",
 ]

nshtrainer/configs/callbacks/early_stopping/__init__.py

@@ -7,9 +7,11 @@ from nshtrainer.callbacks.early_stopping import (
     EarlyStoppingCallbackConfig as EarlyStoppingCallbackConfig,
 )
 from nshtrainer.callbacks.early_stopping import MetricConfig as MetricConfig
+from nshtrainer.callbacks.early_stopping import callback_registry as callback_registry
 
 __all__ = [
     "CallbackConfigBase",
     "EarlyStoppingCallbackConfig",
     "MetricConfig",
+    "callback_registry",
 ]

nshtrainer/configs/callbacks/ema/__init__.py

@@ -4,8 +4,10 @@ __codegen__ = True
 
 from nshtrainer.callbacks.ema import CallbackConfigBase as CallbackConfigBase
 from nshtrainer.callbacks.ema import EMACallbackConfig as EMACallbackConfig
+from nshtrainer.callbacks.ema import callback_registry as callback_registry
 
 __all__ = [
     "CallbackConfigBase",
     "EMACallbackConfig",
+    "callback_registry",
 ]

nshtrainer/configs/callbacks/finite_checks/__init__.py

@@ -6,8 +6,10 @@ from nshtrainer.callbacks.finite_checks import CallbackConfigBase as CallbackCon
 from nshtrainer.callbacks.finite_checks import (
     FiniteChecksCallbackConfig as FiniteChecksCallbackConfig,
 )
+from nshtrainer.callbacks.finite_checks import callback_registry as callback_registry
 
 __all__ = [
     "CallbackConfigBase",
     "FiniteChecksCallbackConfig",
+    "callback_registry",
 ]

nshtrainer/configs/callbacks/gradient_skipping/__init__.py

@@ -8,8 +8,12 @@ from nshtrainer.callbacks.gradient_skipping import (
 from nshtrainer.callbacks.gradient_skipping import (
     GradientSkippingCallbackConfig as GradientSkippingCallbackConfig,
 )
+from nshtrainer.callbacks.gradient_skipping import (
+    callback_registry as callback_registry,
+)
 
 __all__ = [
     "CallbackConfigBase",
     "GradientSkippingCallbackConfig",
+    "callback_registry",
 ]

nshtrainer/configs/callbacks/log_epoch/__init__.py

@@ -6,8 +6,10 @@ from nshtrainer.callbacks.log_epoch import CallbackConfigBase as CallbackConfigB
 from nshtrainer.callbacks.log_epoch import (
     LogEpochCallbackConfig as LogEpochCallbackConfig,
 )
+from nshtrainer.callbacks.log_epoch import callback_registry as callback_registry
 
 __all__ = [
     "CallbackConfigBase",
     "LogEpochCallbackConfig",
+    "callback_registry",
 ]

nshtrainer/configs/callbacks/lr_monitor/__init__.py

@@ -6,8 +6,10 @@ from nshtrainer.callbacks.lr_monitor import CallbackConfigBase as CallbackConfig
 from nshtrainer.callbacks.lr_monitor import (
     LearningRateMonitorConfig as LearningRateMonitorConfig,
 )
+from nshtrainer.callbacks.lr_monitor import callback_registry as callback_registry
 
 __all__ = [
     "CallbackConfigBase",
     "LearningRateMonitorConfig",
+    "callback_registry",
 ]

nshtrainer/configs/callbacks/norm_logging/__init__.py

@@ -6,8 +6,10 @@ from nshtrainer.callbacks.norm_logging import CallbackConfigBase as CallbackConf
 from nshtrainer.callbacks.norm_logging import (
     NormLoggingCallbackConfig as NormLoggingCallbackConfig,
 )
+from nshtrainer.callbacks.norm_logging import callback_registry as callback_registry
 
 __all__ = [
     "CallbackConfigBase",
     "NormLoggingCallbackConfig",
+    "callback_registry",
 ]

nshtrainer/configs/callbacks/print_table/__init__.py

@@ -6,8 +6,10 @@ from nshtrainer.callbacks.print_table import CallbackConfigBase as CallbackConfi
 from nshtrainer.callbacks.print_table import (
     PrintTableMetricsCallbackConfig as PrintTableMetricsCallbackConfig,
 )
+from nshtrainer.callbacks.print_table import callback_registry as callback_registry
 
 __all__ = [
     "CallbackConfigBase",
     "PrintTableMetricsCallbackConfig",
+    "callback_registry",
 ]

nshtrainer/configs/callbacks/rlp_sanity_checks/__init__.py

@@ -8,8 +8,12 @@ from nshtrainer.callbacks.rlp_sanity_checks import (
 from nshtrainer.callbacks.rlp_sanity_checks import (
     RLPSanityChecksCallbackConfig as RLPSanityChecksCallbackConfig,
 )
+from nshtrainer.callbacks.rlp_sanity_checks import (
+    callback_registry as callback_registry,
+)
 
 __all__ = [
     "CallbackConfigBase",
     "RLPSanityChecksCallbackConfig",
+    "callback_registry",
 ]

nshtrainer/configs/callbacks/shared_parameters/__init__.py

@@ -8,8 +8,12 @@ from nshtrainer.callbacks.shared_parameters import (
 from nshtrainer.callbacks.shared_parameters import (
     SharedParametersCallbackConfig as SharedParametersCallbackConfig,
 )
+from nshtrainer.callbacks.shared_parameters import (
+    callback_registry as callback_registry,
+)
 
 __all__ = [
     "CallbackConfigBase",
     "SharedParametersCallbackConfig",
+    "callback_registry",
 ]

nshtrainer/configs/callbacks/timer/__init__.py

@@ -6,8 +6,10 @@ from nshtrainer.callbacks.timer import CallbackConfigBase as CallbackConfigBase
 from nshtrainer.callbacks.timer import (
     EpochTimerCallbackConfig as EpochTimerCallbackConfig,
 )
+from nshtrainer.callbacks.timer import callback_registry as callback_registry
 
 __all__ = [
     "CallbackConfigBase",
     "EpochTimerCallbackConfig",
+    "callback_registry",
 ]