nshtrainer 0.30.1__py3-none-any.whl → 0.32.0__py3-none-any.whl

nshtrainer/ll/__init__.py

@@ -1,3 +1,5 @@
+from typing import TypeAlias
+
 from . import _experimental as _experimental
 from . import actsave as actsave
 from . import callbacks as callbacks
@@ -21,12 +23,9 @@ from .log import init_python_logging as init_python_logging
 from .log import lovely as lovely
 from .log import pretty as pretty
 from .lr_scheduler import LRSchedulerConfig as LRSchedulerConfig
-from .model import Base as Base
 from .model import BaseConfig as BaseConfig
-from .model import BaseProfilerConfig as BaseProfilerConfig
 from .model import CheckpointLoadingConfig as CheckpointLoadingConfig
 from .model import CheckpointSavingConfig as CheckpointSavingConfig
-from .model import ConfigList as ConfigList
 from .model import DirectoryConfig as DirectoryConfig
 from .model import (
     EnvironmentClassInformationConfig as EnvironmentClassInformationConfig,
@@ -43,7 +42,6 @@ from .model import LightningModuleBase as LightningModuleBase
 from .model import LoggingConfig as LoggingConfig
 from .model import MetricConfig as MetricConfig
 from .model import OptimizationConfig as OptimizationConfig
-from .model import PrimaryMetricConfig as PrimaryMetricConfig
 from .model import ReproducibilityConfig as ReproducibilityConfig
 from .model import SanityCheckingConfig as SanityCheckingConfig
 from .model import TrainerConfig as TrainerConfig
@@ -54,3 +52,6 @@ from .runner import Runner as Runner
 from .runner import SnapshotConfig as SnapshotConfig
 from .snoop import snoop as snoop
 from .trainer import Trainer as Trainer
+
+PrimaryMetricConfig: TypeAlias = MetricConfig
+ConfigList: TypeAlias = list[tuple[BaseConfig, type[LightningModuleBase]]]

nshtrainer/ll/model.py

@@ -1,5 +1,12 @@
 from nshtrainer.model import *  # noqa: F403
 
+from ..trainer._config import CheckpointLoadingConfig as CheckpointLoadingConfig
+from ..trainer._config import CheckpointSavingConfig as CheckpointSavingConfig
+from ..trainer._config import GradientClippingConfig as GradientClippingConfig
+from ..trainer._config import LoggingConfig as LoggingConfig
+from ..trainer._config import OptimizationConfig as OptimizationConfig
+from ..trainer._config import ReproducibilityConfig as ReproducibilityConfig
+from ..trainer._config import SanityCheckingConfig as SanityCheckingConfig
 from ..util._environment_info import (
     EnvironmentClassInformationConfig as EnvironmentClassInformationConfig,
 )

nshtrainer/loggers/wandb.py

@@ -129,7 +129,7 @@ class WandbLoggerConfig(CallbackConfigBase, BaseLoggerConfig):
                         "Please either upgrade to a newer version of WandB or disable the `use_wandb_core` option."
                     )
                 else:
-                    wandb.require("core")
+                    wandb.require("core")  # type: ignore
                     log.critical("Using the `wandb-core` backend for WandB.")
             except ImportError:
                 pass

nshtrainer/lr_scheduler/linear_warmup_cosine.py

@@ -6,7 +6,7 @@ from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
 from typing_extensions import override
 
-from ..config import Duration
+from ..util.config import Duration
 from ._base import LRSchedulerConfigBase, LRSchedulerMetadata
 
 

nshtrainer/model/__init__.py

@@ -1,26 +1,5 @@
-from typing_extensions import TypeAlias
-
-from .base import Base as Base
 from .base import LightningModuleBase as LightningModuleBase
 from .config import BaseConfig as BaseConfig
-from .config import BaseProfilerConfig as BaseProfilerConfig
-from .config import BestCheckpointCallbackConfig as BestCheckpointCallbackConfig
-from .config import CheckpointLoadingConfig as CheckpointLoadingConfig
-from .config import CheckpointSavingConfig as CheckpointSavingConfig
 from .config import DirectoryConfig as DirectoryConfig
-from .config import EarlyStoppingConfig as EarlyStoppingConfig
-from .config import GradientClippingConfig as GradientClippingConfig
-from .config import HuggingFaceHubConfig as HuggingFaceHubConfig
-from .config import LastCheckpointCallbackConfig as LastCheckpointCallbackConfig
-from .config import LoggingConfig as LoggingConfig
 from .config import MetricConfig as MetricConfig
-from .config import (
-    OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
-)
-from .config import OptimizationConfig as OptimizationConfig
-from .config import PrimaryMetricConfig as PrimaryMetricConfig
-from .config import ReproducibilityConfig as ReproducibilityConfig
-from .config import SanityCheckingConfig as SanityCheckingConfig
 from .config import TrainerConfig as TrainerConfig
-
-ConfigList: TypeAlias = list[tuple[BaseConfig, type[LightningModuleBase]]]

nshtrainer/model/base.py

@@ -2,39 +2,65 @@ import inspect
 import logging
 from abc import ABC, abstractmethod
 from collections.abc import MutableMapping
-from typing import IO, TYPE_CHECKING, Any, Generic, cast
+from typing import IO, TYPE_CHECKING, Any, Generic, Literal, cast
 
 import torch
+import torch.distributed
 from lightning.fabric.utilities.types import _MAP_LOCATION_TYPE, _PATH
-from lightning.pytorch import LightningModule, Trainer
-from lightning.pytorch.callbacks import Callback
+from lightning.pytorch import LightningModule
+from lightning.pytorch.profilers import PassThroughProfiler, Profiler
 from lightning.pytorch.utilities.types import STEP_OUTPUT
 from typing_extensions import Self, TypeVar, override
 
+from ..callbacks.rlp_sanity_checks import _RLPSanityCheckModuleMixin
 from ..util._environment_info import EnvironmentConfig
 from .config import BaseConfig
-from .modules.callback import CallbackModuleMixin
-from .modules.debug import DebugModuleMixin
-from .modules.distributed import DistributedMixin
-from .modules.logger import LoggerLightningModuleMixin
-from .modules.profiler import ProfilerMixin
-from .modules.rlp_sanity_checks import RLPSanityCheckModuleMixin
-from .modules.shared_parameters import SharedParametersModuleMixin
+from .mixins.logger import LoggerLightningModuleMixin
 
 log = logging.getLogger(__name__)
 
 THparams = TypeVar("THparams", bound=BaseConfig, infer_variance=True)
 
 
-class Base(DebugModuleMixin, Generic[THparams]):
-    @torch.jit.unused
-    @property
-    def config(self) -> THparams:
-        return self.hparams
+T = TypeVar("T", infer_variance=True)
+
+ReduceOpStr = Literal[
+    "avg",
+    "mean",
+    "band",
+    "bor",
+    "bxor",
+    "max",
+    "min",
+    "premul_sum",
+    "product",
+    "sum",
+]
+VALID_REDUCE_OPS = (
+    "avg",
+    "mean",
+    "band",
+    "bor",
+    "bxor",
+    "max",
+    "min",
+    "premul_sum",
+    "product",
+    "sum",
+)
+
 
+class LightningModuleBase(  # pyright: ignore[reportIncompatibleMethodOverride]
+    _RLPSanityCheckModuleMixin,
+    LoggerLightningModuleMixin,
+    LightningModule,
+    ABC,
+    Generic[THparams],
+):
+    # region Config
     @torch.jit.unused
     @property
-    def C(self) -> THparams:
+    def config(self) -> THparams:
         return self.hparams
 
     @property
@@ -43,65 +69,98 @@ class Base(DebugModuleMixin, Generic[THparams]):
             return False
         return self.config.debug
 
-    @property
-    def dev(self) -> bool:
-        if torch.jit.is_scripting():
+    # endregion
+
+    # region Debug
+
+    @torch.jit.unused
+    def breakpoint(self, rank_zero_only: bool = True):
+        if (
+            not rank_zero_only
+            or not torch.distributed.is_initialized()
+            or torch.distributed.get_rank() == 0
+        ):
+            breakpoint()
+
+        if rank_zero_only and torch.distributed.is_initialized():
+            _ = torch.distributed.barrier()
+
+    @torch.jit.unused
+    def ensure_finite(
+        self,
+        tensor: torch.Tensor,
+        name: str | None = None,
+        throw: bool = False,
+    ):
+        name_parts: list[str] = ["Tensor"]
+        if name is not None:
+            name_parts.append(name)
+        name = " ".join(name_parts)
+
+        not_finite = ~torch.isfinite(tensor)
+        if not_finite.any():
+            msg = f"{name} has {not_finite.sum().item()}/{not_finite.numel()} non-finite values."
+            if throw:
+                raise RuntimeError(msg)
+            else:
+                log.warning(msg)
             return False
-        return self.config.debug
+        return True
 
-    @override
-    def __init__(self, hparams: THparams):
-        super().__init__()
+    # endregion
 
-        if not hasattr(self, "hparams"):
-            self.hparams = hparams
+    # region Profiler
+    @property
+    def profiler(self) -> Profiler:
+        if (trainer := self._trainer) is None:
+            raise RuntimeError("trainer is not defined")
 
+        if not hasattr(trainer, "profiler"):
+            raise RuntimeError("trainer does not have profiler")
 
-class DebugFlagCallback(Callback):
-    """
-    Sets the debug flag to true in the following circumstances:
-    - fast_dev_run is enabled
-    - sanity check is running
-    """
+        if (profiler := getattr(trainer, "profiler")) is None:
+            profiler = PassThroughProfiler()
 
-    @override
-    def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str):
-        if not getattr(trainer, "fast_dev_run", False):
-            return
+        return profiler
 
-        hparams = cast(BaseConfig, pl_module.hparams)
-        if not hparams.debug:
-            log.critical("Fast dev run detected, setting debug flag to True.")
-        hparams.debug = True
+    # endregion
 
-    @override
-    def on_sanity_check_start(self, trainer: Trainer, pl_module: LightningModule):
-        hparams = cast(BaseConfig, pl_module.hparams)
-        self._debug = hparams.debug
-        if not self._debug:
-            log.critical("Enabling debug flag during sanity check routine.")
-        hparams.debug = True
+    # region Distributed
+    def all_gather_object(
+        self,
+        object: T,
+        group: torch.distributed.ProcessGroup | None = None,
+    ) -> list[T]:
+        if (
+            not torch.distributed.is_available()
+            or not torch.distributed.is_initialized()
+        ):
+            return [object]
+
+        object_list = [cast(T, None) for _ in range(self.trainer.world_size)]
+        torch.distributed.all_gather_object(object_list, object, group=group)
+        return object_list
+
+    def barrier(self, name: str | None = None):
+        self.trainer.strategy.barrier(name=name)
+
+    def reduce(
+        self,
+        tensor: torch.Tensor,
+        reduce_op: torch.distributed.ReduceOp.RedOpType | ReduceOpStr,
+        group: Any | None = None,
+    ) -> torch.Tensor:
+        if isinstance(reduce_op, str):
+            # validate reduce_op
+            if reduce_op not in VALID_REDUCE_OPS:
+                raise ValueError(
+                    f"reduce_op must be one of {VALID_REDUCE_OPS}, got {reduce_op}"
+                )
 
-    @override
-    def on_sanity_check_end(self, trainer: Trainer, pl_module: LightningModule):
-        hparams = cast(BaseConfig, pl_module.hparams)
-        if not self._debug:
-            log.critical("Sanity check routine complete, disabling debug flag.")
-        hparams.debug = self._debug
+        return self.trainer.strategy.reduce(tensor, group=group, reduce_op=reduce_op)
 
+    # endregion
 
-class LightningModuleBase(  # pyright: ignore[reportIncompatibleMethodOverride]
-    ProfilerMixin,
-    RLPSanityCheckModuleMixin,
-    LoggerLightningModuleMixin,
-    SharedParametersModuleMixin,
-    DistributedMixin,
-    CallbackModuleMixin,
-    Base[THparams],
-    LightningModule,
-    ABC,
-    Generic[THparams],
-):
     # Our own custom __repr__ method.
     # Torch's __repr__ method is too verbose and doesn't provide any useful information.
     @override
@@ -193,10 +252,8 @@ class LightningModuleBase(  # pyright: ignore[reportIncompatibleMethodOverride]
         hparams.environment = EnvironmentConfig.from_current_environment(hparams, self)
         hparams = self.pre_init_update_hparams(hparams)
 
-        super().__init__(hparams)
-
+        super().__init__()
         self.save_hyperparameters(hparams)
-        self.register_callback(lambda: DebugFlagCallback())
 
     def zero_loss(self):
         """