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

nshtrainer/trainer/strategy.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Literal
+
+import nshconfig as C
+from lightning.pytorch.strategies.strategy import Strategy
+from typing_extensions import TypeAliasType
+
+if TYPE_CHECKING:
+    from ._config import TrainerConfig
+
+StrategyLiteral = TypeAliasType(
+    "StrategyLiteral",
+    Literal[
+        "auto",
+        "ddp",
+        "ddp_find_unused_parameters_false",
+        "ddp_find_unused_parameters_true",
+        "ddp_spawn",
+        "ddp_spawn_find_unused_parameters_false",
+        "ddp_spawn_find_unused_parameters_true",
+        "ddp_fork",
+        "ddp_fork_find_unused_parameters_false",
+        "ddp_fork_find_unused_parameters_true",
+        "ddp_notebook",
+        "dp",
+        "deepspeed",
+        "deepspeed_stage_1",
+        "deepspeed_stage_1_offload",
+        "deepspeed_stage_2",
+        "deepspeed_stage_2_offload",
+        "deepspeed_stage_3",
+        "deepspeed_stage_3_offload",
+        "deepspeed_stage_3_offload_nvme",
+        "fsdp",
+        "fsdp_cpu_offload",
+        "single_xla",
+        "xla_fsdp",
+        "xla",
+        "single_tpu",
+    ],
+)
+
+
+class StrategyConfigBase(C.Config, ABC):
+    @abstractmethod
+    def create_strategy(self, trainer_config: "TrainerConfig") -> Strategy: ...
+
+
+StrategyConfig = TypeAliasType("StrategyConfig", StrategyConfigBase)

nshtrainer/trainer/trainer.py

@@ -22,14 +22,12 @@ from .._checkpoint.metadata import _write_checkpoint_metadata
 from ..callbacks.base import resolve_all_callbacks
 from ..util._environment_info import EnvironmentConfig
 from ..util.bf16 import is_bf16_supported_no_emulation
-from ._config import (
-    AcceleratorConfigBase,
-    LightningTrainerKwargs,
-    StrategyConfigBase,
-    TrainerConfig,
-)
+from ._config import LightningTrainerKwargs, TrainerConfig
 from ._runtime_callback import RuntimeTrackerCallback, Stage
+from .accelerator import AcceleratorConfigBase
+from .plugin import PluginConfigBase
 from .signal_connector import _SignalConnector
+from .strategy import StrategyConfigBase
 
 log = logging.getLogger(__name__)
 
@@ -172,12 +170,12 @@ class Trainer(LightningTrainer):
 
         if (accelerator := hparams.accelerator) is not None:
             if isinstance(accelerator, AcceleratorConfigBase):
-                accelerator = accelerator.create_accelerator()
+                accelerator = accelerator.create_accelerator(hparams)
             _update_kwargs(accelerator=accelerator)
 
         if (strategy := hparams.strategy) is not None:
             if isinstance(strategy, StrategyConfigBase):
-                strategy = strategy.create_strategy()
+                strategy = strategy.create_strategy(hparams)
             _update_kwargs(strategy=strategy)
 
         if (precision := hparams.precision) is not None:
@@ -238,7 +236,8 @@ class Trainer(LightningTrainer):
         if plugin_configs := hparams.plugins:
             _update_kwargs(
                 plugins=[
-                    plugin_config.create_plugin() for plugin_config in plugin_configs
+                    plugin_config.create_plugin(hparams)
+                    for plugin_config in plugin_configs
                 ]
             )
 

{nshtrainer-1.0.0b29.dist-info → nshtrainer-1.0.0b30.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 1.0.0b29
+Version: 1.0.0b30
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-1.0.0b29.dist-info → nshtrainer-1.0.0b30.dist-info}/RECORD

@@ -1,5 +1,5 @@
 nshtrainer/.nshconfig.generated.json,sha256=yZd6cn1RhvNNJUgiUTRYut8ofZYvbulnpPG-rZIRhi4,106
-nshtrainer/__init__.py,sha256=VcqBfL8RgCcZDaY645nxeDmOspqerx4x46wggCMnS0E,692
+nshtrainer/__init__.py,sha256=52OB7QRlhrTCIdDecpT7yEZyZM1XvYxywhuORn1eKoY,814
 nshtrainer/_callback.py,sha256=tXQCDzS6CvMTuTY5lQSH5qZs1pXUi-gt9bQdpXMVdEs,12715
 nshtrainer/_checkpoint/metadata.py,sha256=PHy-54Cg-o3OtCffAqrVv6ZVMU7zhRo_-sZiSEEno1Y,5019
 nshtrainer/_checkpoint/saver.py,sha256=LOP8jjKF0Dw9x9H-BKrLMWlEp1XTan2DUK0zQUCWw5U,1360
@@ -31,7 +31,7 @@ nshtrainer/callbacks/shared_parameters.py,sha256=ggMI1krkqN7sGOrjK_I96IsTMYMXHoV
 nshtrainer/callbacks/timer.py,sha256=BB-M7tV4QNYOwY_Su6j9P7IILxVRae_upmDq4qsxiao,4670
 nshtrainer/callbacks/wandb_upload_code.py,sha256=PTqNE1QB5U8NR5zhbiQZrmQuugX2UV7B12UdMpo9aV0,2353
 nshtrainer/callbacks/wandb_watch.py,sha256=tTTcFzxd2Ia9xu8tCogQ5CLJZBq1ne5JlpGVE75vKYs,2976
-nshtrainer/configs/__init__.py,sha256=zyo4lV9ObB3T3_hhBhzWGNb6MRma4h7QHD3OrypxqEw,10582
+nshtrainer/configs/__init__.py,sha256=eS3naq6EG1vCq28G2nAW1CqYFdsrh6ueBlzX_LazgUw,14159
 nshtrainer/configs/_checkpoint/__init__.py,sha256=6s7Y68StboqscY2G4P_QG443jz5aiym5SjOogIljWLg,342
 nshtrainer/configs/_checkpoint/metadata/__init__.py,sha256=oOPfYkXTjKgm6pluGsG6V1TPyCEGjsQpHVL-LffSUFQ,290
 nshtrainer/configs/_directory/__init__.py,sha256=_oO7vM9DhzHSxtZcv86sTi7hZIptnK1gr-AP9mqQ370,386
@@ -81,9 +81,17 @@ nshtrainer/configs/profiler/_base/__init__.py,sha256=ekYfPg-VDhCAFM5nJka2TxUYdRD
 nshtrainer/configs/profiler/advanced/__init__.py,sha256=-ThpUat16Ij_0avkMUVVA8wCWDG_q_tM7KQofnWQCtg,308
 nshtrainer/configs/profiler/pytorch/__init__.py,sha256=soAU1s2_Pa1na4gW8CK-iysJBO5M_7YeZC2_x40iEdg,294
 nshtrainer/configs/profiler/simple/__init__.py,sha256=3Wb11lPuFuyasq8xS1CZ4WLuBCLS_nVSQGVllvOOi0Y,289
-nshtrainer/configs/trainer/__init__.py,sha256=KIDYjJsc-WYXKiH2RNzAZJD5MKOTdO9wdtu_vWDNPxU,3936
-nshtrainer/configs/trainer/_config/__init__.py,sha256=1_Ad5uTvXdVuHMJB3s8s-0EraDwNZssg3sXBmVouF9w,3847
-nshtrainer/configs/trainer/trainer/__init__.py,sha256=DDuBRx0kVNMW0z_sqKTUt8-Ql7bOpargi4KcHHvDu_c,486
+nshtrainer/configs/trainer/__init__.py,sha256=hKMI_2ve5zcsQys2DDQDv7OmshYsIG0uJlCLreVHpF0,7779
+nshtrainer/configs/trainer/_config/__init__.py,sha256=Xw6I_9tUemDbHncpjKHRqye_e1_OyubK_FJcvdcQ0yc,4020
+nshtrainer/configs/trainer/accelerator/__init__.py,sha256=3H6R3wlwbKL1TzDqGCChZk78-BcE2czLouo7Djiq3nA,898
+nshtrainer/configs/trainer/plugin/__init__.py,sha256=NkHQxMPkrtTtdIAO4dQUE9SWEcHRDB0yUXLkTjnl4dA,3332
+nshtrainer/configs/trainer/plugin/base/__init__.py,sha256=GuubcKrbXt4VjJRT8VpNUQqBtyuutre_CfkS0EWZ5_E,368
+nshtrainer/configs/trainer/plugin/environment/__init__.py,sha256=3o16x4qRAOvkJH9Vg4-QwsEODDC6aP_OXRnPPkm_xSo,1376
+nshtrainer/configs/trainer/plugin/io/__init__.py,sha256=W6G67JnigB6d3MiwLrbSKgtIZLUccXznp-IXwkK1J4U,743
+nshtrainer/configs/trainer/plugin/layer_sync/__init__.py,sha256=SYDZk2M6sgpt4sEuoURuS8EKYmaqGcvYxETE9jvTrEE,431
+nshtrainer/configs/trainer/plugin/precision/__init__.py,sha256=szlqSfK2XuWdkf72LQzQFv3SlWfKFdRUpBEYIxQ3TPs,1507
+nshtrainer/configs/trainer/strategy/__init__.py,sha256=50whNloJVBq_bdbLaPQnPBTeS1Rcs8MwxTCYBj1kKa4,273
+nshtrainer/configs/trainer/trainer/__init__.py,sha256=QnuhMQNAa1nSVN2o50_WeKAQG_qkNlkeoq9zTjjwmTI,586
 nshtrainer/configs/util/__init__.py,sha256=qXittS7f7MyaqJnjvFLKnKsyb6bXTD3dEV16jXVDaH4,2104
 nshtrainer/configs/util/_environment_info/__init__.py,sha256=eB4E0Ck7XCeSC5gbUdA5thd7TXnjGCL0t8GZIFj7uCI,1644
 nshtrainer/configs/util/config/__init__.py,sha256=nEFiDG3-dvvTytYn1tEkPFzp7fgaGRp2j7toSN7yRGs,501
@@ -121,11 +129,19 @@ nshtrainer/profiler/_base.py,sha256=kFcSVn9gJuMwgDxbfyHh46CmEAIPZjxw3yjPbKgzvwA,
 nshtrainer/profiler/advanced.py,sha256=XrM3FX0ThCv5UwUrrH0l4Ow4LGAtpiBww2N8QAU5NOQ,1160
 nshtrainer/profiler/pytorch.py,sha256=8K37XvPnCApUpIK8tA2zNMFIaIiTLSoxKQoiyCPBm1Q,2757
 nshtrainer/profiler/simple.py,sha256=PimjqcU-JuS-8C0ZGHAdwCxgNLij4x0FH6WXsjBQzZs,1005
-nshtrainer/trainer/__init__.py,sha256=MmoydVS6aYeav7zgDAUHxAQrV_PMQsbnZTCuPnLH9Wk,128
-nshtrainer/trainer/_config.py,sha256=VD0DfdS-pyQ2nFG83c4u5AUkSAHODmXLX5s2qtvS_to,35400
+nshtrainer/trainer/__init__.py,sha256=ggDHzIUbABezh4BjEwrxyWuXmuDBV-x4jv9gwXgVHU0,250
+nshtrainer/trainer/_config.py,sha256=0GgofvaWf5Vo9REXNJpTvpVVRlFExGTOzcOt4jwJXNk,34129
 nshtrainer/trainer/_runtime_callback.py,sha256=6F2Gq27Q8OFfN3RtdNC6QRA8ac0LC1hh4DUE3V5WgbI,4217
+nshtrainer/trainer/accelerator.py,sha256=Bqq-ry7DeCY4zw9_zBvTZiijpA-uUHrDjtbLV652m4M,2415
+nshtrainer/trainer/plugin/__init__.py,sha256=UM8f70Ml3RGpsXeQr1Yh1yBcxxjFNGFGgJbniCn_rws,366
+nshtrainer/trainer/plugin/base.py,sha256=9-qUHXGpll_yCylun0899sbmJDpyhD9IQcBtVrJx38I,919
+nshtrainer/trainer/plugin/environment.py,sha256=NW0qbsbvDPe59JGOMgPLq1fj7szLucIV1WRTxCrcjF4,4367
+nshtrainer/trainer/plugin/io.py,sha256=nm6YDCVZAhmPvLaLnw6q4BrK2Gj2wvD5ZLDhj1xneEE,2030
+nshtrainer/trainer/plugin/layer_sync.py,sha256=h-ydZwXepnsw5-paLgiDatqPyQ_8C0QEv6DrYFIaXOo,735
+nshtrainer/trainer/plugin/precision.py,sha256=I0QsB1bVxmsFmBOkgrAfGONsuYae_lD9Bz0PfJEQvH4,5598
 nshtrainer/trainer/signal_connector.py,sha256=GhfGcSzfaTNhnj2QFkBDq5aT7FqbLMA7eC8SYQs8_8w,10828
-nshtrainer/trainer/trainer.py,sha256=HHqT83zWtYY9g5yD6X9aWrVh5VSpILW8PhoE6fp4snE,20734
+nshtrainer/trainer/strategy.py,sha256=VPTn5z3zvXTydY8IJchjhjcOfpvtoejnvUkq5E4WTus,1368
+nshtrainer/trainer/trainer.py,sha256=l2kJs27v4IHZnzxExr0zX0sVex0wukgiD2Wn_0wiGJg,20836
 nshtrainer/util/_environment_info.py,sha256=MT8mBe6ZolRfKiwU-les1P-lPNPqXpHQcfADrh_A3uY,24629
 nshtrainer/util/bf16.py,sha256=9QhHZCkYSfYpIcxwAMoXyuh2yTSHBzT-EdLQB297jEs,762
 nshtrainer/util/config/__init__.py,sha256=Z39JJufSb61Lhn2GfVcv3eFW_eorOrN9-9llDWlnZZM,272
@@ -137,6 +153,6 @@ nshtrainer/util/seed.py,sha256=diMV8iwBKN7Xxt5pELmui-gyqyT80_CZzomrWhNss0k,316
 nshtrainer/util/slurm.py,sha256=HflkP5iI_r4UHMyPjw9R4dD5AHsJUpcfJw5PLvGYBRM,1603
 nshtrainer/util/typed.py,sha256=Xt5fUU6zwLKSTLUdenovnKK0N8qUq89Kddz2_XeykVQ,164
 nshtrainer/util/typing_utils.py,sha256=MjY-CUX9R5Tzat-BlFnQjwl1PQ_W2yZQoXhkYHlJ_VA,442
-nshtrainer-1.0.0b29.dist-info/METADATA,sha256=YRehZvU9svmmfAmwFrdmu-Tzxgi_EHbFwrn-ewD8W9c,988
-nshtrainer-1.0.0b29.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-nshtrainer-1.0.0b29.dist-info/RECORD,,
+nshtrainer-1.0.0b30.dist-info/METADATA,sha256=zxFm4X5APkZR6E4E8-jzVghTwYEYCJQzCHpCV_8hWzg,988
+nshtrainer-1.0.0b30.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+nshtrainer-1.0.0b30.dist-info/RECORD,,