nshtrainer 1.1.1b1__py3-none-any.whl → 1.2.0__py3-none-any.whl

nshtrainer/_directory.py

@@ -65,9 +65,9 @@ class DirectoryConfig(C.Config):
     ) -> Path:
         # The subdir will be $CWD/nshtrainer/{id}/{log, stdio, checkpoint, activation}/
         if (subdir := getattr(self, subdirectory, None)) is not None:
-            assert isinstance(
-                subdir, Path
-            ), f"Expected a Path for {subdirectory}, got {type(subdir)}"
+            assert isinstance(subdir, Path), (
+                f"Expected a Path for {subdirectory}, got {type(subdir)}"
+            )
             return subdir
 
         dir = self.resolve_run_root_directory(run_id)

nshtrainer/callbacks/__init__.py

@@ -23,6 +23,12 @@ from .directory_setup import DirectorySetupCallback as DirectorySetupCallback
 from .directory_setup import (
     DirectorySetupCallbackConfig as DirectorySetupCallbackConfig,
 )
+from .distributed_prediction_writer import (
+    DistributedPredictionWriter as DistributedPredictionWriter,
+)
+from .distributed_prediction_writer import (
+    DistributedPredictionWriterConfig as DistributedPredictionWriterConfig,
+)
 from .early_stopping import EarlyStoppingCallback as EarlyStoppingCallback
 from .early_stopping import EarlyStoppingCallbackConfig as EarlyStoppingCallbackConfig
 from .ema import EMACallback as EMACallback

nshtrainer/callbacks/base.py

@@ -23,6 +23,10 @@ class CallbackMetadataConfig(TypedDict, total=False):
     """Priority of the callback. Callbacks with higher priority will be loaded first.
     Default is `0`."""
 
+    enabled_for_barebones: bool
+    """Whether this callback is enabled for barebones mode.
+    Default is `False`."""
+
 
 @dataclass(frozen=True)
 class CallbackWithMetadata:
@@ -91,10 +95,20 @@ def _filter_ignore_if_exists(callbacks: list[CallbackWithMetadata]):
 
 
 def _process_and_filter_callbacks(
+    trainer_config: TrainerConfig,
     callbacks: Iterable[CallbackWithMetadata],
 ) -> list[Callback]:
     callbacks = list(callbacks)
 
+    # If we're in barebones mode, used the callback metadata
+    # to decide to keep/remove the callback.
+    if trainer_config.barebones:
+        callbacks = [
+            callback
+            for callback in callbacks
+            if callback.metadata.get("enabled_for_barebones", False)
+        ]
+
     # Sort by priority (higher priority first)
     callbacks.sort(
         key=lambda callback: callback.metadata.get("priority", 0),
@@ -114,9 +128,14 @@ def resolve_all_callbacks(trainer_config: TrainerConfig):
         if config is not None
     ]
     callbacks = _process_and_filter_callbacks(
-        callback
-        for callback_config in callback_configs
-        for callback in _create_callbacks_with_metadata(callback_config, trainer_config)
+        trainer_config,
+        (
+            callback
+            for callback_config in callback_configs
+            for callback in _create_callbacks_with_metadata(
+                callback_config, trainer_config
+            )
+        ),
     )
     return callbacks
 

nshtrainer/callbacks/distributed_prediction_writer.py

@@ -0,0 +1,166 @@
+from __future__ import annotations
+
+import functools
+import logging
+from collections.abc import Iterator, Sequence
+from pathlib import Path
+from typing import Any, ClassVar, Literal, overload
+
+import torch
+from lightning.fabric.utilities.apply_func import move_data_to_device
+from lightning.pytorch.callbacks import BasePredictionWriter
+from typing_extensions import final, override
+
+from .base import CallbackConfigBase, CallbackMetadataConfig, callback_registry
+
+log = logging.getLogger(__name__)
+
+
+@final
+@callback_registry.register
+class DistributedPredictionWriterConfig(CallbackConfigBase):
+    metadata: ClassVar[CallbackMetadataConfig] = CallbackMetadataConfig(
+        enabled_for_barebones=True
+    )
+    """Metadata for the callback."""
+
+    name: Literal["distributed_prediction_writer"] = "distributed_prediction_writer"
+
+    dirpath: Path | None = None
+    """Directory to save the predictions to. If None, will use the default directory."""
+
+    move_to_cpu_on_save: bool = True
+    """Whether to move the predictions to CPU before saving. Default is True."""
+
+    save_raw: bool = True
+    """Whether to save the raw predictions."""
+
+    save_processed: bool = True
+    """Whether to process and save the predictions.
+
+    "Processing" means that the model's batched predictions are split into individual predictions
+    and saved as a list of tensors.
+    """
+
+    @override
+    def create_callbacks(self, trainer_config):
+        if (dirpath := self.dirpath) is None:
+            dirpath = trainer_config.directory.resolve_subdirectory(
+                trainer_config.id, "predictions"
+            )
+
+        yield DistributedPredictionWriter(self, dirpath)
+
+
+def _move_and_save(data, path: Path, move_to_cpu: bool):
+    if move_to_cpu:
+        data = move_data_to_device(data, "cpu")
+
+    # Save the data to the specified path
+    torch.save(data, path)
+
+
+class DistributedPredictionWriter(BasePredictionWriter):
+    def __init__(
+        self,
+        config: DistributedPredictionWriterConfig,
+        output_dir: Path,
+    ):
+        self.config = config
+
+        super().__init__(write_interval="batch")
+
+        self.output_dir = output_dir
+
+    @override
+    def write_on_batch_end(
+        self,
+        trainer,
+        pl_module,
+        prediction,
+        batch_indices,
+        batch,
+        batch_idx,
+        dataloader_idx,
+    ):
+        save = functools.partial(
+            _move_and_save,
+            move_to_cpu=self.config.move_to_cpu_on_save,
+        )
+
+        # Regular, unstructured writing.
+        if self.config.save_raw:
+            output_dir = (
+                self.output_dir
+                / "raw"
+                / f"dataloader_{dataloader_idx}"
+                / f"rank_{trainer.global_rank}"
+                / f"batch_{batch_idx}"
+            )
+            output_dir.mkdir(parents=True, exist_ok=True)
+            save(prediction, output_dir / "predictions.pt")
+            save(batch, output_dir / "batch.pt")
+            save(batch_indices, output_dir / "batch_indices.pt")
+
+        if self.config.save_processed:
+            # Processed writing.
+            from ..model.base import LightningModuleBase
+
+            if not isinstance(pl_module, LightningModuleBase):
+                raise ValueError(
+                    "The model must be a subclass of LightningModuleBase to use the distributed prediction writer."
+                )
+
+            output_dir = self.output_dir / "processed" / f"dataloader_{dataloader_idx}"
+            output_dir.mkdir(parents=True, exist_ok=True)
+
+            # Split into individual predictions
+            assert batch_indices is not None, (
+                "Batch indices must be provided for processed writing."
+            )
+            for sample in pl_module.split_batched_predictions(
+                batch, prediction, batch_indices
+            ):
+                sample = {
+                    **sample,
+                    "global_rank": trainer.global_rank,
+                    "world_size": trainer.world_size,
+                    "is_global_zero": trainer.is_global_zero,
+                }
+                save(sample, output_dir / f"{sample['index']}.pt")
+
+
+class DistributedPredictionReader(Sequence[tuple[Any, Any]]):
+    def __init__(self, output_dir: Path):
+        self.output_dir = output_dir
+
+    @override
+    def __len__(self) -> int:
+        return len(list(self.output_dir.glob("*.pt")))
+
+    @overload
+    def __getitem__(self, index: int) -> tuple[Any, Any]: ...
+
+    @overload
+    def __getitem__(self, index: slice) -> list[tuple[Any, Any]]: ...
+
+    @override
+    def __getitem__(
+        self, index: int | slice
+    ) -> tuple[Any, Any] | list[tuple[Any, Any]]:
+        if isinstance(index, slice):
+            # Handle slice indexing
+            indices = range(*index.indices(len(self)))
+            return [self.__getitem__(i) for i in indices]
+
+        # Handle integer indexing
+        path = self.output_dir / f"{index}.pt"
+        if not path.exists():
+            raise FileNotFoundError(f"File {path} does not exist.")
+        sample = torch.load(path)
+        return sample["batch"], sample["prediction"]
+
+    @override
+    def __iter__(self) -> Iterator[tuple[Any, Any]]:
+        for i in range(len(self)):
+            yield self[i]

nshtrainer/configs/__init__.py

@@ -21,6 +21,9 @@ from nshtrainer.callbacks import DebugFlagCallbackConfig as DebugFlagCallbackCon
 from nshtrainer.callbacks import (
     DirectorySetupCallbackConfig as DirectorySetupCallbackConfig,
 )
+from nshtrainer.callbacks import (
+    DistributedPredictionWriterConfig as DistributedPredictionWriterConfig,
+)
 from nshtrainer.callbacks import (
     EarlyStoppingCallbackConfig as EarlyStoppingCallbackConfig,
 )
@@ -95,9 +98,21 @@ from nshtrainer.nn.nonlinearity import (
     SwiGLUNonlinearityConfig as SwiGLUNonlinearityConfig,
 )
 from nshtrainer.nn.nonlinearity import nonlinearity_registry as nonlinearity_registry
+from nshtrainer.optimizer import AdadeltaConfig as AdadeltaConfig
+from nshtrainer.optimizer import AdafactorConfig as AdafactorConfig
+from nshtrainer.optimizer import AdagradConfig as AdagradConfig
+from nshtrainer.optimizer import AdamaxConfig as AdamaxConfig
+from nshtrainer.optimizer import AdamConfig as AdamConfig
 from nshtrainer.optimizer import AdamWConfig as AdamWConfig
+from nshtrainer.optimizer import ASGDConfig as ASGDConfig
+from nshtrainer.optimizer import NAdamConfig as NAdamConfig
 from nshtrainer.optimizer import OptimizerConfig as OptimizerConfig
 from nshtrainer.optimizer import OptimizerConfigBase as OptimizerConfigBase
+from nshtrainer.optimizer import RAdamConfig as RAdamConfig
+from nshtrainer.optimizer import RMSpropConfig as RMSpropConfig
+from nshtrainer.optimizer import RpropConfig as RpropConfig
+from nshtrainer.optimizer import SGDConfig as SGDConfig
+from nshtrainer.optimizer import Union as Union
 from nshtrainer.optimizer import optimizer_registry as optimizer_registry
 from nshtrainer.profiler import AdvancedProfilerConfig as AdvancedProfilerConfig
 from nshtrainer.profiler import BaseProfilerConfig as BaseProfilerConfig
@@ -225,11 +240,17 @@ from . import trainer as trainer
 from . import util as util
 
 __all__ = [
+    "ASGDConfig",
     "AcceleratorConfig",
     "AcceleratorConfigBase",
     "ActSaveConfig",
     "ActSaveLoggerConfig",
+    "AdadeltaConfig",
+    "AdafactorConfig",
+    "AdagradConfig",
+    "AdamConfig",
     "AdamWConfig",
+    "AdamaxConfig",
     "AdvancedProfilerConfig",
     "AsyncCheckpointIOPlugin",
     "BaseCheckpointCallbackConfig",
@@ -249,6 +270,7 @@ __all__ = [
     "DeepSpeedPluginConfig",
     "DirectoryConfig",
     "DirectorySetupCallbackConfig",
+    "DistributedPredictionWriterConfig",
     "DoublePrecisionPluginConfig",
     "DurationConfig",
     "ELUNonlinearityConfig",
@@ -294,6 +316,7 @@ __all__ = [
     "MetricValidationCallbackConfig",
     "MishNonlinearityConfig",
     "MixedPrecisionPluginConfig",
+    "NAdamConfig",
     "NonlinearityConfig",
     "NonlinearityConfigBase",
     "NormLoggingCallbackConfig",
@@ -306,10 +329,14 @@ __all__ = [
     "PrintTableMetricsCallbackConfig",
     "ProfilerConfig",
     "PyTorchProfilerConfig",
+    "RAdamConfig",
     "RLPSanityChecksCallbackConfig",
+    "RMSpropConfig",
     "RNGConfig",
     "ReLUNonlinearityConfig",
     "ReduceLROnPlateauConfig",
+    "RpropConfig",
+    "SGDConfig",
     "SLURMEnvironmentPlugin",
     "SanityCheckingConfig",
     "SharedParametersCallbackConfig",
@@ -331,6 +358,7 @@ __all__ = [
     "TorchSyncBatchNormPlugin",
     "TrainerConfig",
     "TransformerEnginePluginConfig",
+    "Union",
     "WandbLoggerConfig",
     "WandbUploadCodeCallbackConfig",
     "WandbWatchCallbackConfig",

nshtrainer/configs/callbacks/__init__.py

@@ -12,6 +12,9 @@ from nshtrainer.callbacks import DebugFlagCallbackConfig as DebugFlagCallbackCon
 from nshtrainer.callbacks import (
     DirectorySetupCallbackConfig as DirectorySetupCallbackConfig,
 )
+from nshtrainer.callbacks import (
+    DistributedPredictionWriterConfig as DistributedPredictionWriterConfig,
+)
 from nshtrainer.callbacks import (
     EarlyStoppingCallbackConfig as EarlyStoppingCallbackConfig,
 )
@@ -62,6 +65,7 @@ from . import base as base
 from . import checkpoint as checkpoint
 from . import debug_flag as debug_flag
 from . import directory_setup as directory_setup
+from . import distributed_prediction_writer as distributed_prediction_writer
 from . import early_stopping as early_stopping
 from . import ema as ema
 from . import finite_checks as finite_checks
@@ -86,6 +90,7 @@ __all__ = [
     "CheckpointMetadata",
     "DebugFlagCallbackConfig",
     "DirectorySetupCallbackConfig",
+    "DistributedPredictionWriterConfig",
     "EMACallbackConfig",
     "EarlyStoppingCallbackConfig",
     "EpochTimerCallbackConfig",
@@ -109,6 +114,7 @@ __all__ = [
     "checkpoint",
     "debug_flag",
     "directory_setup",
+    "distributed_prediction_writer",
     "early_stopping",
     "ema",
     "finite_checks",

nshtrainer/configs/callbacks/distributed_prediction_writer/__init__.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+__codegen__ = True
+
+from nshtrainer.callbacks.distributed_prediction_writer import (
+    CallbackConfigBase as CallbackConfigBase,
+)
+from nshtrainer.callbacks.distributed_prediction_writer import (
+    DistributedPredictionWriterConfig as DistributedPredictionWriterConfig,
+)
+from nshtrainer.callbacks.distributed_prediction_writer import (
+    callback_registry as callback_registry,
+)
+
+__all__ = [
+    "CallbackConfigBase",
+    "DistributedPredictionWriterConfig",
+    "callback_registry",
+]

nshtrainer/configs/optimizer/__init__.py

@@ -2,14 +2,38 @@ from __future__ import annotations
 
 __codegen__ = True
 
+from nshtrainer.optimizer import AdadeltaConfig as AdadeltaConfig
+from nshtrainer.optimizer import AdafactorConfig as AdafactorConfig
+from nshtrainer.optimizer import AdagradConfig as AdagradConfig
+from nshtrainer.optimizer import AdamaxConfig as AdamaxConfig
+from nshtrainer.optimizer import AdamConfig as AdamConfig
 from nshtrainer.optimizer import AdamWConfig as AdamWConfig
+from nshtrainer.optimizer import ASGDConfig as ASGDConfig
+from nshtrainer.optimizer import NAdamConfig as NAdamConfig
 from nshtrainer.optimizer import OptimizerConfig as OptimizerConfig
 from nshtrainer.optimizer import OptimizerConfigBase as OptimizerConfigBase
+from nshtrainer.optimizer import RAdamConfig as RAdamConfig
+from nshtrainer.optimizer import RMSpropConfig as RMSpropConfig
+from nshtrainer.optimizer import RpropConfig as RpropConfig
+from nshtrainer.optimizer import SGDConfig as SGDConfig
+from nshtrainer.optimizer import Union as Union
 from nshtrainer.optimizer import optimizer_registry as optimizer_registry
 
 __all__ = [
+    "ASGDConfig",
+    "AdadeltaConfig",
+    "AdafactorConfig",
+    "AdagradConfig",
+    "AdamConfig",
     "AdamWConfig",
+    "AdamaxConfig",
+    "NAdamConfig",
     "OptimizerConfig",
     "OptimizerConfigBase",
+    "RAdamConfig",
+    "RMSpropConfig",
+    "RpropConfig",
+    "SGDConfig",
+    "Union",
     "optimizer_registry",
 ]

nshtrainer/configs/trainer/__init__.py

@@ -22,6 +22,9 @@ from nshtrainer.trainer._config import (
     DebugFlagCallbackConfig as DebugFlagCallbackConfig,
 )
 from nshtrainer.trainer._config import DirectoryConfig as DirectoryConfig
+from nshtrainer.trainer._config import (
+    DistributedPredictionWriterConfig as DistributedPredictionWriterConfig,
+)
 from nshtrainer.trainer._config import (
     EarlyStoppingCallbackConfig as EarlyStoppingCallbackConfig,
 )
@@ -149,6 +152,7 @@ __all__ = [
     "DebugFlagCallbackConfig",
     "DeepSpeedPluginConfig",
     "DirectoryConfig",
+    "DistributedPredictionWriterConfig",
     "DoublePrecisionPluginConfig",
     "EarlyStoppingCallbackConfig",
     "EnvironmentConfig",

nshtrainer/configs/trainer/_config/__init__.py

@@ -18,6 +18,9 @@ from nshtrainer.trainer._config import (
     DebugFlagCallbackConfig as DebugFlagCallbackConfig,
 )
 from nshtrainer.trainer._config import DirectoryConfig as DirectoryConfig
+from nshtrainer.trainer._config import (
+    DistributedPredictionWriterConfig as DistributedPredictionWriterConfig,
+)
 from nshtrainer.trainer._config import (
     EarlyStoppingCallbackConfig as EarlyStoppingCallbackConfig,
 )
@@ -70,6 +73,7 @@ __all__ = [
     "CheckpointSavingConfig",
     "DebugFlagCallbackConfig",
     "DirectoryConfig",
+    "DistributedPredictionWriterConfig",
     "EarlyStoppingCallbackConfig",
     "EnvironmentConfig",
     "GradientClippingConfig",

nshtrainer/model/base.py

@@ -2,9 +2,9 @@ from __future__ import annotations
 
 import logging
 from abc import ABC, abstractmethod
-from collections.abc import Callable, Mapping
+from collections.abc import Callable, Iterable, Mapping, Sequence
 from pathlib import Path
-from typing import Any, Generic, Literal, cast
+from typing import Any, Generic, Literal, TypedDict, cast
 
 import nshconfig as C
 import torch
@@ -53,6 +53,47 @@ VALID_REDUCE_OPS = (
 )
 
 
+class IndividualSample(TypedDict):
+    """
+    A dictionary that contains the individual sample.
+    This is used to split the batched predictions into individual predictions.
+    """
+
+    index: int
+    """The index of the sample in the batch."""
+
+    batch: Any
+    """The batch to split."""
+
+    prediction: Any
+    """The batched prediction to split."""
+
+
+def default_split_batched_predictions(
+    batch: Any,
+    prediction: Any,
+    batch_indices: Sequence[Any],
+) -> Iterable[IndividualSample]:
+    """
+    Splits the batched predictions into a list of individual predictions.
+    Args:
+        batch: The batch to split.
+        prediction: The batched prediction to split.
+        batch_indices: The indices of the batches.
+    Returns:
+        A tuple of two sequences: the corresponding batches and the individual predictions.
+    """
+    import torch.utils._pytree as tree
+
+    for sample_idx, batch_idx in enumerate(batch_indices):
+        # Create a dictionary for each sample
+        yield IndividualSample(
+            index=batch_idx,
+            batch=tree.tree_map(lambda x: x[sample_idx], batch),
+            prediction=tree.tree_map(lambda x: x[sample_idx], prediction),
+        )
+
+
 class LightningModuleBase(
     DebugModuleMixin,
     RLPSanityCheckModuleMixin,
@@ -171,6 +212,23 @@ class LightningModuleBase(
         loss = cast(torch.Tensor, loss)
         return loss
 
+    def split_batched_predictions(
+        self,
+        batch: Any,
+        prediction: Any,
+        batch_indices: Sequence[Any],
+    ) -> Iterable[IndividualSample]:
+        """
+        Splits the batched predictions into a list of individual predictions.
+        Args:
+            batch: The batch to split.
+            prediction: The batched prediction to split.
+            batch_indices: The indices of the batches.
+        Returns:
+            A tuple of two sequences: the corresponding batches and the individual predictions.
+        """
+        return default_split_batched_predictions(batch, prediction, batch_indices)
+
     @override
     @classmethod
     def load_from_checkpoint(cls, *args, **kwargs) -> Never:

nshtrainer/optimizer.py

@@ -2,10 +2,11 @@ from __future__ import annotations
 
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
-from typing import Annotated, Any, Literal
+from typing import Annotated, Any, Literal, Tuple, Union
 
 import nshconfig as C
 import torch.nn as nn
+from torch import Tensor
 from torch.optim import Optimizer
 from typing_extensions import TypeAliasType, final, override
 
@@ -45,6 +46,18 @@ class AdamWConfig(OptimizerConfigBase):
     amsgrad: bool = False
     """Whether to use the AMSGrad variant of this algorithm."""
 
+    maximize: bool = False
+    """Maximize the objective with respect to the params, instead of minimizing."""
+
+    foreach: bool | None = None
+    """Whether foreach implementation of optimizer is used."""
+
+    capturable: bool = False
+    """Whether this instance is safe to capture in a CUDA graph."""
+
+    differentiable: bool = False
+    """Whether autograd should occur through the optimizer step in training."""
+
     @override
     def create_optimizer(
         self,
@@ -59,6 +72,551 @@ class AdamWConfig(OptimizerConfigBase):
             betas=self.betas,
             eps=self.eps,
             amsgrad=self.amsgrad,
+            maximize=self.maximize,
+            foreach=self.foreach,
+            capturable=self.capturable,
+            differentiable=self.differentiable,
+        )
+
+
+@final
+@optimizer_registry.register
+class AdafactorConfig(OptimizerConfigBase):
+    name: Literal["adafactor"] = "adafactor"
+    lr: float
+    """Learning rate for the optimizer. If None, uses relative step size."""
+
+    eps1: float | None = None
+    """Term added to the denominator to improve numerical stability (default: None)."""
+
+    eps2: float = 1e-3
+    """Term added to the denominator to improve numerical stability (default: 1e-3)."""
+
+    beta2_decay: float = -0.8
+    """Coefficient used for computing running averages of square gradient (default: -0.8)."""
+
+    weight_decay: float = 0.0
+    """Weight decay (L2 penalty) (default: 0.0)."""
+
+    maximize: bool = False
+    """Maximize the params based on the objective, instead of minimizing."""
+
+    @override
+    def create_optimizer(
+        self,
+        parameters: Iterable[nn.Parameter] | Iterable[dict[str, Any]],
+    ):
+        from torch.optim import Adafactor
+
+        return Adafactor(
+            parameters,
+            lr=self.lr,
+            eps=(self.eps1, self.eps2),
+            beta2_decay=self.beta2_decay,
+            weight_decay=self.weight_decay,
+            maximize=self.maximize,
+        )
+
+
+@final
+@optimizer_registry.register
+class AdadeltaConfig(OptimizerConfigBase):
+    name: Literal["adadelta"] = "adadelta"
+
+    lr: float
+    """Learning rate for the optimizer."""
+
+    rho: float = 0.9
+    """Coefficient used for computing a running average of squared gradients."""
+
+    eps: float = 1e-6
+    """Term added to the denominator to improve numerical stability."""
+
+    weight_decay: float = 0.0
+    """Weight decay (L2 penalty) for the optimizer."""
+
+    maximize: bool = False
+    """Maximize the params based on the objective, instead of minimizing."""
+
+    foreach: bool | None = None
+    """Whether foreach implementation of optimizer is used."""
+
+    capturable: bool = False
+    """Whether this instance is safe to capture in a CUDA graph."""
+
+    differentiable: bool = False
+    """Whether autograd should occur through the optimizer step in training."""
+
+    @override
+    def create_optimizer(
+        self,
+        parameters: Iterable[nn.Parameter] | Iterable[dict[str, Any]],
+    ):
+        from torch.optim import Adadelta
+
+        return Adadelta(
+            parameters,
+            lr=self.lr,
+            rho=self.rho,
+            eps=self.eps,
+            weight_decay=self.weight_decay,
+            maximize=self.maximize,
+            foreach=self.foreach,
+            capturable=self.capturable,
+            differentiable=self.differentiable,
+        )
+
+
+@final
+@optimizer_registry.register
+class AdagradConfig(OptimizerConfigBase):
+    name: Literal["adagrad"] = "adagrad"
+
+    lr: float
+    """Learning rate for the optimizer."""
+
+    lr_decay: float = 0.0
+    """Learning rate decay."""
+
+    weight_decay: float = 0.0
+    """Weight decay (L2 penalty) for the optimizer."""
+
+    initial_accumulator_value: float = 0.0
+    """Initial value for the accumulator."""
+
+    eps: float = 1e-10
+    """Term added to the denominator to improve numerical stability."""
+
+    maximize: bool = False
+    """Maximize the params based on the objective, instead of minimizing."""
+
+    foreach: bool | None = None
+    """Whether foreach implementation of optimizer is used."""
+
+    differentiable: bool = False
+    """Whether autograd should occur through the optimizer step in training."""
+
+    fused: bool | None = None
+    """Whether the fused implementation is used."""
+
+    @override
+    def create_optimizer(
+        self,
+        parameters: Iterable[nn.Parameter] | Iterable[dict[str, Any]],
+    ):
+        from torch.optim import Adagrad
+
+        return Adagrad(
+            parameters,
+            lr=self.lr,
+            lr_decay=self.lr_decay,
+            weight_decay=self.weight_decay,
+            initial_accumulator_value=self.initial_accumulator_value,
+            eps=self.eps,
+            maximize=self.maximize,
+            foreach=self.foreach,
+            differentiable=self.differentiable,
+            fused=self.fused,
+        )
+
+
+@final
+@optimizer_registry.register
+class AdamConfig(OptimizerConfigBase):
+    name: Literal["adam"] = "adam"
+
+    lr: float
+    """Learning rate for the optimizer."""
+
+    betas: tuple[float, float] = (0.9, 0.999)
+    """Coefficients used for computing running averages of gradient and its square."""
+
+    eps: float = 1e-8
+    """Term added to the denominator to improve numerical stability."""
+
+    weight_decay: float = 0.0
+    """Weight decay (L2 penalty) for the optimizer."""
+
+    amsgrad: bool = False
+    """Whether to use the AMSGrad variant of this algorithm."""
+
+    maximize: bool = False
+    """Maximize the params based on the objective, instead of minimizing."""
+
+    foreach: bool | None = None
+    """Whether foreach implementation of optimizer is used."""
+
+    capturable: bool = False
+    """Whether this instance is safe to capture in a CUDA graph."""
+
+    differentiable: bool = False
+    """Whether autograd should occur through the optimizer step in training."""
+
+    fused: bool | None = None
+    """Whether the fused implementation is used."""
+
+    @override
+    def create_optimizer(
+        self,
+        parameters: Iterable[nn.Parameter] | Iterable[dict[str, Any]],
+    ):
+        from torch.optim import Adam
+
+        return Adam(
+            parameters,
+            lr=self.lr,
+            betas=self.betas,
+            eps=self.eps,
+            weight_decay=self.weight_decay,
+            amsgrad=self.amsgrad,
+            maximize=self.maximize,
+            foreach=self.foreach,
+            capturable=self.capturable,
+            differentiable=self.differentiable,
+            fused=self.fused,
+        )
+
+
+@final
+@optimizer_registry.register
+class AdamaxConfig(OptimizerConfigBase):
+    name: Literal["adamax"] = "adamax"
+
+    lr: float
+    """Learning rate for the optimizer."""
+
+    betas: tuple[float, float] = (0.9, 0.999)
+    """Coefficients used for computing running averages of gradient and its square."""
+
+    eps: float = 1e-8
+    """Term added to the denominator to improve numerical stability."""
+
+    weight_decay: float = 0.0
+    """Weight decay (L2 penalty) for the optimizer."""
+
+    maximize: bool = False
+    """Maximize the params based on the objective, instead of minimizing."""
+
+    foreach: bool | None = None
+    """Whether foreach implementation of optimizer is used."""
+
+    capturable: bool = False
+    """Whether this instance is safe to capture in a CUDA graph."""
+
+    differentiable: bool = False
+    """Whether autograd should occur through the optimizer step in training."""
+
+    @override
+    def create_optimizer(
+        self,
+        parameters: Iterable[nn.Parameter] | Iterable[dict[str, Any]],
+    ):
+        from torch.optim import Adamax
+
+        return Adamax(
+            parameters,
+            lr=self.lr,
+            betas=self.betas,
+            eps=self.eps,
+            weight_decay=self.weight_decay,
+            maximize=self.maximize,
+            foreach=self.foreach,
+            capturable=self.capturable,
+            differentiable=self.differentiable,
+        )
+
+
+@final
+@optimizer_registry.register
+class ASGDConfig(OptimizerConfigBase):
+    name: Literal["asgd"] = "asgd"
+
+    lr: float
+    """Learning rate for the optimizer."""
+
+    lambd: float = 1e-4
+    """Decay term."""
+
+    alpha: float = 0.75
+    """Power for eta update."""
+
+    t0: float = 1e6
+    """Point at which to start averaging."""
+
+    weight_decay: float = 0.0
+    """Weight decay (L2 penalty) for the optimizer."""
+
+    maximize: bool = False
+    """Maximize the params based on the objective, instead of minimizing."""
+
+    @override
+    def create_optimizer(
+        self,
+        parameters: Iterable[nn.Parameter] | Iterable[dict[str, Any]],
+    ):
+        from torch.optim import ASGD
+
+        return ASGD(
+            parameters,
+            lr=self.lr,
+            lambd=self.lambd,
+            alpha=self.alpha,
+            t0=self.t0,
+            weight_decay=self.weight_decay,
+            maximize=self.maximize,
+        )
+
+
+@final
+@optimizer_registry.register
+class NAdamConfig(OptimizerConfigBase):
+    name: Literal["nadam"] = "nadam"
+
+    lr: float
+    """Learning rate for the optimizer."""
+
+    betas: tuple[float, float] = (0.9, 0.999)
+    """Coefficients used for computing running averages of gradient and its square."""
+
+    eps: float = 1e-8
+    """Term added to the denominator to improve numerical stability."""
+
+    weight_decay: float = 0.0
+    """Weight decay (L2 penalty) for the optimizer."""
+
+    momentum_decay: float = 4e-3
+    """Momentum decay."""
+
+    decoupled_weight_decay: bool = False
+    """Whether to use decoupled weight decay."""
+
+    maximize: bool = False
+    """Maximize the params based on the objective, instead of minimizing."""
+
+    foreach: bool | None = None
+    """Whether foreach implementation of optimizer is used."""
+
+    capturable: bool = False
+    """Whether this instance is safe to capture in a CUDA graph."""
+
+    differentiable: bool = False
+    """Whether autograd should occur through the optimizer step in training."""
+
+    @override
+    def create_optimizer(
+        self,
+        parameters: Iterable[nn.Parameter] | Iterable[dict[str, Any]],
+    ):
+        from torch.optim import NAdam
+
+        return NAdam(
+            parameters,
+            lr=self.lr,
+            betas=self.betas,
+            eps=self.eps,
+            weight_decay=self.weight_decay,
+            momentum_decay=self.momentum_decay,
+            decoupled_weight_decay=self.decoupled_weight_decay,
+            maximize=self.maximize,
+            foreach=self.foreach,
+            capturable=self.capturable,
+            differentiable=self.differentiable,
+        )
+
+
+@final
+@optimizer_registry.register
+class RAdamConfig(OptimizerConfigBase):
+    name: Literal["radam"] = "radam"
+
+    lr: float
+    """Learning rate for the optimizer."""
+
+    betas: tuple[float, float] = (0.9, 0.999)
+    """Coefficients used for computing running averages of gradient and its square."""
+
+    eps: float = 1e-8
+    """Term added to the denominator to improve numerical stability."""
+
+    weight_decay: float = 0.0
+    """Weight decay (L2 penalty) for the optimizer."""
+
+    decoupled_weight_decay: bool = False
+    """Whether to use decoupled weight decay."""
+
+    maximize: bool = False
+    """Maximize the params based on the objective, instead of minimizing."""
+
+    foreach: bool | None = None
+    """Whether foreach implementation of optimizer is used."""
+
+    capturable: bool = False
+    """Whether this instance is safe to capture in a CUDA graph."""
+
+    differentiable: bool = False
+    """Whether autograd should occur through the optimizer step in training."""
+
+    @override
+    def create_optimizer(
+        self,
+        parameters: Iterable[nn.Parameter] | Iterable[dict[str, Any]],
+    ):
+        from torch.optim import RAdam
+
+        return RAdam(
+            parameters,
+            lr=self.lr,
+            betas=self.betas,
+            eps=self.eps,
+            weight_decay=self.weight_decay,
+            decoupled_weight_decay=self.decoupled_weight_decay,
+            maximize=self.maximize,
+            foreach=self.foreach,
+            capturable=self.capturable,
+            differentiable=self.differentiable,
+        )
+
+
+@final
+@optimizer_registry.register
+class RMSpropConfig(OptimizerConfigBase):
+    name: Literal["rmsprop"] = "rmsprop"
+
+    lr: float
+    """Learning rate for the optimizer."""
+
+    alpha: float = 0.99
+    """Smoothing constant."""
+
+    eps: float = 1e-8
+    """Term added to the denominator to improve numerical stability."""
+
+    weight_decay: float = 0.0
+    """Weight decay (L2 penalty) for the optimizer."""
+
+    momentum: float = 0.0
+    """Momentum factor."""
+
+    centered: bool = False
+    """If True, compute the centered RMSProp, the gradient is normalized by an estimation of its variance."""
+
+    maximize: bool = False
+    """Maximize the params based on the objective, instead of minimizing."""
+
+    foreach: bool | None = None
+    """Whether foreach implementation of optimizer is used."""
+
+    capturable: bool = False
+    """Whether this instance is safe to capture in a CUDA graph."""
+
+    differentiable: bool = False
+    """Whether autograd should occur through the optimizer step in training."""
+
+    @override
+    def create_optimizer(
+        self,
+        parameters: Iterable[nn.Parameter] | Iterable[dict[str, Any]],
+    ):
+        from torch.optim import RMSprop
+
+        return RMSprop(
+            parameters,
+            lr=self.lr,
+            alpha=self.alpha,
+            eps=self.eps,
+            weight_decay=self.weight_decay,
+            momentum=self.momentum,
+            centered=self.centered,
+            maximize=self.maximize,
+            foreach=self.foreach,
+            capturable=self.capturable,
+            differentiable=self.differentiable,
+        )
+
+
+@final
+@optimizer_registry.register
+class RpropConfig(OptimizerConfigBase):
+    name: Literal["rprop"] = "rprop"
+
+    lr: float
+    """Learning rate for the optimizer."""
+
+    etas: tuple[float, float] = (0.5, 1.2)
+    """Pair of (etaminus, etaplus), multiplicative increase and decrease factors."""
+
+    step_sizes: tuple[float, float] = (1e-6, 50.0)
+    """Pair of minimal and maximal allowed step sizes."""
+
+    maximize: bool = False
+    """Maximize the params based on the objective, instead of minimizing."""
+
+    @override
+    def create_optimizer(
+        self,
+        parameters: Iterable[nn.Parameter] | Iterable[dict[str, Any]],
+    ):
+        from torch.optim import Rprop
+
+        return Rprop(
+            parameters,
+            lr=self.lr,
+            etas=self.etas,
+            step_sizes=self.step_sizes,
+            maximize=self.maximize,
+        )
+
+
+@final
+@optimizer_registry.register
+class SGDConfig(OptimizerConfigBase):
+    name: Literal["sgd"] = "sgd"
+
+    lr: float
+    """Learning rate for the optimizer."""
+
+    momentum: float = 0.0
+    """Momentum factor."""
+
+    dampening: float = 0.0
+    """Dampening for momentum."""
+
+    weight_decay: float = 0.0
+    """Weight decay (L2 penalty) for the optimizer."""
+
+    nesterov: bool = False
+    """Enables Nesterov momentum."""
+
+    maximize: bool = False
+    """Maximize the params based on the objective, instead of minimizing."""
+
+    foreach: bool | None = None
+    """Whether foreach implementation of optimizer is used."""
+
+    differentiable: bool = False
+    """Whether autograd should occur through the optimizer step in training."""
+
+    fused: bool | None = None
+    """Whether the fused implementation is used."""
+
+    @override
+    def create_optimizer(
+        self,
+        parameters: Iterable[nn.Parameter] | Iterable[dict[str, Any]],
+    ):
+        from torch.optim import SGD
+
+        return SGD(
+            parameters,
+            lr=self.lr,
+            momentum=self.momentum,
+            dampening=self.dampening,
+            weight_decay=self.weight_decay,
+            nesterov=self.nesterov,
+            maximize=self.maximize,
+            foreach=self.foreach,
+            differentiable=self.differentiable,
+            fused=self.fused,
         )
 
 

nshtrainer/trainer/_config.py

@@ -31,6 +31,7 @@ from .._hf_hub import HuggingFaceHubConfig
 from ..callbacks import (
     BestCheckpointCallbackConfig,
     CallbackConfig,
+    DistributedPredictionWriterConfig,
     EarlyStoppingCallbackConfig,
     LastCheckpointCallbackConfig,
     NormLoggingCallbackConfig,
@@ -701,6 +702,14 @@ class TrainerConfig(C.Config):
     auto_validate_metrics: MetricValidationCallbackConfig | None = None
     """If enabled, will automatically validate the metrics before starting the training routine."""
 
+    distributed_predict: DistributedPredictionWriterConfig | None = (
+        DistributedPredictionWriterConfig()
+    )
+    """If enabled, will use a custom BasePredictionWriter callback to automatically
+    handle distributed prediction. This is useful for running prediction on multiple GPUs
+    seamlessly.
+    """
+
     lightning_kwargs: LightningTrainerKwargs = LightningTrainerKwargs()
     """
     Additional keyword arguments to pass to the Lightning `pl.Trainer` constructor.
@@ -752,10 +761,6 @@ class TrainerConfig(C.Config):
         )
 
     def _nshtrainer_all_callback_configs(self) -> Iterable[CallbackConfigBase | None]:
-        # Disable all callbacks if barebones mode is enabled
-        if self.barebones:
-            return
-
         yield self.early_stopping
         yield self.checkpoint_saving
         yield self.lr_monitor
@@ -772,6 +777,7 @@ class TrainerConfig(C.Config):
         yield self.reduce_lr_on_plateau_sanity_checking
         yield self.auto_set_debug_flag
         yield self.auto_validate_metrics
+        yield self.distributed_predict
         yield from self.callbacks
 
     def _nshtrainer_all_logger_configs(self) -> Iterable[LoggerConfigBase | None]:

nshtrainer/trainer/trainer.py

@@ -10,12 +10,16 @@ import torch
 from lightning.fabric.plugins.environments.lsf import LSFEnvironment
 from lightning.fabric.plugins.environments.slurm import SLURMEnvironment
 from lightning.fabric.plugins.precision.precision import _PRECISION_INPUT
-from lightning.pytorch import LightningModule
+from lightning.pytorch import LightningDataModule, LightningModule
 from lightning.pytorch import Trainer as LightningTrainer
 from lightning.pytorch.callbacks import Callback
 from lightning.pytorch.profilers import Profiler
 from lightning.pytorch.trainer.states import TrainerFn
-from lightning.pytorch.utilities.types import _EVALUATE_OUTPUT, _PREDICT_OUTPUT
+from lightning.pytorch.utilities.types import (
+    _EVALUATE_OUTPUT,
+    _PREDICT_OUTPUT,
+    EVAL_DATALOADERS,
+)
 from typing_extensions import Never, Unpack, assert_never, deprecated, override
 
 from .._checkpoint.metadata import write_checkpoint_metadata
@@ -532,3 +536,18 @@ class Trainer(LightningTrainer):
             update_hparams_dict=update_hparams_dict,
         )
         return cls(hparams)
+
+    def distributed_predict(
+        self,
+        model: LightningModule | None = None,
+        dataloaders: EVAL_DATALOADERS | LightningDataModule | None = None,
+        datamodule: LightningDataModule | None = None,
+        ckpt_path: str | Path | None = None,
+    ):
+        self.predict(
+            model,
+            dataloaders,
+            datamodule,
+            return_predictions=False,
+            ckpt_path=ckpt_path,
+        )

{nshtrainer-1.1.1b1.dist-info → nshtrainer-1.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: nshtrainer
-Version: 1.1.1b1
+Version: 1.2.0
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-1.1.1b1.dist-info → nshtrainer-1.2.0.dist-info}/RECORD

@@ -3,12 +3,12 @@ nshtrainer/__init__.py,sha256=VcqBfL8RgCcZDaY645nxeDmOspqerx4x46wggCMnS0E,692
 nshtrainer/_callback.py,sha256=ZDppiJ4d65tRXTEWYPZLH_F1xFizdz1pkWJe_sQ5uII,12564
 nshtrainer/_checkpoint/metadata.py,sha256=Hh5a7OkdknUEbkEwX6vS88-XLEeuVDoR6a3en2uLzQE,5597
 nshtrainer/_checkpoint/saver.py,sha256=utcrYKSosd04N9m2GIylufO5DO05D90qVU3mvadfApU,1658
-nshtrainer/_directory.py,sha256=TJR9ccyuzRlAVfVjGyeQ3E2AFAcz-XbBCxWfiXo2SlY,3191
+nshtrainer/_directory.py,sha256=SuXJe9xJXZkDXWWfeOS9rEDz6vZUA6mpnEdkAW0ZQnY,3193
 nshtrainer/_experimental/__init__.py,sha256=U4S_2y3zgLZVfMenHRaJFBW8yqh2mUBuI291LGQVOJ8,35
 nshtrainer/_hf_hub.py,sha256=4OsCbIITnZk_YLyoMrVyZ0SIN04FBxlC0ig2Et8UAdo,14287
-nshtrainer/callbacks/__init__.py,sha256=w80d6PGNu3wjUj9NiRGMqCX9NnXD5ZlvbY-DIK4zjPE,3766
+nshtrainer/callbacks/__init__.py,sha256=m6eJuprZfBELuKpngKXre33B9yPXkG7jlKVmI-0yXRQ,4000
 nshtrainer/callbacks/actsave.py,sha256=NSXIIu62MNYe5gz479SMW33bdoKYoYtWtd_iTWFpKpc,3881
-nshtrainer/callbacks/base.py,sha256=Alaou1IHAIlMEM7g58d_02ozY2xWlshBN7fsw5Ee21s,3683
+nshtrainer/callbacks/base.py,sha256=K9aom1WVVRYxl-tHWgtmDUQZ1o63NgznvLsjauTKcCc,4225
 nshtrainer/callbacks/checkpoint/__init__.py,sha256=l8tkHc83_mLiU0-wT09SWdRzwpm2ulbkLzcuCmuTwzE,620
 nshtrainer/callbacks/checkpoint/_base.py,sha256=f7lpk8W4xqxk3PolBEU3AWt9VTIpoLW7wMUhC5DNm3c,6345
 nshtrainer/callbacks/checkpoint/best_checkpoint.py,sha256=aCs3E1eucfDlUeW2Iq_Ke7hb96BxHanmvn7PCCbqq0E,2648
@@ -16,6 +16,7 @@ nshtrainer/callbacks/checkpoint/last_checkpoint.py,sha256=vn-as3ex7kaTRcKsIurVtM
 nshtrainer/callbacks/checkpoint/on_exception_checkpoint.py,sha256=nljzETqkHwA-4g8mxaeFK5HxA8My0dlIPzIUscSMWyk,3525
 nshtrainer/callbacks/debug_flag.py,sha256=96fuP0C7C6dSs1GiMeUYzzs0X3Q4Pjt9JVWg3b75fU4,1748
 nshtrainer/callbacks/directory_setup.py,sha256=wPas_Ren8ANejogmIdKhqqgj4ulxz9AS_8xVIAfRXa0,2565
+nshtrainer/callbacks/distributed_prediction_writer.py,sha256=OSh2C6XF7Nki4eFByNVhwlt69izkxnlmfPx54w4rvBo,5274
 nshtrainer/callbacks/early_stopping.py,sha256=rC_qYKCQWjRQJFo0ky46uG0aDJdYP8vsSlKunk0bUVI,4765
 nshtrainer/callbacks/ema.py,sha256=dBFiUXG0xmyCw8-ayuSzJMKqSbepl6Ii5VIbhFlT5ug,12255
 nshtrainer/callbacks/finite_checks.py,sha256=3lZ3kEIjmYQfqTF0DcrgZ9_98ZLQhQj8usH7SgWst3o,2185
@@ -32,12 +33,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=4WNs4Zv4PtHWD0KKH4X7j_zFt-COrEB0KhNIljsA6Rc,14740
+nshtrainer/configs/__init__.py,sha256=KD3uClMwnA4LfQ7rY5phDdUbp3j8NoZfaGbGPbpaJVs,15848
 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=PB3Jg-8_vMhp-mCFw2_Tqt05drKwHK6Ovl9mb8NNiXs,4506
+nshtrainer/configs/callbacks/__init__.py,sha256=tP9urR73NIanyxpbi4EERsxOnGNiptbQpmsj-v53a38,4774
 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
@@ -47,6 +48,7 @@ nshtrainer/configs/callbacks/checkpoint/last_checkpoint/__init__.py,sha256=SIRfz
 nshtrainer/configs/callbacks/checkpoint/on_exception_checkpoint/__init__.py,sha256=VSkO0TYCAYy_9mQuOBoAND7D3Cg6w6nMCpqivQZLPcE,551
 nshtrainer/configs/callbacks/debug_flag/__init__.py,sha256=s_ifB-DbZjar0w11pr2oVAlcMTWWMnK_tCNilfswL04,425
 nshtrainer/configs/callbacks/directory_setup/__init__.py,sha256=e8GCRy2Alds3AXLwp4ieSGtn8S0YjmKJ5khOaQ0zKGs,464
+nshtrainer/configs/callbacks/distributed_prediction_writer/__init__.py,sha256=npO97m5inRgAnGtGBwz_MNJz44B2cG4j9LZFCllQcrk,530
 nshtrainer/configs/callbacks/early_stopping/__init__.py,sha256=m8N6H11PjqcWqXP5ZxWC8L4PHMUI6avYyN5rUNprjuQ,546
 nshtrainer/configs/callbacks/ema/__init__.py,sha256=DUJrbDD8wWX_s0_4dwKpT_IWKSVpBmhe4-1aELq7G6w,377
 nshtrainer/configs/callbacks/finite_checks/__init__.py,sha256=e-vx9Kn-noqw4wPvZw7fDMfb9Tsa6Duk0TIa8ZIgIIE,443
@@ -77,14 +79,14 @@ nshtrainer/configs/nn/__init__.py,sha256=Ms2gIqbRxNVm6GHKCddCJTTqMwUPifjjHD_fCfJ
 nshtrainer/configs/nn/mlp/__init__.py,sha256=O6kQ6utZNJPG9Fax5pRdZcHa3J-XFKKdXcc_PQg0jk0,347
 nshtrainer/configs/nn/nonlinearity/__init__.py,sha256=LCTbTyelCMABVw505CGQ4UpEGlAnIhflSLFqwAQXLQA,2155
 nshtrainer/configs/nn/rng/__init__.py,sha256=4iC6vwxbfNeXyvpwZ1Z5Kcy-he4cu7mg3UpLD-RLrHc,141
-nshtrainer/configs/optimizer/__init__.py,sha256=itIDIHQvGm50eZ7JLyNElahnNUMPJ__4PMmTjc0RQ6o,444
+nshtrainer/configs/optimizer/__init__.py,sha256=8ztp5UD-edfzwF-qdJTeZwlv-YWJ5Sn230b9aWxJyQQ,1398
 nshtrainer/configs/profiler/__init__.py,sha256=2ssaIpfVnvcbfNvZ-JeKp1Cx4NO1LknkVqTm1hu7Lvw,768
 nshtrainer/configs/profiler/_base/__init__.py,sha256=ekYfPg-VDhCAFM5nJka2TxUYdRDm1CKqjwUOQNbQjD4,176
 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=a8pzGVid52abAVARPbgjaN566H1ZM44FH_x95bsBaGE,7880
-nshtrainer/configs/trainer/_config/__init__.py,sha256=6DXdtP-uH11TopQ7kzId9fco-wVkD7ZfevbBqDpN6TE,3817
+nshtrainer/configs/trainer/__init__.py,sha256=PF9rYuVpk0IuhjcxS_hmBTT6A0oq7AWZDcx0Gfqi7MM,8040
+nshtrainer/configs/trainer/_config/__init__.py,sha256=5B8pjyNHfyFJ6p8dD5VSHD1tw2CcZ87Eq2C_Req3t60,3977
 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=slW5z1FZw2qICXO9l9DnLIDB1Yl7KOcxPEZkyYIHrp4,276
@@ -116,7 +118,7 @@ nshtrainer/lr_scheduler/reduce_lr_on_plateau.py,sha256=irPyDjfUX843ze4bJM9sW8WSe
 nshtrainer/metrics/__init__.py,sha256=Nqkn_jsDf3n5WtfMcnaaEftYjIIT2b-S7rmsB1MOMkU,86
 nshtrainer/metrics/_config.py,sha256=ox_ScK6V0J9nzIMhEB0qpToNKpt83VVgOVSRFCV-wBc,595
 nshtrainer/model/__init__.py,sha256=3G-bwPPSRStWdsdwG9-rn0bXcRpEiP1BiQpF_qavtls,97
-nshtrainer/model/base.py,sha256=LsOK5mMhYG5J0eSFKZKdd1fTvr38sgi8LLVSqoW6OCU,8386
+nshtrainer/model/base.py,sha256=Pv3M3QStWQp-DnfGFsLPAmp87HHrX1NrkAa4JcyBoDk,10255
 nshtrainer/model/mixins/callback.py,sha256=0LPgve4VszHbLipid4mpI1qnnmdGS2spivs0dXLvqHw,3154
 nshtrainer/model/mixins/debug.py,sha256=ydLuAAaa7M5bX0gougZ5gWuZnvn4Ra9assal3IZ9hq8,2086
 nshtrainer/model/mixins/logger.py,sha256=7u9fQig-SVFA9RFIB4U0gqJAzruh49mgmXXvZ6VkDUk,11694
@@ -126,14 +128,14 @@ nshtrainer/nn/module_dict.py,sha256=9plb8aQUx5TUEPhX5jI9u8LrpTeKe7jZAHi8iIqcN8w,
 nshtrainer/nn/module_list.py,sha256=UB43pcwD_3nUke_DyLQt-iXKhWdKM6Zjm84lRC1hPYA,1755
 nshtrainer/nn/nonlinearity.py,sha256=xmaL4QCRvCxqmaGIOwetJeKK-6IK4m2OV7D3SjxSwJQ,6322
 nshtrainer/nn/rng.py,sha256=IJGvX9v8qBkfgBrMlNU2aj-MbYTPoncFyJzvPkzCQpM,512
-nshtrainer/optimizer.py,sha256=u968GRNPUNn3f_9BEY2RBNuJq5O3wJWams3NG0dkrOA,1738
+nshtrainer/optimizer.py,sha256=8pjOny7NxIt04PXxn3zOyJ2soL7nmj8yBVV82r_tNsc,17522
 nshtrainer/profiler/__init__.py,sha256=RjaNBoVcTFu8lF0dNlFp-2LaPYdonoIbDy2_KhgF0Ek,594
 nshtrainer/profiler/_base.py,sha256=kFcSVn9gJuMwgDxbfyHh46CmEAIPZjxw3yjPbKgzvwA,950
 nshtrainer/profiler/advanced.py,sha256=XrM3FX0ThCv5UwUrrH0l4Ow4LGAtpiBww2N8QAU5NOQ,1160
 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=s-_XoLc9mbNAdroRJyOKd3dLTyrFLQkPyGJkKDmBYf8,33267
+nshtrainer/trainer/_config.py,sha256=tdWAYh-KGXBpgdY8fwvOejjRZN-AS2Ze0f_9s2VEuZ0,33556
 nshtrainer/trainer/_log_hparams.py,sha256=XH2lZ4U_3AZBhOt91ocsEhdL_NRz35oWvqLCUFDohUs,2389
 nshtrainer/trainer/_runtime_callback.py,sha256=6F2Gq27Q8OFfN3RtdNC6QRA8ac0LC1hh4DUE3V5WgbI,4217
 nshtrainer/trainer/accelerator.py,sha256=Bqq-ry7DeCY4zw9_zBvTZiijpA-uUHrDjtbLV652m4M,2415
@@ -145,7 +147,7 @@ nshtrainer/trainer/plugin/layer_sync.py,sha256=-BbEyWZ063O7tZme7Gdu1lVxK6p1NeuLc
 nshtrainer/trainer/plugin/precision.py,sha256=7lf7KZd_yFyPmhLApjEIv0pkoDB5zdxi-7in0wRj3z8,5436
 nshtrainer/trainer/signal_connector.py,sha256=GhfGcSzfaTNhnj2QFkBDq5aT7FqbLMA7eC8SYQs8_8w,10828
 nshtrainer/trainer/strategy.py,sha256=VPTn5z3zvXTydY8IJchjhjcOfpvtoejnvUkq5E4WTus,1368
-nshtrainer/trainer/trainer.py,sha256=BKRicDlLI7KstzuP0SmzJzp0U4GK5lhZcKHS1IuL5sA,21197
+nshtrainer/trainer/trainer.py,sha256=smoN61iixWYDWGFvxrt8VwryZVy_NzqqjUcgOid0gRA,21696
 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
@@ -157,6 +159,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.1.1b1.dist-info/METADATA,sha256=wdOIQ91eUgWrIHfPLP06FD4uMkyyIfToR3VhBY-BXsE,962
-nshtrainer-1.1.1b1.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
-nshtrainer-1.1.1b1.dist-info/RECORD,,
+nshtrainer-1.2.0.dist-info/METADATA,sha256=HkNLruaJJuf3ijnGe7NqNd9emBR6QHMRh2-taC5wTrU,960
+nshtrainer-1.2.0.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+nshtrainer-1.2.0.dist-info/RECORD,,

{nshtrainer-1.1.1b1.dist-info → nshtrainer-1.2.0.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: poetry-core 2.1.1
+Generator: poetry-core 2.1.2
 Root-Is-Purelib: true
 Tag: py3-none-any