kostyl-toolkit 0.1.38__tar.gz → 0.1.39__tar.gz

{kostyl_toolkit-0.1.38 → kostyl_toolkit-0.1.39}/PKG-INFO

@@ -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 → kostyl_toolkit-0.1.39}/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_toolkit-0.1.39/kostyl/ml/configs/hyperparams.py

@@ -0,0 +1,120 @@
+from typing import Literal
+
+from pydantic import BaseModel
+from pydantic import Field
+from pydantic import model_validator
+
+from kostyl.utils.logging import setup_logger
+
+
+logger = setup_logger(fmt="only_message")
+
+
+class AdamConfig(BaseModel):
+    """Adam optimizer hyperparameters configuration."""
+
+    type: Literal["AdamW", "Adam"] = "AdamW"
+    betas: tuple[float, float] = (0.9, 0.999)
+
+
+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):
+    """Adam optimizer with low-precision hyperparameters configuration."""
+
+    type: Literal["Adam8bit", "Adam4bit", "AdamFp8"]
+    betas: tuple[float, float] = (0.9, 0.999)
+    block_size: int
+    bf16_stochastic_round: bool = False
+    is_adamw: bool = True
+
+
+OPTIMIZER = AdamConfig | AdamWithPrecisionConfig | MuonConfig
+SCHEDULER = Literal[
+    "linear",
+    "cosine",
+    "plateau-with-cosine-annealing",
+    "plateau-with-linear-annealing",
+]
+
+
+class Lr(BaseModel):
+    """Learning rate hyperparameters configuration."""
+
+    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_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(
+                "scheduler_type is None, warmup_value and warmup_ratio will be ignored."
+            )
+            self.warmup_value = 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":
+        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
+
+    @model_validator(mode="after")
+    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
+    lr: Lr
+    weight_decay: WeightDecay

{kostyl_toolkit-0.1.38 → kostyl_toolkit-0.1.39}/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_toolkit-0.1.38 → kostyl_toolkit-0.1.39}/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_toolkit-0.1.38 → kostyl_toolkit-0.1.39}/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_toolkit-0.1.38 → kostyl_toolkit-0.1.39}/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 → kostyl_toolkit-0.1.39}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "kostyl-toolkit"
-version = "0.1.38"
+version = "0.1.39"
 description = "Kickass Orchestration System for Training, Yielding & Logging "
 readme = "README.md"
 requires-python = ">=3.12"
@@ -10,10 +10,8 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
-ml-core = [
+ml = [
     "case-converter>=1.2.0",
-    "clearml[s3]>=2.0.2",
-    "lightning>=2.5.6",
     "pydantic>=2.12.4",
     "torch>=2.9.1",
     "transformers>=4.57.1",
@@ -30,7 +28,7 @@ dev = [
     "pyarrow>=22.0.0",
 ]
 
-ml-core = [
+ml = [
     "case-converter>=1.2.0",
     "clearml[s3]>=2.0.2",
     "lightning>=2.5.6",

kostyl_toolkit-0.1.38/kostyl/ml/configs/hyperparams.py

@@ -1,94 +0,0 @@
-from typing import Literal
-
-from pydantic import BaseModel
-from pydantic import Field
-from pydantic import model_validator
-
-from kostyl.utils.logging import setup_logger
-
-
-logger = setup_logger(fmt="only_message")
-
-
-class AdamConfig(BaseModel):
-    """AdamW optimizer hyperparameters configuration."""
-
-    type: Literal["AdamW"] = "AdamW"
-    betas: tuple[float, float] = (0.9, 0.999)
-    is_adamw: bool = True
-
-
-class AdamWithPrecisionConfig(BaseModel):
-    """Adam optimizer with low-precision hyperparameters configuration."""
-
-    type: Literal["Adam8bit", "Adam4bit", "AdamFp8"]
-    betas: tuple[float, float] = (0.9, 0.999)
-    block_size: int
-    bf16_stochastic_round: bool = False
-    is_adamw: bool = True
-
-
-Optimizer = AdamConfig | AdamWithPrecisionConfig
-
-
-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
-    )
-    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)
-
-    @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
-            logger.warning(
-                "use_scheduler is False, warmup_value and warmup_iters_ratio will be ignored."
-            )
-            self.warmup_value = None
-            self.warmup_iters_ratio = None
-        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):
-            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.")
-        return self
-
-
-class HyperparamsConfig(BaseModel):
-    """Model training hyperparameters configuration."""
-
-    grad_clip_val: float | None = Field(default=None, gt=0, validate_default=False)
-    optimizer: Optimizer
-    lr: Lr
-    weight_decay: WeightDecay