nshtrainer 0.30.0__py3-none-any.whl → 0.31.0__py3-none-any.whl

nshtrainer/trainer/trainer.py

@@ -18,10 +18,8 @@ from typing_extensions import Unpack, assert_never, override
 
 from .._checkpoint.metadata import _write_checkpoint_metadata
 from ..callbacks.base import resolve_all_callbacks
-from ..model.config import (
+from ._config import (
     AcceleratorConfigProtocol,
-    BaseConfig,
-    BaseProfilerConfig,
     LightningTrainerKwargs,
     StrategyConfigProtocol,
 )
@@ -29,6 +27,9 @@ from ._runtime_callback import RuntimeTrackerCallback, Stage
 from .checkpoint_connector import _CheckpointConnector
 from .signal_connector import _SignalConnector
 
+if TYPE_CHECKING:
+    from ..model.config import BaseConfig
+
 log = logging.getLogger(__name__)
 
 
@@ -58,14 +59,14 @@ def _is_bf16_supported_no_emulation():
 
 class Trainer(LightningTrainer):
     @classmethod
-    def _pre_init(cls, config: BaseConfig):
+    def _pre_init(cls, config: "BaseConfig"):
         if (precision := config.trainer.set_float32_matmul_precision) is not None:
             torch.set_float32_matmul_precision(precision)
 
     @classmethod
     def _update_kwargs(
         cls,
-        config: BaseConfig,
+        config: "BaseConfig",
         kwargs_ctor: LightningTrainerKwargs,
     ):
         kwargs: LightningTrainerKwargs = {
@@ -217,18 +218,16 @@ class Trainer(LightningTrainer):
                 gradient_clip_val=grad_clip_config.value,
             )
 
-        if profiler := config.trainer.profiler:
-            # If the profiler is an ProfilerConfig instance, then we instantiate it.
-            if isinstance(profiler, BaseProfilerConfig):
-                profiler = profiler.create_profiler(config)
-                # Make sure that the profiler is an instance of `Profiler`.
-                if not isinstance(profiler, Profiler):
-                    raise ValueError(f"{profiler=} is not an instance of `{Profiler}`.")
-
+        if profiler_config := config.trainer.profiler:
+            if (profiler := profiler_config.create_profiler(config)) is None:
+                log.warning(f"Profiler config {profiler_config=} returned None.")
+            # Make sure that the profiler is an instance of `Profiler`.
+            elif not isinstance(profiler, Profiler):
+                raise ValueError(f"{profiler=} is not an instance of `{Profiler}`.")
             # Otherwise, if the profiler is a string (e.g., "simpe", "advanced", "pytorch"),
             #   then we just pass it through.
-            # kwargs["profiler"] = profiler
-            _update_kwargs(profiler=profiler)
+            else:
+                _update_kwargs(profiler=profiler)
 
         if callbacks := resolve_all_callbacks(config):
             _update_kwargs(callbacks=callbacks)
@@ -281,7 +280,7 @@ class Trainer(LightningTrainer):
     @override
     def __init__(
         self,
-        config: BaseConfig,
+        config: "BaseConfig",
         /,
         **kwargs: Unpack[LightningTrainerKwargs],
     ):
@@ -424,7 +423,7 @@ class Trainer(LightningTrainer):
         # Save the checkpoint metadata
         metadata_path = None
         lm = self._base_module
-        root_config = cast(BaseConfig, lm.hparams)
+        root_config = cast("BaseConfig", lm.hparams)
         if root_config.trainer.save_checkpoint_metadata and self.is_global_zero:
             # Generate the metadata and write to disk
             if (

nshtrainer/{config → util/config}/__init__.py

@@ -1,3 +1,4 @@
+from . import duration as duration
 from .duration import Duration as Duration
 from .duration import Epochs as Epochs
 from .duration import Steps as Steps

{nshtrainer-0.30.0.dist-info → nshtrainer-0.31.0.dist-info}/METADATA

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

{nshtrainer-0.30.0.dist-info → nshtrainer-0.31.0.dist-info}/RECORD

@@ -1,11 +1,12 @@
-nshtrainer/__init__.py,sha256=sUb2yNdkHHhrKWCeWA5QKIA1Xx3jkO1QGD5Pa-HvgbA,614
+nshtrainer/__init__.py,sha256=flMI50Hj1Ie8c1YMSUQ759AqtNBQLT_zHaV2J9EUmOs,573
 nshtrainer/_callback.py,sha256=A1zLsTy4b_wOYnInLLXGSRdHzT2yNa6mPEql-ozm0u0,1013
 nshtrainer/_checkpoint/loader.py,sha256=5vjg-OFChXJjgiOVv8vnV8nwTscfdDtEdxQRz6uPfDE,14158
 nshtrainer/_checkpoint/metadata.py,sha256=5D4PgKodzhLsmQvuF3xxkH49epKaegxi4wh_ImDTtns,4737
 nshtrainer/_checkpoint/saver.py,sha256=MbX_WjkDtHHAf9Ms-KXDlknkjiPXVoGIe2ciO28AdZ0,1264
+nshtrainer/_directory.py,sha256=RjnW6vKTeKlz2vQWT3cG0Jje5BkFXA7HpUubDhcSiq4,2993
 nshtrainer/_experimental/__init__.py,sha256=pEXPyI184UuDHvfh4p9Kg9nQZQZI41e4_HvNd4BK-yg,81
 nshtrainer/_hf_hub.py,sha256=0bkXkqhve5D1onMW-fCfuvVKlTn0i6jv_6uMNgZ7OHQ,12974
-nshtrainer/callbacks/__init__.py,sha256=4qocBDzQbLLhhbIEfvbA3SQB_Dy9ZJH7keMwPay-ZS8,2359
+nshtrainer/callbacks/__init__.py,sha256=NpEV8bMU12ClFN2sLKLBDXnuwIHYyZOCNxDZgjrV104,2892
 nshtrainer/callbacks/_throughput_monitor_callback.py,sha256=aJo_11rc4lo0IYOd-kHmPDtzdC4ctgXyRudkRJqH4m4,23184
 nshtrainer/callbacks/actsave.py,sha256=qbnaKts4_dvjPeAaPtv7Ds12_vEWzaHUfg_--49NB9I,4041
 nshtrainer/callbacks/base.py,sha256=NpjeKmonJ1Kaz5_39XSn3LlDwvbGjk6WV8BpHSNCvI4,3508
@@ -14,6 +15,7 @@ nshtrainer/callbacks/checkpoint/_base.py,sha256=vvlwuD-20NozYVIolGGShmUdkkNYeuwN
 nshtrainer/callbacks/checkpoint/best_checkpoint.py,sha256=8BHgLAd3Tuzf5sup0guEAKF1jJiAwYsjdKBFYZw98ac,2171
 nshtrainer/callbacks/checkpoint/last_checkpoint.py,sha256=CWWv0cSwQ1VAX26N7hAyMxbNCk26Keh39oQguBEK5To,1102
 nshtrainer/callbacks/checkpoint/on_exception_checkpoint.py,sha256=ctT88EGT22_t_6tr5r7Sfo43cuve6XeroBnBYRMPOus,3372
+nshtrainer/callbacks/directory_setup.py,sha256=c0uY0oTqLcQ3egInHO7G6BeQQgk_xvOLoHH8FR-9U0U,2629
 nshtrainer/callbacks/early_stopping.py,sha256=VWuJz0oN87b6SwBeVc32YNpeJr1wts8K45k8JJJmG9I,4617
 nshtrainer/callbacks/ema.py,sha256=8-WHmKFP3VfnzMviJaIFmVD9xHPqIPmq9NRF5xdu3c8,12131
 nshtrainer/callbacks/finite_checks.py,sha256=gJC_RUr3ais3FJI0uB6wUZnDdE3WRwCix3ppA3PwQXA,2077
@@ -22,15 +24,16 @@ nshtrainer/callbacks/interval.py,sha256=smz5Zl8cN6X6yHKVsMRS2e3SEkzRCP3LvwE1ONvL
 nshtrainer/callbacks/log_epoch.py,sha256=fTa_K_Y8A7g09630cG4YkDE6AzSMPkjb9bpPm4gtqos,1120
 nshtrainer/callbacks/norm_logging.py,sha256=T2psu8mYsw9iahPKT6aUPjkGrZ4TIzm6_UUUmE09GJs,6274
 nshtrainer/callbacks/print_table.py,sha256=_FdAHhqylWGk4Z0c2FrLFeiMA4jhfA_beZRK_BHpzmE,2837
+nshtrainer/callbacks/rlp_sanity_checks.py,sha256=c30G9jAu42QLLIS5LnusdSnI3wqyIHgOUFDRcKESuNI,9935
+nshtrainer/callbacks/shared_parameters.py,sha256=fqlDweFDXPV_bfcAWpRgaJIad9i5AehYDtuJjDtUum4,2922
 nshtrainer/callbacks/throughput_monitor.py,sha256=H_ocXErZxUO3dxFk8Tx_VQdpI9E_Ztvqof5WtFevLyQ,1838
 nshtrainer/callbacks/timer.py,sha256=quS79oYClDUvQxJkNWmDMe0hwRUkkREgTgqzVrnom50,4607
 nshtrainer/callbacks/wandb_watch.py,sha256=Y6SEXfIx3kDDQbI5zpP53BVq0FBLJbLd3RJsiHZk1-Y,2921
-nshtrainer/config/__init__.py,sha256=v9RtlM1Pqj_4fCDfskgxEtiGtbWH3Tj7lqNsKCDQ4gk,119
-nshtrainer/config/duration.py,sha256=f_obz0eorkktI3HzAuIawABDkvuL4lDqCxcPb3UW7Q4,692
+nshtrainer/config.py,sha256=W6nAmn5Y1GVZto9vkx4v8i5XdikMSdVYDiq7kbDEWAg,5900
 nshtrainer/data/__init__.py,sha256=7mk1tr7SWUZ7ySbsf0y0ZPszk7u4QznPhQ-7wnpH9ec,149
 nshtrainer/data/balanced_batch_sampler.py,sha256=dGBTDDtlBU6c-ZlVQOCnTW7SjTB5hczWsOWEdUWjvkA,4385
 nshtrainer/data/transform.py,sha256=6SNs3_TpNpfhcwTwvPKyEJ3opM1OT7LmMEYQNHKgRl8,2227
-nshtrainer/ll/__init__.py,sha256=6UTt2apSD8tOZw3M7hyd-33v4RKSpNNATlWFbW4cNnU,2523
+nshtrainer/ll/__init__.py,sha256=L-aTi1V1bbvnZjOro8NvI393zbHQSFR9movWSRK9Mds,2477
 nshtrainer/ll/_experimental.py,sha256=oBQCKOEVYoxuUU9eLb-Fg2B2mzZD7SA0zfAO6lmWZ88,53
 nshtrainer/ll/actsave.py,sha256=2lbiseSrjcwFT6AiyLNWarTWl1bnzliVWlu1iOfnP30,209
 nshtrainer/ll/callbacks.py,sha256=AxyUmc8aGRSjx6WwwgXYCmdJ73rwLuEAEH0AGRosojQ,49
@@ -38,7 +41,7 @@ nshtrainer/ll/config.py,sha256=fKumJf42HY2FITX1QUM1OTXkYD6U2np2ciyd4PFRPZ8,145
 nshtrainer/ll/data.py,sha256=zRG0FRje-jtSHximVzkHIHzpwsyQxpHCoACFihNKLPM,44
 nshtrainer/ll/log.py,sha256=d4BB3TyM8imK65EXOiOeUTF0zFM1ropbe7Vq3DeB0xU,140
 nshtrainer/ll/lr_scheduler.py,sha256=7xjhN6L69BCUzFhcy33NtMtPuCzHiB611zVWFg92lQ0,52
-nshtrainer/ll/model.py,sha256=cxFQfFc-2mAYBGwDpP8m5tjQBs7M47cZ6JoPXksPaoI,473
+nshtrainer/ll/model.py,sha256=Cw8Vq8IUL6YU1fTUcOIZsXcNJ3XyKgQY4YENIsL9H7c,996
 nshtrainer/ll/nn.py,sha256=8qiRDFwojIxkB7-LtNWk4mLL2tJbaskHYofDsOIHiNg,42
 nshtrainer/ll/optimizer.py,sha256=3T-VZtT73jVvwCNJGDjgGEbzs-1LFTzMQH-SB_58mSo,49
 nshtrainer/ll/runner.py,sha256=B0m5VEhNKIjF1aFmqPkonkQxDoRL2jeHZGsV3zwhSVE,117
@@ -51,44 +54,47 @@ nshtrainer/loggers/__init__.py,sha256=C_xk0A3_qKbNdTmzK85AgjRHFD3w-jPRS2ig-iPhfE
 nshtrainer/loggers/_base.py,sha256=xiZKEK0ALJkcqf4OpVNRY0QbZsamR_WR7x7m_68YHXQ,705
 nshtrainer/loggers/csv.py,sha256=D_lYyd94bZ8jAgnRo-ARtFgVcInaD9zktxtsUD9RWCI,1052
 nshtrainer/loggers/tensorboard.py,sha256=wL2amRSdP68zbslZvBeM0ZQBnjF3hIKsz-_lBbdomaM,2216
-nshtrainer/loggers/wandb.py,sha256=FPwbf618AYmuPzHdhd1ZFhJ8qDjwTUiSe7cm7g3KCyM,5112
+nshtrainer/loggers/wandb.py,sha256=8B2BMMzILRSUEiCkmp_fBpcXs69euRKViTiaV__DJZk,5128
 nshtrainer/lr_scheduler/__init__.py,sha256=uEvgaFAs-4s_bAEMaildy0GT6OvgpgOEKTuzqutESHE,736
 nshtrainer/lr_scheduler/_base.py,sha256=7xOIuxQ86YHbFWG5a3gX46emQj1WN_LaY4-i0Q1TDBg,3659
-nshtrainer/lr_scheduler/linear_warmup_cosine.py,sha256=pmX5n7mmhSqPTz4Nu9g_JTsE9gzCkuU4V3GuAHUsDoA,5451
+nshtrainer/lr_scheduler/linear_warmup_cosine.py,sha256=YQm84Sb4SWrofpBwa39DCslJvu2uorjbpWaGWyys1l4,5352
 nshtrainer/lr_scheduler/reduce_lr_on_plateau.py,sha256=h76oTHYpMxauV_l6lviya5DW-WKArwxxf7ZQizhmbCw,2782
 nshtrainer/metrics/__init__.py,sha256=ObLIELGguIEcUpRsUkqh1ltrvZii6vglTpJGrPvoy00,50
 nshtrainer/metrics/_config.py,sha256=jgRBfDAQLFTW7AiUY7CRtdfts6CR6keeuqm0FFMWCzQ,1288
-nshtrainer/model/__init__.py,sha256=VyRziPT3YilP6xjLi_StsSqtlvn7N4LOMzgukRsOnF8,1380
-nshtrainer/model/base.py,sha256=oQVolDk81acy4OlckwQEBHuX2gCaVSYiIA0JaDIfhQ4,17517
-nshtrainer/model/config.py,sha256=zcCLcqvg4u7Zg6SLtCnqdIfiW8I0eART47lf1LCYl-A,43326
-nshtrainer/model/modules/callback.py,sha256=1z6gUDBd35KG3phGzRekgZM6SIk-wj5Uo6APN4YhRR0,8549
-nshtrainer/model/modules/debug.py,sha256=Yy7XEdPou9BkCsD5hJchwJGmCVGrfUru5g9VjPM4uAw,1120
-nshtrainer/model/modules/distributed.py,sha256=ABpR9d-3uBS_fivfy_WYW-dExW6vp5BPaoPQnOudHng,1725
-nshtrainer/model/modules/logger.py,sha256=CJWSmNT8SV5GLtfml-qGYenqRPXcNOMsJRGEavAd8Hw,5464
-nshtrainer/model/modules/profiler.py,sha256=rQ_jRMcM1Z2AIROZlRnBRHM5rkTpq67afZPD6CIRfXs,825
-nshtrainer/model/modules/rlp_sanity_checks.py,sha256=I_ralr2ThQ-D_FkVQTwbdXLLlgHJEr7-s01I5wSDjps,8893
-nshtrainer/model/modules/shared_parameters.py,sha256=ZiRKkZXr6RwdwLCdZCJPl3dXe7bnT8Z9yTeRK5bXBGk,2687
+nshtrainer/model/__init__.py,sha256=2i_VEy6u_Y1LUGKljHXWeekvhnUcanZM2QyaaBM1Bmw,261
+nshtrainer/model/base.py,sha256=hT27FtzwKQiEL0C8RcaTKYXlanfvzTxHOJpHUcWiItk,19891
+nshtrainer/model/config.py,sha256=Q4Wong6w3cp_Sq7s8iZdABKF-LZBbSCFn_TQPYkhkrI,6572
+nshtrainer/model/mixins/callback.py,sha256=lh3imlw1H3ESIG4WFA5frooSlWi6-RPUUDRFGRzEg4A,8571
+nshtrainer/model/mixins/logger.py,sha256=xOymSTofukEYZGkGojXsMEO__ZlBI5lIPZVmlotMEX8,5291
 nshtrainer/nn/__init__.py,sha256=0QPFl02a71WZQjLMGOlFNMmsYP5aa1q3eABHmnWH58Q,1427
 nshtrainer/nn/mlp.py,sha256=V0FrScpIUdg_IgIO8GMtIsGEtmHjwF14i2IWxmZrsqg,5952
 nshtrainer/nn/module_dict.py,sha256=NOY0B6WDTnktyWH4GthsprMQo0bpehC-hCq9SfD8paE,2329
 nshtrainer/nn/module_list.py,sha256=fb2u5Rqdjff8Pekyr9hkCPkBorQ-fldzzFAjsgWAm30,1719
 nshtrainer/nn/nonlinearity.py,sha256=4sYE4MN5zojc-go1k0PYtqssVRuXrM7D4tbpIXp5K-E,6078
 nshtrainer/optimizer.py,sha256=kuJEA1pvB3y1FcsfhAoOJujVqEZqFHlmYO8GW6JeA1g,1527
+nshtrainer/profiler/__init__.py,sha256=RQYkqQBVWuVvfdtAJIk2x5bNsXownklT87Mr_j-uXjw,474
+nshtrainer/profiler/_base.py,sha256=YF5lsJBIl9qts9GLW5Z62JuYeo4SnIArhlFwTGkfTb4,897
+nshtrainer/profiler/advanced.py,sha256=44asloha0aGUW8YwjQt3lm3ve8H-N6mM4QgseUSLT30,1112
+nshtrainer/profiler/pytorch.py,sha256=tGeRvoPP5ulWX2RkfXrQvMBoki1T95dpz5p8mwyon1I,2709
+nshtrainer/profiler/simple.py,sha256=MbMfsJvligd0mtGiltxJ0T8MQVDP9T9BzQZFwswl66Y,957
 nshtrainer/runner.py,sha256=USAjrExHkN5oVNVunsoPnLxfQrEHSaa54S3RipOe544,3605
 nshtrainer/scripts/find_packages.py,sha256=ixYivZobumyyGsf2B9oYMLyLTRcBzY_vUv-u3bNW-hs,1424
 nshtrainer/trainer/__init__.py,sha256=P2rmr8oBVTHk-HJHYPcUwWqDEArMbPR4_rPpATbWK3E,40
+nshtrainer/trainer/_config.py,sha256=2SgO1L5VBfWQ5g7Dg2dTx_vq2_Wo7dTqt2A4GlQaGo0,28673
 nshtrainer/trainer/_runtime_callback.py,sha256=sd2cUdRJG-UCdQr9ruZvEYpNGNF1t2W2fuxwwVlQD9E,4164
 nshtrainer/trainer/checkpoint_connector.py,sha256=r0ir4xYSdf_jebM0x09qaO6nJsvsiRQDyM0fs80ppOQ,2347
 nshtrainer/trainer/signal_connector.py,sha256=2EzkVktlasl8PgWAKNLDZRUMY__gRlDy1HdinAU-tfU,10740
-nshtrainer/trainer/trainer.py,sha256=L4nYXq6Gts2sS9CQGenwEcvMET4L5vO5c60KM5Hm8Do,17544
+nshtrainer/trainer/trainer.py,sha256=iYueHW-m8fHyC8SQuXmpgxq_-GUa7pAJik7rDFPXmy0,17499
 nshtrainer/util/_environment_info.py,sha256=CFUUZYjXhBLWGc0jtPNOaZgYMueUDEHpEaWFA1f3GoY,24213
 nshtrainer/util/_useful_types.py,sha256=dwZokFkIe7M5i2GR3nQ9A1lhGw06DMAFfH5atyquqSA,8000
+nshtrainer/util/config/__init__.py,sha256=6iCFLhujhbOi7Q694e--Sq-ethiGoGHShm699GPV8Zg,154
+nshtrainer/util/config/duration.py,sha256=f_obz0eorkktI3HzAuIawABDkvuL4lDqCxcPb3UW7Q4,692
 nshtrainer/util/environment.py,sha256=AeW_kLl-N70wmb6L_JLz1wRj0kA70xs6RCmc9iUqczE,4159
 nshtrainer/util/path.py,sha256=VkpuhR4GaZtSFBVqbGAvfjcrU-PR8xwiGzzwFNOWP9c,2995
 nshtrainer/util/seed.py,sha256=Or2wMPsnQxfnZ2xfBiyMcHFIUt3tGTNeMMyOEanCkqs,280
 nshtrainer/util/slurm.py,sha256=rofIU26z3SdL79SF45tNez6juou1cyDLz07oXEZb9Hg,1566
 nshtrainer/util/typed.py,sha256=NGuDkDzFlc1fAoaXjOFZVbmj0mRFjsQi1E_hPa7Bn5U,128
 nshtrainer/util/typing_utils.py,sha256=8ptjSSLZxlmy4FY6lzzkoGoF5fGNClo8-B_c0XHQaNU,385
-nshtrainer-0.30.0.dist-info/METADATA,sha256=lDudS-lD7exw8lNe_3vT13ysnk491QCkObXGLQtjhMk,916
-nshtrainer-0.30.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.30.0.dist-info/RECORD,,
+nshtrainer-0.31.0.dist-info/METADATA,sha256=99b-8IvPlMmTrjyb5EK1kKsgKj8lWhGw4gZvM5sKyzc,916
+nshtrainer-0.31.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.31.0.dist-info/RECORD,,

nshtrainer/model/modules/debug.py

@@ -1,42 +0,0 @@
-import logging
-
-import torch
-import torch.distributed
-
-log = logging.getLogger(__name__)
-
-
-class DebugModuleMixin:
-    @torch.jit.unused
-    def breakpoint(self, rank_zero_only: bool = True):
-        if (
-            not rank_zero_only
-            or not torch.distributed.is_initialized()
-            or torch.distributed.get_rank() == 0
-        ):
-            breakpoint()
-
-        if rank_zero_only and torch.distributed.is_initialized():
-            _ = torch.distributed.barrier()
-
-    @torch.jit.unused
-    def ensure_finite(
-        self,
-        tensor: torch.Tensor,
-        name: str | None = None,
-        throw: bool = False,
-    ):
-        name_parts: list[str] = ["Tensor"]
-        if name is not None:
-            name_parts.append(name)
-        name = " ".join(name_parts)
-
-        not_finite = ~torch.isfinite(tensor)
-        if not_finite.any():
-            msg = f"{name} has {not_finite.sum().item()}/{not_finite.numel()} non-finite values."
-            if throw:
-                raise RuntimeError(msg)
-            else:
-                log.warning(msg)
-            return False
-        return True

nshtrainer/model/modules/distributed.py

@@ -1,70 +0,0 @@
-from typing import Any, Literal, cast
-
-import torch.distributed
-from lightning.pytorch import LightningModule
-from torch.distributed import ReduceOp
-from typing_extensions import TypeVar
-
-from ...util.typing_utils import mixin_base_type
-
-T = TypeVar("T", infer_variance=True)
-
-ReduceOpStr = Literal[
-    "avg",
-    "mean",
-    "band",
-    "bor",
-    "bxor",
-    "max",
-    "min",
-    "premul_sum",
-    "product",
-    "sum",
-]
-VALID_REDUCE_OPS = (
-    "avg",
-    "mean",
-    "band",
-    "bor",
-    "bxor",
-    "max",
-    "min",
-    "premul_sum",
-    "product",
-    "sum",
-)
-
-
-class DistributedMixin(mixin_base_type(LightningModule)):
-    def all_gather_object(
-        self,
-        object: T,
-        group: torch.distributed.ProcessGroup | None = None,
-    ) -> list[T]:
-        if (
-            not torch.distributed.is_available()
-            or not torch.distributed.is_initialized()
-        ):
-            return [object]
-
-        object_list = [cast(T, None) for _ in range(self.trainer.world_size)]
-        torch.distributed.all_gather_object(object_list, object, group=group)
-        return object_list
-
-    def barrier(self, name: str | None = None):
-        self.trainer.strategy.barrier(name=name)
-
-    def reduce(
-        self,
-        tensor: torch.Tensor,
-        reduce_op: ReduceOp.RedOpType | ReduceOpStr,
-        group: Any | None = None,
-    ) -> torch.Tensor:
-        if isinstance(reduce_op, str):
-            # validate reduce_op
-            if reduce_op not in VALID_REDUCE_OPS:
-                raise ValueError(
-                    f"reduce_op must be one of {VALID_REDUCE_OPS}, got {reduce_op}"
-                )
-
-        return self.trainer.strategy.reduce(tensor, group=group, reduce_op=reduce_op)

nshtrainer/model/modules/profiler.py

@@ -1,24 +0,0 @@
-from lightning.pytorch import LightningDataModule, LightningModule
-from lightning.pytorch.profilers import PassThroughProfiler
-
-from ...util.typing_utils import mixin_base_type
-
-
-class ProfilerMixin(mixin_base_type(LightningModule)):
-    @property
-    def profiler(self):
-        if not isinstance(self, (LightningModule, LightningDataModule)):
-            raise TypeError(
-                "`profiler` can only be used on LightningModule or LightningDataModule"
-            )
-
-        if (trainer := self.trainer) is None:
-            raise RuntimeError("trainer is not defined")
-
-        if not hasattr(trainer, "profiler"):
-            raise RuntimeError("trainer does not have profiler")
-
-        if (profiler := getattr(trainer, "profiler")) is None:
-            profiler = PassThroughProfiler()
-
-        return profiler

nshtrainer/model/modules/rlp_sanity_checks.py

@@ -1,202 +0,0 @@
-import logging
-from collections.abc import Mapping
-from typing import cast
-
-import torch
-from lightning.pytorch import LightningModule, Trainer
-from lightning.pytorch.utilities.types import (
-    LRSchedulerConfigType,
-    LRSchedulerTypeUnion,
-)
-from typing_extensions import Protocol, override, runtime_checkable
-
-from ...util.typing_utils import mixin_base_type
-from ..config import BaseConfig
-from .callback import CallbackModuleMixin
-
-log = logging.getLogger(__name__)
-
-
-def _on_train_start_callback(trainer: Trainer, pl_module: LightningModule):
-    # If we're in PL's "sanity check" mode, we don't need to run this check
-    if trainer.sanity_checking:
-        return
-
-    config = cast(BaseConfig, pl_module.hparams)
-    if config.trainer.sanity_checking.reduce_lr_on_plateau == "disable":
-        return
-
-    # if no lr schedulers, return
-    if not trainer.lr_scheduler_configs:
-        return
-
-    errors: list[str] = []
-    disable_message = (
-        "Otherwise, set `config.trainer.sanity_checking.reduce_lr_on_plateau='disable'` "
-        "to disable this sanity check."
-    )
-
-    for lr_scheduler_config in trainer.lr_scheduler_configs:
-        if not lr_scheduler_config.reduce_on_plateau:
-            continue
-
-        match lr_scheduler_config.interval:
-            case "epoch":
-                # we need to make sure that the trainer runs val every `frequency` epochs
-
-                # If `trainer.check_val_every_n_epoch` is None, then Lightning
-                # will run val every `int(trainer.val_check_interval)` steps.
-                # So, first we need to make sure that `trainer.val_check_interval` is not None first.
-                if trainer.check_val_every_n_epoch is None:
-                    errors.append(
-                        "Trainer is not running validation at epoch intervals "
-                        "(i.e., `trainer.check_val_every_n_epoch` is None) but "
-                        f"a ReduceLRPlateau scheduler with interval={lr_scheduler_config.interval} is used."
-                        f"Please set `config.trainer.check_val_every_n_epoch={lr_scheduler_config.frequency}`. "
-                        + disable_message
-                    )
-
-                # Second, we make sure that the trainer runs val at least every `frequency` epochs
-                if (
-                    trainer.check_val_every_n_epoch is not None
-                    and lr_scheduler_config.frequency % trainer.check_val_every_n_epoch
-                    != 0
-                ):
-                    errors.append(
-                        f"Trainer is not running validation every {lr_scheduler_config.frequency} epochs but "
-                        f"a ReduceLRPlateau scheduler with interval={lr_scheduler_config.interval} and frequency={lr_scheduler_config.frequency} is used."
-                        f"Please set `config.trainer.check_val_every_n_epoch` to a multiple of {lr_scheduler_config.frequency}. "
-                        + disable_message
-                    )
-
-            case "step":
-                # In this case, we need to make sure that the trainer runs val at step intervals
-                # that are multiples of `frequency`.
-
-                # First, we make sure that validation is run at step intervals
-                if trainer.check_val_every_n_epoch is not None:
-                    errors.append(
-                        "Trainer is running validation at epoch intervals "
-                        "(i.e., `trainer.check_val_every_n_epoch` is not None) but "
-                        f"a ReduceLRPlateau scheduler with interval={lr_scheduler_config.interval} is used."
-                        "Please set `config.trainer.check_val_every_n_epoch=None` "
-                        f"and `config.trainer.val_check_interval={lr_scheduler_config.frequency}`. "
-                        + disable_message
-                    )
-
-                # Second, we make sure `trainer.val_check_interval` is an integer
-                if not isinstance(trainer.val_check_interval, int):
-                    errors.append(
-                        f"Trainer is not running validation at step intervals "
-                        f"(i.e., `trainer.val_check_interval` is not an integer) but "
-                        f"a ReduceLRPlateau scheduler with interval={lr_scheduler_config.interval} is used."
-                        "Please set `config.trainer.val_check_interval=None` "
-                        f"and `config.trainer.val_check_interval={lr_scheduler_config.frequency}`. "
-                        + disable_message
-                    )
-
-                # Third, we make sure that the trainer runs val at least every `frequency` steps
-                if (
-                    isinstance(trainer.val_check_interval, int)
-                    and trainer.val_check_interval % lr_scheduler_config.frequency != 0
-                ):
-                    errors.append(
-                        f"Trainer is not running validation every {lr_scheduler_config.frequency} steps but "
-                        f"a ReduceLRPlateau scheduler with interval={lr_scheduler_config.interval} and frequency={lr_scheduler_config.frequency} is used."
-                        "Please set `config.trainer.val_check_interval` "
-                        f"to a multiple of {lr_scheduler_config.frequency}. "
-                        + disable_message
-                    )
-
-            case _:
-                pass
-
-    if not errors:
-        return
-
-    message = (
-        "ReduceLRPlateau sanity checks failed with the following errors:\n"
-        + "\n".join(errors)
-    )
-    match config.trainer.sanity_checking.reduce_lr_on_plateau:
-        case "warn":
-            log.warning(message)
-        case "error":
-            raise ValueError(message)
-        case _:
-            pass
-
-
-@runtime_checkable
-class CustomRLPImplementation(Protocol):
-    __reduce_lr_on_plateau__: bool
-
-
-class RLPSanityCheckModuleMixin(mixin_base_type(CallbackModuleMixin)):
-    @override
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-        global _on_train_start_callback
-        self.register_callback(on_train_start=_on_train_start_callback)
-
-    def reduce_lr_on_plateau_config(
-        self,
-        lr_scheduler: LRSchedulerTypeUnion | LRSchedulerConfigType,
-    ) -> LRSchedulerConfigType:
-        if (trainer := self._trainer) is None:
-            raise RuntimeError(
-                "Could not determine the frequency of ReduceLRPlateau scheduler "
-                "because `self.trainer` is None."
-            )
-
-        # First, resolve the LR scheduler from the provided config.
-        lr_scheduler_config: LRSchedulerConfigType
-        match lr_scheduler:
-            case Mapping():
-                lr_scheduler_config = cast(LRSchedulerConfigType, lr_scheduler)
-            case _:
-                lr_scheduler_config = {"scheduler": lr_scheduler}
-
-        # Make sure the scheduler is a ReduceLRPlateau scheduler. Otherwise, warn the user.
-        if (
-            not isinstance(
-                lr_scheduler_config["scheduler"],
-                torch.optim.lr_scheduler.ReduceLROnPlateau,
-            )
-        ) and (
-            not isinstance(lr_scheduler_config["scheduler"], CustomRLPImplementation)
-            or not lr_scheduler_config["scheduler"].__reduce_lr_on_plateau__
-        ):
-            log.warning(
-                "`reduce_lr_on_plateau_config` should only be used with a ReduceLRPlateau scheduler. "
-                f"The provided scheduler, {lr_scheduler_config['scheduler']}, does not subclass "
-                "`torch.optim.lr_scheduler.ReduceLROnPlateau`. "
-                "Please ensure that the scheduler is a ReduceLRPlateau scheduler. "
-                "If you are using a custom ReduceLRPlateau scheduler implementation, "
-                "please either (1) make sure that it subclasses `torch.optim.lr_scheduler.ReduceLROnPlateau`, "
-                "or (2) set the scheduler's `__reduce_lr_on_plateau__` attribute to `True`."
-            )
-
-        # If trainer.check_val_every_n_epoch is an integer, then we run val at epoch intervals.
-        if trainer.check_val_every_n_epoch is not None:
-            return {
-                "reduce_on_plateau": True,
-                "interval": "epoch",
-                "frequency": trainer.check_val_every_n_epoch,
-                **lr_scheduler_config,
-            }
-
-        # Otherwise, we run val at step intervals.
-        if not isinstance(trainer.val_check_batch, int):
-            raise ValueError(
-                "Could not determine the frequency of ReduceLRPlateau scheduler "
-                f"because {trainer.val_check_batch=} is not an integer."
-            )
-
-        return {
-            "reduce_on_plateau": True,
-            "interval": "step",
-            "frequency": trainer.val_check_batch,
-            **lr_scheduler_config,
-        }

nshtrainer/model/modules/shared_parameters.py

@@ -1,72 +0,0 @@
-import logging
-from collections.abc import Sequence
-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 = logging.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")