nshtrainer 0.33.2__py3-none-any.whl → 0.34.0__py3-none-any.whl

nshtrainer/config.py

@@ -1,9 +1,21 @@
 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 (
@@ -53,13 +65,13 @@ from nshtrainer.callbacks.throughput_monitor import (
 )
 from nshtrainer.callbacks.timer import EpochTimerConfig as EpochTimerConfig
 from nshtrainer.callbacks.wandb_watch import WandbWatchConfig as WandbWatchConfig
+from nshtrainer.config import LRSchedulerConfig as LRSchedulerConfig
 from nshtrainer.loggers._base import BaseLoggerConfig as BaseLoggerConfig
 from nshtrainer.loggers.csv import CSVLoggerConfig as CSVLoggerConfig
 from nshtrainer.loggers.tensorboard import (
     TensorboardLoggerConfig as TensorboardLoggerConfig,
 )
 from nshtrainer.loggers.wandb import WandbLoggerConfig as WandbLoggerConfig
-from nshtrainer.lr_scheduler import LRSchedulerConfig as LRSchedulerConfig
 from nshtrainer.lr_scheduler._base import LRSchedulerConfigBase as LRSchedulerConfigBase
 from nshtrainer.lr_scheduler.linear_warmup_cosine import (
     DurationConfig as DurationConfig,
@@ -129,9 +141,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/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/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/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/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/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

{nshtrainer-0.33.2.dist-info → nshtrainer-0.34.0.dist-info}/METADATA

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

{nshtrainer-0.33.2.dist-info → nshtrainer-0.34.0.dist-info}/RECORD

@@ -30,9 +30,9 @@ nshtrainer/callbacks/shared_parameters.py,sha256=fqlDweFDXPV_bfcAWpRgaJIad9i5Aeh
 nshtrainer/callbacks/throughput_monitor.py,sha256=H_ocXErZxUO3dxFk8Tx_VQdpI9E_Ztvqof5WtFevLyQ,1838
 nshtrainer/callbacks/timer.py,sha256=quS79oYClDUvQxJkNWmDMe0hwRUkkREgTgqzVrnom50,4607
 nshtrainer/callbacks/wandb_watch.py,sha256=Y6SEXfIx3kDDQbI5zpP53BVq0FBLJbLd3RJsiHZk1-Y,2921
-nshtrainer/config.py,sha256=EqvSp06RmSkCvo13-5bkecoCcE1nVViwvFIivTZOXoI,6883
+nshtrainer/config.py,sha256=skar_Wfz50_sU2NZS8PEjqofWeon4g4cyIgby3Da81g,8308
 nshtrainer/data/__init__.py,sha256=7mk1tr7SWUZ7ySbsf0y0ZPszk7u4QznPhQ-7wnpH9ec,149
-nshtrainer/data/balanced_batch_sampler.py,sha256=dGBTDDtlBU6c-ZlVQOCnTW7SjTB5hczWsOWEdUWjvkA,4385
+nshtrainer/data/balanced_batch_sampler.py,sha256=WAjhbO9EsZ_UadhdW3obBsjvEDMc2V-irpjegqIb7AI,4791
 nshtrainer/data/transform.py,sha256=6SNs3_TpNpfhcwTwvPKyEJ3opM1OT7LmMEYQNHKgRl8,2227
 nshtrainer/ll/__init__.py,sha256=L-aTi1V1bbvnZjOro8NvI393zbHQSFR9movWSRK9Mds,2477
 nshtrainer/ll/_experimental.py,sha256=oBQCKOEVYoxuUU9eLb-Fg2B2mzZD7SA0zfAO6lmWZ88,53
@@ -84,10 +84,12 @@ nshtrainer/trainer/_config.py,sha256=ZIodM5Ek1lpkWFhQ_VfmKR7q1mZFFwtjfx8FH72H8WM
 nshtrainer/trainer/_runtime_callback.py,sha256=sd2cUdRJG-UCdQr9ruZvEYpNGNF1t2W2fuxwwVlQD9E,4164
 nshtrainer/trainer/checkpoint_connector.py,sha256=r0ir4xYSdf_jebM0x09qaO6nJsvsiRQDyM0fs80ppOQ,2347
 nshtrainer/trainer/signal_connector.py,sha256=2EzkVktlasl8PgWAKNLDZRUMY__gRlDy1HdinAU-tfU,10740
-nshtrainer/trainer/trainer.py,sha256=iYueHW-m8fHyC8SQuXmpgxq_-GUa7pAJik7rDFPXmy0,17499
+nshtrainer/trainer/trainer.py,sha256=8T4LB31ygXXS3DECkvD2uqgElAxkulacYvZyL_-imJs,16839
 nshtrainer/util/_environment_info.py,sha256=CFUUZYjXhBLWGc0jtPNOaZgYMueUDEHpEaWFA1f3GoY,24213
 nshtrainer/util/_useful_types.py,sha256=dwZokFkIe7M5i2GR3nQ9A1lhGw06DMAFfH5atyquqSA,8000
-nshtrainer/util/config/__init__.py,sha256=N2AOhaZC93DszvCdwvNL9KgnzJ2M3P-esFBY6VGih6Y,190
+nshtrainer/util/bf16.py,sha256=VUnIG6aA4XtZscZc_dxv5ln_jlEbdU3eMFwDb5SEWSI,726
+nshtrainer/util/config/__init__.py,sha256=o8fwPf_dctE_7CAkT0wNOBkvmxnzYzXeHpLedrZLt54,236
+nshtrainer/util/config/dtype.py,sha256=JtYjrcBFNBlziJnLAE6QS0QV4PUXhGspYH1hNFrB3ks,1965
 nshtrainer/util/config/duration.py,sha256=pgIKQ88Dg8y1YAKUvUsNWu9hc9O79kdYBfgmC3a_-kQ,728
 nshtrainer/util/environment.py,sha256=AeW_kLl-N70wmb6L_JLz1wRj0kA70xs6RCmc9iUqczE,4159
 nshtrainer/util/path.py,sha256=VkpuhR4GaZtSFBVqbGAvfjcrU-PR8xwiGzzwFNOWP9c,2995
@@ -95,6 +97,6 @@ nshtrainer/util/seed.py,sha256=Or2wMPsnQxfnZ2xfBiyMcHFIUt3tGTNeMMyOEanCkqs,280
 nshtrainer/util/slurm.py,sha256=rofIU26z3SdL79SF45tNez6juou1cyDLz07oXEZb9Hg,1566
 nshtrainer/util/typed.py,sha256=NGuDkDzFlc1fAoaXjOFZVbmj0mRFjsQi1E_hPa7Bn5U,128
 nshtrainer/util/typing_utils.py,sha256=8ptjSSLZxlmy4FY6lzzkoGoF5fGNClo8-B_c0XHQaNU,385
-nshtrainer-0.33.2.dist-info/METADATA,sha256=zDC_xehJGE3RlCACScFpu64qL1TKd_D8VyhjmRxNDkw,916
-nshtrainer-0.33.2.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.33.2.dist-info/RECORD,,
+nshtrainer-0.34.0.dist-info/METADATA,sha256=GYC9ejdKV3MCyOFhJcFjI-uedTWLGWj-SE5S79ruug4,916
+nshtrainer-0.34.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.34.0.dist-info/RECORD,,