nshtrainer 1.0.0b54__tar.gz → 1.0.0b56__tar.gz

{nshtrainer-1.0.0b54 → nshtrainer-1.0.0b56}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: nshtrainer
-Version: 1.0.0b54
+Version: 1.0.0b56
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-1.0.0b54 → nshtrainer-1.0.0b56}/pyproject.toml

@@ -1,13 +1,13 @@
 [tool.poetry]
 name = "nshtrainer"
-version = "1.0.0-beta54"
+version = "1.0.0-beta56"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"
 
 [tool.poetry.dependencies]
 python = "^3.10"
-nshrunner = { version = "*", allow-prereleases = true }
+nshrunner = { version = "*" }
 nshconfig = "*"
 nshutils = { version = "*", optional = true }
 psutil = "*"
@@ -25,7 +25,7 @@ tensorboard = { version = "*", optional = true }
 huggingface-hub = { version = "*", optional = true }
 
 [tool.poetry.group.dev.dependencies]
-pyright = "*"
+basedpyright = "*"
 ruff = "*"
 ipykernel = "*"
 ipywidgets = "*"
@@ -36,7 +36,7 @@ pytest-cov = "^6.0.0"
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
 
-[tool.pyright]
+[tool.basedpyright]
 typeCheckingMode = "standard"
 deprecateTypingAliases = true
 strictListInference = true

{nshtrainer-1.0.0b54 → nshtrainer-1.0.0b56}/src/nshtrainer/configs/__init__.py

@@ -85,6 +85,7 @@ from nshtrainer.nn import NonlinearityConfig as NonlinearityConfig
 from nshtrainer.nn import NonlinearityConfigBase as NonlinearityConfigBase
 from nshtrainer.nn import PReLUConfig as PReLUConfig
 from nshtrainer.nn import ReLUNonlinearityConfig as ReLUNonlinearityConfig
+from nshtrainer.nn import RNGConfig as RNGConfig
 from nshtrainer.nn import SigmoidNonlinearityConfig as SigmoidNonlinearityConfig
 from nshtrainer.nn import SiLUNonlinearityConfig as SiLUNonlinearityConfig
 from nshtrainer.nn import SoftmaxNonlinearityConfig as SoftmaxNonlinearityConfig
@@ -306,6 +307,7 @@ __all__ = [
     "ProfilerConfig",
     "PyTorchProfilerConfig",
     "RLPSanityChecksCallbackConfig",
+    "RNGConfig",
     "ReLUNonlinearityConfig",
     "ReduceLROnPlateauConfig",
     "SLURMEnvironmentPlugin",

{nshtrainer-1.0.0b54 → nshtrainer-1.0.0b56}/src/nshtrainer/configs/lr_scheduler/__init__.py

@@ -12,6 +12,7 @@ from nshtrainer.lr_scheduler.base import lr_scheduler_registry as lr_scheduler_r
 from nshtrainer.lr_scheduler.linear_warmup_cosine import (
     DurationConfig as DurationConfig,
 )
+from nshtrainer.lr_scheduler.reduce_lr_on_plateau import EpochsConfig as EpochsConfig
 from nshtrainer.lr_scheduler.reduce_lr_on_plateau import MetricConfig as MetricConfig
 
 from . import base as base
@@ -20,6 +21,7 @@ from . import reduce_lr_on_plateau as reduce_lr_on_plateau
 
 __all__ = [
     "DurationConfig",
+    "EpochsConfig",
     "LRSchedulerConfig",
     "LRSchedulerConfigBase",
     "LinearWarmupCosineDecayLRSchedulerConfig",

{nshtrainer-1.0.0b54 → nshtrainer-1.0.0b56}/src/nshtrainer/configs/lr_scheduler/reduce_lr_on_plateau/__init__.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 __codegen__ = True
 
+from nshtrainer.lr_scheduler.reduce_lr_on_plateau import EpochsConfig as EpochsConfig
 from nshtrainer.lr_scheduler.reduce_lr_on_plateau import (
     LRSchedulerConfigBase as LRSchedulerConfigBase,
 )
@@ -14,6 +15,7 @@ from nshtrainer.lr_scheduler.reduce_lr_on_plateau import (
 )
 
 __all__ = [
+    "EpochsConfig",
     "LRSchedulerConfigBase",
     "MetricConfig",
     "ReduceLROnPlateauConfig",

{nshtrainer-1.0.0b54 → nshtrainer-1.0.0b56}/src/nshtrainer/configs/nn/__init__.py

@@ -11,6 +11,7 @@ from nshtrainer.nn import NonlinearityConfig as NonlinearityConfig
 from nshtrainer.nn import NonlinearityConfigBase as NonlinearityConfigBase
 from nshtrainer.nn import PReLUConfig as PReLUConfig
 from nshtrainer.nn import ReLUNonlinearityConfig as ReLUNonlinearityConfig
+from nshtrainer.nn import RNGConfig as RNGConfig
 from nshtrainer.nn import SigmoidNonlinearityConfig as SigmoidNonlinearityConfig
 from nshtrainer.nn import SiLUNonlinearityConfig as SiLUNonlinearityConfig
 from nshtrainer.nn import SoftmaxNonlinearityConfig as SoftmaxNonlinearityConfig
@@ -25,6 +26,7 @@ from nshtrainer.nn.nonlinearity import nonlinearity_registry as nonlinearity_reg
 
 from . import mlp as mlp
 from . import nonlinearity as nonlinearity
+from . import rng as rng
 
 __all__ = [
     "ELUNonlinearityConfig",
@@ -35,6 +37,7 @@ __all__ = [
     "NonlinearityConfig",
     "NonlinearityConfigBase",
     "PReLUConfig",
+    "RNGConfig",
     "ReLUNonlinearityConfig",
     "SiLUNonlinearityConfig",
     "SigmoidNonlinearityConfig",
@@ -47,4 +50,5 @@ __all__ = [
     "mlp",
     "nonlinearity",
     "nonlinearity_registry",
+    "rng",
 ]

nshtrainer-1.0.0b56/src/nshtrainer/configs/nn/rng/__init__.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+__codegen__ = True
+
+from nshtrainer.nn.rng import RNGConfig as RNGConfig
+
+__all__ = [
+    "RNGConfig",
+]

{nshtrainer-1.0.0b54 → nshtrainer-1.0.0b56}/src/nshtrainer/configs/trainer/plugin/base/__init__.py

@@ -2,12 +2,10 @@ from __future__ import annotations
 
 __codegen__ = True
 
-from nshtrainer.trainer.plugin.base import PluginConfig as PluginConfig
 from nshtrainer.trainer.plugin.base import PluginConfigBase as PluginConfigBase
 from nshtrainer.trainer.plugin.base import plugin_registry as plugin_registry
 
 __all__ = [
-    "PluginConfig",
     "PluginConfigBase",
     "plugin_registry",
 ]

{nshtrainer-1.0.0b54 → nshtrainer-1.0.0b56}/src/nshtrainer/configs/trainer/plugin/environment/__init__.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 __codegen__ = True
 
-from nshtrainer.trainer.plugin.environment import DTypeConfig as DTypeConfig
 from nshtrainer.trainer.plugin.environment import (
     KubeflowEnvironmentPlugin as KubeflowEnvironmentPlugin,
 )
@@ -28,7 +27,6 @@ from nshtrainer.trainer.plugin.environment import (
 from nshtrainer.trainer.plugin.environment import plugin_registry as plugin_registry
 
 __all__ = [
-    "DTypeConfig",
     "KubeflowEnvironmentPlugin",
     "LSFEnvironmentPlugin",
     "LightningEnvironmentPlugin",

{nshtrainer-1.0.0b54 → nshtrainer-1.0.0b56}/src/nshtrainer/configs/trainer/plugin/io/__init__.py

@@ -5,7 +5,6 @@ __codegen__ = True
 from nshtrainer.trainer.plugin.io import (
     AsyncCheckpointIOPlugin as AsyncCheckpointIOPlugin,
 )
-from nshtrainer.trainer.plugin.io import PluginConfig as PluginConfig
 from nshtrainer.trainer.plugin.io import PluginConfigBase as PluginConfigBase
 from nshtrainer.trainer.plugin.io import (
     TorchCheckpointIOPlugin as TorchCheckpointIOPlugin,
@@ -15,7 +14,6 @@ from nshtrainer.trainer.plugin.io import plugin_registry as plugin_registry
 
 __all__ = [
     "AsyncCheckpointIOPlugin",
-    "PluginConfig",
     "PluginConfigBase",
     "TorchCheckpointIOPlugin",
     "XLACheckpointIOPlugin",

{nshtrainer-1.0.0b54 → nshtrainer-1.0.0b56}/src/nshtrainer/nn/__init__.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 from .mlp import MLP as MLP
 from .mlp import MLPConfig as MLPConfig
 from .mlp import ResidualSequential as ResidualSequential
-from .mlp import custom_seed_context as custom_seed_context
 from .module_dict import TypedModuleDict as TypedModuleDict
 from .module_list import TypedModuleList as TypedModuleList
 from .nonlinearity import ELUNonlinearityConfig as ELUNonlinearityConfig
@@ -21,3 +20,5 @@ from .nonlinearity import SoftplusNonlinearityConfig as SoftplusNonlinearityConf
 from .nonlinearity import SoftsignNonlinearityConfig as SoftsignNonlinearityConfig
 from .nonlinearity import SwishNonlinearityConfig as SwishNonlinearityConfig
 from .nonlinearity import TanhNonlinearityConfig as TanhNonlinearityConfig
+from .rng import RNGConfig as RNGConfig
+from .rng import rng_context as rng_context

nshtrainer-1.0.0b56/src/nshtrainer/nn/rng.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+import contextlib
+
+import nshconfig as C
+import torch
+
+
+@contextlib.contextmanager
+def rng_context(config: RNGConfig | None):
+    with contextlib.ExitStack() as stack:
+        if config is not None:
+            stack.enter_context(
+                torch.random.fork_rng(devices=range(torch.cuda.device_count()))
+            )
+            torch.manual_seed(config.seed)
+
+        yield
+
+
+class RNGConfig(C.Config):
+    seed: int
+    """Random seed to use for initialization."""

{nshtrainer-1.0.0b54 → nshtrainer-1.0.0b56}/src/nshtrainer/trainer/plugin/__init__.py

@@ -1,10 +1,18 @@
 from __future__ import annotations
 
+from typing import Annotated
+
+from typing_extensions import TypeAliasType
+
 from . import environment as environment
 from . import io as io
 from . import layer_sync as layer_sync
 from . import precision as precision
 from .base import Plugin as Plugin
-from .base import PluginConfig as PluginConfig
 from .base import PluginConfigBase as PluginConfigBase
 from .base import plugin_registry as plugin_registry
+
+PluginConfig = TypeAliasType(
+    "PluginConfig",
+    Annotated[PluginConfigBase, plugin_registry.DynamicResolution()],
+)

{nshtrainer-1.0.0b54 → nshtrainer-1.0.0b56}/src/nshtrainer/trainer/plugin/base.py

@@ -1,8 +1,7 @@
 from __future__ import annotations
 
-import logging
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Annotated
+from typing import TYPE_CHECKING
 
 import nshconfig as C
 from lightning.fabric.plugins import CheckpointIO, ClusterEnvironment
@@ -12,7 +11,6 @@ from typing_extensions import TypeAliasType
 
 if TYPE_CHECKING:
     from .._config import TrainerConfig
-log = logging.getLogger(__name__)
 
 
 Plugin = TypeAliasType(
@@ -26,8 +24,3 @@ class PluginConfigBase(C.Config, ABC):
 
 
 plugin_registry = C.Registry(PluginConfigBase, discriminator="name")
-
-PluginConfig = TypeAliasType(
-    "PluginConfig",
-    Annotated[PluginConfigBase, plugin_registry.DynamicResolution()],
-)

{nshtrainer-1.0.0b54 → nshtrainer-1.0.0b56}/src/nshtrainer/trainer/plugin/environment.py

@@ -1,27 +1,25 @@
 from __future__ import annotations
 
 import signal
-from typing import Any, Literal
+from typing import Literal
 
-from lightning.pytorch.plugins.environments import ClusterEnvironment
-from typing_extensions import override
+from typing_extensions import TypeAliasType, override
 
-from ...util.config.dtype import DTypeConfig
 from .base import PluginConfigBase, plugin_registry
 
 
 @plugin_registry.register
 class KubeflowEnvironmentPlugin(PluginConfigBase):
-    name: Literal["kubeflow_environment"] = "kubeflow_environment"
-
     """Environment for distributed training using the PyTorchJob operator from Kubeflow.
 
     This environment, unlike others, does not get auto-detected and needs to be passed
     to the Fabric/Trainer constructor manually.
     """
 
+    name: Literal["kubeflow_environment"] = "kubeflow_environment"
+
     @override
-    def create_plugin(self, trainer_config) -> ClusterEnvironment:
+    def create_plugin(self, trainer_config):
         from lightning.fabric.plugins.environments.kubeflow import KubeflowEnvironment
 
         return KubeflowEnvironment()
@@ -29,8 +27,6 @@ class KubeflowEnvironmentPlugin(PluginConfigBase):
 
 @plugin_registry.register
 class LightningEnvironmentPlugin(PluginConfigBase):
-    name: Literal["lightning_environment"] = "lightning_environment"
-
     """The default environment used by Lightning for a single node or free cluster (not managed).
 
     There are two modes the Lightning environment can operate with:
@@ -40,8 +36,10 @@ class LightningEnvironmentPlugin(PluginConfigBase):
        The appropriate environment variables need to be set, and at minimum `LOCAL_RANK`.
     """
 
+    name: Literal["lightning_environment"] = "lightning_environment"
+
     @override
-    def create_plugin(self, trainer_config) -> ClusterEnvironment:
+    def create_plugin(self, trainer_config):
         from lightning.fabric.plugins.environments.lightning import LightningEnvironment
 
         return LightningEnvironment()
@@ -49,16 +47,16 @@ class LightningEnvironmentPlugin(PluginConfigBase):
 
 @plugin_registry.register
 class LSFEnvironmentPlugin(PluginConfigBase):
-    name: Literal["lsf_environment"] = "lsf_environment"
-
     """An environment for running on clusters managed by the LSF resource manager.
 
     It is expected that any execution using this ClusterEnvironment was executed
     using the Job Step Manager i.e. `jsrun`.
     """
 
+    name: Literal["lsf_environment"] = "lsf_environment"
+
     @override
-    def create_plugin(self, trainer_config) -> ClusterEnvironment:
+    def create_plugin(self, trainer_config):
         from lightning.fabric.plugins.environments.lsf import LSFEnvironment
 
         return LSFEnvironment()
@@ -66,48 +64,108 @@ class LSFEnvironmentPlugin(PluginConfigBase):
 
 @plugin_registry.register
 class MPIEnvironmentPlugin(PluginConfigBase):
-    name: Literal["mpi_environment"] = "mpi_environment"
-
     """An environment for running on clusters with processes created through MPI.
 
     Requires the installation of the `mpi4py` package.
     """
 
+    name: Literal["mpi_environment"] = "mpi_environment"
+
     @override
-    def create_plugin(self, trainer_config) -> ClusterEnvironment:
+    def create_plugin(self, trainer_config):
         from lightning.fabric.plugins.environments.mpi import MPIEnvironment
 
         return MPIEnvironment()
 
 
+SignalAlias = TypeAliasType(
+    "SignalAlias",
+    Literal[
+        "SIGABRT",
+        "SIGFPE",
+        "SIGILL",
+        "SIGINT",
+        "SIGSEGV",
+        "SIGTERM",
+        "SIGBREAK",
+        "CTRL_C_EVENT",
+        "CTRL_BREAK_EVENT",
+        "SIGALRM",
+        "SIGBUS",
+        "SIGCHLD",
+        "SIGCONT",
+        "SIGHUP",
+        "SIGIO",
+        "SIGIOT",
+        "SIGKILL",
+        "SIGPIPE",
+        "SIGPROF",
+        "SIGQUIT",
+        "SIGSTOP",
+        "SIGSYS",
+        "SIGTRAP",
+        "SIGTSTP",
+        "SIGTTIN",
+        "SIGTTOU",
+        "SIGURG",
+        "SIGUSR1",
+        "SIGUSR2",
+        "SIGVTALRM",
+        "SIGWINCH",
+        "SIGXCPU",
+        "SIGXFSZ",
+        "SIGEMT",
+        "SIGINFO",
+        "SIGCLD",
+        "SIGPOLL",
+        "SIGPWR",
+        "SIGRTMAX",
+        "SIGRTMIN",
+        "SIGSTKFLT",
+    ],
+)
+
+
 @plugin_registry.register
 class SLURMEnvironmentPlugin(PluginConfigBase):
+    """An environment for running on clusters managed by the SLURM resource manager."""
+
     name: Literal["slurm_environment"] = "slurm_environment"
 
     auto_requeue: bool = True
     """Whether automatic job resubmission is enabled or not."""
 
-    requeue_signal: signal.Signals | None = None
+    requeue_signal: SignalAlias | None = None
     """The signal that SLURM will send to indicate that the job should be requeued."""
 
     @override
-    def create_plugin(self, trainer_config) -> ClusterEnvironment:
+    def create_plugin(self, trainer_config):
         from lightning.fabric.plugins.environments.slurm import SLURMEnvironment
 
+        requeue_signal = None
+        if self.requeue_signal is not None:
+            try:
+                requeue_signal = signal.Signals[self.requeue_signal]
+            except KeyError:
+                raise ValueError(
+                    f"Invalid signal name: {self.requeue_signal}. "
+                    "Please provide a valid signal name from the signal module."
+                )
+
         return SLURMEnvironment(
             auto_requeue=self.auto_requeue,
-            requeue_signal=self.requeue_signal,
+            requeue_signal=requeue_signal,
         )
 
 
 @plugin_registry.register
 class TorchElasticEnvironmentPlugin(PluginConfigBase):
-    name: Literal["torchelastic_environment"] = "torchelastic_environment"
-
     """Environment for fault-tolerant and elastic training with torchelastic."""
 
+    name: Literal["torchelastic_environment"] = "torchelastic_environment"
+
     @override
-    def create_plugin(self, trainer_config) -> ClusterEnvironment:
+    def create_plugin(self, trainer_config):
         from lightning.fabric.plugins.environments.torchelastic import (
             TorchElasticEnvironment,
         )
@@ -117,12 +175,12 @@ class TorchElasticEnvironmentPlugin(PluginConfigBase):
 
 @plugin_registry.register
 class XLAEnvironmentPlugin(PluginConfigBase):
-    name: Literal["xla_environment"] = "xla_environment"
-
     """Cluster environment for training on a TPU Pod with the PyTorch/XLA library."""
 
+    name: Literal["xla_environment"] = "xla_environment"
+
     @override
-    def create_plugin(self, trainer_config) -> ClusterEnvironment:
+    def create_plugin(self, trainer_config):
         from lightning.fabric.plugins.environments.xla import XLAEnvironment
 
         return XLAEnvironment()

{nshtrainer-1.0.0b54 → nshtrainer-1.0.0b56}/src/nshtrainer/trainer/plugin/io.py

@@ -2,27 +2,52 @@ from __future__ import annotations
 
 from typing import Literal
 
-from lightning.pytorch.plugins.io import CheckpointIO
 from typing_extensions import override
 
-from .base import PluginConfig, PluginConfigBase, plugin_registry
+from .base import PluginConfigBase, plugin_registry
 
 
 @plugin_registry.register
-class AsyncCheckpointIOPlugin(PluginConfigBase):
-    name: Literal["async_checkpoint"] = "async_checkpoint"
+class TorchCheckpointIOPlugin(PluginConfigBase):
+    """CheckpointIO that utilizes torch.save and torch.load to save and load checkpoints respectively."""
+
+    name: Literal["torch_checkpoint"] = "torch_checkpoint"
+
+    @override
+    def create_plugin(self, trainer_config):
+        from lightning.fabric.plugins.io.torch_io import TorchCheckpointIO
+
+        return TorchCheckpointIO()
+
+
+@plugin_registry.register
+class XLACheckpointIOPlugin(PluginConfigBase):
+    """CheckpointIO that utilizes xm.save to save checkpoints for TPU training strategies."""
 
+    name: Literal["xla_checkpoint"] = "xla_checkpoint"
+
+    @override
+    def create_plugin(self, trainer_config):
+        from lightning.fabric.plugins.io.xla import XLACheckpointIO
+
+        return XLACheckpointIO()
+
+
+@plugin_registry.register
+class AsyncCheckpointIOPlugin(PluginConfigBase):
     """Enables saving the checkpoints asynchronously in a thread.
 
     .. warning::  This is an experimental feature.
     """
 
-    checkpoint_io: PluginConfig | None = None
+    name: Literal["async_checkpoint"] = "async_checkpoint"
+
+    checkpoint_io: TorchCheckpointIOPlugin | None = None
     """A checkpoint IO plugin that is used as the basis for async checkpointing."""
 
     @override
-    def create_plugin(self, trainer_config) -> CheckpointIO:
-        from lightning.pytorch.plugins.io.async_plugin import AsyncCheckpointIO
+    def create_plugin(self, trainer_config):
+        from lightning.pytorch.plugins.io import AsyncCheckpointIO, CheckpointIO
 
         base_io = (
             self.checkpoint_io.create_plugin(trainer_config)
@@ -34,29 +59,3 @@ class AsyncCheckpointIOPlugin(PluginConfigBase):
                 f"Expected `checkpoint_io` to be a `CheckpointIO` instance, but got {type(base_io)}."
             )
         return AsyncCheckpointIO(checkpoint_io=base_io)
-
-
-@plugin_registry.register
-class TorchCheckpointIOPlugin(PluginConfigBase):
-    name: Literal["torch_checkpoint"] = "torch_checkpoint"
-
-    """CheckpointIO that utilizes torch.save and torch.load to save and load checkpoints respectively."""
-
-    @override
-    def create_plugin(self, trainer_config) -> CheckpointIO:
-        from lightning.fabric.plugins.io.torch_io import TorchCheckpointIO
-
-        return TorchCheckpointIO()
-
-
-@plugin_registry.register
-class XLACheckpointIOPlugin(PluginConfigBase):
-    name: Literal["xla_checkpoint"] = "xla_checkpoint"
-
-    """CheckpointIO that utilizes xm.save to save checkpoints for TPU training strategies."""
-
-    @override
-    def create_plugin(self, trainer_config) -> CheckpointIO:
-        from lightning.fabric.plugins.io.xla import XLACheckpointIO
-
-        return XLACheckpointIO()

{nshtrainer-1.0.0b54 → nshtrainer-1.0.0b56}/src/nshtrainer/trainer/plugin/layer_sync.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 from typing import Literal
 
-from lightning.pytorch.plugins.layer_sync import LayerSync
 from typing_extensions import override
 
 from .base import PluginConfigBase, plugin_registry
@@ -10,16 +9,16 @@ from .base import PluginConfigBase, plugin_registry
 
 @plugin_registry.register
 class TorchSyncBatchNormPlugin(PluginConfigBase):
-    name: Literal["torch_sync_batchnorm"] = "torch_sync_batchnorm"
-
     """A plugin that wraps all batch normalization layers of a model with synchronization
     logic for multiprocessing.
 
     This plugin has no effect in single-device operation.
     """
 
+    name: Literal["torch_sync_batchnorm"] = "torch_sync_batchnorm"
+
     @override
-    def create_plugin(self, trainer_config) -> LayerSync:
+    def create_plugin(self, trainer_config):
         from lightning.pytorch.plugins.layer_sync import TorchSyncBatchNorm
 
         return TorchSyncBatchNorm()

{nshtrainer-1.0.0b54 → nshtrainer-1.0.0b56}/src/nshtrainer/trainer/plugin/precision.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 from typing import Any, Literal
 
-from lightning.pytorch.plugins.precision import Precision
 from typing_extensions import override
 
 from ...util.config.dtype import DTypeConfig
@@ -20,7 +19,7 @@ class MixedPrecisionPluginConfig(PluginConfigBase):
     """The device for ``torch.autocast``."""
 
     @override
-    def create_plugin(self, trainer_config) -> Precision:
+    def create_plugin(self, trainer_config):
         from lightning.pytorch.plugins.precision.amp import MixedPrecision
 
         return MixedPrecision(self.precision, self.device)
@@ -45,7 +44,7 @@ class BitsandbytesPluginConfig(PluginConfigBase):
     """
 
     @override
-    def create_plugin(self, trainer_config) -> Precision:
+    def create_plugin(self, trainer_config):
         from lightning.pytorch.plugins.precision.bitsandbytes import (
             BitsandbytesPrecision,
         )
@@ -66,7 +65,7 @@ class DeepSpeedPluginConfig(PluginConfigBase):
     mixed precision (16-mixed, bf16-mixed)."""
 
     @override
-    def create_plugin(self, trainer_config) -> Precision:
+    def create_plugin(self, trainer_config):
         from lightning.pytorch.plugins.precision.deepspeed import DeepSpeedPrecision
 
         return DeepSpeedPrecision(precision=self.precision)
@@ -80,7 +79,7 @@ class DoublePrecisionPluginConfig(PluginConfigBase):
     """Plugin for training with double (``torch.float64``) precision."""
 
     @override
-    def create_plugin(self, trainer_config) -> Precision:
+    def create_plugin(self, trainer_config):
         from lightning.pytorch.plugins.precision.double import DoublePrecision
 
         return DoublePrecision()
@@ -95,7 +94,7 @@ class FSDPPrecisionPluginConfig(PluginConfigBase):
     mixed precision (16-mixed, bf16-mixed)."""
 
     @override
-    def create_plugin(self, trainer_config) -> Precision:
+    def create_plugin(self, trainer_config):
         from lightning.pytorch.plugins.precision.fsdp import FSDPPrecision
 
         return FSDPPrecision(precision=self.precision)
@@ -109,7 +108,7 @@ class HalfPrecisionPluginConfig(PluginConfigBase):
     """Whether to use ``torch.float16`` (``'16-true'``) or ``torch.bfloat16`` (``'bf16-true'``)."""
 
     @override
-    def create_plugin(self, trainer_config) -> Precision:
+    def create_plugin(self, trainer_config):
         from lightning.pytorch.plugins.precision.half import HalfPrecision
 
         return HalfPrecision(precision=self.precision)
@@ -134,7 +133,7 @@ class TransformerEnginePluginConfig(PluginConfigBase):
     Defaults to the same as weights_dtype."""
 
     @override
-    def create_plugin(self, trainer_config) -> Precision:
+    def create_plugin(self, trainer_config):
         from lightning.pytorch.plugins.precision.transformer_engine import (
             TransformerEnginePrecision,
         )
@@ -157,7 +156,7 @@ class XLAPluginConfig(PluginConfigBase):
     """Full precision (32-true) or half precision (16-true, bf16-true)."""
 
     @override
-    def create_plugin(self, trainer_config) -> Precision:
+    def create_plugin(self, trainer_config):
         from lightning.pytorch.plugins.precision.xla import XLAPrecision
 
         return XLAPrecision(precision=self.precision)