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

nshtrainer/__init__.py

@@ -7,10 +7,9 @@ from . import metrics as metrics
 from . import model as model
 from . import nn as nn
 from . import optimizer as optimizer
+from . import profiler as profiler
 from .metrics import MetricConfig as MetricConfig
-from .model import Base as Base
 from .model import BaseConfig as BaseConfig
-from .model import ConfigList as ConfigList
 from .model import LightningModuleBase as LightningModuleBase
 from .runner import Runner as Runner
 from .trainer import Trainer as Trainer

nshtrainer/_directory.py

@@ -0,0 +1,85 @@
+import logging
+from pathlib import Path
+
+import nshconfig as C
+
+from .callbacks.directory_setup import DirectorySetupConfig
+from .loggers import LoggerConfig
+
+log = logging.getLogger(__name__)
+
+
+class DirectoryConfig(C.Config):
+    project_root: Path | None = None
+    """
+    Root directory for this project.
+
+    This isn't specific to the run; it is the parent directory of all runs.
+    """
+
+    log: Path | None = None
+    """Base directory for all experiment tracking (e.g., WandB, Tensorboard, etc.) files. If None, will use nshtrainer/{id}/log/."""
+
+    stdio: Path | None = None
+    """stdout/stderr log directory to use for the trainer. If None, will use nshtrainer/{id}/stdio/."""
+
+    checkpoint: Path | None = None
+    """Checkpoint directory to use for the trainer. If None, will use nshtrainer/{id}/checkpoint/."""
+
+    activation: Path | None = None
+    """Activation directory to use for the trainer. If None, will use nshtrainer/{id}/activation/."""
+
+    profile: Path | None = None
+    """Directory to save profiling information to. If None, will use nshtrainer/{id}/profile/."""
+
+    setup_callback: DirectorySetupConfig = DirectorySetupConfig()
+    """Configuration for the directory setup PyTorch Lightning callback."""
+
+    def resolve_run_root_directory(self, run_id: str) -> Path:
+        if (project_root_dir := self.project_root) is None:
+            project_root_dir = Path.cwd()
+
+        # The default base dir is $CWD/nshtrainer/{id}/
+        base_dir = project_root_dir / "nshtrainer"
+        base_dir.mkdir(exist_ok=True)
+
+        # Add a .gitignore file to the nshtrainer directory
+        #   which will ignore all files except for the .gitignore file itself
+        gitignore_path = base_dir / ".gitignore"
+        if not gitignore_path.exists():
+            gitignore_path.touch()
+            gitignore_path.write_text("*\n")
+
+        base_dir = base_dir / run_id
+        base_dir.mkdir(exist_ok=True)
+
+        return base_dir
+
+    def resolve_subdirectory(
+        self,
+        run_id: str,
+        # subdirectory: Literal["log", "stdio", "checkpoint", "activation", "profile"],
+        subdirectory: str,
+    ) -> 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)}"
+            return subdir
+
+        dir = self.resolve_run_root_directory(run_id)
+        dir = dir / subdirectory
+        dir.mkdir(exist_ok=True)
+        return dir
+
+    def _resolve_log_directory_for_logger(self, run_id: str, logger: LoggerConfig):
+        if (log_dir := logger.log_dir) is not None:
+            return log_dir
+
+        # Save to nshtrainer/{id}/log/{logger name}
+        log_dir = self.resolve_subdirectory(run_id, "log")
+        log_dir = log_dir / logger.name
+        log_dir.mkdir(exist_ok=True)
+
+        return log_dir

nshtrainer/callbacks/__init__.py

@@ -12,6 +12,8 @@ from .checkpoint import OnExceptionCheckpoint as OnExceptionCheckpoint
 from .checkpoint import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
 )
+from .directory_setup import DirectorySetupCallback as DirectorySetupCallback
+from .directory_setup import DirectorySetupConfig as DirectorySetupConfig
 from .early_stopping import EarlyStopping as EarlyStopping
 from .early_stopping import EarlyStoppingConfig as EarlyStoppingConfig
 from .ema import EMA as EMA
@@ -28,6 +30,10 @@ from .norm_logging import NormLoggingCallback as NormLoggingCallback
 from .norm_logging import NormLoggingConfig as NormLoggingConfig
 from .print_table import PrintTableMetricsCallback as PrintTableMetricsCallback
 from .print_table import PrintTableMetricsConfig as PrintTableMetricsConfig
+from .rlp_sanity_checks import RLPSanityChecksCallback as RLPSanityChecksCallback
+from .rlp_sanity_checks import RLPSanityChecksConfig as RLPSanityChecksConfig
+from .shared_parameters import SharedParametersCallback as SharedParametersCallback
+from .shared_parameters import SharedParametersConfig as SharedParametersConfig
 from .throughput_monitor import ThroughputMonitorConfig as ThroughputMonitorConfig
 from .timer import EpochTimer as EpochTimer
 from .timer import EpochTimerConfig as EpochTimerConfig
@@ -46,6 +52,8 @@ CallbackConfig = Annotated[
     | BestCheckpointCallbackConfig
     | LastCheckpointCallbackConfig
     | OnExceptionCheckpointCallbackConfig
+    | SharedParametersConfig
+    | RLPSanityChecksConfig
     | WandbWatchConfig,
     C.Field(discriminator="name"),
 ]

nshtrainer/callbacks/directory_setup.py

@@ -0,0 +1,85 @@
+import logging
+import os
+from pathlib import Path
+from typing import Literal
+
+from lightning.pytorch import Callback
+from typing_extensions import override
+
+from .base import CallbackConfigBase
+
+log = logging.getLogger(__name__)
+
+
+def _create_symlink_to_nshrunner(base_dir: Path):
+    # Resolve the current nshrunner session directory
+    if not (session_dir := os.environ.get("NSHRUNNER_SESSION_DIR")):
+        log.warning("NSHRUNNER_SESSION_DIR is not set. Skipping symlink creation.")
+        return
+    session_dir = Path(session_dir)
+    if not session_dir.exists() or not session_dir.is_dir():
+        log.warning(
+            f"NSHRUNNER_SESSION_DIR is not a valid directory: {session_dir}. "
+            "Skipping symlink creation."
+        )
+        return
+
+    # Create the symlink
+    symlink_path = base_dir / "nshrunner"
+    if symlink_path.exists():
+        # If it already points to the correct directory, we're done
+        if symlink_path.resolve() == session_dir.resolve():
+            return
+
+        # Otherwise, we should log a warning and remove the existing symlink
+        log.warning(
+            f"A symlink pointing to {symlink_path.resolve()} already exists at {symlink_path}. "
+            "Removing the existing symlink."
+        )
+        symlink_path.unlink()
+
+    symlink_path.symlink_to(session_dir)
+
+
+class DirectorySetupConfig(CallbackConfigBase):
+    name: Literal["directory_setup"] = "directory_setup"
+
+    enabled: bool = True
+    """Whether to enable the directory setup callback."""
+
+    create_symlink_to_nshrunner_root: bool = True
+    """Should we create a symlink to the root folder for the Runner (if we're in one)?"""
+
+    def __bool__(self):
+        return self.enabled
+
+    def create_callbacks(self, root_config):
+        if not self:
+            return
+
+        yield DirectorySetupCallback(self)
+
+
+class DirectorySetupCallback(Callback):
+    @override
+    def __init__(self, config: DirectorySetupConfig):
+        super().__init__()
+
+        self.config = config
+        del config
+
+    @override
+    def setup(self, trainer, pl_module, stage):
+        super().setup(trainer, pl_module, stage)
+
+        # Create a symlink to the root folder for the Runner
+        if self.config.create_symlink_to_nshrunner_root:
+            # Resolve the base dir
+            from ..model.config import BaseConfig
+
+            assert isinstance(
+                config := pl_module.hparams, BaseConfig
+            ), f"Expected a BaseConfig, got {type(config)}"
+
+            base_dir = config.directory.resolve_run_root_directory(config.id)
+            _create_symlink_to_nshrunner(base_dir)

nshtrainer/callbacks/rlp_sanity_checks.py

@@ -0,0 +1,230 @@
+import logging
+from collections.abc import Mapping
+from typing import Literal, cast
+
+import torch
+from lightning.pytorch import LightningModule
+from lightning.pytorch.callbacks import Callback
+from lightning.pytorch.utilities.types import (
+    LRSchedulerConfigType,
+    LRSchedulerTypeUnion,
+)
+from typing_extensions import Protocol, override, runtime_checkable
+
+from .base import CallbackConfigBase
+
+log = logging.getLogger(__name__)
+
+
+class RLPSanityChecksConfig(CallbackConfigBase):
+    """
+    If enabled, will do some sanity checks if the `ReduceLROnPlateau` scheduler is used:
+        - If the ``interval`` is step, it makes sure that validation is called every ``frequency`` steps.
+        - If the ``interval`` is epoch, it makes sure that validation is called every ``frequency`` epochs.
+    """
+
+    name: Literal["rlp_sanity_checks"] = "rlp_sanity_checks"
+
+    enabled: bool = True
+    """Whether to enable ReduceLRPlateau sanity checks."""
+
+    on_error: Literal["warn", "error"] = "error"
+    """What to do when a sanity check fails."""
+
+    def __bool__(self):
+        return self.enabled
+
+    def create_callbacks(self, root_config):
+        if not self:
+            return
+
+        yield RLPSanityChecksCallback(self)
+
+
+class RLPSanityChecksCallback(Callback):
+    @override
+    def __init__(self, config: RLPSanityChecksConfig):
+        super().__init__()
+
+        self.config = config
+        del config
+
+    @override
+    def on_train_start(self, trainer, pl_module):
+        # If we're in PL's "sanity check" mode, we don't need to run this check
+        if trainer.sanity_checking:
+            return
+
+        # If the sanity check is disabled, return.
+        if not self.config:
+            return
+
+        # If no lr schedulers, return.
+        if not trainer.lr_scheduler_configs:
+            return
+
+        errors: list[str] = []
+        disable_message = (
+            "Otherwise, set `config.trainer.sanity_checking.reduce_lr_on_plateau = None` "
+            "to disable this sanity check."
+        )
+
+        for lr_scheduler_config in trainer.lr_scheduler_configs:
+            if not lr_scheduler_config.reduce_on_plateau:
+                continue
+
+            match lr_scheduler_config.interval:
+                case "epoch":
+                    # we need to make sure that the trainer runs val every `frequency` epochs
+
+                    # If `trainer.check_val_every_n_epoch` is None, then Lightning
+                    # will run val every `int(trainer.val_check_interval)` steps.
+                    # So, first we need to make sure that `trainer.val_check_interval` is not None first.
+                    if trainer.check_val_every_n_epoch is None:
+                        errors.append(
+                            "Trainer is not running validation at epoch intervals "
+                            "(i.e., `trainer.check_val_every_n_epoch` is None) but "
+                            f"a ReduceLRPlateau scheduler with interval={lr_scheduler_config.interval} is used."
+                            f"Please set `config.trainer.check_val_every_n_epoch={lr_scheduler_config.frequency}`. "
+                            + disable_message
+                        )
+
+                    # Second, we make sure that the trainer runs val at least every `frequency` epochs
+                    if (
+                        trainer.check_val_every_n_epoch is not None
+                        and lr_scheduler_config.frequency
+                        % trainer.check_val_every_n_epoch
+                        != 0
+                    ):
+                        errors.append(
+                            f"Trainer is not running validation every {lr_scheduler_config.frequency} epochs but "
+                            f"a ReduceLRPlateau scheduler with interval={lr_scheduler_config.interval} and frequency={lr_scheduler_config.frequency} is used."
+                            f"Please set `config.trainer.check_val_every_n_epoch` to a multiple of {lr_scheduler_config.frequency}. "
+                            + disable_message
+                        )
+
+                case "step":
+                    # In this case, we need to make sure that the trainer runs val at step intervals
+                    # that are multiples of `frequency`.
+
+                    # First, we make sure that validation is run at step intervals
+                    if trainer.check_val_every_n_epoch is not None:
+                        errors.append(
+                            "Trainer is running validation at epoch intervals "
+                            "(i.e., `trainer.check_val_every_n_epoch` is not None) but "
+                            f"a ReduceLRPlateau scheduler with interval={lr_scheduler_config.interval} is used."
+                            "Please set `config.trainer.check_val_every_n_epoch=None` "
+                            f"and `config.trainer.val_check_interval={lr_scheduler_config.frequency}`. "
+                            + disable_message
+                        )
+
+                    # Second, we make sure `trainer.val_check_interval` is an integer
+                    if not isinstance(trainer.val_check_interval, int):
+                        errors.append(
+                            f"Trainer is not running validation at step intervals "
+                            f"(i.e., `trainer.val_check_interval` is not an integer) but "
+                            f"a ReduceLRPlateau scheduler with interval={lr_scheduler_config.interval} is used."
+                            "Please set `config.trainer.val_check_interval=None` "
+                            f"and `config.trainer.val_check_interval={lr_scheduler_config.frequency}`. "
+                            + disable_message
+                        )
+
+                    # Third, we make sure that the trainer runs val at least every `frequency` steps
+                    if (
+                        isinstance(trainer.val_check_interval, int)
+                        and trainer.val_check_interval % lr_scheduler_config.frequency
+                        != 0
+                    ):
+                        errors.append(
+                            f"Trainer is not running validation every {lr_scheduler_config.frequency} steps but "
+                            f"a ReduceLRPlateau scheduler with interval={lr_scheduler_config.interval} and frequency={lr_scheduler_config.frequency} is used."
+                            "Please set `config.trainer.val_check_interval` "
+                            f"to a multiple of {lr_scheduler_config.frequency}. "
+                            + disable_message
+                        )
+
+                case _:
+                    pass
+
+        if not errors:
+            return
+
+        message = (
+            "ReduceLRPlateau sanity checks failed with the following errors:\n"
+            + "\n".join(errors)
+        )
+        match self.config.on_error:
+            case "warn":
+                log.warning(message)
+            case "error":
+                raise ValueError(message)
+            case _:
+                pass
+
+
+@runtime_checkable
+class CustomRLPImplementation(Protocol):
+    __reduce_lr_on_plateau__: bool
+
+
+class _RLPSanityCheckModuleMixin(LightningModule):
+    def reduce_lr_on_plateau_config(
+        self,
+        lr_scheduler: LRSchedulerTypeUnion | LRSchedulerConfigType,
+    ) -> LRSchedulerConfigType:
+        if (trainer := self._trainer) is None:
+            raise RuntimeError(
+                "Could not determine the frequency of ReduceLRPlateau scheduler "
+                "because `self.trainer` is None."
+            )
+
+        # First, resolve the LR scheduler from the provided config.
+        lr_scheduler_config: LRSchedulerConfigType
+        match lr_scheduler:
+            case Mapping():
+                lr_scheduler_config = cast(LRSchedulerConfigType, lr_scheduler)
+            case _:
+                lr_scheduler_config = {"scheduler": lr_scheduler}
+
+        # Make sure the scheduler is a ReduceLRPlateau scheduler. Otherwise, warn the user.
+        if (
+            not isinstance(
+                lr_scheduler_config["scheduler"],
+                torch.optim.lr_scheduler.ReduceLROnPlateau,
+            )
+        ) and (
+            not isinstance(lr_scheduler_config["scheduler"], CustomRLPImplementation)
+            or not lr_scheduler_config["scheduler"].__reduce_lr_on_plateau__
+        ):
+            log.warning(
+                "`reduce_lr_on_plateau_config` should only be used with a ReduceLRPlateau scheduler. "
+                f"The provided scheduler, {lr_scheduler_config['scheduler']}, does not subclass "
+                "`torch.optim.lr_scheduler.ReduceLROnPlateau`. "
+                "Please ensure that the scheduler is a ReduceLRPlateau scheduler. "
+                "If you are using a custom ReduceLRPlateau scheduler implementation, "
+                "please either (1) make sure that it subclasses `torch.optim.lr_scheduler.ReduceLROnPlateau`, "
+                "or (2) set the scheduler's `__reduce_lr_on_plateau__` attribute to `True`."
+            )
+
+        # If trainer.check_val_every_n_epoch is an integer, then we run val at epoch intervals.
+        if trainer.check_val_every_n_epoch is not None:
+            return {
+                "reduce_on_plateau": True,
+                "interval": "epoch",
+                "frequency": trainer.check_val_every_n_epoch,
+                **lr_scheduler_config,
+            }
+
+        # Otherwise, we run val at step intervals.
+        if not isinstance(trainer.val_check_batch, int):
+            raise ValueError(
+                "Could not determine the frequency of ReduceLRPlateau scheduler "
+                f"because {trainer.val_check_batch=} is not an integer."
+            )
+
+        return {
+            "reduce_on_plateau": True,
+            "interval": "step",
+            "frequency": trainer.val_check_batch,
+            **lr_scheduler_config,
+        }

nshtrainer/callbacks/shared_parameters.py

@@ -0,0 +1,87 @@
+import logging
+from collections.abc import Iterable
+from typing import Literal, Protocol, TypeAlias, runtime_checkable
+
+import torch.nn as nn
+from lightning.pytorch import LightningModule, Trainer
+from lightning.pytorch.callbacks import Callback
+from typing_extensions import override
+
+from .base import CallbackConfigBase
+
+log = logging.getLogger(__name__)
+
+
+def _parameters_to_names(parameters: Iterable[nn.Parameter], model: nn.Module):
+    mapping = {id(p): n for n, p in model.named_parameters()}
+    return [mapping[id(p)] for p in parameters]
+
+
+class SharedParametersConfig(CallbackConfigBase):
+    """A callback that allows scaling the gradients of shared parameters that
+    are registered in the ``self.shared_parameters`` list of the root module.
+
+    This is useful for models that share parameters across multiple modules and
+    want to downscale the gradients of these parameters to avoid overfitting.
+    """
+
+    name: Literal["shared_parameters"] = "shared_parameters"
+
+    @override
+    def create_callbacks(self, root_config):
+        yield SharedParametersCallback(self)
+
+
+SharedParametersList: TypeAlias = list[tuple[nn.Parameter, int | float]]
+
+
+@runtime_checkable
+class ModuleWithSharedParameters(Protocol):
+    @property
+    def shared_parameters(self) -> SharedParametersList: ...
+
+
+class SharedParametersCallback(Callback):
+    @override
+    def __init__(self, config: SharedParametersConfig):
+        super().__init__()
+
+        self.config = config
+        del config
+
+        self._warned_shared_parameters = False
+
+    def _shared_parameters(self, pl_module: LightningModule) -> SharedParametersList:
+        if not isinstance(pl_module, ModuleWithSharedParameters):
+            return []
+
+        return pl_module.shared_parameters
+
+    @override
+    def on_after_backward(self, trainer: Trainer, pl_module: LightningModule):
+        if not (shared_parameters := self._shared_parameters(pl_module)):
+            log.debug(
+                "No shared parameters to scale, skipping SharedParametersCallback"
+            )
+            return
+
+        log.debug(f"Scaling {len(shared_parameters)} shared parameters...")
+        no_grad_parameters: list[nn.Parameter] = []
+        for p, factor in shared_parameters:
+            if not hasattr(p, "grad") or p.grad is None:
+                no_grad_parameters.append(p)
+                continue
+
+            _ = p.grad.data.div_(factor)
+
+        if no_grad_parameters and not self._warned_shared_parameters:
+            no_grad_parameters_str = ", ".join(
+                _parameters_to_names(no_grad_parameters, pl_module)
+            )
+            log.warning(
+                "The following parameters were marked as shared, but had no gradients: "
+                f"{no_grad_parameters_str}"
+            )
+            self._warned_shared_parameters = True
+
+        log.debug(f"Done scaling shared parameters. (len={len(shared_parameters)})")

nshtrainer/config.py

@@ -0,0 +1,67 @@
+from nshconfig._config import Config as Config
+from nshsnap._config import SnapshotConfig as SnapshotConfig
+from nshtrainer._checkpoint.loader import CheckpointLoadingConfig as CheckpointLoadingConfig
+from nshtrainer._checkpoint.metadata import CheckpointMetadata as CheckpointMetadata
+from nshtrainer._directory import DirectoryConfig as DirectoryConfig
+from nshtrainer._hf_hub import HuggingFaceHubAutoCreateConfig as HuggingFaceHubAutoCreateConfig
+from nshtrainer._hf_hub import HuggingFaceHubConfig as HuggingFaceHubConfig
+from nshtrainer.callbacks.actsave import ActSaveConfig as ActSaveConfig
+from nshtrainer.callbacks.base import CallbackConfigBase as CallbackConfigBase
+from nshtrainer.callbacks.checkpoint._base import BaseCheckpointCallbackConfig as BaseCheckpointCallbackConfig
+from nshtrainer.callbacks.checkpoint.best_checkpoint import BestCheckpointCallbackConfig as BestCheckpointCallbackConfig
+from nshtrainer.callbacks.checkpoint.last_checkpoint import LastCheckpointCallbackConfig as LastCheckpointCallbackConfig
+from nshtrainer.callbacks.checkpoint.on_exception_checkpoint import OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig
+from nshtrainer.callbacks.directory_setup import DirectorySetupConfig as DirectorySetupConfig
+from nshtrainer.callbacks.early_stopping import EarlyStoppingConfig as EarlyStoppingConfig
+from nshtrainer.callbacks.ema import EMAConfig as EMAConfig
+from nshtrainer.callbacks.finite_checks import FiniteChecksConfig as FiniteChecksConfig
+from nshtrainer.callbacks.gradient_skipping import GradientSkippingConfig as GradientSkippingConfig
+from nshtrainer.callbacks.norm_logging import NormLoggingConfig as NormLoggingConfig
+from nshtrainer.callbacks.print_table import PrintTableMetricsConfig as PrintTableMetricsConfig
+from nshtrainer.callbacks.rlp_sanity_checks import RLPSanityChecksConfig as RLPSanityChecksConfig
+from nshtrainer.callbacks.shared_parameters import SharedParametersConfig as SharedParametersConfig
+from nshtrainer.callbacks.throughput_monitor import ThroughputMonitorConfig as ThroughputMonitorConfig
+from nshtrainer.callbacks.timer import EpochTimerConfig as EpochTimerConfig
+from nshtrainer.callbacks.wandb_watch import WandbWatchConfig as WandbWatchConfig
+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._base import LRSchedulerConfigBase as LRSchedulerConfigBase
+from nshtrainer.lr_scheduler.linear_warmup_cosine import LinearWarmupCosineDecayLRSchedulerConfig as LinearWarmupCosineDecayLRSchedulerConfig
+from nshtrainer.lr_scheduler.reduce_lr_on_plateau import ReduceLROnPlateauConfig as ReduceLROnPlateauConfig
+from nshtrainer.metrics._config import MetricConfig as MetricConfig
+from nshtrainer.model.config import BaseConfig as BaseConfig
+from nshtrainer.nn.mlp import MLPConfig as MLPConfig
+from nshtrainer.nn.nonlinearity import BaseNonlinearityConfig as BaseNonlinearityConfig
+from nshtrainer.nn.nonlinearity import ELUNonlinearityConfig as ELUNonlinearityConfig
+from nshtrainer.nn.nonlinearity import GELUNonlinearityConfig as GELUNonlinearityConfig
+from nshtrainer.nn.nonlinearity import LeakyReLUNonlinearityConfig as LeakyReLUNonlinearityConfig
+from nshtrainer.nn.nonlinearity import MishNonlinearityConfig as MishNonlinearityConfig
+from nshtrainer.nn.nonlinearity import PReLUConfig as PReLUConfig
+from nshtrainer.nn.nonlinearity import ReLUNonlinearityConfig as ReLUNonlinearityConfig
+from nshtrainer.nn.nonlinearity import SiLUNonlinearityConfig as SiLUNonlinearityConfig
+from nshtrainer.nn.nonlinearity import SigmoidNonlinearityConfig as SigmoidNonlinearityConfig
+from nshtrainer.nn.nonlinearity import SoftmaxNonlinearityConfig as SoftmaxNonlinearityConfig
+from nshtrainer.nn.nonlinearity import SoftplusNonlinearityConfig as SoftplusNonlinearityConfig
+from nshtrainer.nn.nonlinearity import SoftsignNonlinearityConfig as SoftsignNonlinearityConfig
+from nshtrainer.nn.nonlinearity import SwiGLUNonlinearityConfig as SwiGLUNonlinearityConfig
+from nshtrainer.nn.nonlinearity import SwishNonlinearityConfig as SwishNonlinearityConfig
+from nshtrainer.nn.nonlinearity import TanhNonlinearityConfig as TanhNonlinearityConfig
+from nshtrainer.optimizer import AdamWConfig as AdamWConfig
+from nshtrainer.optimizer import OptimizerConfigBase as OptimizerConfigBase
+from nshtrainer.profiler._base import BaseProfilerConfig as BaseProfilerConfig
+from nshtrainer.profiler.advanced import AdvancedProfilerConfig as AdvancedProfilerConfig
+from nshtrainer.profiler.pytorch import PyTorchProfilerConfig as PyTorchProfilerConfig
+from nshtrainer.profiler.simple import SimpleProfilerConfig as SimpleProfilerConfig
+from nshtrainer.trainer._config import CheckpointSavingConfig as CheckpointSavingConfig
+from nshtrainer.trainer._config import GradientClippingConfig as GradientClippingConfig
+from nshtrainer.trainer._config import LoggingConfig as LoggingConfig
+from nshtrainer.trainer._config import OptimizationConfig as OptimizationConfig
+from nshtrainer.trainer._config import ReproducibilityConfig as ReproducibilityConfig
+from nshtrainer.trainer._config import SanityCheckingConfig as SanityCheckingConfig
+from nshtrainer.trainer._config import TrainerConfig as TrainerConfig
+from nshtrainer.util._environment_info import EnvironmentClassInformationConfig as EnvironmentClassInformationConfig
+from nshtrainer.util._environment_info import EnvironmentConfig as EnvironmentConfig
+from nshtrainer.util._environment_info import EnvironmentLinuxEnvironmentConfig as EnvironmentLinuxEnvironmentConfig
+from nshtrainer.util._environment_info import EnvironmentSLURMInformationConfig as EnvironmentSLURMInformationConfig

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