nshtrainer 1.0.0b12__tar.gz → 1.0.0b13__tar.gz

{nshtrainer-1.0.0b12 → nshtrainer-1.0.0b13}/PKG-INFO

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

{nshtrainer-1.0.0b12 → nshtrainer-1.0.0b13}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshtrainer"
-version = "1.0.0-beta12"
+version = "1.0.0-beta13"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"

nshtrainer-1.0.0b13/src/nshtrainer/callbacks/lr_monitor.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from typing import Literal
+
+from lightning.pytorch.callbacks import LearningRateMonitor
+
+from .base import CallbackConfigBase
+
+
+class LearningRateMonitorConfig(CallbackConfigBase):
+    logging_interval: Literal["step", "epoch"] | None = None
+    """
+    Set to 'epoch' or 'step' to log 'lr' of all optimizers at the same interval, set to None to log at individual interval according to the 'interval' key of each scheduler. Defaults to None.
+    """
+
+    log_momentum: bool = False
+    """
+    Option to also log the momentum values of the optimizer, if the optimizer has the 'momentum' or 'betas' attribute. Defaults to False.
+    """
+
+    log_weight_decay: bool = False
+    """
+    Option to also log the weight decay values of the optimizer. Defaults to False.
+    """
+
+    def create_callbacks(self, trainer_config):
+        yield LearningRateMonitor(
+            logging_interval=self.logging_interval,
+            log_momentum=self.log_momentum,
+            log_weight_decay=self.log_weight_decay,
+        )

{nshtrainer-1.0.0b12 → nshtrainer-1.0.0b13}/src/nshtrainer/configs/__init__.py

@@ -132,10 +132,8 @@ if TYPE_CHECKING:
     from nshtrainer.trainer._config import (
         GradientClippingConfig as GradientClippingConfig,
     )
-    from nshtrainer.trainer._config import LoggingConfig as LoggingConfig
-    from nshtrainer.trainer._config import OptimizationConfig as OptimizationConfig
     from nshtrainer.trainer._config import (
-        ReproducibilityConfig as ReproducibilityConfig,
+        LearningRateMonitorConfig as LearningRateMonitorConfig,
     )
     from nshtrainer.trainer._config import SanityCheckingConfig as SanityCheckingConfig
     from nshtrainer.util._environment_info import (
@@ -325,6 +323,10 @@ else:
             ).LastCheckpointStrategyConfig
         if name == "LeakyReLUNonlinearityConfig":
             return importlib.import_module("nshtrainer.nn").LeakyReLUNonlinearityConfig
+        if name == "LearningRateMonitorConfig":
+            return importlib.import_module(
+                "nshtrainer.trainer._config"
+            ).LearningRateMonitorConfig
         if name == "LinearWarmupCosineDecayLRSchedulerConfig":
             return importlib.import_module(
                 "nshtrainer.lr_scheduler"
@@ -333,8 +335,6 @@ else:
             return importlib.import_module(
                 "nshtrainer.callbacks"
             ).LogEpochCallbackConfig
-        if name == "LoggingConfig":
-            return importlib.import_module("nshtrainer.trainer._config").LoggingConfig
         if name == "MLPConfig":
             return importlib.import_module("nshtrainer.nn").MLPConfig
         if name == "MetricConfig":
@@ -349,10 +349,6 @@ else:
             return importlib.import_module(
                 "nshtrainer.callbacks"
             ).OnExceptionCheckpointCallbackConfig
-        if name == "OptimizationConfig":
-            return importlib.import_module(
-                "nshtrainer.trainer._config"
-            ).OptimizationConfig
         if name == "OptimizerConfigBase":
             return importlib.import_module("nshtrainer.optimizer").OptimizerConfigBase
         if name == "PReLUConfig":
@@ -373,10 +369,6 @@ else:
             return importlib.import_module(
                 "nshtrainer.lr_scheduler"
             ).ReduceLROnPlateauConfig
-        if name == "ReproducibilityConfig":
-            return importlib.import_module(
-                "nshtrainer.trainer._config"
-            ).ReproducibilityConfig
         if name == "SanityCheckingConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"

{nshtrainer-1.0.0b12 → nshtrainer-1.0.0b13}/src/nshtrainer/configs/callbacks/__init__.py

@@ -62,6 +62,9 @@ if TYPE_CHECKING:
         CheckpointMetadata as CheckpointMetadata,
     )
     from nshtrainer.callbacks.early_stopping import MetricConfig as MetricConfig
+    from nshtrainer.callbacks.lr_monitor import (
+        LearningRateMonitorConfig as LearningRateMonitorConfig,
+    )
 else:
 
     def __getattr__(name):
@@ -115,6 +118,10 @@ else:
             return importlib.import_module(
                 "nshtrainer.callbacks"
             ).LastCheckpointCallbackConfig
+        if name == "LearningRateMonitorConfig":
+            return importlib.import_module(
+                "nshtrainer.callbacks.lr_monitor"
+            ).LearningRateMonitorConfig
         if name == "LogEpochCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks"
@@ -167,6 +174,7 @@ from . import ema as ema
 from . import finite_checks as finite_checks
 from . import gradient_skipping as gradient_skipping
 from . import log_epoch as log_epoch
+from . import lr_monitor as lr_monitor
 from . import norm_logging as norm_logging
 from . import print_table as print_table
 from . import rlp_sanity_checks as rlp_sanity_checks

nshtrainer-1.0.0b13/src/nshtrainer/configs/callbacks/lr_monitor/__init__.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+__codegen__ = True
+
+from typing import TYPE_CHECKING
+
+# Config/alias imports
+
+if TYPE_CHECKING:
+    from nshtrainer.callbacks.lr_monitor import CallbackConfigBase as CallbackConfigBase
+    from nshtrainer.callbacks.lr_monitor import (
+        LearningRateMonitorConfig as LearningRateMonitorConfig,
+    )
+else:
+
+    def __getattr__(name):
+        import importlib
+
+        if name in globals():
+            return globals()[name]
+        if name == "CallbackConfigBase":
+            return importlib.import_module(
+                "nshtrainer.callbacks.lr_monitor"
+            ).CallbackConfigBase
+        if name == "LearningRateMonitorConfig":
+            return importlib.import_module(
+                "nshtrainer.callbacks.lr_monitor"
+            ).LearningRateMonitorConfig
+        raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+
+# Submodule exports

{nshtrainer-1.0.0b12 → nshtrainer-1.0.0b13}/src/nshtrainer/configs/trainer/__init__.py

@@ -9,6 +9,7 @@ from typing import TYPE_CHECKING
 if TYPE_CHECKING:
     from nshtrainer.trainer import TrainerConfig as TrainerConfig
     from nshtrainer.trainer._config import ActSaveLoggerConfig as ActSaveLoggerConfig
+    from nshtrainer.trainer._config import BaseLoggerConfig as BaseLoggerConfig
     from nshtrainer.trainer._config import (
         BestCheckpointCallbackConfig as BestCheckpointCallbackConfig,
     )
@@ -39,20 +40,21 @@ if TYPE_CHECKING:
     from nshtrainer.trainer._config import (
         LastCheckpointCallbackConfig as LastCheckpointCallbackConfig,
     )
+    from nshtrainer.trainer._config import (
+        LearningRateMonitorConfig as LearningRateMonitorConfig,
+    )
     from nshtrainer.trainer._config import (
         LogEpochCallbackConfig as LogEpochCallbackConfig,
     )
     from nshtrainer.trainer._config import LoggerConfig as LoggerConfig
-    from nshtrainer.trainer._config import LoggingConfig as LoggingConfig
     from nshtrainer.trainer._config import MetricConfig as MetricConfig
     from nshtrainer.trainer._config import (
-        OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
+        NormLoggingCallbackConfig as NormLoggingCallbackConfig,
     )
-    from nshtrainer.trainer._config import OptimizationConfig as OptimizationConfig
-    from nshtrainer.trainer._config import ProfilerConfig as ProfilerConfig
     from nshtrainer.trainer._config import (
-        ReproducibilityConfig as ReproducibilityConfig,
+        OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
     )
+    from nshtrainer.trainer._config import ProfilerConfig as ProfilerConfig
     from nshtrainer.trainer._config import (
         RLPSanityChecksCallbackConfig as RLPSanityChecksCallbackConfig,
     )
@@ -75,6 +77,10 @@ else:
             return importlib.import_module(
                 "nshtrainer.trainer._config"
             ).ActSaveLoggerConfig
+        if name == "BaseLoggerConfig":
+            return importlib.import_module(
+                "nshtrainer.trainer._config"
+            ).BaseLoggerConfig
         if name == "BestCheckpointCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
@@ -119,30 +125,28 @@ else:
             return importlib.import_module(
                 "nshtrainer.trainer._config"
             ).LastCheckpointCallbackConfig
+        if name == "LearningRateMonitorConfig":
+            return importlib.import_module(
+                "nshtrainer.trainer._config"
+            ).LearningRateMonitorConfig
         if name == "LogEpochCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
             ).LogEpochCallbackConfig
-        if name == "LoggingConfig":
-            return importlib.import_module("nshtrainer.trainer._config").LoggingConfig
         if name == "MetricConfig":
             return importlib.import_module("nshtrainer.trainer._config").MetricConfig
-        if name == "OnExceptionCheckpointCallbackConfig":
+        if name == "NormLoggingCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
-            ).OnExceptionCheckpointCallbackConfig
-        if name == "OptimizationConfig":
+            ).NormLoggingCallbackConfig
+        if name == "OnExceptionCheckpointCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
-            ).OptimizationConfig
+            ).OnExceptionCheckpointCallbackConfig
         if name == "RLPSanityChecksCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
             ).RLPSanityChecksCallbackConfig
-        if name == "ReproducibilityConfig":
-            return importlib.import_module(
-                "nshtrainer.trainer._config"
-            ).ReproducibilityConfig
         if name == "SanityCheckingConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"

{nshtrainer-1.0.0b12 → nshtrainer-1.0.0b13}/src/nshtrainer/configs/trainer/_config/__init__.py

@@ -8,6 +8,7 @@ from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from nshtrainer.trainer._config import ActSaveLoggerConfig as ActSaveLoggerConfig
+    from nshtrainer.trainer._config import BaseLoggerConfig as BaseLoggerConfig
     from nshtrainer.trainer._config import (
         BestCheckpointCallbackConfig as BestCheckpointCallbackConfig,
     )
@@ -38,20 +39,21 @@ if TYPE_CHECKING:
     from nshtrainer.trainer._config import (
         LastCheckpointCallbackConfig as LastCheckpointCallbackConfig,
     )
+    from nshtrainer.trainer._config import (
+        LearningRateMonitorConfig as LearningRateMonitorConfig,
+    )
     from nshtrainer.trainer._config import (
         LogEpochCallbackConfig as LogEpochCallbackConfig,
     )
     from nshtrainer.trainer._config import LoggerConfig as LoggerConfig
-    from nshtrainer.trainer._config import LoggingConfig as LoggingConfig
     from nshtrainer.trainer._config import MetricConfig as MetricConfig
     from nshtrainer.trainer._config import (
-        OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
+        NormLoggingCallbackConfig as NormLoggingCallbackConfig,
     )
-    from nshtrainer.trainer._config import OptimizationConfig as OptimizationConfig
-    from nshtrainer.trainer._config import ProfilerConfig as ProfilerConfig
     from nshtrainer.trainer._config import (
-        ReproducibilityConfig as ReproducibilityConfig,
+        OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
     )
+    from nshtrainer.trainer._config import ProfilerConfig as ProfilerConfig
     from nshtrainer.trainer._config import (
         RLPSanityChecksCallbackConfig as RLPSanityChecksCallbackConfig,
     )
@@ -75,6 +77,10 @@ else:
             return importlib.import_module(
                 "nshtrainer.trainer._config"
             ).ActSaveLoggerConfig
+        if name == "BaseLoggerConfig":
+            return importlib.import_module(
+                "nshtrainer.trainer._config"
+            ).BaseLoggerConfig
         if name == "BestCheckpointCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
@@ -119,30 +125,28 @@ else:
             return importlib.import_module(
                 "nshtrainer.trainer._config"
             ).LastCheckpointCallbackConfig
+        if name == "LearningRateMonitorConfig":
+            return importlib.import_module(
+                "nshtrainer.trainer._config"
+            ).LearningRateMonitorConfig
         if name == "LogEpochCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
             ).LogEpochCallbackConfig
-        if name == "LoggingConfig":
-            return importlib.import_module("nshtrainer.trainer._config").LoggingConfig
         if name == "MetricConfig":
             return importlib.import_module("nshtrainer.trainer._config").MetricConfig
-        if name == "OnExceptionCheckpointCallbackConfig":
+        if name == "NormLoggingCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
-            ).OnExceptionCheckpointCallbackConfig
-        if name == "OptimizationConfig":
+            ).NormLoggingCallbackConfig
+        if name == "OnExceptionCheckpointCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
-            ).OptimizationConfig
+            ).OnExceptionCheckpointCallbackConfig
         if name == "RLPSanityChecksCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
             ).RLPSanityChecksCallbackConfig
-        if name == "ReproducibilityConfig":
-            return importlib.import_module(
-                "nshtrainer.trainer._config"
-            ).ReproducibilityConfig
         if name == "SanityCheckingConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"

{nshtrainer-1.0.0b12 → nshtrainer-1.0.0b13}/src/nshtrainer/data/datamodule.py

@@ -8,8 +8,6 @@ from typing import Any, Generic, cast
 import nshconfig as C
 import torch
 from lightning.pytorch import LightningDataModule
-from lightning.pytorch.utilities.model_helpers import is_overridden
-from lightning.pytorch.utilities.rank_zero import rank_zero_warn
 from typing_extensions import Never, TypeVar, deprecated, override
 
 from ..model.mixins.callback import CallbackRegistrarModuleMixin

{nshtrainer-1.0.0b12 → nshtrainer-1.0.0b13}/src/nshtrainer/trainer/_config.py

@@ -40,11 +40,13 @@ from ..callbacks import (
     CallbackConfig,
     EarlyStoppingCallbackConfig,
     LastCheckpointCallbackConfig,
+    NormLoggingCallbackConfig,
     OnExceptionCheckpointCallbackConfig,
 )
 from ..callbacks.base import CallbackConfigBase
 from ..callbacks.debug_flag import DebugFlagCallbackConfig
 from ..callbacks.log_epoch import LogEpochCallbackConfig
+from ..callbacks.lr_monitor import LearningRateMonitorConfig
 from ..callbacks.rlp_sanity_checks import RLPSanityChecksCallbackConfig
 from ..callbacks.shared_parameters import SharedParametersCallbackConfig
 from ..loggers import (
@@ -53,6 +55,7 @@ from ..loggers import (
     TensorboardLoggerConfig,
     WandbLoggerConfig,
 )
+from ..loggers._base import BaseLoggerConfig
 from ..loggers.actsave import ActSaveLoggerConfig
 from ..metrics._config import MetricConfig
 from ..profiler import ProfilerConfig
@@ -61,103 +64,6 @@ from ..util._environment_info import EnvironmentConfig
 log = logging.getLogger(__name__)
 
 
-class LoggingConfig(CallbackConfigBase):
-    enabled: bool = True
-    """Enable experiment tracking."""
-
-    loggers: Sequence[LoggerConfig] = [
-        WandbLoggerConfig(),
-        CSVLoggerConfig(),
-        TensorboardLoggerConfig(),
-    ]
-    """Loggers to use for experiment tracking."""
-
-    log_lr: bool | Literal["step", "epoch"] = True
-    """If enabled, will register a `LearningRateMonitor` callback to log the learning rate to the logger."""
-    log_epoch: LogEpochCallbackConfig | None = LogEpochCallbackConfig()
-    """If enabled, will log the fractional epoch number to the logger."""
-
-    actsave_logger: ActSaveLoggerConfig | None = None
-    """If enabled, will automatically save logged metrics using ActSave (if nshutils is installed)."""
-
-    @property
-    def wandb(self):
-        return next(
-            (
-                logger
-                for logger in self.loggers
-                if isinstance(logger, WandbLoggerConfig)
-            ),
-            None,
-        )
-
-    @property
-    def csv(self):
-        return next(
-            (logger for logger in self.loggers if isinstance(logger, CSVLoggerConfig)),
-            None,
-        )
-
-    @property
-    def tensorboard(self):
-        return next(
-            (
-                logger
-                for logger in self.loggers
-                if isinstance(logger, TensorboardLoggerConfig)
-            ),
-            None,
-        )
-
-    def create_loggers(self, trainer_config: TrainerConfig):
-        """
-        Constructs and returns a list of loggers based on the provided root configuration.
-
-        Args:
-            trainer_config (TrainerConfig): The root configuration object.
-
-        Returns:
-            list[Logger]: A list of constructed loggers.
-        """
-        if not self.enabled:
-            return
-
-        for logger_config in sorted(
-            self.loggers,
-            key=lambda x: x.priority,
-            reverse=True,
-        ):
-            if not logger_config.enabled:
-                continue
-            if (logger := logger_config.create_logger(trainer_config)) is None:
-                continue
-            yield logger
-
-        # If the actsave_metrics is enabled, add the ActSave logger
-        if self.actsave_logger:
-            yield self.actsave_logger.create_logger(trainer_config)
-
-    @override
-    def create_callbacks(self, trainer_config):
-        if self.log_lr:
-            from lightning.pytorch.callbacks import LearningRateMonitor
-
-            logging_interval: str | None = None
-            if isinstance(self.log_lr, str):
-                logging_interval = self.log_lr
-
-            yield LearningRateMonitor(logging_interval=logging_interval)
-
-        if self.log_epoch:
-            yield from self.log_epoch.create_callbacks(trainer_config)
-
-        for logger in self.loggers:
-            if not logger or not isinstance(logger, CallbackConfigBase):
-                continue
-
-            yield from logger.create_callbacks(trainer_config)
-
-
 class GradientClippingConfig(C.Config):
     enabled: bool = True
     """Enable gradient clipping."""
@@ -167,32 +73,6 @@ class GradientClippingConfig(C.Config):
     """Norm type to use for gradient clipping."""
 
 
-class OptimizationConfig(CallbackConfigBase):
-    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."""
-
-    gradient_clipping: GradientClippingConfig | None = None
-    """Gradient clipping configuration, or None to disable gradient clipping."""
-
-    @override
-    def create_callbacks(self, trainer_config):
-        from ..callbacks.norm_logging import NormLoggingCallbackConfig
-
-        yield from NormLoggingCallbackConfig(
-            log_grad_norm=self.log_grad_norm,
-            log_grad_norm_per_param=self.log_grad_norm_per_param,
-            log_param_norm=self.log_param_norm,
-            log_param_norm_per_param=self.log_param_norm_per_param,
-        ).create_callbacks(trainer_config)
-
-
 TPlugin = TypeVar(
     "TPlugin",
     Precision,
@@ -252,15 +132,6 @@ StrategyLiteral: TypeAlias = Literal[
 ]
 
 
-class ReproducibilityConfig(C.Config):
-    deterministic: bool | Literal["warn"] | None = None
-    """
-    If ``True``, sets whether PyTorch operations must use deterministic algorithms.
-        Set to ``"warn"`` to use deterministic algorithms whenever possible, throwing warnings on operations
-        that don't support deterministic mode. If not set, defaults to ``False``. Default: ``None``.
-    """
-
-
 CheckpointCallbackConfig: TypeAlias = Annotated[
     BestCheckpointCallbackConfig
     | LastCheckpointCallbackConfig
@@ -634,14 +505,34 @@ class TrainerConfig(C.Config):
     hf_hub: HuggingFaceHubConfig = HuggingFaceHubConfig()
     """Hugging Face Hub configuration options."""
 
-    logging: LoggingConfig = LoggingConfig()
-    """Logging/experiment tracking (e.g., WandB) configuration options."""
+    loggers: Sequence[LoggerConfig] = [
+        WandbLoggerConfig(),
+        CSVLoggerConfig(),
+        TensorboardLoggerConfig(),
+    ]
+    """Loggers to use for experiment tracking."""
 
-    optimizer: OptimizationConfig = OptimizationConfig()
-    """Optimization configuration options."""
+    actsave_logger: ActSaveLoggerConfig | None = None
+    """If enabled, will automatically save logged metrics using ActSave (if nshutils is installed)."""
 
-    reproducibility: ReproducibilityConfig = ReproducibilityConfig()
-    """Reproducibility configuration options."""
+    lr_monitor: LearningRateMonitorConfig | None = LearningRateMonitorConfig()
+    """Learning rate monitoring configuration options."""
+
+    log_epoch: LogEpochCallbackConfig | None = LogEpochCallbackConfig()
+    """If enabled, will log the fractional epoch number to the logger."""
+
+    gradient_clipping: GradientClippingConfig | None = None
+    """Gradient clipping configuration, or None to disable gradient clipping."""
+
+    log_norms: NormLoggingCallbackConfig | None = None
+    """Norm logging configuration options."""
+
+    deterministic: bool | Literal["warn"] | None = None
+    """
+    If ``True``, sets whether PyTorch operations must use deterministic algorithms.
+        Set to ``"warn"`` to use deterministic algorithms whenever possible, throwing warnings on operations
+        that don't support deterministic mode. If not set, defaults to ``False``. Default: ``None``.
+    """
 
     reduce_lr_on_plateau_sanity_checking: RLPSanityChecksCallbackConfig | None = (
         RLPSanityChecksCallbackConfig()
@@ -856,27 +747,87 @@ class TrainerConfig(C.Config):
     set_float32_matmul_precision: Literal["medium", "high", "highest"] | None = None
     """If enabled, will set the torch float32 matmul precision to the specified value. Useful for faster training on Ampere+ GPUs."""
 
+    @property
+    def wandb_logger(self):
+        return next(
+            (
+                logger
+                for logger in self.loggers
+                if isinstance(logger, WandbLoggerConfig)
+            ),
+            None,
+        )
+
+    @property
+    def csv_logger(self):
+        return next(
+            (logger for logger in self.loggers if isinstance(logger, CSVLoggerConfig)),
+            None,
+        )
+
+    @property
+    def tensorboard_logger(self):
+        return next(
+            (
+                logger
+                for logger in self.loggers
+                if isinstance(logger, TensorboardLoggerConfig)
+            ),
+            None,
+        )
+
     def _nshtrainer_all_callback_configs(self) -> Iterable[CallbackConfigBase | None]:
         yield self.early_stopping
         yield self.checkpoint_saving
-        yield self.logging
-        yield self.optimizer
+        yield self.lr_monitor
+        yield from (
+            logger_config
+            for logger_config in self.loggers
+            if logger_config is not None
+            and isinstance(logger_config, CallbackConfigBase)
+        )
+        yield self.log_epoch
+        yield self.log_norms
         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
 
+    def _nshtrainer_all_logger_configs(self) -> Iterable[BaseLoggerConfig | None]:
+        yield from self.loggers
+        yield self.actsave_logger
+
     # region Helper Methods
+    def fast_dev_run_(self, value: int | bool = True, /):
+        """
+        Enables fast_dev_run mode for the trainer.
+        This will run the training loop for a specified number of batches,
+        if an integer is provided, or for a single batch if True is provided.
+        """
+        self.fast_dev_run = value
+        return self
+
     def with_fast_dev_run(self, value: int | bool = True, /):
         """
         Enables fast_dev_run mode for the trainer.
         This will run the training loop for a specified number of batches,
         if an integer is provided, or for a single batch if True is provided.
         """
-        config = copy.deepcopy(self)
-        config.fast_dev_run = value
-        return config
+        return copy.deepcopy(self).fast_dev_run_(value)
+
+    def project_root_(self, project_root: str | Path | os.PathLike):
+        """
+        Set the project root directory for the trainer.
+
+        Args:
+            project_root (Path): The base directory to use.
+
+        Returns:
+            self: The current instance of the class.
+        """
+        self.directory.project_root = Path(project_root)
+        return self
 
     def with_project_root(self, project_root: str | Path | os.PathLike):
         """
@@ -888,9 +839,7 @@ class TrainerConfig(C.Config):
         Returns:
             self: The current instance of the class.
         """
-        config = copy.deepcopy(self)
-        config.directory.project_root = Path(project_root)
-        return config
+        return copy.deepcopy(self).project_root_(project_root)
 
     def reset_run(
         self,