nshtrainer 0.33.2__tar.gz → 0.34.1__tar.gz

{nshtrainer-0.33.2 → nshtrainer-0.34.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.33.2
+Version: 0.34.1
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-0.33.2 → nshtrainer-0.34.1}/pyproject.toml

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

{nshtrainer-0.33.2 → nshtrainer-0.34.1}/src/nshtrainer/config.py

@@ -1,9 +1,18 @@
-from nshconfig._config import Config as Config
-from nshsnap._config import SnapshotConfig as SnapshotConfig
-
+from nshtrainer._checkpoint.loader import (
+    BestCheckpointStrategyConfig as BestCheckpointStrategyConfig,
+)
 from nshtrainer._checkpoint.loader import (
     CheckpointLoadingConfig as CheckpointLoadingConfig,
 )
+from nshtrainer._checkpoint.loader import (
+    CheckpointLoadingStrategyConfig as CheckpointLoadingStrategyConfig,
+)
+from nshtrainer._checkpoint.loader import (
+    LastCheckpointStrategyConfig as LastCheckpointStrategyConfig,
+)
+from nshtrainer._checkpoint.loader import (
+    UserProvidedPathCheckpointStrategyConfig as UserProvidedPathCheckpointStrategyConfig,
+)
 from nshtrainer._checkpoint.metadata import CheckpointMetadata as CheckpointMetadata
 from nshtrainer._directory import DirectoryConfig as DirectoryConfig
 from nshtrainer._hf_hub import (
@@ -129,9 +138,31 @@ from nshtrainer.util._environment_info import (
     EnvironmentClassInformationConfig as EnvironmentClassInformationConfig,
 )
 from nshtrainer.util._environment_info import EnvironmentConfig as EnvironmentConfig
+from nshtrainer.util._environment_info import (
+    EnvironmentCUDAConfig as EnvironmentCUDAConfig,
+)
+from nshtrainer.util._environment_info import (
+    EnvironmentGPUConfig as EnvironmentGPUConfig,
+)
+from nshtrainer.util._environment_info import (
+    EnvironmentHardwareConfig as EnvironmentHardwareConfig,
+)
 from nshtrainer.util._environment_info import (
     EnvironmentLinuxEnvironmentConfig as EnvironmentLinuxEnvironmentConfig,
 )
+from nshtrainer.util._environment_info import (
+    EnvironmentLSFInformationConfig as EnvironmentLSFInformationConfig,
+)
+from nshtrainer.util._environment_info import (
+    EnvironmentPackageConfig as EnvironmentPackageConfig,
+)
 from nshtrainer.util._environment_info import (
     EnvironmentSLURMInformationConfig as EnvironmentSLURMInformationConfig,
 )
+from nshtrainer.util._environment_info import (
+    EnvironmentSnapshotConfig as EnvironmentSnapshotConfig,
+)
+from nshtrainer.util._environment_info import GitRepositoryConfig as GitRepositoryConfig
+from nshtrainer.util.config.dtype import DTypeConfig as DTypeConfig
+from nshtrainer.util.config.duration import EpochsConfig as EpochsConfig
+from nshtrainer.util.config.duration import StepsConfig as StepsConfig

{nshtrainer-0.33.2 → nshtrainer-0.34.1}/src/nshtrainer/data/balanced_batch_sampler.py

@@ -1,13 +1,12 @@
 import heapq
 import logging
-from functools import cached_property
 from typing import Any, Protocol, runtime_checkable
 
 import numpy as np
 import torch
 import torch.distributed
 from lightning_fabric.utilities.distributed import _DatasetSamplerWrapper
-from torch.utils.data import BatchSampler, Dataset, DistributedSampler
+from torch.utils.data import BatchSampler, DistributedSampler
 from typing_extensions import override
 
 log = logging.getLogger(__name__)
@@ -47,24 +46,16 @@ class DatasetWithSizes(Protocol):
     def data_sizes(self, indices: list[int]) -> np.ndarray: ...
 
 
-class BalancedBatchSampler(BatchSampler):
-    @staticmethod
-    def _ensure_supported(dataset: Any):
-        if not isinstance(dataset, Dataset):
-            raise ValueError(
-                "BalancedBatchSampler requires a dataset that implements `__getitem__`"
-            )
-
-        if not isinstance(dataset, DatasetWithSizes):
-            raise ValueError(
-                "BalancedBatchSampler requires a dataset that implements `data_sizes`"
-            )
+@runtime_checkable
+class DataSizesFunction(Protocol):
+    def __call__(self, dataset: Any, indices: list[int]) -> np.ndarray: ...
 
-        log.critical(f"BalancedBatchSampler: Resolved dataset to {type(dataset)}")
-        return dataset
 
+class BalancedBatchSampler(BatchSampler):
     @staticmethod
-    def _unwrap_dataset(dataset: Dataset) -> Dataset:
+    def _unwrap_dataset(dataset: Any):
+        # Lightning's DistributedSampler wraps the dataset in a _DatasetSamplerWrapper,
+        # so we need to unwrap it to get the actual dataset.
         if isinstance(dataset, _DatasetSamplerWrapper):
             if (data_source := getattr(dataset._sampler, "data_source", None)) is None:
                 raise ValueError("Could not unwrap dataset from _DatasetSamplerWrapper")
@@ -79,12 +70,6 @@ class BalancedBatchSampler(BatchSampler):
             )
         return self.sampler
 
-    @cached_property
-    def dataset(self):
-        return self._ensure_supported(
-            self._unwrap_dataset(self.distributed_sampler.dataset)
-        )
-
     def __init__(
         self,
         sampler: DistributedSampler,
@@ -92,10 +77,12 @@ class BalancedBatchSampler(BatchSampler):
         batch_size: int,
         device: torch.device,
         drop_last: bool = False,
+        data_sizes_fn: DataSizesFunction | None = None,
     ):
         super().__init__(sampler, batch_size, drop_last=drop_last)
 
         self._device = device
+        self._data_sizes_fn = data_sizes_fn
 
         log.info(
             f"Created BalancedBatchSampler with {sampler=}, {batch_size=}, {drop_last=}"
@@ -105,17 +92,34 @@ class BalancedBatchSampler(BatchSampler):
     def _dist_enabled():
         return torch.distributed.is_available() and torch.distributed.is_initialized()
 
+    def _dataset_sizes(self, indices: list[int]) -> np.ndarray:
+        dataset = self._unwrap_dataset(self.distributed_sampler.dataset)
+        # Dataset much either implement `data_sizes`, or we need to provide a custom
+        # implementation of the dataset sizes function.
+        if isinstance(dataset, DatasetWithSizes):
+            log.critical(f"BalancedBatchSampler: Resolved dataset to {type(dataset)}")
+            return dataset.data_sizes(indices)
+
+        if (data_sizes_fn := self._data_sizes_fn) is not None:
+            return data_sizes_fn(dataset, indices)
+
+        raise ValueError(
+            "Dataset must implement the `data_sizes` method, "
+            "or a custom data_sizes_fn must be provided "
+            "to the BalancedBatchSampler."
+        )
+
     @override
     def __iter__(self):
         if not self._dist_enabled():
             yield from super().__iter__()
             return
 
-        for batch_idx in super().__iter__():
-            sizes = self.dataset.data_sizes(batch_idx)
+        for batch_idxs in super().__iter__():
+            sizes = self._dataset_sizes(batch_idxs)
             idx_sizes = torch.stack(
                 [
-                    torch.tensor(batch_idx, device=self._device),
+                    torch.tensor(batch_idxs, device=self._device),
                     torch.tensor(sizes, device=self._device),
                 ]
             )

{nshtrainer-0.33.2 → nshtrainer-0.34.1}/src/nshtrainer/trainer/trainer.py

@@ -18,6 +18,7 @@ from typing_extensions import Unpack, assert_never, override
 
 from .._checkpoint.metadata import _write_checkpoint_metadata
 from ..callbacks.base import resolve_all_callbacks
+from ..util.bf16 import is_bf16_supported_no_emulation
 from ._config import (
     AcceleratorConfigProtocol,
     LightningTrainerKwargs,
@@ -33,30 +34,6 @@ if TYPE_CHECKING:
 log = logging.getLogger(__name__)
 
 
-def _is_bf16_supported_no_emulation():
-    r"""Return a bool indicating if the current CUDA/ROCm device supports dtype bfloat16."""
-    version = getattr(torch, "version")
-
-    # Check for ROCm, if true return true, no ROCM_VERSION check required,
-    # since it is supported on AMD GPU archs.
-    if version.hip:
-        return True
-
-    device = torch.cuda.current_device()
-
-    # Check for CUDA version and device compute capability.
-    # This is a fast way to check for it.
-    cuda_version = version.cuda
-    if (
-        cuda_version is not None
-        and int(cuda_version.split(".")[0]) >= 11
-        and torch.cuda.get_device_properties(device).major >= 8
-    ):
-        return True
-
-    return False
-
-
 class Trainer(LightningTrainer):
     @classmethod
     def _pre_init(cls, config: "BaseConfig"):
@@ -188,7 +165,7 @@ class Trainer(LightningTrainer):
                     try:
                         resolved_precision = (
                             "bf16-mixed"
-                            if _is_bf16_supported_no_emulation()
+                            if is_bf16_supported_no_emulation()
                             else "16-mixed"
                         )
                     except BaseException:

nshtrainer-0.34.1/src/nshtrainer/util/bf16.py

@@ -0,0 +1,25 @@
+import torch
+
+
+def is_bf16_supported_no_emulation():
+    r"""Return a bool indicating if the current CUDA/ROCm device supports dtype bfloat16."""
+    version = getattr(torch, "version")
+
+    # Check for ROCm, if true return true, no ROCM_VERSION check required,
+    # since it is supported on AMD GPU archs.
+    if version.hip:
+        return True
+
+    device = torch.cuda.current_device()
+
+    # Check for CUDA version and device compute capability.
+    # This is a fast way to check for it.
+    cuda_version = version.cuda
+    if (
+        cuda_version is not None
+        and int(cuda_version.split(".")[0]) >= 11
+        and torch.cuda.get_device_properties(device).major >= 8
+    ):
+        return True
+
+    return False

{nshtrainer-0.33.2 → nshtrainer-0.34.1}/src/nshtrainer/util/config/__init__.py

@@ -1,4 +1,5 @@
 from . import duration as duration
+from .dtype import DTypeConfig as DTypeConfig
 from .duration import DurationConfig as DurationConfig
 from .duration import EpochsConfig as EpochsConfig
 from .duration import StepsConfig as StepsConfig

nshtrainer-0.34.1/src/nshtrainer/util/config/dtype.py

@@ -0,0 +1,89 @@
+from typing import TYPE_CHECKING, Literal, TypeAlias
+
+import nshconfig as C
+import torch
+from typing_extensions import assert_never
+
+from ..bf16 import is_bf16_supported_no_emulation
+
+if TYPE_CHECKING:
+    from ...model.base import BaseConfig
+
+DTypeName: TypeAlias = Literal[
+    "float32",
+    "float",
+    "float64",
+    "double",
+    "float16",
+    "bfloat16",
+    "float8_e4m3fn",
+    "float8_e4m3fnuz",
+    "float8_e5m2",
+    "float8_e5m2fnuz",
+    "half",
+    "uint8",
+    "uint16",
+    "uint32",
+    "uint64",
+    "int8",
+    "int16",
+    "short",
+    "int32",
+    "int",
+    "int64",
+    "long",
+    "complex32",
+    "complex64",
+    "chalf",
+    "cfloat",
+    "complex128",
+    "cdouble",
+    "quint8",
+    "qint8",
+    "qint32",
+    "bool",
+    "quint4x2",
+    "quint2x4",
+    "bits1x8",
+    "bits2x4",
+    "bits4x2",
+    "bits8",
+    "bits16",
+]
+
+
+class DTypeConfig(C.Config):
+    name: DTypeName
+    """The name of the dtype."""
+
+    @classmethod
+    def from_base_config(cls, config: "BaseConfig"):
+        if (precision := config.trainer.precision) is None:
+            precision = "32-true"
+
+        match precision:
+            case "16-mixed-auto":
+                return (
+                    cls(name="bfloat16")
+                    if is_bf16_supported_no_emulation()
+                    else cls(name="float16")
+                )
+            case "fp16-mixed":
+                return cls(name="float16")
+            case "bf16-mixed":
+                return cls(name="bfloat16")
+            case "32-true":
+                return cls(name="float32")
+            case "64-true":
+                return cls(name="float64")
+            case _:
+                assert_never(config.trainer.precision)
+
+    @property
+    def torch_dtype(self):
+        if ((dtype := getattr(torch, self.name, None)) is None) or not isinstance(
+            dtype, torch.dtype
+        ):
+            raise ValueError(f"Unknown dtype {self.name}")
+
+        return dtype