nshtrainer 1.0.0b44__tar.gz → 1.0.0b46__tar.gz

{nshtrainer-1.0.0b44 → nshtrainer-1.0.0b46}/PKG-INFO

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

{nshtrainer-1.0.0b44 → nshtrainer-1.0.0b46}/pyproject.toml

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

{nshtrainer-1.0.0b44 → nshtrainer-1.0.0b46}/src/nshtrainer/_checkpoint/metadata.py

@@ -139,15 +139,30 @@ def remove_checkpoint_metadata(checkpoint_path: Path):
         log.debug(f"Removed {path}")
 
 
+def remove_checkpoint_metadata_link(ckpt_link_path: Path):
+    path = _metadata_path(ckpt_link_path)
+    # If the metadata does not exist, we can safely ignore this
+    if not path.exists(follow_symlinks=False):
+        # This is EXTREMELY important here
+        # Otherwise, we've already deleted the file that the symlink
+        # used to point to, so this always returns False
+        log.debug(f"Metadata file does not exist: {path}")
+        return
+
+    # If the metadata exists, we can remove it
+    try:
+        path.unlink(missing_ok=True)
+    except Exception:
+        log.warning(f"Failed to remove {path}", exc_info=True)
+    else:
+        log.debug(f"Removed {path}")
+
+
 def link_checkpoint_metadata(checkpoint_path: Path, linked_checkpoint_path: Path):
     # First, remove any existing metadata files
-    remove_checkpoint_metadata(linked_checkpoint_path)
+    remove_checkpoint_metadata_link(linked_checkpoint_path)
 
     # Link the metadata files to the new checkpoint
     path = _metadata_path(checkpoint_path)
     linked_path = _metadata_path(linked_checkpoint_path)
-
-    if not path.exists():
-        raise FileNotFoundError(f"Checkpoint path does not exist: {checkpoint_path}")
-
     try_symlink_or_copy(path, linked_path)

{nshtrainer-1.0.0b44 → nshtrainer-1.0.0b46}/src/nshtrainer/_checkpoint/saver.py

@@ -8,7 +8,11 @@ from pathlib import Path
 from lightning.pytorch import Trainer
 
 from ..util.path import try_symlink_or_copy
-from .metadata import link_checkpoint_metadata, remove_checkpoint_metadata
+from .metadata import (
+    link_checkpoint_metadata,
+    remove_checkpoint_metadata,
+    remove_checkpoint_metadata_link,
+)
 
 log = logging.getLogger(__name__)
 
@@ -39,7 +43,7 @@ def link_checkpoint(
             log.debug(f"Removed {linkpath=}")
 
         if metadata:
-            remove_checkpoint_metadata(linkpath)
+            remove_checkpoint_metadata_link(linkpath)
 
     try_symlink_or_copy(filepath, linkpath)
     if metadata:

{nshtrainer-1.0.0b44 → nshtrainer-1.0.0b46}/src/nshtrainer/callbacks/checkpoint/_base.py

@@ -160,7 +160,7 @@ class CheckpointBase(Checkpoint, ABC, Generic[TConfig]):
         filepath = self.resolve_checkpoint_path(self.current_metrics(trainer))
         trainer.save_checkpoint(filepath, self.config.save_weights_only)
 
-        if trainer.is_global_zero:
+        if trainer.hparams.save_checkpoint_metadata and trainer.is_global_zero:
             # Remove old checkpoints
             self.remove_old_checkpoints(trainer)
 

{nshtrainer-1.0.0b44 → nshtrainer-1.0.0b46}/src/nshtrainer/callbacks/metric_validation.py

@@ -5,8 +5,8 @@ from typing import Literal
 
 from lightning.pytorch.utilities.exceptions import MisconfigurationException
 from typing_extensions import final, override, assert_never
-
-from .._callback import NTCallbackBase
+from lightning.pytorch import Trainer
+from lightning.pytorch.callbacks import Callback
 from ..metrics import MetricConfig
 from .base import CallbackConfigBase, callback_registry
 
@@ -43,33 +43,48 @@ class MetricValidationCallbackConfig(CallbackConfigBase):
         yield MetricValidationCallback(self, metrics)
 
 
-class MetricValidationCallback(NTCallbackBase):
+class MetricValidationCallback(Callback):
     def __init__(
-        self, config: MetricValidationCallbackConfig, metrics: list[MetricConfig]
+        self,
+        config: MetricValidationCallbackConfig,
+        metrics: list[MetricConfig],
     ):
         super().__init__()
 
         self.config = config
         self.metrics = metrics
 
-    @override
-    def on_sanity_check_end(self, trainer, pl_module):
-        super().on_sanity_check_end(trainer, pl_module)
-
-        log.debug("Validating 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())
-        for metric in self.metrics:
-            if metric.validation_monitor in logged_metrics:
-                continue
 
+        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(
-                        f"Metric '{metric.validation_monitor}' not found in logged metrics."
-                    )
+                    raise MisconfigurationException(msg)
                 case "warn":
-                    log.warning(
-                        f"Metric '{metric.validation_monitor}' not found in logged metrics."
-                    )
+                    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-1.0.0b44 → nshtrainer-1.0.0b46}/src/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-1.0.0b44 → nshtrainer-1.0.0b46}/src/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-1.0.0b44 → nshtrainer-1.0.0b46}/src/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-1.0.0b44 → nshtrainer-1.0.0b46}/src/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.0b44 → nshtrainer-1.0.0b46}/src/nshtrainer/util/path.py

@@ -120,7 +120,8 @@ def try_symlink_or_copy(
                 shutil.copy(file_path, link_path)
         else:
             link_path.symlink_to(
-                symlink_target, target_is_directory=target_is_directory
+                symlink_target,
+                target_is_directory=target_is_directory,
             )
     except Exception:
         log.warning(