nshtrainer 1.0.0b29__py3-none-any.whl → 1.0.0b30__py3-none-any.whl

nshtrainer/configs/trainer/plugin/precision/__init__.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+__codegen__ = True
+
+from nshtrainer.trainer.plugin.precision import (
+    BitsandbytesPluginConfig as BitsandbytesPluginConfig,
+)
+from nshtrainer.trainer.plugin.precision import (
+    DeepSpeedPluginConfig as DeepSpeedPluginConfig,
+)
+from nshtrainer.trainer.plugin.precision import (
+    DoublePrecisionPluginConfig as DoublePrecisionPluginConfig,
+)
+from nshtrainer.trainer.plugin.precision import DTypeConfig as DTypeConfig
+from nshtrainer.trainer.plugin.precision import (
+    FSDPPrecisionPluginConfig as FSDPPrecisionPluginConfig,
+)
+from nshtrainer.trainer.plugin.precision import (
+    HalfPrecisionPluginConfig as HalfPrecisionPluginConfig,
+)
+from nshtrainer.trainer.plugin.precision import (
+    MixedPrecisionPluginConfig as MixedPrecisionPluginConfig,
+)
+from nshtrainer.trainer.plugin.precision import PluginConfigBase as PluginConfigBase
+from nshtrainer.trainer.plugin.precision import (
+    TransformerEnginePluginConfig as TransformerEnginePluginConfig,
+)
+from nshtrainer.trainer.plugin.precision import XLAPluginConfig as XLAPluginConfig
+from nshtrainer.trainer.plugin.precision import plugin_registry as plugin_registry
+
+__all__ = [
+    "BitsandbytesPluginConfig",
+    "DTypeConfig",
+    "DeepSpeedPluginConfig",
+    "DoublePrecisionPluginConfig",
+    "FSDPPrecisionPluginConfig",
+    "HalfPrecisionPluginConfig",
+    "MixedPrecisionPluginConfig",
+    "PluginConfigBase",
+    "TransformerEnginePluginConfig",
+    "XLAPluginConfig",
+    "plugin_registry",
+]

nshtrainer/configs/trainer/strategy/__init__.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+__codegen__ = True
+
+from nshtrainer.trainer.strategy import StrategyConfig as StrategyConfig
+from nshtrainer.trainer.strategy import StrategyConfigBase as StrategyConfigBase
+
+__all__ = [
+    "StrategyConfig",
+    "StrategyConfigBase",
+]

nshtrainer/configs/trainer/trainer/__init__.py

@@ -4,12 +4,14 @@ __codegen__ = True
 
 from nshtrainer.trainer.trainer import AcceleratorConfigBase as AcceleratorConfigBase
 from nshtrainer.trainer.trainer import EnvironmentConfig as EnvironmentConfig
+from nshtrainer.trainer.trainer import PluginConfigBase as PluginConfigBase
 from nshtrainer.trainer.trainer import StrategyConfigBase as StrategyConfigBase
 from nshtrainer.trainer.trainer import TrainerConfig as TrainerConfig
 
 __all__ = [
     "AcceleratorConfigBase",
     "EnvironmentConfig",
+    "PluginConfigBase",
     "StrategyConfigBase",
     "TrainerConfig",
 ]

nshtrainer/trainer/__init__.py

@@ -1,4 +1,6 @@
 from __future__ import annotations
 
 from ._config import TrainerConfig as TrainerConfig
+from ._config import accelerator_registry as accelerator_registry
+from ._config import plugin_registry as plugin_registry
 from .trainer import Trainer as Trainer

nshtrainer/trainer/_config.py

@@ -5,7 +5,6 @@ import logging
 import os
 import string
 import time
-from abc import ABC, abstractmethod
 from collections.abc import Iterable, Sequence
 from datetime import timedelta
 from pathlib import Path
@@ -18,14 +17,11 @@ from typing import (
 
 import nshconfig as C
 import numpy as np
-from lightning.fabric.plugins import CheckpointIO, ClusterEnvironment
 from lightning.fabric.plugins.precision.precision import _PRECISION_INPUT
 from lightning.pytorch.accelerators import Accelerator
 from lightning.pytorch.callbacks.callback import Callback
 from lightning.pytorch.loggers import Logger
 from lightning.pytorch.plugins import _PLUGIN_INPUT
-from lightning.pytorch.plugins.layer_sync import LayerSync
-from lightning.pytorch.plugins.precision.precision import Precision
 from lightning.pytorch.profilers import Profiler
 from lightning.pytorch.strategies.strategy import Strategy
 from typing_extensions import TypeAliasType, TypedDict, override
@@ -58,6 +54,9 @@ from ..loggers.actsave import ActSaveLoggerConfig
 from ..metrics._config import MetricConfig
 from ..profiler import ProfilerConfig
 from ..util._environment_info import EnvironmentConfig
+from .accelerator import AcceleratorConfig, AcceleratorLiteral, accelerator_registry
+from .plugin import PluginConfig, plugin_registry
+from .strategy import StrategyConfig
 
 log = logging.getLogger(__name__)
 
@@ -71,37 +70,6 @@ class GradientClippingConfig(C.Config):
     """Norm type to use for gradient clipping."""
 
 
-Plugin = TypeAliasType(
-    "Plugin", Precision | ClusterEnvironment | CheckpointIO | LayerSync
-)
-
-
-class PluginConfigBase(C.Config, ABC):
-    @abstractmethod
-    def create_plugin(self) -> Plugin: ...
-
-
-plugin_registry = C.Registry(PluginConfigBase, discriminator="name")
-PluginConfig = TypeAliasType(
-    "PluginConfig", Annotated[PluginConfigBase, plugin_registry.DynamicResolution()]
-)
-
-AcceleratorLiteral = TypeAliasType(
-    "AcceleratorLiteral", Literal["cpu", "gpu", "tpu", "ipu", "hpu", "mps", "auto"]
-)
-
-
-class AcceleratorConfigBase(C.Config, ABC):
-    @abstractmethod
-    def create_accelerator(self) -> Accelerator: ...
-
-
-accelerator_registry = C.Registry(AcceleratorConfigBase, discriminator="name")
-AcceleratorConfig = TypeAliasType(
-    "AcceleratorConfig",
-    Annotated[AcceleratorConfigBase, accelerator_registry.DynamicResolution()],
-)
-
 StrategyLiteral = TypeAliasType(
     "StrategyLiteral",
     Literal[
@@ -135,17 +103,6 @@ StrategyLiteral = TypeAliasType(
 )
 
 
-class StrategyConfigBase(C.Config, ABC):
-    @abstractmethod
-    def create_strategy(self) -> Strategy: ...
-
-
-strategy_registry = C.Registry(StrategyConfigBase, discriminator="name")
-StrategyConfig = TypeAliasType(
-    "StrategyConfig",
-    Annotated[StrategyConfigBase, strategy_registry.DynamicResolution()],
-)
-
 CheckpointCallbackConfig = TypeAliasType(
     "CheckpointCallbackConfig",
     Annotated[
@@ -441,7 +398,6 @@ class SanityCheckingConfig(C.Config):
 
 
 @plugin_registry.rebuild_on_registers
-@strategy_registry.rebuild_on_registers
 @accelerator_registry.rebuild_on_registers
 class TrainerConfig(C.Config):
     # region Active Run Configuration

nshtrainer/trainer/accelerator.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Annotated, Literal
+
+import nshconfig as C
+from lightning.pytorch.accelerators import Accelerator
+from typing_extensions import TypeAliasType, override
+
+if TYPE_CHECKING:
+    from ._config import TrainerConfig
+
+AcceleratorLiteral = TypeAliasType(
+    "AcceleratorLiteral", Literal["cpu", "gpu", "tpu", "ipu", "hpu", "mps", "auto"]
+)
+
+
+class AcceleratorConfigBase(C.Config, ABC):
+    @abstractmethod
+    def create_accelerator(self, trainer_config: "TrainerConfig") -> Accelerator: ...
+
+
+accelerator_registry = C.Registry(AcceleratorConfigBase, discriminator="name")
+
+AcceleratorConfig = TypeAliasType(
+    "AcceleratorConfig",
+    Annotated[AcceleratorConfigBase, accelerator_registry.DynamicResolution()],
+)
+
+
+@accelerator_registry.register
+class CPUAcceleratorConfig(AcceleratorConfigBase):
+    name: Literal["cpu"] = "cpu"
+
+    """Accelerator for CPU devices."""
+
+    @override
+    def create_accelerator(self, trainer_config) -> Accelerator:
+        from lightning.pytorch.accelerators.cpu import CPUAccelerator
+
+        return CPUAccelerator()
+
+
+@accelerator_registry.register
+class CUDAAcceleratorConfig(AcceleratorConfigBase):
+    name: Literal["gpu"] = "gpu"
+
+    """Accelerator for NVIDIA CUDA devices."""
+
+    @override
+    def create_accelerator(self, trainer_config) -> Accelerator:
+        from lightning.pytorch.accelerators.cuda import CUDAAccelerator
+
+        return CUDAAccelerator()
+
+
+@accelerator_registry.register
+class MPSAcceleratorConfig(AcceleratorConfigBase):
+    name: Literal["mps"] = "mps"
+
+    """Accelerator for Metal Apple Silicon GPU devices.
+
+    .. warning::  Use of this accelerator beyond import and instantiation is experimental.
+    """
+
+    @override
+    def create_accelerator(self, trainer_config) -> Accelerator:
+        from lightning.pytorch.accelerators.mps import MPSAccelerator
+
+        return MPSAccelerator()
+
+
+@accelerator_registry.register
+class XLAAcceleratorConfig(AcceleratorConfigBase):
+    name: Literal["tpu"] = "tpu"
+
+    """Accelerator for XLA devices, normally TPUs.
+
+    .. warning::  Use of this accelerator beyond import and instantiation is experimental.
+    """
+
+    @override
+    def create_accelerator(self, trainer_config) -> Accelerator:
+        from lightning.pytorch.accelerators.xla import XLAAccelerator
+
+        return XLAAccelerator()

nshtrainer/trainer/plugin/__init__.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+
+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

nshtrainer/trainer/plugin/base.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+import logging
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Annotated
+
+import nshconfig as C
+from lightning.fabric.plugins import CheckpointIO, ClusterEnvironment
+from lightning.pytorch.plugins.layer_sync import LayerSync
+from lightning.pytorch.plugins.precision.precision import Precision
+from typing_extensions import TypeAliasType
+
+if TYPE_CHECKING:
+    from .._config import TrainerConfig
+log = logging.getLogger(__name__)
+
+
+Plugin = TypeAliasType(
+    "Plugin", Precision | ClusterEnvironment | CheckpointIO | LayerSync
+)
+
+
+class PluginConfigBase(C.Config, ABC):
+    @abstractmethod
+    def create_plugin(self, trainer_config: "TrainerConfig") -> Plugin: ...
+
+
+plugin_registry = C.Registry(PluginConfigBase, discriminator="name")
+
+PluginConfig = TypeAliasType(
+    "PluginConfig",
+    Annotated[PluginConfigBase, plugin_registry.DynamicResolution()],
+)

nshtrainer/trainer/plugin/environment.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+import signal
+from typing import Any, Literal
+
+from lightning.pytorch.plugins.environments import ClusterEnvironment
+from typing_extensions import 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.
+    """
+
+    @override
+    def create_plugin(self, trainer_config) -> ClusterEnvironment:
+        from lightning.fabric.plugins.environments.kubeflow import KubeflowEnvironment
+
+        return KubeflowEnvironment()
+
+
+@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:
+    1. User launches main process by `python train.py ...` with no additional environment variables.
+       Lightning will spawn new worker processes for distributed training in the current node.
+    2. User launches all processes manually or with utilities like `torch.distributed.launch`.
+       The appropriate environment variables need to be set, and at minimum `LOCAL_RANK`.
+    """
+
+    @override
+    def create_plugin(self, trainer_config) -> ClusterEnvironment:
+        from lightning.fabric.plugins.environments.lightning import LightningEnvironment
+
+        return LightningEnvironment()
+
+
+@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`.
+    """
+
+    @override
+    def create_plugin(self, trainer_config) -> ClusterEnvironment:
+        from lightning.fabric.plugins.environments.lsf import LSFEnvironment
+
+        return LSFEnvironment()
+
+
+@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.
+    """
+
+    @override
+    def create_plugin(self, trainer_config) -> ClusterEnvironment:
+        from lightning.fabric.plugins.environments.mpi import MPIEnvironment
+
+        return MPIEnvironment()
+
+
+@plugin_registry.register
+class SLURMEnvironmentPlugin(PluginConfigBase):
+    name: Literal["slurm_environment"] = "slurm_environment"
+
+    auto_requeue: bool = True
+    """Whether automatic job resubmission is enabled or not."""
+
+    requeue_signal: signal.Signals | None = None
+    """The signal that SLURM will send to indicate that the job should be requeued."""
+
+    @override
+    def create_plugin(self, trainer_config) -> ClusterEnvironment:
+        from lightning.fabric.plugins.environments.slurm import SLURMEnvironment
+
+        return SLURMEnvironment(
+            auto_requeue=self.auto_requeue,
+            requeue_signal=self.requeue_signal,
+        )
+
+
+@plugin_registry.register
+class TorchElasticEnvironmentPlugin(PluginConfigBase):
+    name: Literal["torchelastic_environment"] = "torchelastic_environment"
+
+    """Environment for fault-tolerant and elastic training with torchelastic."""
+
+    @override
+    def create_plugin(self, trainer_config) -> ClusterEnvironment:
+        from lightning.fabric.plugins.environments.torchelastic import (
+            TorchElasticEnvironment,
+        )
+
+        return TorchElasticEnvironment()
+
+
+@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."""
+
+    @override
+    def create_plugin(self, trainer_config) -> ClusterEnvironment:
+        from lightning.fabric.plugins.environments.xla import XLAEnvironment
+
+        return XLAEnvironment()

nshtrainer/trainer/plugin/io.py

@@ -0,0 +1,62 @@
+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
+
+
+@plugin_registry.register
+class AsyncCheckpointIOPlugin(PluginConfigBase):
+    name: Literal["async_checkpoint"] = "async_checkpoint"
+
+    """Enables saving the checkpoints asynchronously in a thread.
+
+    .. warning::  This is an experimental feature.
+    """
+
+    checkpoint_io: PluginConfig | 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
+
+        base_io = (
+            self.checkpoint_io.create_plugin(trainer_config)
+            if self.checkpoint_io
+            else None
+        )
+        if base_io is not None and not isinstance(base_io, CheckpointIO):
+            raise TypeError(
+                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/trainer/plugin/layer_sync.py

@@ -0,0 +1,25 @@
+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
+
+
+@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.
+    """
+
+    @override
+    def create_plugin(self, trainer_config) -> LayerSync:
+        from lightning.pytorch.plugins.layer_sync import TorchSyncBatchNorm
+
+        return TorchSyncBatchNorm()

nshtrainer/trainer/plugin/precision.py

@@ -0,0 +1,163 @@
+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
+from .base import PluginConfigBase, plugin_registry
+
+
+@plugin_registry.register
+class MixedPrecisionPluginConfig(PluginConfigBase):
+    name: Literal["mixed_precision"] = "mixed_precision"
+
+    precision: Literal["16-mixed", "bf16-mixed"]
+    """Whether to use ``torch.float16`` (``'16-mixed'``) or ``torch.bfloat16`` (``'bf16-mixed'``)."""
+
+    device: str
+    """The device for ``torch.autocast``."""
+
+    @override
+    def create_plugin(self, trainer_config) -> Precision:
+        from lightning.pytorch.plugins.precision.amp import MixedPrecision
+
+        return MixedPrecision(self.precision, self.device)
+
+
+@plugin_registry.register
+class BitsandbytesPluginConfig(PluginConfigBase):
+    name: Literal["bitsandbytes_precision"] = "bitsandbytes_precision"
+
+    mode: Literal["nf4", "nf4-dq", "fp4", "fp4-dq", "int8", "int8-training"]
+    """The quantization mode to use."""
+
+    dtype: DTypeConfig | None = None
+    """The compute dtype to use."""
+
+    ignore_modules: set[str] | None = None
+    """The submodules whose Linear layers should not be replaced.
+
+    This might be desirable for numerical stability. The string will be checked
+    as a prefix, so a value like "transformer.blocks" will ignore all linear
+    layers in all of the transformer blocks.
+    """
+
+    @override
+    def create_plugin(self, trainer_config) -> Precision:
+        from lightning.pytorch.plugins.precision.bitsandbytes import (
+            BitsandbytesPrecision,
+        )
+
+        return BitsandbytesPrecision(
+            mode=self.mode,
+            dtype=self.dtype.torch_dtype if self.dtype is not None else None,
+            ignore_modules=self.ignore_modules,
+        )
+
+
+@plugin_registry.register
+class DeepSpeedPluginConfig(PluginConfigBase):
+    name: Literal["deepspeed_precision"] = "deepspeed_precision"
+
+    precision: Literal["16-true", "bf16-true", "16-mixed", "bf16-mixed", "32-true"]
+    """Full precision (32-true), half precision (16-true, bf16-true) or
+    mixed precision (16-mixed, bf16-mixed)."""
+
+    @override
+    def create_plugin(self, trainer_config) -> Precision:
+        from lightning.pytorch.plugins.precision.deepspeed import DeepSpeedPrecision
+
+        return DeepSpeedPrecision(precision=self.precision)
+
+
+@plugin_registry.register
+class DoublePrecisionPluginConfig(PluginConfigBase):
+    name: Literal["double_precision"] = "double_precision"
+
+    precision: Literal["64-true"] = "64-true"
+    """Plugin for training with double (``torch.float64``) precision."""
+
+    @override
+    def create_plugin(self, trainer_config) -> Precision:
+        from lightning.pytorch.plugins.precision.double import DoublePrecision
+
+        return DoublePrecision()
+
+
+@plugin_registry.register
+class FSDPPrecisionPluginConfig(PluginConfigBase):
+    name: Literal["fsdp_precision"] = "fsdp_precision"
+
+    precision: Literal["16-true", "bf16-true", "16-mixed", "bf16-mixed", "32-true"]
+    """Full precision (32-true), half precision (16-true, bf16-true) or
+    mixed precision (16-mixed, bf16-mixed)."""
+
+    @override
+    def create_plugin(self, trainer_config) -> Precision:
+        from lightning.pytorch.plugins.precision.fsdp import FSDPPrecision
+
+        return FSDPPrecision(precision=self.precision)
+
+
+@plugin_registry.register
+class HalfPrecisionPluginConfig(PluginConfigBase):
+    name: Literal["half_precision"] = "half_precision"
+
+    precision: Literal["bf16-true", "16-true"]
+    """Whether to use ``torch.float16`` (``'16-true'``) or ``torch.bfloat16`` (``'bf16-true'``)."""
+
+    @override
+    def create_plugin(self, trainer_config) -> Precision:
+        from lightning.pytorch.plugins.precision.half import HalfPrecision
+
+        return HalfPrecision(precision=self.precision)
+
+
+@plugin_registry.register
+class TransformerEnginePluginConfig(PluginConfigBase):
+    name: Literal["transformer_engine_precision"] = "transformer_engine_precision"
+
+    weights_dtype: DTypeConfig
+    """The weights dtype to use."""
+
+    recipe: dict[str, Any] | None = None
+    """Recipe for the DelayedScaling configuration in dict format."""
+
+    replace_layers: bool | None = None
+    """Whether to replace ``Linear`` and ``LayerNorm`` layers automatically with their
+    Transformer Engine alternatives."""
+
+    fallback_compute_dtype: DTypeConfig | None = None
+    """The compute dtype to use for operations that don't support fp8 autocast.
+    Defaults to the same as weights_dtype."""
+
+    @override
+    def create_plugin(self, trainer_config) -> Precision:
+        from lightning.pytorch.plugins.precision.transformer_engine import (
+            TransformerEnginePrecision,
+        )
+
+        return TransformerEnginePrecision(
+            weights_dtype=self.weights_dtype.torch_dtype,
+            recipe=self.recipe,
+            replace_layers=self.replace_layers,
+            fallback_compute_dtype=self.fallback_compute_dtype.torch_dtype
+            if self.fallback_compute_dtype
+            else None,
+        )
+
+
+@plugin_registry.register
+class XLAPluginConfig(PluginConfigBase):
+    name: Literal["xla_precision"] = "xla_precision"
+
+    precision: Literal["32-true", "16-true", "bf16-true"]
+    """Full precision (32-true) or half precision (16-true, bf16-true)."""
+
+    @override
+    def create_plugin(self, trainer_config) -> Precision:
+        from lightning.pytorch.plugins.precision.xla import XLAPrecision
+
+        return XLAPrecision(precision=self.precision)