nshtrainer 0.1.0__py3-none-any.whl

nshtrainer/model/modules/shared_parameters.py

@@ -0,0 +1,72 @@
+from collections.abc import Sequence
+from logging import getLogger
+from typing import cast
+
+import torch.nn as nn
+from lightning.pytorch import LightningModule, Trainer
+from typing_extensions import override
+
+from ...util.typing_utils import mixin_base_type
+from ..config import BaseConfig
+from .callback import CallbackRegistrarModuleMixin
+
+log = getLogger(__name__)
+
+
+def _parameters_to_names(parameters: Sequence[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 SharedParametersModuleMixin(mixin_base_type(CallbackRegistrarModuleMixin)):
+    @override
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.shared_parameters: list[tuple[nn.Parameter, int | float]] = []
+        self._warned_shared_parameters = False
+
+        def on_after_backward(_trainer: Trainer, pl_module: LightningModule):
+            nonlocal self
+
+            config = cast(BaseConfig, pl_module.hparams)
+            if not config.trainer.supports_shared_parameters:
+                return
+
+            log.debug(f"Scaling {len(self.shared_parameters)} shared parameters...")
+            no_grad_parameters: list[nn.Parameter] = []
+            for p, factor in self.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(self.shared_parameters)})"
+            )
+
+        self.register_callback(on_after_backward=on_after_backward)
+
+    def register_shared_parameters(
+        self, parameters: list[tuple[nn.Parameter, int | float]]
+    ):
+        for parameter, factor in parameters:
+            if not isinstance(parameter, nn.Parameter):
+                raise ValueError("Shared parameters must be PyTorch parameters")
+            if not isinstance(factor, (int, float)):
+                raise ValueError("Factor must be an integer or float")
+
+            self.shared_parameters.append((parameter, factor))
+
+        log.info(f"Registered {len(parameters)} shared parameters")

nshtrainer/nn/__init__.py

@@ -0,0 +1,19 @@
+from .mlp import MLP as MLP
+from .mlp import ResidualSequential as ResidualSequential
+from .module_dict import TypedModuleDict as TypedModuleDict
+from .module_list import TypedModuleList as TypedModuleList
+from .nonlinearity import BaseNonlinearityConfig as BaseNonlinearityConfig
+from .nonlinearity import ELUNonlinearityConfig as ELUNonlinearityConfig
+from .nonlinearity import GELUNonlinearityConfig as GELUNonlinearityConfig
+from .nonlinearity import LeakyReLUNonlinearityConfig as LeakyReLUNonlinearityConfig
+from .nonlinearity import MishNonlinearityConfig as MishNonlinearityConfig
+from .nonlinearity import NonlinearityConfig as NonlinearityConfig
+from .nonlinearity import PReLUConfig as PReLUConfig
+from .nonlinearity import ReLUNonlinearityConfig as ReLUNonlinearityConfig
+from .nonlinearity import SigmoidNonlinearityConfig as SigmoidNonlinearityConfig
+from .nonlinearity import SiLUNonlinearityConfig as SiLUNonlinearityConfig
+from .nonlinearity import SoftmaxNonlinearityConfig as SoftmaxNonlinearityConfig
+from .nonlinearity import SoftplusNonlinearityConfig as SoftplusNonlinearityConfig
+from .nonlinearity import SoftsignNonlinearityConfig as SoftsignNonlinearityConfig
+from .nonlinearity import SwishNonlinearityConfig as SwishNonlinearityConfig
+from .nonlinearity import TanhNonlinearityConfig as TanhNonlinearityConfig

nshtrainer/nn/mlp.py

@@ -0,0 +1,106 @@
+import copy
+from collections.abc import Callable, Sequence
+from typing import Literal, Protocol, runtime_checkable
+
+import torch
+import torch.nn as nn
+from typing_extensions import override
+
+from .nonlinearity import BaseNonlinearityConfig
+
+
+@runtime_checkable
+class LinearModuleConstructor(Protocol):
+    def __call__(
+        self, in_features: int, out_features: int, bias: bool = True
+    ) -> nn.Module: ...
+
+
+class ResidualSequential(nn.Sequential):
+    @override
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return input + super().forward(input)
+
+
+def MLP(
+    dims: Sequence[int],
+    activation: BaseNonlinearityConfig
+    | nn.Module
+    | Callable[[], nn.Module]
+    | None = None,
+    nonlinearity: BaseNonlinearityConfig
+    | nn.Module
+    | Callable[[], nn.Module]
+    | None = None,
+    bias: bool = True,
+    no_bias_scalar: bool = True,
+    ln: bool | Literal["pre", "post"] = False,
+    dropout: float | None = None,
+    residual: bool = False,
+    pre_layers: Sequence[nn.Module] = [],
+    post_layers: Sequence[nn.Module] = [],
+    linear_cls: LinearModuleConstructor = nn.Linear,
+):
+    """
+    Constructs a multi-layer perceptron (MLP) with the given dimensions and activation function.
+
+    Args:
+        dims (Sequence[int]): List of integers representing the dimensions of the MLP.
+        nonlinearity (Callable[[], nn.Module]): Activation function to use between layers.
+        activation (Callable[[], nn.Module]): Activation function to use between layers.
+        bias (bool, optional): Whether to include bias terms in the linear layers. Defaults to True.
+        no_bias_scalar (bool, optional): Whether to exclude bias terms when the output dimension is 1. Defaults to True.
+        ln (bool | Literal["pre", "post"], optional): Whether to apply layer normalization before or after the linear layers. Defaults to False.
+        dropout (float | None, optional): Dropout probability to apply between layers. Defaults to None.
+        residual (bool, optional): Whether to use residual connections between layers. Defaults to False.
+        pre_layers (Sequence[nn.Module], optional): List of layers to insert before the linear layers. Defaults to [].
+        post_layers (Sequence[nn.Module], optional): List of layers to insert after the linear layers. Defaults to [].
+
+    Returns:
+        nn.Sequential: The constructed MLP.
+    """
+
+    if activation is None:
+        activation = nonlinearity
+
+    if len(dims) < 2:
+        raise ValueError("mlp requires at least 2 dimensions")
+    if ln is True:
+        ln = "pre"
+    elif isinstance(ln, str) and ln not in ("pre", "post"):
+        raise ValueError("ln must be a boolean or 'pre' or 'post'")
+
+    layers: list[nn.Module] = []
+    if ln == "pre":
+        layers.append(nn.LayerNorm(dims[0]))
+
+    layers.extend(pre_layers)
+
+    for i in range(len(dims) - 1):
+        in_features = dims[i]
+        out_features = dims[i + 1]
+        bias_ = bias and not (no_bias_scalar and out_features == 1)
+        layers.append(linear_cls(in_features, out_features, bias=bias_))
+        if dropout is not None:
+            layers.append(nn.Dropout(dropout))
+        if i < len(dims) - 2:
+            match activation:
+                case BaseNonlinearityConfig():
+                    layers.append(activation.create_module())
+                case nn.Module():
+                    # In this case, we create a deep copy of the module to avoid sharing parameters (if any).
+                    layers.append(copy.deepcopy(activation))
+                case Callable():
+                    layers.append(activation())
+                case _:
+                    raise ValueError(
+                        "Either `nonlinearity` or `activation` must be provided"
+                    )
+
+    layers.extend(post_layers)
+
+    if ln == "post":
+        layers.append(nn.LayerNorm(dims[-1]))
+
+    cls = ResidualSequential if residual else nn.Sequential
+    return cls(*layers)

nshtrainer/nn/module_dict.py

@@ -0,0 +1,66 @@
+from collections.abc import Iterable, Mapping
+from typing import Generic, cast
+
+import torch.nn as nn
+from typing_extensions import TypeVar
+
+TModule = TypeVar("TModule", bound=nn.Module, infer_variance=True)
+
+
+class TypedModuleDict(nn.Module, Generic[TModule]):
+    def __init__(
+        self,
+        modules: Mapping[str, TModule],
+        key_prefix: str = "_typed_moduledict_",
+        # we use a key prefix to avoid attribute name collisions
+        # (which is a common issue in nn.ModuleDict as it uses `__setattr__` to set the modules)
+    ):
+        super().__init__()
+
+        self.key_prefix = key_prefix
+        self._module_dict = nn.ModuleDict(
+            {self._with_prefix(k): v for k, v in modules.items()}
+        )
+
+    def _with_prefix(self, key: str) -> str:
+        return f"{self.key_prefix}{key}"
+
+    def _remove_prefix(self, key: str) -> str:
+        assert key.startswith(
+            self.key_prefix
+        ), f"{key} does not start with {self.key_prefix}"
+        return key[len(self.key_prefix) :]
+
+    def __setitem__(self, key: str, module: TModule) -> None:
+        key = self._with_prefix(key)
+        return self._module_dict.__setitem__(key, module)
+
+    def __getitem__(self, key: str) -> TModule:
+        key = self._with_prefix(key)
+        return self._module_dict.__getitem__(key)  # type: ignore
+
+    def update(self, modules: Mapping[str, TModule]) -> None:
+        return self._module_dict.update(
+            {self._with_prefix(k): v for k, v in modules.items()}
+        )
+
+    def get(self, key: str) -> TModule | None:
+        key = self._with_prefix(key)
+        if (value := self._module_dict._modules.get(key)) is None:
+            return None
+        return cast(TModule, value)
+
+    def keys(self) -> Iterable[str]:
+        r"""Return an iterable of the ModuleDict keys."""
+        return [self._remove_prefix(k) for k in self._module_dict.keys()]
+
+    def items(self) -> Iterable[tuple[str, TModule]]:
+        r"""Return an iterable of the ModuleDict key/value pairs."""
+        return [
+            (self._remove_prefix(k), cast(TModule, v))
+            for k, v in self._module_dict.items()
+        ]
+
+    def values(self) -> Iterable[TModule]:
+        r"""Return an iterable of the ModuleDict values."""
+        return cast(Iterable[TModule], self._module_dict.values())

nshtrainer/nn/module_list.py

@@ -0,0 +1,50 @@
+from collections.abc import Iterable, Iterator
+from typing import Generic, TypeVar, overload
+
+import torch.nn as nn
+from typing_extensions import override
+
+TModule = TypeVar("TModule", bound=nn.Module)
+
+
+class TypedModuleList(nn.ModuleList, Generic[TModule]):
+    def __init__(self, modules: Iterable[TModule] | None = None) -> None:
+        super().__init__(modules)
+
+    @overload
+    def __getitem__(self, idx: int) -> TModule: ...
+
+    @overload
+    def __getitem__(self, idx: slice) -> "TypedModuleList[TModule]": ...
+
+    @override
+    def __getitem__(self, idx: int | slice) -> TModule | "TypedModuleList[TModule]":
+        return super().__getitem__(idx)  # type: ignore
+
+    @override
+    def __setitem__(self, idx: int, module: TModule) -> None:  # type: ignore
+        return super().__setitem__(idx, module)
+
+    @override
+    def __iter__(self) -> Iterator[TModule]:
+        return super().__iter__()  # type: ignore
+
+    @override
+    def __iadd__(self, modules: Iterable[TModule]) -> "TypedModuleList[TModule]":  # type: ignore
+        return super().__iadd__(modules)  # type: ignore
+
+    @override
+    def __add__(self, modules: Iterable[TModule]) -> "TypedModuleList[TModule]":  # type: ignore
+        return super().__add__(modules)  # type: ignore
+
+    @override
+    def insert(self, idx: int, module: TModule) -> None:  # type: ignore
+        return super().insert(idx, module)
+
+    @override
+    def append(self, module: TModule) -> "TypedModuleList[TModule]":  # type: ignore
+        return super().append(module)  # type: ignore
+
+    @override
+    def extend(self, modules: Iterable[TModule]) -> "TypedModuleList[TModule]":  # type: ignore
+        return super().extend(modules)  # type: ignore

nshtrainer/nn/nonlinearity.py

@@ -0,0 +1,157 @@
+from abc import ABC, abstractmethod
+from typing import Annotated, Literal
+
+import torch
+import torch.nn as nn
+from typing_extensions import override
+
+from ..config import Field, TypedConfig
+
+
+class BaseNonlinearityConfig(TypedConfig, ABC):
+    @abstractmethod
+    def create_module(self) -> nn.Module:
+        pass
+
+
+class ReLUNonlinearityConfig(BaseNonlinearityConfig):
+    name: Literal["relu"] = "relu"
+
+    @override
+    def create_module(self) -> nn.Module:
+        return nn.ReLU()
+
+
+class SigmoidNonlinearityConfig(BaseNonlinearityConfig):
+    name: Literal["sigmoid"] = "sigmoid"
+
+    @override
+    def create_module(self) -> nn.Module:
+        return nn.Sigmoid()
+
+
+class TanhNonlinearityConfig(BaseNonlinearityConfig):
+    name: Literal["tanh"] = "tanh"
+
+    @override
+    def create_module(self) -> nn.Module:
+        return nn.Tanh()
+
+
+class SoftmaxNonlinearityConfig(BaseNonlinearityConfig):
+    name: Literal["softmax"] = "softmax"
+
+    @override
+    def create_module(self) -> nn.Module:
+        return nn.Softmax(dim=1)
+
+
+class SoftplusNonlinearityConfig(BaseNonlinearityConfig):
+    name: Literal["softplus"] = "softplus"
+
+    @override
+    def create_module(self) -> nn.Module:
+        return nn.Softplus()
+
+
+class SoftsignNonlinearityConfig(BaseNonlinearityConfig):
+    name: Literal["softsign"] = "softsign"
+
+    @override
+    def create_module(self) -> nn.Module:
+        return nn.Softsign()
+
+
+class ELUNonlinearityConfig(BaseNonlinearityConfig):
+    name: Literal["elu"] = "elu"
+
+    @override
+    def create_module(self) -> nn.Module:
+        return nn.ELU()
+
+
+class LeakyReLUNonlinearityConfig(BaseNonlinearityConfig):
+    name: Literal["leaky_relu"] = "leaky_relu"
+
+    negative_slope: float | None = None
+
+    @override
+    def create_module(self) -> nn.Module:
+        kwargs = {}
+        if self.negative_slope is not None:
+            kwargs["negative_slope"] = self.negative_slope
+        return nn.LeakyReLU(**kwargs)
+
+
+class PReLUConfig(BaseNonlinearityConfig):
+    name: Literal["prelu"] = "prelu"
+
+    @override
+    def create_module(self) -> nn.Module:
+        return nn.PReLU()
+
+
+class GELUNonlinearityConfig(BaseNonlinearityConfig):
+    name: Literal["gelu"] = "gelu"
+
+    @override
+    def create_module(self) -> nn.Module:
+        return nn.GELU()
+
+
+class SwishNonlinearityConfig(BaseNonlinearityConfig):
+    name: Literal["swish"] = "swish"
+
+    @override
+    def create_module(self) -> nn.Module:
+        return nn.SiLU()
+
+
+class SiLUNonlinearityConfig(BaseNonlinearityConfig):
+    name: Literal["silu"] = "silu"
+
+    @override
+    def create_module(self) -> nn.Module:
+        return nn.SiLU()
+
+
+class MishNonlinearityConfig(BaseNonlinearityConfig):
+    name: Literal["mish"] = "mish"
+
+    @override
+    def create_module(self) -> nn.Module:
+        return nn.Mish()
+
+
+class SwiGLU(nn.SiLU):
+    @override
+    def forward(self, input: torch.Tensor):
+        input, gate = input.chunk(2, dim=-1)
+        return input * super().forward(gate)
+
+
+class SwiGLUNonlinearityConfig(BaseNonlinearityConfig):
+    name: Literal["swiglu"] = "swiglu"
+
+    @override
+    def create_module(self) -> nn.Module:
+        return SwiGLU()
+
+
+NonlinearityConfig = Annotated[
+    ReLUNonlinearityConfig
+    | SigmoidNonlinearityConfig
+    | TanhNonlinearityConfig
+    | SoftmaxNonlinearityConfig
+    | SoftplusNonlinearityConfig
+    | SoftsignNonlinearityConfig
+    | ELUNonlinearityConfig
+    | LeakyReLUNonlinearityConfig
+    | PReLUConfig
+    | GELUNonlinearityConfig
+    | SwishNonlinearityConfig
+    | SiLUNonlinearityConfig
+    | MishNonlinearityConfig
+    | SwiGLUNonlinearityConfig,
+    Field(discriminator="name"),
+]

nshtrainer/optimizer.py

@@ -0,0 +1,62 @@
+from abc import ABC, abstractmethod
+from collections.abc import Iterable
+from typing import Annotated, Any, Literal, TypeAlias
+
+import torch.nn as nn
+from torch.optim import Optimizer
+from typing_extensions import override
+
+from .config import Field, TypedConfig
+
+
+class OptimizerConfigBase(TypedConfig, ABC):
+    @abstractmethod
+    def create_optimizer(
+        self,
+        parameters: Iterable[nn.Parameter] | Iterable[dict[str, Any]],
+    ) -> Optimizer: ...
+
+
+class AdamWConfig(OptimizerConfigBase):
+    name: Literal["adamw"] = "adamw"
+
+    lr: float
+    """Learning rate for the optimizer."""
+
+    weight_decay: float = 1.0e-2
+    """Weight decay (L2 penalty) for the optimizer."""
+
+    betas: tuple[float, float] = (0.9, 0.999)
+    """
+    Betas for the optimizer:
+    (beta1, beta2) are the coefficients used for computing running averages of
+    gradient and its square.
+    """
+
+    eps: float = 1e-8
+    """Term added to the denominator to improve numerical stability."""
+
+    amsgrad: bool = False
+    """Whether to use the AMSGrad variant of this algorithm."""
+
+    @override
+    def create_optimizer(
+        self,
+        parameters: Iterable[nn.Parameter] | Iterable[dict[str, Any]],
+    ):
+        from torch.optim import AdamW
+
+        return AdamW(
+            parameters,
+            lr=self.lr,
+            weight_decay=self.weight_decay,
+            betas=self.betas,
+            eps=self.eps,
+            amsgrad=self.amsgrad,
+        )
+
+
+OptimizerConfig: TypeAlias = Annotated[
+    AdamWConfig,
+    Field(discriminator="name"),
+]

nshtrainer/runner.py

@@ -0,0 +1,21 @@
+from dataclasses import dataclass
+from typing import Generic
+
+from nshrunner import Runner as _Runner
+from typing_extensions import Concatenate, TypeVar, TypeVarTuple, Unpack, override
+
+from .model.config import BaseConfig
+
+TConfig = TypeVar("TConfig", bound=BaseConfig, infer_variance=True)
+TArguments = TypeVarTuple("TArguments")
+TReturn = TypeVar("TReturn", infer_variance=True)
+
+
+@dataclass(frozen=True)
+class Runner(
+    _Runner[Unpack[tuple[TConfig, Unpack[TArguments]]], TReturn],
+    Generic[TConfig, Unpack[TArguments], TReturn],
+):
+    @override
+    def default_validate_fn():
+        pass

nshtrainer/scripts/check_env.py

@@ -0,0 +1,41 @@
+REQUIRED_PACKAGES = [
+    "beartype",
+    "cloudpickle",
+    "jaxtyping",
+    "lightning",
+    "lightning_fabric",
+    "lightning_utilities",
+    "lovely_numpy",
+    "lovely_tensors",
+    "numpy",
+    "psutil",
+    "pydantic",
+    "pydantic_core",
+    "pysnooper",
+    "rich",
+    "tabulate",
+    "torch",
+    "torchmetrics",
+    "tqdm",
+    "typing_extensions",
+    "wrapt",
+    "yaml",
+]
+
+
+def main():
+    import importlib.util
+    import sys
+
+    missing_packages: list[str] = []
+    for package_name in REQUIRED_PACKAGES:
+        spec = importlib.util.find_spec(package_name)
+        if spec is None:
+            missing_packages.append(package_name)
+
+    if missing_packages:
+        sys.exit(f"Error: Missing required packages: {', '.join(missing_packages)}")
+
+
+if __name__ == "__main__":
+    main()

nshtrainer/scripts/find_packages.py

@@ -0,0 +1,51 @@
+import argparse
+import ast
+import glob
+import sys
+from pathlib import Path
+
+
+def get_imports(file_path: Path):
+    with open(file_path, "r") as file:
+        try:
+            tree = ast.parse(file.read())
+        except SyntaxError:
+            print(f"Syntax error in file: {file_path}", file=sys.stderr)
+            return set()
+
+    imports = set()
+    for node in ast.walk(tree):
+        if isinstance(node, ast.Import):
+            for alias in node.names:
+                imports.add(alias.name.split(".")[0])
+        elif isinstance(node, ast.ImportFrom):
+            if node.level == 0 and node.module:  # Absolute import
+                imports.add(node.module.split(".")[0])
+    return imports
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Find unique Python packages used in files."
+    )
+    parser.add_argument("glob_pattern", help="Glob pattern to match files")
+    parser.add_argument(
+        "--exclude-std", action="store_true", help="Exclude Python standard libraries"
+    )
+    args = parser.parse_args()
+
+    all_imports = set()
+    for file_path in glob.glob(args.glob_pattern, recursive=True):
+        all_imports.update(get_imports(Path(file_path)))
+
+    if args.exclude_std:
+        std_libs = set(sys.stdlib_module_names)
+        std_libs.update({"pkg_resources"})
+        all_imports = all_imports - std_libs
+
+    for package in sorted(all_imports):
+        print(package)
+
+
+if __name__ == "__main__":
+    main()

nshtrainer/trainer/__init__.py

@@ -0,0 +1 @@
+from .trainer import Trainer as Trainer