nshtrainer 1.0.0b33__py3-none-any.whl → 1.0.0b37__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,36 +52,34 @@ 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,
 )
 from nshtrainer.loggers import ActSaveLoggerConfig as ActSaveLoggerConfig
-from nshtrainer.loggers import BaseLoggerConfig as BaseLoggerConfig
 from nshtrainer.loggers import CSVLoggerConfig as CSVLoggerConfig
 from nshtrainer.loggers import LoggerConfig as LoggerConfig
+from nshtrainer.loggers import LoggerConfigBase as LoggerConfigBase
 from nshtrainer.loggers import TensorboardLoggerConfig as TensorboardLoggerConfig
 from nshtrainer.loggers import WandbLoggerConfig as WandbLoggerConfig
+from nshtrainer.loggers import logger_registry as logger_registry
 from nshtrainer.lr_scheduler import (
     LinearWarmupCosineDecayLRSchedulerConfig as LinearWarmupCosineDecayLRSchedulerConfig,
 )
 from nshtrainer.lr_scheduler import LRSchedulerConfig as LRSchedulerConfig
 from nshtrainer.lr_scheduler import LRSchedulerConfigBase as LRSchedulerConfigBase
 from nshtrainer.lr_scheduler import ReduceLROnPlateauConfig as ReduceLROnPlateauConfig
-from nshtrainer.nn import BaseNonlinearityConfig as BaseNonlinearityConfig
+from nshtrainer.lr_scheduler.base import lr_scheduler_registry as lr_scheduler_registry
 from nshtrainer.nn import ELUNonlinearityConfig as ELUNonlinearityConfig
 from nshtrainer.nn import GELUNonlinearityConfig as GELUNonlinearityConfig
 from nshtrainer.nn import LeakyReLUNonlinearityConfig as LeakyReLUNonlinearityConfig
 from nshtrainer.nn import MishNonlinearityConfig as MishNonlinearityConfig
 from nshtrainer.nn import MLPConfig as MLPConfig
 from nshtrainer.nn import NonlinearityConfig as NonlinearityConfig
+from nshtrainer.nn import NonlinearityConfigBase as NonlinearityConfigBase
 from nshtrainer.nn import PReLUConfig as PReLUConfig
 from nshtrainer.nn import ReLUNonlinearityConfig as ReLUNonlinearityConfig
 from nshtrainer.nn import SigmoidNonlinearityConfig as SigmoidNonlinearityConfig
@@ -91,9 +92,11 @@ from nshtrainer.nn import TanhNonlinearityConfig as TanhNonlinearityConfig
 from nshtrainer.nn.nonlinearity import (
     SwiGLUNonlinearityConfig as SwiGLUNonlinearityConfig,
 )
+from nshtrainer.nn.nonlinearity import nonlinearity_registry as nonlinearity_registry
 from nshtrainer.optimizer import AdamWConfig as AdamWConfig
 from nshtrainer.optimizer import OptimizerConfig as OptimizerConfig
 from nshtrainer.optimizer import OptimizerConfigBase as OptimizerConfigBase
+from nshtrainer.optimizer import optimizer_registry as optimizer_registry
 from nshtrainer.profiler import AdvancedProfilerConfig as AdvancedProfilerConfig
 from nshtrainer.profiler import BaseProfilerConfig as BaseProfilerConfig
 from nshtrainer.profiler import ProfilerConfig as ProfilerConfig
@@ -106,9 +109,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
@@ -229,8 +229,6 @@ __all__ = [
     "AdvancedProfilerConfig",
     "AsyncCheckpointIOPlugin",
     "BaseCheckpointCallbackConfig",
-    "BaseLoggerConfig",
-    "BaseNonlinearityConfig",
     "BaseProfilerConfig",
     "BestCheckpointCallbackConfig",
     "BitsandbytesPluginConfig",
@@ -284,6 +282,7 @@ __all__ = [
     "LinearWarmupCosineDecayLRSchedulerConfig",
     "LogEpochCallbackConfig",
     "LoggerConfig",
+    "LoggerConfigBase",
     "MLPConfig",
     "MPIEnvironmentPlugin",
     "MPSAcceleratorConfig",
@@ -291,6 +290,7 @@ __all__ = [
     "MishNonlinearityConfig",
     "MixedPrecisionPluginConfig",
     "NonlinearityConfig",
+    "NonlinearityConfigBase",
     "NormLoggingCallbackConfig",
     "OnExceptionCheckpointCallbackConfig",
     "OptimizerConfig",
@@ -320,7 +320,6 @@ __all__ = [
     "SwishNonlinearityConfig",
     "TanhNonlinearityConfig",
     "TensorboardLoggerConfig",
-    "TimeCheckpointCallbackConfig",
     "TorchCheckpointIOPlugin",
     "TorchElasticEnvironmentPlugin",
     "TorchSyncBatchNormPlugin",
@@ -337,12 +336,17 @@ __all__ = [
     "_directory",
     "_hf_hub",
     "accelerator_registry",
+    "callback_registry",
     "callbacks",
+    "logger_registry",
     "loggers",
     "lr_scheduler",
+    "lr_scheduler_registry",
     "metrics",
     "nn",
+    "nonlinearity_registry",
     "optimizer",
+    "optimizer_registry",
     "plugin_registry",
     "profiler",
     "trainer",

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",
 ]