kostyl-toolkit 0.1.38__py3-none-any.whl → 0.1.39__py3-none-any.whl

kostyl/ml/configs/__init__.py

@@ -1,6 +1,6 @@
+from .hyperparams import OPTIMIZER
 from .hyperparams import HyperparamsConfig
 from .hyperparams import Lr
-from .hyperparams import Optimizer
 from .hyperparams import WeightDecay
 from .mixins import ConfigLoadingMixin
 from .training_settings import CheckpointConfig
@@ -14,6 +14,7 @@ from .training_settings import TrainingSettings
 
 
 __all__ = [
+    "OPTIMIZER",
     "CheckpointConfig",
     "ConfigLoadingMixin",
     "DDPStrategyConfig",
@@ -23,7 +24,6 @@ __all__ = [
     "HyperparamsConfig",
     "LightningTrainerParameters",
     "Lr",
-    "Optimizer",
     "SingleDeviceStrategyConfig",
     "TrainingSettings",
     "WeightDecay",

kostyl/ml/configs/hyperparams.py

@@ -11,11 +11,19 @@ logger = setup_logger(fmt="only_message")
 
 
 class AdamConfig(BaseModel):
-    """AdamW optimizer hyperparameters configuration."""
+    """Adam optimizer hyperparameters configuration."""
 
-    type: Literal["AdamW"] = "AdamW"
+    type: Literal["AdamW", "Adam"] = "AdamW"
     betas: tuple[float, float] = (0.9, 0.999)
-    is_adamw: bool = True
+
+
+class MuonConfig(BaseModel):
+    """Muon optimizer hyperparameters configuration."""
+
+    type: Literal["Muon"]
+    nesterov: bool = True
+    ns_coefficients: tuple[float, float, float] = (3.4445, -4.7750, 2.0315)
+    ns_steps: int = 5
 
 
 class AdamWithPrecisionConfig(BaseModel):
@@ -28,67 +36,85 @@ class AdamWithPrecisionConfig(BaseModel):
     is_adamw: bool = True
 
 
-Optimizer = AdamConfig | AdamWithPrecisionConfig
+OPTIMIZER = AdamConfig | AdamWithPrecisionConfig | MuonConfig
+SCHEDULER = Literal[
+    "linear",
+    "cosine",
+    "plateau-with-cosine-annealing",
+    "plateau-with-linear-annealing",
+]
 
 
 class Lr(BaseModel):
     """Learning rate hyperparameters configuration."""
 
-    use_scheduler: bool = False
-    warmup_iters_ratio: float | None = Field(
-        default=None, gt=0, lt=1, validate_default=False
-    )
+    scheduler_type: SCHEDULER | None = None
+
+    freeze_ratio: float | None = Field(default=None, ge=0, le=1)
+    warmup_ratio: float | None = Field(default=None, gt=0, lt=1, validate_default=False)
     warmup_value: float | None = Field(default=None, gt=0, validate_default=False)
     base_value: float
     final_value: float | None = Field(default=None, gt=0, validate_default=False)
+    plateau_ratio: float | None = Field(
+        default=None, gt=0, lt=1, validate_default=False
+    )
 
     @model_validator(mode="after")
-    def validate_warmup(self) -> "Lr":
-        """Validates the warmup parameters based on use_scheduler."""
-        if (self.warmup_value is None) != (self.warmup_iters_ratio is None):  # fmt: skip
-            raise ValueError(
-                "Both warmup_value and warmup_iters_ratio must be provided or neither"
-            )
-        if ((self.warmup_value is not None) or (self.warmup_iters_ratio is not None)) and not self.use_scheduler:  # fmt: skip
+    def _validate_freeze_ratio(self) -> "Lr":
+        if self.scheduler_type is None and self.freeze_ratio is not None:
+            logger.warning("use_scheduler is False, freeze_ratio will be ignored.")
+            self.freeze_ratio = None
+        return self
+
+    @model_validator(mode="after")
+    def _validate_warmup(self) -> "Lr":
+        if ((self.warmup_value is not None) or (self.warmup_ratio is not None)) and self.scheduler_type is None:  # fmt: skip
             logger.warning(
-                "use_scheduler is False, warmup_value and warmup_iters_ratio will be ignored."
+                "scheduler_type is None, warmup_value and warmup_ratio will be ignored."
             )
             self.warmup_value = None
-            self.warmup_iters_ratio = None
+            self.warmup_ratio = None
+        if (self.warmup_value is None) != (self.warmup_ratio is None):  # fmt: skip
+            raise ValueError(
+                "Both warmup_value and warmup_ratio must be provided or neither"
+            )
         return self
 
     @model_validator(mode="after")
-    def validate_final_value(self) -> "Lr":
-        """Validates the final_value based on use_scheduler."""
-        if self.use_scheduler and (self.final_value is None):
-            raise ValueError("If use_scheduler is True, final_value must be provided.")
-        if (not self.use_scheduler) and (self.final_value is not None):
+    def _validate_final_value(self) -> "Lr":
+        if (self.scheduler_type in {"linear"}) and (self.final_value is not None):
+            raise ValueError("If scheduler_type is 'linear', final_value must be None.")
+        if (self.scheduler_type is None) and (self.final_value is not None):
             logger.warning("use_scheduler is False, final_value will be ignored.")
             self.final_value = None
         return self
 
-
-class WeightDecay(BaseModel):
-    """Weight decay hyperparameters configuration."""
-
-    use_scheduler: bool = False
-    base_value: float
-    final_value: float | None = None
-
     @model_validator(mode="after")
-    def validate_final_value(self) -> "WeightDecay":
-        """Validates the final_value based on use_scheduler."""
-        if self.use_scheduler and self.final_value is None:
-            raise ValueError("If use_scheduler is True, final_value must be provided.")
-        if not self.use_scheduler and self.final_value is not None:
-            logger.warning("use_scheduler is False, final_value will be ignored.")
+    def _validate_plateau_ratio(self) -> "Lr":
+        if self.scheduler_type is not None:
+            if self.scheduler_type.startswith("plateau") and self.plateau_ratio is None:
+                raise ValueError(
+                    "If scheduler_type is 'plateau-with-*', plateau_ratio must be provided."
+                )
+            if (
+                not self.scheduler_type.startswith("plateau")
+                and self.plateau_ratio is not None
+            ):
+                logger.warning(
+                    "scheduler_type is not 'plateau-with-*', plateau_ratio will be ignored."
+                )
+                self.plateau_ratio = None
         return self
 
 
+class WeightDecay(Lr):
+    """Weight decay hyperparameters configuration."""
+
+
 class HyperparamsConfig(BaseModel):
     """Model training hyperparameters configuration."""
 
     grad_clip_val: float | None = Field(default=None, gt=0, validate_default=False)
-    optimizer: Optimizer
+    optimizer: OPTIMIZER
     lr: Lr
     weight_decay: WeightDecay

kostyl/ml/integrations/lightning/callbacks/checkpoint.py

@@ -286,7 +286,7 @@ def setup_checkpoint_callback(
     ckpt_cfg: CheckpointConfig,
     checkpoint_uploader: ModelCheckpointUploader | None = None,
     upload_strategy: Literal["only-best", "every-checkpoint"] | None = None,
-    remove_folder_if_exists: bool = True,
+    remove_folder_if_exists: bool = False,
 ) -> ModelCheckpointWithCheckpointUploader | ModelCheckpoint:
     """
     Create and configure a checkpoint callback for model saving.

kostyl/ml/integrations/lightning/loggers/tb_logger.py

@@ -11,14 +11,17 @@ logger = setup_logger()
 
 
 def setup_tb_logger(
-    runs_dir: Path,
+    runs_dir: Path, remove_folder_if_exists: bool = False
 ) -> TensorBoardLogger:
     """Sets up a TensorBoardLogger for PyTorch Lightning."""
     if runs_dir.exists():
         if is_local_zero_rank():
             logger.warning(f"TensorBoard log directory {runs_dir} already exists.")
-            rmtree(runs_dir)
-            logger.warning(f"Removed existing TensorBoard log directory {runs_dir}.")
+            if remove_folder_if_exists:
+                rmtree(runs_dir)
+                logger.warning(
+                    f"Removed existing TensorBoard log directory {runs_dir}."
+                )
     else:
         logger.info(f"Creating TensorBoard log directory {runs_dir}.")
         runs_dir.mkdir(parents=True, exist_ok=True)

kostyl/ml/integrations/lightning/utils.py

@@ -4,16 +4,9 @@ import lightning as L
 import torch.distributed as dist
 from torch.distributed import ProcessGroup
 
-from kostyl.ml.configs import DDPStrategyConfig
-from kostyl.ml.configs import FSDP1StrategyConfig
-from kostyl.ml.configs import SingleDeviceStrategyConfig
 from kostyl.utils.logging import setup_logger
 
 
-TRAINING_STRATEGIES = (
-    FSDP1StrategyConfig | DDPStrategyConfig | SingleDeviceStrategyConfig
-)
-
 logger = setup_logger()
 
 

kostyl/ml/schedulers/linear.py

@@ -13,21 +13,21 @@ class _LinearScheduleBase(BaseScheduler):
         self,
         param_name: str,
         num_iters: int,
-        start_value: float,
+        base_value: float,
         final_value: float,
     ) -> None:
         self.param_name = param_name
         self.num_iters = num_iters
-        self.start_value = start_value
+        self.base_value = base_value
         self.final_value = final_value
 
         self.scheduled_values: npt.NDArray[np.float64] = np.array([], dtype=np.float64)
-        self.current_value_ = self.start_value
+        self.current_value_ = self.base_value
         return
 
     def _create_scheduler(self) -> None:
         self.scheduled_values = np.linspace(
-            self.start_value, self.final_value, num=self.num_iters, dtype=np.float64
+            self.base_value, self.final_value, num=self.num_iters, dtype=np.float64
         )
         self._verify()
         return
@@ -68,7 +68,7 @@ class LinearScheduler(_LinearScheduleBase):
         optimizer: torch.optim.Optimizer,
         param_group_field: str,
         num_iters: int,
-        start_value: float,
+        base_value: float,
         final_value: float,
         multiplier_field: str | None = None,
         skip_if_zero: bool = False,
@@ -82,7 +82,7 @@ class LinearScheduler(_LinearScheduleBase):
             optimizer: Optimizer whose param groups are updated in-place.
             param_group_field: Name of the field that receives the scheduled value.
             num_iters: Number of scheduler iterations before clamping at ``final_value``.
-            start_value: Value used on the first iteration.
+            base_value: Value used on the first iteration.
             final_value: Value used once ``num_iters`` iterations are consumed.
             multiplier_field: Optional per-group multiplier applied to the scheduled value.
             skip_if_zero: Leave groups untouched when their target field equals zero.
@@ -98,7 +98,7 @@ class LinearScheduler(_LinearScheduleBase):
         super().__init__(
             param_name=param_group_field,
             num_iters=num_iters,
-            start_value=start_value,
+            base_value=base_value,
             final_value=final_value,
         )
         self.param_group_field = param_group_field

{kostyl_toolkit-0.1.38.dist-info → kostyl_toolkit-0.1.39.dist-info}/METADATA

@@ -1,17 +1,15 @@
 Metadata-Version: 2.3
 Name: kostyl-toolkit
-Version: 0.1.38
+Version: 0.1.39
 Summary: Kickass Orchestration System for Training, Yielding & Logging 
 Requires-Dist: case-converter>=1.2.0
 Requires-Dist: loguru>=0.7.3
-Requires-Dist: case-converter>=1.2.0 ; extra == 'ml-core'
-Requires-Dist: clearml[s3]>=2.0.2 ; extra == 'ml-core'
-Requires-Dist: lightning>=2.5.6 ; extra == 'ml-core'
-Requires-Dist: pydantic>=2.12.4 ; extra == 'ml-core'
-Requires-Dist: torch>=2.9.1 ; extra == 'ml-core'
-Requires-Dist: transformers>=4.57.1 ; extra == 'ml-core'
+Requires-Dist: case-converter>=1.2.0 ; extra == 'ml'
+Requires-Dist: pydantic>=2.12.4 ; extra == 'ml'
+Requires-Dist: torch>=2.9.1 ; extra == 'ml'
+Requires-Dist: transformers>=4.57.1 ; extra == 'ml'
 Requires-Python: >=3.12
-Provides-Extra: ml-core
+Provides-Extra: ml
 Description-Content-Type: text/markdown
 
 # Kostyl Toolkit

{kostyl_toolkit-0.1.38.dist-info → kostyl_toolkit-0.1.39.dist-info}/RECORD

@@ -1,8 +1,8 @@
 kostyl/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 kostyl/ml/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 kostyl/ml/base_uploader.py,sha256=KxHuohCcNK18kTVFBBqDu_IOQefluhSXOzwC56O66wc,484
-kostyl/ml/configs/__init__.py,sha256=djYjLxA7riFcSibAKfWHns-BCESEPrqSz_ZY2rJO-cc,913
-kostyl/ml/configs/hyperparams.py,sha256=lvtbvOFEoTBAJug7FR35xMQdPLgDQjRoP2fyDP-jD7E,3305
+kostyl/ml/configs/__init__.py,sha256=9AbM3ugsvc0bRKHA99DTRuUc2N7iSV9BH5FTDOp0cpw,913
+kostyl/ml/configs/hyperparams.py,sha256=wU6NoZoTG9WlnU4z3g4ViUfjjp8fb8YWUS_2gAd5wsY,4131
 kostyl/ml/configs/mixins.py,sha256=xHHAoRoPbzP9ECFP9duzg6SzegHcoLI8Pr9NrLoWNHs,1411
 kostyl/ml/configs/training_settings.py,sha256=wT9CHuLaKrLwonsc87Ee421EyFis_c9fqOgn9bSClm8,2747
 kostyl/ml/data_collator.py,sha256=kxiaMDKwSKXGBtrF8yXxHcypf7t_6syU-NwO1LcX50k,4062
@@ -16,25 +16,25 @@ kostyl/ml/integrations/clearml/loading_utils.py,sha256=NAMmB9NTGCXCHh-bR_nrQZyqI
 kostyl/ml/integrations/clearml/version_utils.py,sha256=GBjIIZbH_itd5sj7XpvxjkyZwxxGOpEcQ3BiWaJTyq8,1210
 kostyl/ml/integrations/lightning/__init__.py,sha256=r96os8kTuKIAymx3k9Td1JBrO2PH7nQAWUC54NsY5yY,392
 kostyl/ml/integrations/lightning/callbacks/__init__.py,sha256=EnKkNwwNDZnEqKRlpY4FVrqP88ECPF6nlT2bSLUIKRk,194
-kostyl/ml/integrations/lightning/callbacks/checkpoint.py,sha256=SfcaQRkXviMUej0UgrfXcqMDlRKYaAN3rgYCMKI97Os,18433
+kostyl/ml/integrations/lightning/callbacks/checkpoint.py,sha256=iJQKaQUzCMNza7SUKtCVDol_Cy3rZy9KfBKQ6kT6Swg,18434
 kostyl/ml/integrations/lightning/callbacks/early_stopping.py,sha256=D5nyjktCJ9XYAf28-kgXG8jORvXLl1N3nbDQnvValPM,615
 kostyl/ml/integrations/lightning/loggers/__init__.py,sha256=e51dszaoJbuzwBkbdugmuDsPldoSO4yaRgmZUg1Bdy0,71
-kostyl/ml/integrations/lightning/loggers/tb_logger.py,sha256=CpjlcEIT187cJXJgRYafqfzvcnwPgPaVZ0vLUflIr7k,899
+kostyl/ml/integrations/lightning/loggers/tb_logger.py,sha256=jVzEArHr1bpO-HRdKjlLJ4BJnxomHPAuBKCdYQRHRn4,1023
 kostyl/ml/integrations/lightning/metrics_formatting.py,sha256=U6vdNENZLvp2dT1L3HqFKtXrHwGKoDXN93hvamPGHjM,1341
 kostyl/ml/integrations/lightning/mixins.py,sha256=hVIsIUu6Iryrz6S7GQTqog9vNq8LQyjJd2aoJ5Ws6KU,5253
 kostyl/ml/integrations/lightning/module.py,sha256=39hcVNZSGyj5tLpXyX8IoqMGWt5vf6-Bx5JnNJ2-Wag,5218
-kostyl/ml/integrations/lightning/utils.py,sha256=DhLy_3JA5VyMQkB1v6xxRxDNHfisjXFYVjuIKPpO81M,1967
+kostyl/ml/integrations/lightning/utils.py,sha256=QhbK5iTv07xkRKuyXiK05aoY_ObSGorzGu1WcFvvFtI,1712
 kostyl/ml/params_groups.py,sha256=nUyw5d06Pvy9QPiYtZzLYR87xwXqJLxbHthgQH8oSCM,3583
 kostyl/ml/schedulers/__init__.py,sha256=VIo8MOP4w5Ll24XqFb3QGi2rKvys6c0dEFYPIdDoPlw,526
 kostyl/ml/schedulers/base.py,sha256=bjmwgdZpnSqpCnHPnKC6MEiRO79cwxMJpZq-eQVNs2M,1353
 kostyl/ml/schedulers/composite.py,sha256=ee4xlMDMMtjKPkbTF2ue9GTr9DuGCGjZWf11mHbi6aE,2387
 kostyl/ml/schedulers/cosine.py,sha256=y8ylrgVOkVcr2-ExoqqNW--tdDX88TBYPQCOppIf2_M,8685
-kostyl/ml/schedulers/linear.py,sha256=RnnnblRuRXP3LT03QVIHUaK2kNsiMP1AedrMoeyh3qk,5843
+kostyl/ml/schedulers/linear.py,sha256=a-Dq5VYoM1Z7vBb4_2Np0MAOfRmO1QVZnOzEoY8nM6k,5834
 kostyl/ml/schedulers/plateau.py,sha256=N-hiostPtTR0W4xnEJYB_1dv0DRx39iufLkGUrSIoWE,11235
 kostyl/utils/__init__.py,sha256=hkpmB6c5pr4Ti5BshOROebb7cvjDZfNCw83qZ_FFKMM,240
 kostyl/utils/dict_manipulations.py,sha256=e3vBicID74nYP8lHkVTQc4-IQwoJimrbFELy5uSF6Gk,1073
 kostyl/utils/fs.py,sha256=gAQNIU4R_2DhwjgzOS8BOMe0gZymtY1eZwmdgOdDgqo,510
 kostyl/utils/logging.py,sha256=CgNFNogcK0hoZmygvBWlTcq5A3m2Pfv9eOAP_gwx0pM,6633
-kostyl_toolkit-0.1.38.dist-info/WHEEL,sha256=e_m4S054HL0hyR3CpOk-b7Q7fDX6BuFkgL5OjAExXas,80
-kostyl_toolkit-0.1.38.dist-info/METADATA,sha256=nz5AzlWjKBqh7OZCklk-efWZ1jVDihw3YrrpLyoII3k,4269
-kostyl_toolkit-0.1.38.dist-info/RECORD,,
+kostyl_toolkit-0.1.39.dist-info/WHEEL,sha256=e_m4S054HL0hyR3CpOk-b7Q7fDX6BuFkgL5OjAExXas,80
+kostyl_toolkit-0.1.39.dist-info/METADATA,sha256=9AA-uH9-jSciHx3xh4WGxHkzTyjEESi0AFRitTkIM5w,4136
+kostyl_toolkit-0.1.39.dist-info/RECORD,,