nshtrainer 1.0.0b37__tar.gz → 1.0.0b40__tar.gz

{nshtrainer-1.0.0b37 → nshtrainer-1.0.0b40}/PKG-INFO

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

{nshtrainer-1.0.0b37 → nshtrainer-1.0.0b40}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshtrainer"
-version = "1.0.0-beta37"
+version = "1.0.0-beta40"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"
@@ -29,6 +29,8 @@ pyright = "*"
 ruff = "*"
 ipykernel = "*"
 ipywidgets = "*"
+pytest = "^8.3.5"
+pytest-cov = "^6.0.0"
 
 [build-system]
 requires = ["poetry-core"]
@@ -58,3 +60,19 @@ extra = [
     "huggingface-hub",
     "nshutils",
 ]
+
+[tool.pytest]
+testpaths = ["tests"]
+python_files = "test_*.py"
+python_functions = "test_*"
+python_classes = "Test*"
+
+[tool.pytest.ini_options]
+minversion = "6.0"
+# addopts = "--cov=src/nshtrainer --cov-report=term-missing --cov-report=xml --"
+addopts = [
+    "--import-mode=importlib",
+    "--cov=src/nshtrainer",
+    "--cov-report=term-missing",
+    "--cov-report=xml",
+]

{nshtrainer-1.0.0b37 → nshtrainer-1.0.0b40}/src/nshtrainer/_directory.py

@@ -81,7 +81,7 @@ class DirectoryConfig(C.Config):
 
         # Save to nshtrainer/{id}/log/{logger name}
         log_dir = self.resolve_subdirectory(run_id, "log")
-        log_dir = log_dir / getattr(logger, "name")
+        log_dir = log_dir / logger.resolve_logger_dirname()
         # ^ NOTE: Logger must have a `name` attribute, as this is
         # the discriminator for the logger registry
         log_dir.mkdir(exist_ok=True)

{nshtrainer-1.0.0b37 → nshtrainer-1.0.0b40}/src/nshtrainer/callbacks/checkpoint/last_checkpoint.py

@@ -21,11 +21,8 @@ log = logging.getLogger(__name__)
 class LastCheckpointCallbackConfig(BaseCheckpointCallbackConfig):
     name: Literal["last_checkpoint"] = "last_checkpoint"
 
-    save_on_time_interval: bool = True
-    """Whether to save checkpoints based on time interval."""
-
-    interval: timedelta = timedelta(hours=12)
-    """Time interval between checkpoints when save_on_time_interval is True."""
+    save_on_time_interval: timedelta | None = None
+    """Save a checkpoint every `save_on_time_interval` seconds. If `None`, this feature is disabled."""
 
     @override
     def create_checkpoint(self, trainer_config, dirpath):
@@ -38,8 +35,6 @@ class LastCheckpointCallback(CheckpointBase[LastCheckpointCallbackConfig]):
         super().__init__(config, dirpath)
         self.start_time = time.time()
         self.last_checkpoint_time = self.start_time
-        self.interval_seconds = config.interval.total_seconds()
-        self.save_on_time_interval = config.save_on_time_interval
 
     @override
     def name(self):
@@ -57,12 +52,18 @@ class LastCheckpointCallback(CheckpointBase[LastCheckpointCallbackConfig]):
     def topk_sort_reverse(self):
         return True
 
-    def _should_checkpoint(self) -> bool:
-        if not self.save_on_time_interval:
+    def _local_should_checkpoint(self) -> bool:
+        if (interval := self.config.save_on_time_interval) is None:
             return False
+
         current_time = time.time()
         elapsed_time = current_time - self.last_checkpoint_time
-        return elapsed_time >= self.interval_seconds
+        return elapsed_time >= interval.total_seconds()
+
+    def _should_checkpoint(self, trainer: Trainer):
+        if self.config.save_on_time_interval is None:
+            return False
+        return trainer.strategy.broadcast(self._local_should_checkpoint(), src=0)
 
     def _format_duration(self, seconds: float) -> str:
         """Format duration in seconds to a human-readable string."""
@@ -98,7 +99,7 @@ class LastCheckpointCallback(CheckpointBase[LastCheckpointCallbackConfig]):
         *args,
         **kwargs,
     ):
-        if not self._should_checkpoint():
+        if not self._should_checkpoint(trainer):
             return
         self.save_checkpoints(trainer)
 
@@ -110,5 +111,5 @@ class LastCheckpointCallback(CheckpointBase[LastCheckpointCallbackConfig]):
     def save_checkpoints(self, trainer):
         super().save_checkpoints(trainer)
 
-        if self.save_on_time_interval:
+        if self.config.save_on_time_interval is not None:
             self.last_checkpoint_time = time.time()

nshtrainer-1.0.0b40/src/nshtrainer/configs/.gitattributes

@@ -0,0 +1 @@
+* linguist-generated=true

{nshtrainer-1.0.0b37 → nshtrainer-1.0.0b40}/src/nshtrainer/loggers/base.py

@@ -30,5 +30,14 @@ class LoggerConfigBase(C.Config, ABC):
     def __bool__(self):
         return self.enabled
 
+    def resolve_logger_dirname(self) -> str:
+        if not (name := getattr(self, "name", None)):
+            raise ValueError(
+                "Logger must have a name attribute to resolve the directory name.\n"
+                "Otherwise, you must override `resolve_logger_dirname`."
+            )
+
+        return name
+
 
 logger_registry = C.Registry(LoggerConfigBase, discriminator="name")

{nshtrainer-1.0.0b37 → nshtrainer-1.0.0b40}/src/nshtrainer/nn/mlp.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import contextlib
 import copy
 from collections.abc import Callable, Sequence
 from typing import Literal, Protocol, runtime_checkable
@@ -44,6 +45,9 @@ class MLPConfigDict(TypedDict):
     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
@@ -64,15 +68,20 @@ class MLPConfig(C.Config):
     residual: bool = False
     """Whether to use residual connections between layers."""
 
+    seed: int | None = None
+    """Random seed to use for initialization. If None, the default Torch behavior is used."""
+
     def to_kwargs(self) -> MLPConfigDict:
-        return {
+        kwargs: MLPConfigDict = {
             "bias": self.bias,
             "no_bias_scalar": self.no_bias_scalar,
             "nonlinearity": self.nonlinearity,
             "ln": self.ln,
             "dropout": self.dropout,
             "residual": self.residual,
+            "seed": self.seed,
         }
+        return kwargs
 
     def create_module(
         self,
@@ -108,6 +117,7 @@ def MLP(
     pre_layers: Sequence[nn.Module] = [],
     post_layers: Sequence[nn.Module] = [],
     linear_cls: LinearModuleConstructor = nn.Linear,
+    seed: int | None = None,
 ):
     """
     Constructs a multi-layer perceptron (MLP) with the given dimensions and activation function.
@@ -123,52 +133,61 @@ def MLP(
         residual (bool, optional): Whether to use residual connections between layers. Defaults to False.
         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.
 
     Returns:
         nn.Sequential: The constructed MLP.
     """
 
-    if activation is None:
-        activation = 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"):
-        raise ValueError("ln must be a boolean or 'pre' or 'post'")
-
-    layers: list[nn.Module] = []
-    if ln == "pre":
-        layers.append(nn.LayerNorm(dims[0]))
-
-    layers.extend(pre_layers)
-
-    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)
-        layers.append(linear_cls(in_features, out_features, bias=bias_))
-        if dropout is not None:
-            layers.append(nn.Dropout(dropout))
-        if i < len(dims) - 2:
-            match activation:
-                case NonlinearityConfigBase():
-                    layers.append(activation.create_module())
-                case nn.Module():
-                    # In this case, we create a deep copy of the module to avoid sharing parameters (if any).
-                    layers.append(copy.deepcopy(activation))
-                case Callable():
-                    layers.append(activation())
-                case _:
-                    raise ValueError(
-                        "Either `nonlinearity` or `activation` must be provided"
-                    )
-
-    layers.extend(post_layers)
-
-    if ln == "post":
-        layers.append(nn.LayerNorm(dims[-1]))
-
-    cls = ResidualSequential if residual else nn.Sequential
-    return cls(*layers)
+    with contextlib.ExitStack() as stack:
+        if seed is not None:
+            stack.enter_context(
+                torch.random.fork_rng(devices=range(torch.cuda.device_count()))
+            )
+            torch.manual_seed(seed)
+
+        if activation is None:
+            activation = 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"):
+            raise ValueError("ln must be a boolean or 'pre' or 'post'")
+
+        layers: list[nn.Module] = []
+        if ln == "pre":
+            layers.append(nn.LayerNorm(dims[0]))
+
+        layers.extend(pre_layers)
+
+        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)
+            layers.append(linear_cls(in_features, out_features, bias=bias_))
+            if dropout is not None:
+                layers.append(nn.Dropout(dropout))
+            if i < len(dims) - 2:
+                match activation:
+                    case NonlinearityConfigBase():
+                        layers.append(activation.create_module())
+                    case nn.Module():
+                        # In this case, we create a deep copy of the module to avoid sharing parameters (if any).
+                        layers.append(copy.deepcopy(activation))
+                    case Callable():
+                        layers.append(activation())
+                    case _:
+                        raise ValueError(
+                            "Either `nonlinearity` or `activation` must be provided"
+                        )
+
+        layers.extend(post_layers)
+
+        if ln == "post":
+            layers.append(nn.LayerNorm(dims[-1]))
+
+        cls = ResidualSequential if residual else nn.Sequential
+        return cls(*layers)

nshtrainer-1.0.0b40/src/nshtrainer/nn/tests/test_mlp.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+from typing import cast
+
+import pytest
+import torch
+
+from nshtrainer.nn.mlp import MLP
+
+
+def test_mlp_seed_reproducibility():
+    """Test that the seed parameter in MLP ensures reproducible weights."""
+
+    # Test dimensions
+    dims = [10, 20, 5]
+
+    # Create two MLPs with the same seed
+    seed1 = 42
+    mlp1 = MLP(dims, activation=torch.nn.ReLU(), seed=seed1)
+    mlp2 = MLP(dims, activation=torch.nn.ReLU(), seed=seed1)
+
+    # Create an MLP with a different seed
+    seed2 = 123
+    mlp3 = MLP(dims, activation=torch.nn.ReLU(), seed=seed2)
+
+    # Check first layer weights
+    layer1_weights1 = cast(torch.Tensor, mlp1[0].weight)
+    layer1_weights2 = cast(torch.Tensor, mlp2[0].weight)
+    layer1_weights3 = cast(torch.Tensor, mlp3[0].weight)
+
+    # Same seed should produce identical weights
+    assert torch.allclose(layer1_weights1, layer1_weights2)
+
+    # Different seeds should produce different weights
+    assert not torch.allclose(layer1_weights1, layer1_weights3)
+
+    # Check second layer weights
+    layer2_weights1 = cast(torch.Tensor, mlp1[2].weight)
+    layer2_weights2 = cast(torch.Tensor, mlp2[2].weight)
+    layer2_weights3 = cast(torch.Tensor, mlp3[2].weight)
+
+    # Same seed should produce identical weights for all layers
+    assert torch.allclose(layer2_weights1, layer2_weights2)
+
+    # Different seeds should produce different weights for all layers
+    assert not torch.allclose(layer2_weights1, layer2_weights3)
+
+    # Test that not providing a seed gives different results each time
+    mlp4 = MLP(dims, activation=torch.nn.ReLU(), seed=None)
+    mlp5 = MLP(dims, activation=torch.nn.ReLU(), seed=None)
+
+    # Without seeds, weights should be different
+    assert not torch.allclose(
+        cast(torch.Tensor, mlp4[0].weight), cast(torch.Tensor, mlp5[0].weight)
+    )