nshtrainer 1.1.1b1__tar.gz → 1.2.0__tar.gz

{nshtrainer-1.1.1b1 → nshtrainer-1.2.0}/PKG-INFO

@@ -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 → nshtrainer-1.2.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "nshtrainer"
-version = "1.1.1-beta.1"
+version = "1.2.0"
 description = ""
 authors = [{ name = "Nima Shoghi", email = "nimashoghi@gmail.com" }]
 requires-python = ">=3.10,<4.0"
@@ -33,8 +33,9 @@ basedpyright = "*"
 ruff = "*"
 ipykernel = "*"
 ipywidgets = "*"
-pytest = "^8.3.5"
-pytest-cov = "^6.0.0"
+pytest = "*"
+pytest-cov = "*"
+pytest-forked = "*"
 
 [build-system]
 requires = ["poetry-core"]

{nshtrainer-1.1.1b1 → nshtrainer-1.2.0}/src/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-1.1.1b1 → nshtrainer-1.2.0}/src/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-1.1.1b1 → nshtrainer-1.2.0}/src/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-1.2.0/src/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-1.1.1b1 → nshtrainer-1.2.0}/src/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-1.1.1b1 → nshtrainer-1.2.0}/src/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-1.2.0/src/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-1.2.0/src/nshtrainer/configs/optimizer/__init__.py

@@ -0,0 +1,39 @@
+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-1.1.1b1 → nshtrainer-1.2.0}/src/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-1.1.1b1 → nshtrainer-1.2.0}/src/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-1.1.1b1 → nshtrainer-1.2.0}/src/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: