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

nshtrainer/__init__.py

@@ -14,6 +14,8 @@ from .metrics import MetricConfig as MetricConfig
 from .model import LightningModuleBase as LightningModuleBase
 from .trainer import Trainer as Trainer
 from .trainer import TrainerConfig as TrainerConfig
+from .trainer import accelerator_registry as accelerator_registry
+from .trainer import plugin_registry as plugin_registry
 
 try:
     from . import configs as configs

nshtrainer/configs/__init__.py

@@ -4,6 +4,8 @@ __codegen__ = True
 
 from nshtrainer import MetricConfig as MetricConfig
 from nshtrainer import TrainerConfig as TrainerConfig
+from nshtrainer import accelerator_registry as accelerator_registry
+from nshtrainer import plugin_registry as plugin_registry
 from nshtrainer._checkpoint.metadata import CheckpointMetadata as CheckpointMetadata
 from nshtrainer._directory import DirectoryConfig as DirectoryConfig
 from nshtrainer._hf_hub import CallbackConfigBase as CallbackConfigBase
@@ -97,7 +99,7 @@ from nshtrainer.profiler import BaseProfilerConfig as BaseProfilerConfig
 from nshtrainer.profiler import ProfilerConfig as ProfilerConfig
 from nshtrainer.profiler import PyTorchProfilerConfig as PyTorchProfilerConfig
 from nshtrainer.profiler import SimpleProfilerConfig as SimpleProfilerConfig
-from nshtrainer.trainer._config import AcceleratorConfigBase as AcceleratorConfigBase
+from nshtrainer.trainer._config import AcceleratorConfig as AcceleratorConfig
 from nshtrainer.trainer._config import (
     CheckpointCallbackConfig as CheckpointCallbackConfig,
 )
@@ -107,9 +109,71 @@ from nshtrainer.trainer._config import GradientClippingConfig as GradientClippin
 from nshtrainer.trainer._config import (
     LearningRateMonitorConfig as LearningRateMonitorConfig,
 )
-from nshtrainer.trainer._config import PluginConfigBase as PluginConfigBase
 from nshtrainer.trainer._config import SanityCheckingConfig as SanityCheckingConfig
-from nshtrainer.trainer._config import StrategyConfigBase as StrategyConfigBase
+from nshtrainer.trainer._config import StrategyConfig as StrategyConfig
+from nshtrainer.trainer.accelerator import CPUAcceleratorConfig as CPUAcceleratorConfig
+from nshtrainer.trainer.accelerator import (
+    CUDAAcceleratorConfig as CUDAAcceleratorConfig,
+)
+from nshtrainer.trainer.accelerator import MPSAcceleratorConfig as MPSAcceleratorConfig
+from nshtrainer.trainer.accelerator import XLAAcceleratorConfig as XLAAcceleratorConfig
+from nshtrainer.trainer.plugin import PluginConfig as PluginConfig
+from nshtrainer.trainer.plugin import PluginConfigBase as PluginConfigBase
+from nshtrainer.trainer.plugin.environment import (
+    KubeflowEnvironmentPlugin as KubeflowEnvironmentPlugin,
+)
+from nshtrainer.trainer.plugin.environment import (
+    LightningEnvironmentPlugin as LightningEnvironmentPlugin,
+)
+from nshtrainer.trainer.plugin.environment import (
+    LSFEnvironmentPlugin as LSFEnvironmentPlugin,
+)
+from nshtrainer.trainer.plugin.environment import (
+    MPIEnvironmentPlugin as MPIEnvironmentPlugin,
+)
+from nshtrainer.trainer.plugin.environment import (
+    SLURMEnvironmentPlugin as SLURMEnvironmentPlugin,
+)
+from nshtrainer.trainer.plugin.environment import (
+    TorchElasticEnvironmentPlugin as TorchElasticEnvironmentPlugin,
+)
+from nshtrainer.trainer.plugin.environment import (
+    XLAEnvironmentPlugin as XLAEnvironmentPlugin,
+)
+from nshtrainer.trainer.plugin.io import (
+    AsyncCheckpointIOPlugin as AsyncCheckpointIOPlugin,
+)
+from nshtrainer.trainer.plugin.io import (
+    TorchCheckpointIOPlugin as TorchCheckpointIOPlugin,
+)
+from nshtrainer.trainer.plugin.io import XLACheckpointIOPlugin as XLACheckpointIOPlugin
+from nshtrainer.trainer.plugin.layer_sync import (
+    TorchSyncBatchNormPlugin as TorchSyncBatchNormPlugin,
+)
+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 (
+    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 (
+    TransformerEnginePluginConfig as TransformerEnginePluginConfig,
+)
+from nshtrainer.trainer.plugin.precision import XLAPluginConfig as XLAPluginConfig
+from nshtrainer.trainer.trainer import AcceleratorConfigBase as AcceleratorConfigBase
+from nshtrainer.trainer.trainer import StrategyConfigBase as StrategyConfigBase
 from nshtrainer.util._environment_info import (
     EnvironmentClassInformationConfig as EnvironmentClassInformationConfig,
 )
@@ -157,17 +221,22 @@ from . import trainer as trainer
 from . import util as util
 
 __all__ = [
+    "AcceleratorConfig",
     "AcceleratorConfigBase",
     "ActSaveConfig",
     "ActSaveLoggerConfig",
     "AdamWConfig",
     "AdvancedProfilerConfig",
+    "AsyncCheckpointIOPlugin",
     "BaseCheckpointCallbackConfig",
     "BaseLoggerConfig",
     "BaseNonlinearityConfig",
     "BaseProfilerConfig",
     "BestCheckpointCallbackConfig",
+    "BitsandbytesPluginConfig",
+    "CPUAcceleratorConfig",
     "CSVLoggerConfig",
+    "CUDAAcceleratorConfig",
     "CallbackConfig",
     "CallbackConfigBase",
     "CheckpointCallbackConfig",
@@ -175,8 +244,10 @@ __all__ = [
     "CheckpointSavingConfig",
     "DTypeConfig",
     "DebugFlagCallbackConfig",
+    "DeepSpeedPluginConfig",
     "DirectoryConfig",
     "DirectorySetupCallbackConfig",
+    "DoublePrecisionPluginConfig",
     "DurationConfig",
     "ELUNonlinearityConfig",
     "EMACallbackConfig",
@@ -193,30 +264,39 @@ __all__ = [
     "EnvironmentSnapshotConfig",
     "EpochTimerCallbackConfig",
     "EpochsConfig",
+    "FSDPPrecisionPluginConfig",
     "FiniteChecksCallbackConfig",
     "GELUNonlinearityConfig",
     "GitRepositoryConfig",
     "GradientClippingConfig",
     "GradientSkippingCallbackConfig",
+    "HalfPrecisionPluginConfig",
     "HuggingFaceHubAutoCreateConfig",
     "HuggingFaceHubConfig",
+    "KubeflowEnvironmentPlugin",
     "LRSchedulerConfig",
     "LRSchedulerConfigBase",
+    "LSFEnvironmentPlugin",
     "LastCheckpointCallbackConfig",
     "LeakyReLUNonlinearityConfig",
     "LearningRateMonitorConfig",
+    "LightningEnvironmentPlugin",
     "LinearWarmupCosineDecayLRSchedulerConfig",
     "LogEpochCallbackConfig",
     "LoggerConfig",
     "MLPConfig",
+    "MPIEnvironmentPlugin",
+    "MPSAcceleratorConfig",
     "MetricConfig",
     "MishNonlinearityConfig",
+    "MixedPrecisionPluginConfig",
     "NonlinearityConfig",
     "NormLoggingCallbackConfig",
     "OnExceptionCheckpointCallbackConfig",
     "OptimizerConfig",
     "OptimizerConfigBase",
     "PReLUConfig",
+    "PluginConfig",
     "PluginConfigBase",
     "PrintTableMetricsCallbackConfig",
     "ProfilerConfig",
@@ -224,6 +304,7 @@ __all__ = [
     "RLPSanityChecksCallbackConfig",
     "ReLUNonlinearityConfig",
     "ReduceLROnPlateauConfig",
+    "SLURMEnvironmentPlugin",
     "SanityCheckingConfig",
     "SharedParametersCallbackConfig",
     "SiLUNonlinearityConfig",
@@ -233,25 +314,36 @@ __all__ = [
     "SoftplusNonlinearityConfig",
     "SoftsignNonlinearityConfig",
     "StepsConfig",
+    "StrategyConfig",
     "StrategyConfigBase",
     "SwiGLUNonlinearityConfig",
     "SwishNonlinearityConfig",
     "TanhNonlinearityConfig",
     "TensorboardLoggerConfig",
     "TimeCheckpointCallbackConfig",
+    "TorchCheckpointIOPlugin",
+    "TorchElasticEnvironmentPlugin",
+    "TorchSyncBatchNormPlugin",
     "TrainerConfig",
+    "TransformerEnginePluginConfig",
     "WandbLoggerConfig",
     "WandbUploadCodeCallbackConfig",
     "WandbWatchCallbackConfig",
+    "XLAAcceleratorConfig",
+    "XLACheckpointIOPlugin",
+    "XLAEnvironmentPlugin",
+    "XLAPluginConfig",
     "_checkpoint",
     "_directory",
     "_hf_hub",
+    "accelerator_registry",
     "callbacks",
     "loggers",
     "lr_scheduler",
     "metrics",
     "nn",
     "optimizer",
+    "plugin_registry",
     "profiler",
     "trainer",
     "util",

nshtrainer/configs/trainer/__init__.py

@@ -3,7 +3,9 @@ from __future__ import annotations
 __codegen__ = True
 
 from nshtrainer.trainer import TrainerConfig as TrainerConfig
-from nshtrainer.trainer._config import AcceleratorConfigBase as AcceleratorConfigBase
+from nshtrainer.trainer import accelerator_registry as accelerator_registry
+from nshtrainer.trainer import plugin_registry as plugin_registry
+from nshtrainer.trainer._config import AcceleratorConfig as AcceleratorConfig
 from nshtrainer.trainer._config import ActSaveLoggerConfig as ActSaveLoggerConfig
 from nshtrainer.trainer._config import BaseLoggerConfig as BaseLoggerConfig
 from nshtrainer.trainer._config import (
@@ -41,7 +43,6 @@ from nshtrainer.trainer._config import (
 from nshtrainer.trainer._config import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
 )
-from nshtrainer.trainer._config import PluginConfigBase as PluginConfigBase
 from nshtrainer.trainer._config import ProfilerConfig as ProfilerConfig
 from nshtrainer.trainer._config import (
     RLPSanityChecksCallbackConfig as RLPSanityChecksCallbackConfig,
@@ -50,7 +51,7 @@ from nshtrainer.trainer._config import SanityCheckingConfig as SanityCheckingCon
 from nshtrainer.trainer._config import (
     SharedParametersCallbackConfig as SharedParametersCallbackConfig,
 )
-from nshtrainer.trainer._config import StrategyConfigBase as StrategyConfigBase
+from nshtrainer.trainer._config import StrategyConfig as StrategyConfig
 from nshtrainer.trainer._config import (
     TensorboardLoggerConfig as TensorboardLoggerConfig,
 )
@@ -58,43 +59,142 @@ from nshtrainer.trainer._config import (
     TimeCheckpointCallbackConfig as TimeCheckpointCallbackConfig,
 )
 from nshtrainer.trainer._config import WandbLoggerConfig as WandbLoggerConfig
+from nshtrainer.trainer.accelerator import CPUAcceleratorConfig as CPUAcceleratorConfig
+from nshtrainer.trainer.accelerator import (
+    CUDAAcceleratorConfig as CUDAAcceleratorConfig,
+)
+from nshtrainer.trainer.accelerator import MPSAcceleratorConfig as MPSAcceleratorConfig
+from nshtrainer.trainer.accelerator import XLAAcceleratorConfig as XLAAcceleratorConfig
+from nshtrainer.trainer.plugin import PluginConfig as PluginConfig
+from nshtrainer.trainer.plugin import PluginConfigBase as PluginConfigBase
+from nshtrainer.trainer.plugin.environment import (
+    KubeflowEnvironmentPlugin as KubeflowEnvironmentPlugin,
+)
+from nshtrainer.trainer.plugin.environment import (
+    LightningEnvironmentPlugin as LightningEnvironmentPlugin,
+)
+from nshtrainer.trainer.plugin.environment import (
+    LSFEnvironmentPlugin as LSFEnvironmentPlugin,
+)
+from nshtrainer.trainer.plugin.environment import (
+    MPIEnvironmentPlugin as MPIEnvironmentPlugin,
+)
+from nshtrainer.trainer.plugin.environment import (
+    SLURMEnvironmentPlugin as SLURMEnvironmentPlugin,
+)
+from nshtrainer.trainer.plugin.environment import (
+    TorchElasticEnvironmentPlugin as TorchElasticEnvironmentPlugin,
+)
+from nshtrainer.trainer.plugin.environment import (
+    XLAEnvironmentPlugin as XLAEnvironmentPlugin,
+)
+from nshtrainer.trainer.plugin.io import (
+    AsyncCheckpointIOPlugin as AsyncCheckpointIOPlugin,
+)
+from nshtrainer.trainer.plugin.io import (
+    TorchCheckpointIOPlugin as TorchCheckpointIOPlugin,
+)
+from nshtrainer.trainer.plugin.io import XLACheckpointIOPlugin as XLACheckpointIOPlugin
+from nshtrainer.trainer.plugin.layer_sync import (
+    TorchSyncBatchNormPlugin as TorchSyncBatchNormPlugin,
+)
+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 (
+    TransformerEnginePluginConfig as TransformerEnginePluginConfig,
+)
+from nshtrainer.trainer.plugin.precision import XLAPluginConfig as XLAPluginConfig
+from nshtrainer.trainer.trainer import AcceleratorConfigBase as AcceleratorConfigBase
+from nshtrainer.trainer.trainer import StrategyConfigBase as StrategyConfigBase
 
 from . import _config as _config
+from . import accelerator as accelerator
+from . import plugin as plugin
+from . import strategy as strategy
 from . import trainer as trainer
 
 __all__ = [
+    "AcceleratorConfig",
     "AcceleratorConfigBase",
     "ActSaveLoggerConfig",
+    "AsyncCheckpointIOPlugin",
     "BaseLoggerConfig",
     "BestCheckpointCallbackConfig",
+    "BitsandbytesPluginConfig",
+    "CPUAcceleratorConfig",
     "CSVLoggerConfig",
+    "CUDAAcceleratorConfig",
     "CallbackConfig",
     "CallbackConfigBase",
     "CheckpointCallbackConfig",
     "CheckpointSavingConfig",
+    "DTypeConfig",
     "DebugFlagCallbackConfig",
+    "DeepSpeedPluginConfig",
     "DirectoryConfig",
+    "DoublePrecisionPluginConfig",
     "EarlyStoppingCallbackConfig",
     "EnvironmentConfig",
+    "FSDPPrecisionPluginConfig",
     "GradientClippingConfig",
+    "HalfPrecisionPluginConfig",
     "HuggingFaceHubConfig",
+    "KubeflowEnvironmentPlugin",
+    "LSFEnvironmentPlugin",
     "LastCheckpointCallbackConfig",
     "LearningRateMonitorConfig",
+    "LightningEnvironmentPlugin",
     "LogEpochCallbackConfig",
     "LoggerConfig",
+    "MPIEnvironmentPlugin",
+    "MPSAcceleratorConfig",
     "MetricConfig",
+    "MixedPrecisionPluginConfig",
     "NormLoggingCallbackConfig",
     "OnExceptionCheckpointCallbackConfig",
+    "PluginConfig",
     "PluginConfigBase",
     "ProfilerConfig",
     "RLPSanityChecksCallbackConfig",
+    "SLURMEnvironmentPlugin",
     "SanityCheckingConfig",
     "SharedParametersCallbackConfig",
+    "StrategyConfig",
     "StrategyConfigBase",
     "TensorboardLoggerConfig",
     "TimeCheckpointCallbackConfig",
+    "TorchCheckpointIOPlugin",
+    "TorchElasticEnvironmentPlugin",
+    "TorchSyncBatchNormPlugin",
     "TrainerConfig",
+    "TransformerEnginePluginConfig",
     "WandbLoggerConfig",
+    "XLAAcceleratorConfig",
+    "XLACheckpointIOPlugin",
+    "XLAEnvironmentPlugin",
+    "XLAPluginConfig",
     "_config",
+    "accelerator",
+    "accelerator_registry",
+    "plugin",
+    "plugin_registry",
+    "strategy",
     "trainer",
 ]

nshtrainer/configs/trainer/_config/__init__.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 __codegen__ = True
 
-from nshtrainer.trainer._config import AcceleratorConfigBase as AcceleratorConfigBase
+from nshtrainer.trainer._config import AcceleratorConfig as AcceleratorConfig
 from nshtrainer.trainer._config import ActSaveLoggerConfig as ActSaveLoggerConfig
 from nshtrainer.trainer._config import BaseLoggerConfig as BaseLoggerConfig
 from nshtrainer.trainer._config import (
@@ -40,7 +40,7 @@ from nshtrainer.trainer._config import (
 from nshtrainer.trainer._config import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
 )
-from nshtrainer.trainer._config import PluginConfigBase as PluginConfigBase
+from nshtrainer.trainer._config import PluginConfig as PluginConfig
 from nshtrainer.trainer._config import ProfilerConfig as ProfilerConfig
 from nshtrainer.trainer._config import (
     RLPSanityChecksCallbackConfig as RLPSanityChecksCallbackConfig,
@@ -49,7 +49,7 @@ from nshtrainer.trainer._config import SanityCheckingConfig as SanityCheckingCon
 from nshtrainer.trainer._config import (
     SharedParametersCallbackConfig as SharedParametersCallbackConfig,
 )
-from nshtrainer.trainer._config import StrategyConfigBase as StrategyConfigBase
+from nshtrainer.trainer._config import StrategyConfig as StrategyConfig
 from nshtrainer.trainer._config import (
     TensorboardLoggerConfig as TensorboardLoggerConfig,
 )
@@ -58,9 +58,11 @@ from nshtrainer.trainer._config import (
 )
 from nshtrainer.trainer._config import TrainerConfig as TrainerConfig
 from nshtrainer.trainer._config import WandbLoggerConfig as WandbLoggerConfig
+from nshtrainer.trainer._config import accelerator_registry as accelerator_registry
+from nshtrainer.trainer._config import plugin_registry as plugin_registry
 
 __all__ = [
-    "AcceleratorConfigBase",
+    "AcceleratorConfig",
     "ActSaveLoggerConfig",
     "BaseLoggerConfig",
     "BestCheckpointCallbackConfig",
@@ -82,14 +84,16 @@ __all__ = [
     "MetricConfig",
     "NormLoggingCallbackConfig",
     "OnExceptionCheckpointCallbackConfig",
-    "PluginConfigBase",
+    "PluginConfig",
     "ProfilerConfig",
     "RLPSanityChecksCallbackConfig",
     "SanityCheckingConfig",
     "SharedParametersCallbackConfig",
-    "StrategyConfigBase",
+    "StrategyConfig",
     "TensorboardLoggerConfig",
     "TimeCheckpointCallbackConfig",
     "TrainerConfig",
     "WandbLoggerConfig",
+    "accelerator_registry",
+    "plugin_registry",
 ]

nshtrainer/configs/trainer/accelerator/__init__.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+__codegen__ = True
+
+from nshtrainer.trainer.accelerator import AcceleratorConfig as AcceleratorConfig
+from nshtrainer.trainer.accelerator import (
+    AcceleratorConfigBase as AcceleratorConfigBase,
+)
+from nshtrainer.trainer.accelerator import CPUAcceleratorConfig as CPUAcceleratorConfig
+from nshtrainer.trainer.accelerator import (
+    CUDAAcceleratorConfig as CUDAAcceleratorConfig,
+)
+from nshtrainer.trainer.accelerator import MPSAcceleratorConfig as MPSAcceleratorConfig
+from nshtrainer.trainer.accelerator import XLAAcceleratorConfig as XLAAcceleratorConfig
+from nshtrainer.trainer.accelerator import accelerator_registry as accelerator_registry
+
+__all__ = [
+    "AcceleratorConfig",
+    "AcceleratorConfigBase",
+    "CPUAcceleratorConfig",
+    "CUDAAcceleratorConfig",
+    "MPSAcceleratorConfig",
+    "XLAAcceleratorConfig",
+    "accelerator_registry",
+]

nshtrainer/configs/trainer/plugin/__init__.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+__codegen__ = True
+
+from nshtrainer.trainer.plugin import PluginConfig as PluginConfig
+from nshtrainer.trainer.plugin import PluginConfigBase as PluginConfigBase
+from nshtrainer.trainer.plugin import plugin_registry as plugin_registry
+from nshtrainer.trainer.plugin.environment import (
+    KubeflowEnvironmentPlugin as KubeflowEnvironmentPlugin,
+)
+from nshtrainer.trainer.plugin.environment import (
+    LightningEnvironmentPlugin as LightningEnvironmentPlugin,
+)
+from nshtrainer.trainer.plugin.environment import (
+    LSFEnvironmentPlugin as LSFEnvironmentPlugin,
+)
+from nshtrainer.trainer.plugin.environment import (
+    MPIEnvironmentPlugin as MPIEnvironmentPlugin,
+)
+from nshtrainer.trainer.plugin.environment import (
+    SLURMEnvironmentPlugin as SLURMEnvironmentPlugin,
+)
+from nshtrainer.trainer.plugin.environment import (
+    TorchElasticEnvironmentPlugin as TorchElasticEnvironmentPlugin,
+)
+from nshtrainer.trainer.plugin.environment import (
+    XLAEnvironmentPlugin as XLAEnvironmentPlugin,
+)
+from nshtrainer.trainer.plugin.io import (
+    AsyncCheckpointIOPlugin as AsyncCheckpointIOPlugin,
+)
+from nshtrainer.trainer.plugin.io import (
+    TorchCheckpointIOPlugin as TorchCheckpointIOPlugin,
+)
+from nshtrainer.trainer.plugin.io import XLACheckpointIOPlugin as XLACheckpointIOPlugin
+from nshtrainer.trainer.plugin.layer_sync import (
+    TorchSyncBatchNormPlugin as TorchSyncBatchNormPlugin,
+)
+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 (
+    TransformerEnginePluginConfig as TransformerEnginePluginConfig,
+)
+from nshtrainer.trainer.plugin.precision import XLAPluginConfig as XLAPluginConfig
+
+from . import base as base
+from . import environment as environment
+from . import io as io
+from . import layer_sync as layer_sync
+from . import precision as precision
+
+__all__ = [
+    "AsyncCheckpointIOPlugin",
+    "BitsandbytesPluginConfig",
+    "DTypeConfig",
+    "DeepSpeedPluginConfig",
+    "DoublePrecisionPluginConfig",
+    "FSDPPrecisionPluginConfig",
+    "HalfPrecisionPluginConfig",
+    "KubeflowEnvironmentPlugin",
+    "LSFEnvironmentPlugin",
+    "LightningEnvironmentPlugin",
+    "MPIEnvironmentPlugin",
+    "MixedPrecisionPluginConfig",
+    "PluginConfig",
+    "PluginConfigBase",
+    "SLURMEnvironmentPlugin",
+    "TorchCheckpointIOPlugin",
+    "TorchElasticEnvironmentPlugin",
+    "TorchSyncBatchNormPlugin",
+    "TransformerEnginePluginConfig",
+    "XLACheckpointIOPlugin",
+    "XLAEnvironmentPlugin",
+    "XLAPluginConfig",
+    "base",
+    "environment",
+    "io",
+    "layer_sync",
+    "plugin_registry",
+    "precision",
+]

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

@@ -0,0 +1,13 @@
+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/configs/trainer/plugin/environment/__init__.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+__codegen__ = True
+
+from nshtrainer.trainer.plugin.environment import DTypeConfig as DTypeConfig
+from nshtrainer.trainer.plugin.environment import (
+    KubeflowEnvironmentPlugin as KubeflowEnvironmentPlugin,
+)
+from nshtrainer.trainer.plugin.environment import (
+    LightningEnvironmentPlugin as LightningEnvironmentPlugin,
+)
+from nshtrainer.trainer.plugin.environment import (
+    LSFEnvironmentPlugin as LSFEnvironmentPlugin,
+)
+from nshtrainer.trainer.plugin.environment import (
+    MPIEnvironmentPlugin as MPIEnvironmentPlugin,
+)
+from nshtrainer.trainer.plugin.environment import PluginConfigBase as PluginConfigBase
+from nshtrainer.trainer.plugin.environment import (
+    SLURMEnvironmentPlugin as SLURMEnvironmentPlugin,
+)
+from nshtrainer.trainer.plugin.environment import (
+    TorchElasticEnvironmentPlugin as TorchElasticEnvironmentPlugin,
+)
+from nshtrainer.trainer.plugin.environment import (
+    XLAEnvironmentPlugin as XLAEnvironmentPlugin,
+)
+from nshtrainer.trainer.plugin.environment import plugin_registry as plugin_registry
+
+__all__ = [
+    "DTypeConfig",
+    "KubeflowEnvironmentPlugin",
+    "LSFEnvironmentPlugin",
+    "LightningEnvironmentPlugin",
+    "MPIEnvironmentPlugin",
+    "PluginConfigBase",
+    "SLURMEnvironmentPlugin",
+    "TorchElasticEnvironmentPlugin",
+    "XLAEnvironmentPlugin",
+    "plugin_registry",
+]

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

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+__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,
+)
+from nshtrainer.trainer.plugin.io import XLACheckpointIOPlugin as XLACheckpointIOPlugin
+from nshtrainer.trainer.plugin.io import plugin_registry as plugin_registry
+
+__all__ = [
+    "AsyncCheckpointIOPlugin",
+    "PluginConfig",
+    "PluginConfigBase",
+    "TorchCheckpointIOPlugin",
+    "XLACheckpointIOPlugin",
+    "plugin_registry",
+]

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

@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+__codegen__ = True
+
+from nshtrainer.trainer.plugin.layer_sync import PluginConfigBase as PluginConfigBase
+from nshtrainer.trainer.plugin.layer_sync import (
+    TorchSyncBatchNormPlugin as TorchSyncBatchNormPlugin,
+)
+from nshtrainer.trainer.plugin.layer_sync import plugin_registry as plugin_registry
+
+__all__ = [
+    "PluginConfigBase",
+    "TorchSyncBatchNormPlugin",
+    "plugin_registry",
+]