nshtrainer 1.0.0b43__py3-none-any.whl → 1.0.0b45__py3-none-any.whl

nshtrainer/callbacks/__init__.py

@@ -40,6 +40,10 @@ from .log_epoch import LogEpochCallback as LogEpochCallback
 from .log_epoch import LogEpochCallbackConfig as LogEpochCallbackConfig
 from .lr_monitor import LearningRateMonitor as LearningRateMonitor
 from .lr_monitor import LearningRateMonitorConfig as LearningRateMonitorConfig
+from .metric_validation import MetricValidationCallback as MetricValidationCallback
+from .metric_validation import (
+    MetricValidationCallbackConfig as MetricValidationCallbackConfig,
+)
 from .norm_logging import NormLoggingCallback as NormLoggingCallback
 from .norm_logging import NormLoggingCallbackConfig as NormLoggingCallbackConfig
 from .print_table import PrintTableMetricsCallback as PrintTableMetricsCallback

nshtrainer/callbacks/metric_validation.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+import logging
+from typing import Literal
+
+from lightning.pytorch.utilities.exceptions import MisconfigurationException
+from typing_extensions import final, override, assert_never
+from lightning.pytorch import Trainer
+from lightning.pytorch.callbacks import Callback
+from ..metrics import MetricConfig
+from .base import CallbackConfigBase, callback_registry
+
+log = logging.getLogger(__name__)
+
+
+@final
+@callback_registry.register
+class MetricValidationCallbackConfig(CallbackConfigBase):
+    name: Literal["metric_validation"] = "metric_validation"
+
+    error_behavior: Literal["raise", "warn"] = "raise"
+    """
+    Behavior when an error occurs during validation:
+    - "raise": Raise an error and stop the training.
+    - "warn": Log a warning and continue the training.
+    """
+
+    validate_default_metric: bool = True
+    """Whether to validate the default metric from the root config."""
+
+    metrics: list[MetricConfig] = []
+    """List of metrics to validate."""
+
+    @override
+    def create_callbacks(self, trainer_config):
+        metrics = self.metrics.copy()
+        if (
+            self.validate_default_metric
+            and (default_metric := trainer_config.primary_metric) is not None
+        ):
+            metrics.append(default_metric)
+
+        yield MetricValidationCallback(self, metrics)
+
+
+class MetricValidationCallback(Callback):
+    def __init__(
+        self,
+        config: MetricValidationCallbackConfig,
+        metrics: list[MetricConfig],
+    ):
+        super().__init__()
+
+        self.config = config
+        self.metrics = metrics
+
+    def _check_metrics(self, trainer: Trainer):
+        metric_names = ", ".join(metric.validation_monitor for metric in self.metrics)
+        log.info(f"Validating metrics: {metric_names}...")
+        logged_metrics = set(trainer.logged_metrics.keys())
+
+        invalid_metrics: list[str] = []
+        for metric in self.metrics:
+            if metric.validation_monitor not in logged_metrics:
+                invalid_metrics.append(metric.validation_monitor)
+
+        if invalid_metrics:
+            msg = (
+                f"The following metrics were not found in logged metrics: {invalid_metrics}\n"
+                f"List of logged metrics: {list(trainer.logged_metrics.keys())}"
+            )
+            match self.config.error_behavior:
+                case "raise":
+                    raise MisconfigurationException(msg)
+                case "warn":
+                    log.warning(msg)
+                case _:
+                    assert_never(self.config.error_behavior)
+
+    @override
+    def on_sanity_check_end(self, trainer, pl_module):
+        super().on_sanity_check_end(trainer, pl_module)
+
+        self._check_metrics(trainer)
+
+    @override
+    def on_validation_end(self, trainer, pl_module):
+        super().on_validation_end(trainer, pl_module)
+
+        self._check_metrics(trainer)

nshtrainer/configs/__init__.py

@@ -39,6 +39,9 @@ from nshtrainer.callbacks import (
 )
 from nshtrainer.callbacks import LearningRateMonitorConfig as LearningRateMonitorConfig
 from nshtrainer.callbacks import LogEpochCallbackConfig as LogEpochCallbackConfig
+from nshtrainer.callbacks import (
+    MetricValidationCallbackConfig as MetricValidationCallbackConfig,
+)
 from nshtrainer.callbacks import NormLoggingCallbackConfig as NormLoggingCallbackConfig
 from nshtrainer.callbacks import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
@@ -287,6 +290,7 @@ __all__ = [
     "MPIEnvironmentPlugin",
     "MPSAcceleratorConfig",
     "MetricConfig",
+    "MetricValidationCallbackConfig",
     "MishNonlinearityConfig",
     "MixedPrecisionPluginConfig",
     "NonlinearityConfig",

nshtrainer/configs/callbacks/__init__.py

@@ -28,6 +28,9 @@ from nshtrainer.callbacks import (
 )
 from nshtrainer.callbacks import LearningRateMonitorConfig as LearningRateMonitorConfig
 from nshtrainer.callbacks import LogEpochCallbackConfig as LogEpochCallbackConfig
+from nshtrainer.callbacks import (
+    MetricValidationCallbackConfig as MetricValidationCallbackConfig,
+)
 from nshtrainer.callbacks import NormLoggingCallbackConfig as NormLoggingCallbackConfig
 from nshtrainer.callbacks import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
@@ -65,6 +68,7 @@ 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 metric_validation as metric_validation
 from . import norm_logging as norm_logging
 from . import print_table as print_table
 from . import rlp_sanity_checks as rlp_sanity_checks
@@ -91,6 +95,7 @@ __all__ = [
     "LearningRateMonitorConfig",
     "LogEpochCallbackConfig",
     "MetricConfig",
+    "MetricValidationCallbackConfig",
     "NormLoggingCallbackConfig",
     "OnExceptionCheckpointCallbackConfig",
     "PrintTableMetricsCallbackConfig",
@@ -110,6 +115,7 @@ __all__ = [
     "gradient_skipping",
     "log_epoch",
     "lr_monitor",
+    "metric_validation",
     "norm_logging",
     "print_table",
     "rlp_sanity_checks",

nshtrainer/configs/callbacks/metric_validation/__init__.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+__codegen__ = True
+
+from nshtrainer.callbacks.metric_validation import (
+    CallbackConfigBase as CallbackConfigBase,
+)
+from nshtrainer.callbacks.metric_validation import MetricConfig as MetricConfig
+from nshtrainer.callbacks.metric_validation import (
+    MetricValidationCallbackConfig as MetricValidationCallbackConfig,
+)
+from nshtrainer.callbacks.metric_validation import (
+    callback_registry as callback_registry,
+)
+
+__all__ = [
+    "CallbackConfigBase",
+    "MetricConfig",
+    "MetricValidationCallbackConfig",
+    "callback_registry",
+]

nshtrainer/configs/trainer/__init__.py

@@ -38,6 +38,9 @@ from nshtrainer.trainer._config import LogEpochCallbackConfig as LogEpochCallbac
 from nshtrainer.trainer._config import LoggerConfig as LoggerConfig
 from nshtrainer.trainer._config import LoggerConfigBase as LoggerConfigBase
 from nshtrainer.trainer._config import MetricConfig as MetricConfig
+from nshtrainer.trainer._config import (
+    MetricValidationCallbackConfig as MetricValidationCallbackConfig,
+)
 from nshtrainer.trainer._config import (
     NormLoggingCallbackConfig as NormLoggingCallbackConfig,
 )
@@ -164,6 +167,7 @@ __all__ = [
     "MPIEnvironmentPlugin",
     "MPSAcceleratorConfig",
     "MetricConfig",
+    "MetricValidationCallbackConfig",
     "MixedPrecisionPluginConfig",
     "NormLoggingCallbackConfig",
     "OnExceptionCheckpointCallbackConfig",

nshtrainer/configs/trainer/_config/__init__.py

@@ -34,6 +34,9 @@ from nshtrainer.trainer._config import LogEpochCallbackConfig as LogEpochCallbac
 from nshtrainer.trainer._config import LoggerConfig as LoggerConfig
 from nshtrainer.trainer._config import LoggerConfigBase as LoggerConfigBase
 from nshtrainer.trainer._config import MetricConfig as MetricConfig
+from nshtrainer.trainer._config import (
+    MetricValidationCallbackConfig as MetricValidationCallbackConfig,
+)
 from nshtrainer.trainer._config import (
     NormLoggingCallbackConfig as NormLoggingCallbackConfig,
 )
@@ -77,6 +80,7 @@ __all__ = [
     "LoggerConfig",
     "LoggerConfigBase",
     "MetricConfig",
+    "MetricValidationCallbackConfig",
     "NormLoggingCallbackConfig",
     "OnExceptionCheckpointCallbackConfig",
     "PluginConfig",

nshtrainer/configs/trainer/trainer/__init__.py

@@ -4,14 +4,12 @@ __codegen__ = True
 
 from nshtrainer.trainer.trainer import AcceleratorConfigBase as AcceleratorConfigBase
 from nshtrainer.trainer.trainer import EnvironmentConfig as EnvironmentConfig
-from nshtrainer.trainer.trainer import PluginConfigBase as PluginConfigBase
 from nshtrainer.trainer.trainer import StrategyConfigBase as StrategyConfigBase
 from nshtrainer.trainer.trainer import TrainerConfig as TrainerConfig
 
 __all__ = [
     "AcceleratorConfigBase",
     "EnvironmentConfig",
-    "PluginConfigBase",
     "StrategyConfigBase",
     "TrainerConfig",
 ]

nshtrainer/nn/__init__.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 from .mlp import MLP as MLP
 from .mlp import MLPConfig as MLPConfig
-from .mlp import MLPConfigDict as MLPConfigDict
 from .mlp import ResidualSequential as ResidualSequential
 from .mlp import custom_seed_context as custom_seed_context
 from .module_dict import TypedModuleDict as TypedModuleDict

nshtrainer/nn/mlp.py

@@ -3,12 +3,12 @@ from __future__ import annotations
 import contextlib
 import copy
 from collections.abc import Callable, Sequence
-from typing import Literal, Protocol, runtime_checkable
+from typing import Any, Literal, Protocol, runtime_checkable
 
 import nshconfig as C
 import torch
 import torch.nn as nn
-from typing_extensions import TypedDict, override
+from typing_extensions import deprecated, override
 
 from .nonlinearity import NonlinearityConfig, NonlinearityConfigBase
 
@@ -26,29 +26,6 @@ class ResidualSequential(nn.Sequential):
         return input + super().forward(input)
 
 
-class MLPConfigDict(TypedDict):
-    bias: bool
-    """Whether to include bias terms in the linear layers."""
-
-    no_bias_scalar: bool
-    """Whether to exclude bias terms when the output dimension is 1."""
-
-    nonlinearity: NonlinearityConfig | None
-    """Activation function to use between layers."""
-
-    ln: bool | Literal["pre", "post"]
-    """Whether to apply layer normalization before or after the linear layers."""
-
-    dropout: float | None
-    """Dropout probability to apply between layers."""
-
-    residual: bool
-    """Whether to use residual connections between layers."""
-
-    seed: int | None
-    """Random seed to use for initialization. If None, the default Torch behavior is used."""
-
-
 class MLPConfig(C.Config):
     bias: bool = True
     """Whether to include bias terms in the linear layers."""
@@ -71,8 +48,15 @@ class MLPConfig(C.Config):
     seed: int | None = None
     """Random seed to use for initialization. If None, the default Torch behavior is used."""
 
-    def to_kwargs(self) -> MLPConfigDict:
-        kwargs: MLPConfigDict = {
+    @deprecated("Use `nt.nn.MLP(config=...)` instead.")
+    def create_module(
+        self,
+        dims: Sequence[int],
+        pre_layers: Sequence[nn.Module] = [],
+        post_layers: Sequence[nn.Module] = [],
+        linear_cls: LinearModuleConstructor = nn.Linear,
+    ):
+        kwargs: dict[str, Any] = {
             "bias": self.bias,
             "no_bias_scalar": self.no_bias_scalar,
             "nonlinearity": self.nonlinearity,
@@ -81,18 +65,9 @@ class MLPConfig(C.Config):
             "residual": self.residual,
             "seed": self.seed,
         }
-        return kwargs
-
-    def create_module(
-        self,
-        dims: Sequence[int],
-        pre_layers: Sequence[nn.Module] = [],
-        post_layers: Sequence[nn.Module] = [],
-        linear_cls: LinearModuleConstructor = nn.Linear,
-    ):
         return MLP(
             dims,
-            **self.to_kwargs(),
+            **kwargs,
             pre_layers=pre_layers,
             post_layers=post_layers,
             linear_cls=linear_cls,
@@ -121,50 +96,73 @@ def MLP(
     | nn.Module
     | Callable[[], nn.Module]
     | None = None,
-    bias: bool = True,
-    no_bias_scalar: bool = True,
-    ln: bool | Literal["pre", "post"] = False,
+    bias: bool | None = None,
+    no_bias_scalar: bool | None = None,
+    ln: bool | Literal["pre", "post"] | None = None,
     dropout: float | None = None,
-    residual: bool = False,
+    residual: bool | None = None,
     pre_layers: Sequence[nn.Module] = [],
     post_layers: Sequence[nn.Module] = [],
     linear_cls: LinearModuleConstructor = nn.Linear,
     seed: int | None = None,
+    config: MLPConfig | None = None,
 ):
     """
     Constructs a multi-layer perceptron (MLP) with the given dimensions and activation function.
 
     Args:
         dims (Sequence[int]): List of integers representing the dimensions of the MLP.
-        nonlinearity (Callable[[], nn.Module]): Activation function to use between layers.
-        activation (Callable[[], nn.Module]): Activation function to use between layers.
-        bias (bool, optional): Whether to include bias terms in the linear layers. Defaults to True.
-        no_bias_scalar (bool, optional): Whether to exclude bias terms when the output dimension is 1. Defaults to True.
-        ln (bool | Literal["pre", "post"], optional): Whether to apply layer normalization before or after the linear layers. Defaults to False.
-        dropout (float | None, optional): Dropout probability to apply between layers. Defaults to None.
-        residual (bool, optional): Whether to use residual connections between layers. Defaults to False.
+        nonlinearity (Callable[[], nn.Module] | None, optional): Activation function to use between layers.
+        activation (Callable[[], nn.Module] | None, optional): Activation function to use between layers.
+        bias (bool | None, optional): Whether to include bias terms in the linear layers.
+        no_bias_scalar (bool | None, optional): Whether to exclude bias terms when the output dimension is 1.
+        ln (bool | Literal["pre", "post"] | None, optional): Whether to apply layer normalization before or after the linear layers.
+        dropout (float | None, optional): Dropout probability to apply between layers.
+        residual (bool | None, optional): Whether to use residual connections between layers.
         pre_layers (Sequence[nn.Module], optional): List of layers to insert before the linear layers. Defaults to [].
         post_layers (Sequence[nn.Module], optional): List of layers to insert after the linear layers. Defaults to [].
         linear_cls (LinearModuleConstructor, optional): Linear module constructor to use. Defaults to nn.Linear.
-        seed (int | None, optional): Random seed to use for initialization. If None, the default Torch behavior is used. Defaults to None.
+        seed (int | None, optional): Random seed to use for initialization. If None, the default Torch behavior is used.
+        config (MLPConfig | None, optional): Configuration object for the MLP. Parameters specified directly take precedence.
 
     Returns:
         nn.Sequential: The constructed MLP.
     """
 
-    with custom_seed_context(seed):
+    # Resolve parameters: arg if not None, otherwise config value if config exists, otherwise default
+    resolved_bias = bias if bias is not None else (config.bias if config else True)
+    resolved_no_bias_scalar = (
+        no_bias_scalar
+        if no_bias_scalar is not None
+        else (config.no_bias_scalar if config else True)
+    )
+    resolved_nonlinearity = (
+        nonlinearity
+        if nonlinearity is not None
+        else (config.nonlinearity if config else None)
+    )
+    resolved_ln = ln if ln is not None else (config.ln if config else False)
+    resolved_dropout = (
+        dropout if dropout is not None else (config.dropout if config else None)
+    )
+    resolved_residual = (
+        residual if residual is not None else (config.residual if config else False)
+    )
+    resolved_seed = seed if seed is not None else (config.seed if config else None)
+
+    with custom_seed_context(resolved_seed):
         if activation is None:
-            activation = nonlinearity
+            activation = resolved_nonlinearity
 
         if len(dims) < 2:
             raise ValueError("mlp requires at least 2 dimensions")
-        if ln is True:
-            ln = "pre"
-        elif isinstance(ln, str) and ln not in ("pre", "post"):
+        if resolved_ln is True:
+            resolved_ln = "pre"
+        elif isinstance(resolved_ln, str) and resolved_ln not in ("pre", "post"):
             raise ValueError("ln must be a boolean or 'pre' or 'post'")
 
         layers: list[nn.Module] = []
-        if ln == "pre":
+        if resolved_ln == "pre":
             layers.append(nn.LayerNorm(dims[0]))
 
         layers.extend(pre_layers)
@@ -172,10 +170,12 @@ def MLP(
         for i in range(len(dims) - 1):
             in_features = dims[i]
             out_features = dims[i + 1]
-            bias_ = bias and not (no_bias_scalar and out_features == 1)
+            bias_ = resolved_bias and not (
+                resolved_no_bias_scalar and out_features == 1
+            )
             layers.append(linear_cls(in_features, out_features, bias=bias_))
-            if dropout is not None:
-                layers.append(nn.Dropout(dropout))
+            if resolved_dropout is not None:
+                layers.append(nn.Dropout(resolved_dropout))
             if i < len(dims) - 2:
                 match activation:
                     case NonlinearityConfigBase():
@@ -192,8 +192,8 @@ def MLP(
 
         layers.extend(post_layers)
 
-        if ln == "post":
+        if resolved_ln == "post":
             layers.append(nn.LayerNorm(dims[-1]))
 
-        cls = ResidualSequential if residual else nn.Sequential
+        cls = ResidualSequential if resolved_residual else nn.Sequential
         return cls(*layers)

nshtrainer/trainer/_config.py

@@ -40,6 +40,7 @@ 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.metric_validation import MetricValidationCallbackConfig
 from ..callbacks.rlp_sanity_checks import RLPSanityChecksCallbackConfig
 from ..callbacks.shared_parameters import SharedParametersCallbackConfig
 from ..loggers import (
@@ -697,6 +698,10 @@ class TrainerConfig(C.Config):
     - The trainer is running in fast_dev_run mode.
     - The trainer is running a sanity check (which happens before starting the training routine).
     """
+    auto_validate_metrics: MetricValidationCallbackConfig | None = (
+        MetricValidationCallbackConfig()
+    )
+    """If enabled, will automatically validate the metrics before starting the training routine."""
 
     lightning_kwargs: LightningTrainerKwargs = LightningTrainerKwargs()
     """
@@ -768,6 +773,7 @@ class TrainerConfig(C.Config):
         yield self.shared_parameters
         yield self.reduce_lr_on_plateau_sanity_checking
         yield self.auto_set_debug_flag
+        yield self.auto_validate_metrics
         yield from self.callbacks
 
     def _nshtrainer_all_logger_configs(self) -> Iterable[LoggerConfigBase | None]:

nshtrainer/trainer/trainer.py

@@ -25,7 +25,6 @@ from ..util.bf16 import is_bf16_supported_no_emulation
 from ._config import LightningTrainerKwargs, TrainerConfig
 from ._runtime_callback import RuntimeTrackerCallback, Stage
 from .accelerator import AcceleratorConfigBase
-from .plugin import PluginConfigBase
 from .signal_connector import _SignalConnector
 from .strategy import StrategyConfigBase
 

{nshtrainer-1.0.0b43.dist-info → nshtrainer-1.0.0b45.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: nshtrainer
-Version: 1.0.0b43
+Version: 1.0.0b45
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-1.0.0b43.dist-info → nshtrainer-1.0.0b45.dist-info}/RECORD

@@ -6,7 +6,7 @@ nshtrainer/_checkpoint/saver.py,sha256=rWl4d2lCTMU4_wt8yZFL2pFQaP9hj5sPgqHMPQ4zu
 nshtrainer/_directory.py,sha256=TJR9ccyuzRlAVfVjGyeQ3E2AFAcz-XbBCxWfiXo2SlY,3191
 nshtrainer/_experimental/__init__.py,sha256=U4S_2y3zgLZVfMenHRaJFBW8yqh2mUBuI291LGQVOJ8,35
 nshtrainer/_hf_hub.py,sha256=4OsCbIITnZk_YLyoMrVyZ0SIN04FBxlC0ig2Et8UAdo,14287
-nshtrainer/callbacks/__init__.py,sha256=4giOYT8A709UOLRtQEt16QbOAFUHCjJ_aLB7ITTwXJI,3577
+nshtrainer/callbacks/__init__.py,sha256=w80d6PGNu3wjUj9NiRGMqCX9NnXD5ZlvbY-DIK4zjPE,3766
 nshtrainer/callbacks/actsave.py,sha256=NSXIIu62MNYe5gz479SMW33bdoKYoYtWtd_iTWFpKpc,3881
 nshtrainer/callbacks/base.py,sha256=Alaou1IHAIlMEM7g58d_02ozY2xWlshBN7fsw5Ee21s,3683
 nshtrainer/callbacks/checkpoint/__init__.py,sha256=l8tkHc83_mLiU0-wT09SWdRzwpm2ulbkLzcuCmuTwzE,620
@@ -23,6 +23,7 @@ nshtrainer/callbacks/gradient_skipping.py,sha256=8g7oC7PF0LTAEzwiNoaS5tWOnkjk_EB
 nshtrainer/callbacks/interval.py,sha256=UCzUzt3XCFVyQyCWL9lOrStkkxesvduNOYk8yMrGTTk,8116
 nshtrainer/callbacks/log_epoch.py,sha256=B5Dm8XVZwCzKUhUWfT_5PDdDac993191OsbcxxuSVJE,1457
 nshtrainer/callbacks/lr_monitor.py,sha256=qy_C0R40J0hBAukzBwng5FI2jJUpWuXOi5N6FU6ym3I,1210
+nshtrainer/callbacks/metric_validation.py,sha256=tqUVS2n9QRT3v1_8jAGlYBFhLpA6Bm9pxOsfWhD3yZQ,2915
 nshtrainer/callbacks/norm_logging.py,sha256=nVIDWe-ASl5zN830-ODR8QMCqI1ma-QPCIwoy0Wb-Nk,6390
 nshtrainer/callbacks/print_table.py,sha256=VaS4JgI963do79laXK4lUkFQx8v6aRSy22W0zyal_LA,3035
 nshtrainer/callbacks/rlp_sanity_checks.py,sha256=74BZvV2HLO__ucQXsLXb8eJLUZgRFUNJZ6TL9efMp74,10051
@@ -31,12 +32,12 @@ nshtrainer/callbacks/timer.py,sha256=gDcw_K_ikf0bkVgxQ0cDhvvNvz6GLZVLcatuKfh0ORU
 nshtrainer/callbacks/wandb_upload_code.py,sha256=shV7UtnXgY2bUlXdVrXiaDs0PNLlIt7TzNJkJPkzvzI,2414
 nshtrainer/callbacks/wandb_watch.py,sha256=VB14Dy5ZRXQ3di0fPv0K_DFJurLhroLPytnuwQBiJFg,3037
 nshtrainer/configs/.gitattributes,sha256=VeZmarvNEqiRBOHGcllpKm90nL6C8u4tBu7SEm7fj-E,26
-nshtrainer/configs/__init__.py,sha256=MZfcSKhnjtVObBvVv9lu8L2cFTLINP5zcTQvWnz8jdk,14505
+nshtrainer/configs/__init__.py,sha256=0BzCgE1iEJ0Ywmy__mqJZipLQtwZVdz6XK-gHbkA7GY,14650
 nshtrainer/configs/_checkpoint/__init__.py,sha256=6s7Y68StboqscY2G4P_QG443jz5aiym5SjOogIljWLg,342
 nshtrainer/configs/_checkpoint/metadata/__init__.py,sha256=oOPfYkXTjKgm6pluGsG6V1TPyCEGjsQpHVL-LffSUFQ,290
 nshtrainer/configs/_directory/__init__.py,sha256=_oO7vM9DhzHSxtZcv86sTi7hZIptnK1gr-AP9mqQ370,386
 nshtrainer/configs/_hf_hub/__init__.py,sha256=ciFLbV-JV8SVzqo2SyythEuDMnk7gGfdIacB18QYnkY,511
-nshtrainer/configs/callbacks/__init__.py,sha256=jSWkbsdiu9vdGWTzqkDf-Bo9dXr9RengeNZLzWUhi7Y,4283
+nshtrainer/configs/callbacks/__init__.py,sha256=PB3Jg-8_vMhp-mCFw2_Tqt05drKwHK6Ovl9mb8NNiXs,4506
 nshtrainer/configs/callbacks/actsave/__init__.py,sha256=JvjSZtEoA28FC4u-QT3skQzBDVbN9eq07rn4u2ydW-E,377
 nshtrainer/configs/callbacks/base/__init__.py,sha256=wT3RhXttLyf6RFWCIvsoiXcPdfGx5W309WBI18AI5os,278
 nshtrainer/configs/callbacks/checkpoint/__init__.py,sha256=aGJ7vX14YamkMdwYAdPv6XrRnP0aZd5uZ5X0nSLc6IU,1475
@@ -52,6 +53,7 @@ nshtrainer/configs/callbacks/finite_checks/__init__.py,sha256=e-vx9Kn-noqw4wPvZw
 nshtrainer/configs/callbacks/gradient_skipping/__init__.py,sha256=T3eVxxJfnYBrO9WfLiycn4TyWP4vaqJ57yp7Epkg7B4,485
 nshtrainer/configs/callbacks/log_epoch/__init__.py,sha256=IQ5owYYvyk7fiQP1QXYtncRRJrESuq3rRFhab-II2uE,419
 nshtrainer/configs/callbacks/lr_monitor/__init__.py,sha256=qejy1AnXNDHmsFuXRAXQQ5B0TcbKzvpaw-I4dv2AXIs,431
+nshtrainer/configs/callbacks/metric_validation/__init__.py,sha256=_YV0EbISkforE_GDlTTVA6Nn2_l13zX3m1ggcbhnAvs,585
 nshtrainer/configs/callbacks/norm_logging/__init__.py,sha256=j2LrnYEbDGLwJR2lk-jmh-4J_iLEs2HNEoepvJSFLAg,437
 nshtrainer/configs/callbacks/print_table/__init__.py,sha256=t6fA_dBkUCszUXDJKEdnlBH4oEpfAQqcmAlatTFYIyQ,452
 nshtrainer/configs/callbacks/rlp_sanity_checks/__init__.py,sha256=dlP14Wh-w8zG_B4EtNmCIFzVMhf6bXCJ1O9cJWmEFnA,482
@@ -80,8 +82,8 @@ nshtrainer/configs/profiler/_base/__init__.py,sha256=ekYfPg-VDhCAFM5nJka2TxUYdRD
 nshtrainer/configs/profiler/advanced/__init__.py,sha256=-ThpUat16Ij_0avkMUVVA8wCWDG_q_tM7KQofnWQCtg,308
 nshtrainer/configs/profiler/pytorch/__init__.py,sha256=soAU1s2_Pa1na4gW8CK-iysJBO5M_7YeZC2_x40iEdg,294
 nshtrainer/configs/profiler/simple/__init__.py,sha256=3Wb11lPuFuyasq8xS1CZ4WLuBCLS_nVSQGVllvOOi0Y,289
-nshtrainer/configs/trainer/__init__.py,sha256=jYCp4Q9uvutA6NYqfthbREMg09-obD3gHtzEI2Ta-hU,7729
-nshtrainer/configs/trainer/_config/__init__.py,sha256=uof_oJfhwjB1pft7KsRdk_RvNj-tE8wcDBEM7X5qtNc,3666
+nshtrainer/configs/trainer/__init__.py,sha256=a8pzGVid52abAVARPbgjaN566H1ZM44FH_x95bsBaGE,7880
+nshtrainer/configs/trainer/_config/__init__.py,sha256=6DXdtP-uH11TopQ7kzId9fco-wVkD7ZfevbBqDpN6TE,3817
 nshtrainer/configs/trainer/accelerator/__init__.py,sha256=3H6R3wlwbKL1TzDqGCChZk78-BcE2czLouo7Djiq3nA,898
 nshtrainer/configs/trainer/plugin/__init__.py,sha256=NkHQxMPkrtTtdIAO4dQUE9SWEcHRDB0yUXLkTjnl4dA,3332
 nshtrainer/configs/trainer/plugin/base/__init__.py,sha256=GuubcKrbXt4VjJRT8VpNUQqBtyuutre_CfkS0EWZ5_E,368
@@ -90,7 +92,7 @@ nshtrainer/configs/trainer/plugin/io/__init__.py,sha256=W6G67JnigB6d3MiwLrbSKgtI
 nshtrainer/configs/trainer/plugin/layer_sync/__init__.py,sha256=SYDZk2M6sgpt4sEuoURuS8EKYmaqGcvYxETE9jvTrEE,431
 nshtrainer/configs/trainer/plugin/precision/__init__.py,sha256=szlqSfK2XuWdkf72LQzQFv3SlWfKFdRUpBEYIxQ3TPs,1507
 nshtrainer/configs/trainer/strategy/__init__.py,sha256=50whNloJVBq_bdbLaPQnPBTeS1Rcs8MwxTCYBj1kKa4,273
-nshtrainer/configs/trainer/trainer/__init__.py,sha256=QnuhMQNAa1nSVN2o50_WeKAQG_qkNlkeoq9zTjjwmTI,586
+nshtrainer/configs/trainer/trainer/__init__.py,sha256=DDuBRx0kVNMW0z_sqKTUt8-Ql7bOpargi4KcHHvDu_c,486
 nshtrainer/configs/util/__init__.py,sha256=qXittS7f7MyaqJnjvFLKnKsyb6bXTD3dEV16jXVDaH4,2104
 nshtrainer/configs/util/_environment_info/__init__.py,sha256=eB4E0Ck7XCeSC5gbUdA5thd7TXnjGCL0t8GZIFj7uCI,1644
 nshtrainer/configs/util/config/__init__.py,sha256=nEFiDG3-dvvTytYn1tEkPFzp7fgaGRp2j7toSN7yRGs,501
@@ -117,8 +119,8 @@ nshtrainer/model/base.py,sha256=JL3AmH17GQjQIoMrZl3O0vUI7dj5ZsO5iEJgoLPyzHw,1035
 nshtrainer/model/mixins/callback.py,sha256=0LPgve4VszHbLipid4mpI1qnnmdGS2spivs0dXLvqHw,3154
 nshtrainer/model/mixins/debug.py,sha256=1LX9KzeFX9JDPs_a6YCdYDZXLhEk_5rBO2aCqlfBy7w,2087
 nshtrainer/model/mixins/logger.py,sha256=27H99FuLaxc6_dDLG2pid4E_5E0-eLGnc2Ifpt0HYIM,6066
-nshtrainer/nn/__init__.py,sha256=0FgeoaLYtRiSLT8fdPigLD8t-d8DKR8IQDw16JA9lT4,1523
-nshtrainer/nn/mlp.py,sha256=_a8rJJniSCvM08gyQGO-5MUoO18U9_FSGGn3tZL2_U4,7101
+nshtrainer/nn/__init__.py,sha256=5Gg3nieGSC5_dXaI9KUVUUbM13hHexH9831m4hcf6no,1475
+nshtrainer/nn/mlp.py,sha256=nYUgAISzuhC8sav6PloAdyz0PdEoikwppiXIuToEVdE,7550
 nshtrainer/nn/module_dict.py,sha256=9plb8aQUx5TUEPhX5jI9u8LrpTeKe7jZAHi8iIqcN8w,2365
 nshtrainer/nn/module_list.py,sha256=UB43pcwD_3nUke_DyLQt-iXKhWdKM6Zjm84lRC1hPYA,1755
 nshtrainer/nn/nonlinearity.py,sha256=xmaL4QCRvCxqmaGIOwetJeKK-6IK4m2OV7D3SjxSwJQ,6322
@@ -129,7 +131,7 @@ nshtrainer/profiler/advanced.py,sha256=XrM3FX0ThCv5UwUrrH0l4Ow4LGAtpiBww2N8QAU5N
 nshtrainer/profiler/pytorch.py,sha256=8K37XvPnCApUpIK8tA2zNMFIaIiTLSoxKQoiyCPBm1Q,2757
 nshtrainer/profiler/simple.py,sha256=PimjqcU-JuS-8C0ZGHAdwCxgNLij4x0FH6WXsjBQzZs,1005
 nshtrainer/trainer/__init__.py,sha256=fQ7gQRlGWX-90TYT0rttkQyvXDCzo7DAvJgr-jX1zsY,316
-nshtrainer/trainer/_config.py,sha256=QDy6sINVDGEqfHfPTWXSN-06EoEuMSVscHn8fCRTvr0,32981
+nshtrainer/trainer/_config.py,sha256=pCBRtqIC_BzNPqthsDhd7L5_7DG5y8_uVq19lj1mtOM,33311
 nshtrainer/trainer/_runtime_callback.py,sha256=6F2Gq27Q8OFfN3RtdNC6QRA8ac0LC1hh4DUE3V5WgbI,4217
 nshtrainer/trainer/accelerator.py,sha256=Bqq-ry7DeCY4zw9_zBvTZiijpA-uUHrDjtbLV652m4M,2415
 nshtrainer/trainer/plugin/__init__.py,sha256=UM8f70Ml3RGpsXeQr1Yh1yBcxxjFNGFGgJbniCn_rws,366
@@ -140,7 +142,7 @@ nshtrainer/trainer/plugin/layer_sync.py,sha256=h-ydZwXepnsw5-paLgiDatqPyQ_8C0QEv
 nshtrainer/trainer/plugin/precision.py,sha256=I0QsB1bVxmsFmBOkgrAfGONsuYae_lD9Bz0PfJEQvH4,5598
 nshtrainer/trainer/signal_connector.py,sha256=GhfGcSzfaTNhnj2QFkBDq5aT7FqbLMA7eC8SYQs8_8w,10828
 nshtrainer/trainer/strategy.py,sha256=VPTn5z3zvXTydY8IJchjhjcOfpvtoejnvUkq5E4WTus,1368
-nshtrainer/trainer/trainer.py,sha256=ed_Pn-yQCb9BqaHXo2wVhkt2CSfGNEzMAM6RsDoTo-I,20834
+nshtrainer/trainer/trainer.py,sha256=8wMe0qArbDfStS4UdmuKSC2aiAImR3mhj14_kCJiNSM,20797
 nshtrainer/util/_environment_info.py,sha256=MT8mBe6ZolRfKiwU-les1P-lPNPqXpHQcfADrh_A3uY,24629
 nshtrainer/util/bf16.py,sha256=9QhHZCkYSfYpIcxwAMoXyuh2yTSHBzT-EdLQB297jEs,762
 nshtrainer/util/config/__init__.py,sha256=Z39JJufSb61Lhn2GfVcv3eFW_eorOrN9-9llDWlnZZM,272
@@ -152,6 +154,6 @@ nshtrainer/util/seed.py,sha256=diMV8iwBKN7Xxt5pELmui-gyqyT80_CZzomrWhNss0k,316
 nshtrainer/util/slurm.py,sha256=HflkP5iI_r4UHMyPjw9R4dD5AHsJUpcfJw5PLvGYBRM,1603
 nshtrainer/util/typed.py,sha256=Xt5fUU6zwLKSTLUdenovnKK0N8qUq89Kddz2_XeykVQ,164
 nshtrainer/util/typing_utils.py,sha256=MjY-CUX9R5Tzat-BlFnQjwl1PQ_W2yZQoXhkYHlJ_VA,442
-nshtrainer-1.0.0b43.dist-info/METADATA,sha256=ZE3l6CN34ptFgx3SDPfKIgjdV2s3J8qdP729eb58vzo,988
-nshtrainer-1.0.0b43.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
-nshtrainer-1.0.0b43.dist-info/RECORD,,
+nshtrainer-1.0.0b45.dist-info/METADATA,sha256=_RPpe6F7DXpsQSmBF1GTc-E5VUfaC69fIYfoFhsip2s,988
+nshtrainer-1.0.0b45.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
+nshtrainer-1.0.0b45.dist-info/RECORD,,