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

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

@@ -1,17 +1,15 @@
 Metadata-Version: 2.3
 Name: kostyl-toolkit
-Version: 0.1.38
+Version: 0.1.40
 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.40}/kostyl/ml/configs/__init__.py

@@ -1,6 +1,10 @@
+from .hyperparams import OPTIMIZER_CONFIG
+from .hyperparams import AdamConfig
+from .hyperparams import AdamWithPrecisionConfig
 from .hyperparams import HyperparamsConfig
 from .hyperparams import Lr
-from .hyperparams import Optimizer
+from .hyperparams import MuonConfig
+from .hyperparams import ScheduledParamConfig
 from .hyperparams import WeightDecay
 from .mixins import ConfigLoadingMixin
 from .training_settings import CheckpointConfig
@@ -14,6 +18,9 @@ from .training_settings import TrainingSettings
 
 
 __all__ = [
+    "OPTIMIZER_CONFIG",
+    "AdamConfig",
+    "AdamWithPrecisionConfig",
     "CheckpointConfig",
     "ConfigLoadingMixin",
     "DDPStrategyConfig",
@@ -23,7 +30,8 @@ __all__ = [
     "HyperparamsConfig",
     "LightningTrainerParameters",
     "Lr",
-    "Optimizer",
+    "MuonConfig",
+    "ScheduledParamConfig",
     "SingleDeviceStrategyConfig",
     "TrainingSettings",
     "WeightDecay",

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

@@ -0,0 +1,126 @@
+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"]
+    momentum: float = 0.95
+    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", "AdamW8bit", "AdamW4bit", "AdamWFp8"
+    ]
+    betas: tuple[float, float] = (0.9, 0.999)
+    block_size: int
+    bf16_stochastic_round: bool = False
+
+
+OPTIMIZER_CONFIG = AdamConfig | AdamWithPrecisionConfig | MuonConfig
+SCHEDULER = Literal[
+    "linear",
+    "cosine",
+    "plateau-with-cosine-annealing",
+    "plateau-with-linear-annealing",
+]
+
+
+class ScheduledParamConfig(BaseModel):
+    """Base configuration for a scheduled hyperparameter."""
+
+    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) -> "ScheduledParamConfig":
+        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) -> "ScheduledParamConfig":
+        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) -> "ScheduledParamConfig":
+        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) -> "ScheduledParamConfig":
+        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 Lr(ScheduledParamConfig):
+    """Learning rate hyperparameters configuration."""
+
+
+class WeightDecay(ScheduledParamConfig):
+    """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_CONFIG
+    lr: Lr
+    weight_decay: WeightDecay

kostyl_toolkit-0.1.40/kostyl/ml/integrations/clearml/__init__.py

@@ -0,0 +1,29 @@
+try:
+    import clearml  # noqa: F401
+except ImportError as e:
+    raise ImportError(
+        "ClearML integration requires the 'clearml' package. "
+        "Please install it via 'pip install clearml'."
+    ) from e
+from .checkpoint_uploader import ClearMLCheckpointUploader
+from .config_mixin import ConfigSyncingClearmlMixin
+from .dataset_utils import collect_clearml_datasets
+from .dataset_utils import download_clearml_datasets
+from .dataset_utils import get_datasets_paths
+from .loading_utils import load_model_from_clearml
+from .loading_utils import load_tokenizer_from_clearml
+from .version_utils import find_version_in_tags
+from .version_utils import increment_version
+
+
+__all__ = [
+    "ClearMLCheckpointUploader",
+    "ConfigSyncingClearmlMixin",
+    "collect_clearml_datasets",
+    "download_clearml_datasets",
+    "find_version_in_tags",
+    "get_datasets_paths",
+    "increment_version",
+    "load_model_from_clearml",
+    "load_tokenizer_from_clearml",
+]

{kostyl_toolkit-0.1.38 → kostyl_toolkit-0.1.40}/kostyl/ml/integrations/clearml/checkpoint_uploader.py

@@ -1,5 +1,5 @@
 from collections.abc import Callable
-from functools import partial
+from datetime import datetime
 from pathlib import Path
 from typing import override
 
@@ -24,7 +24,7 @@ class ClearMLCheckpointUploader(ModelCheckpointUploader):
         comment: str | None = None,
         framework: str | None = None,
         base_model_id: str | None = None,
-        new_model_per_upload: bool = True,
+        upload_as_new_model: bool = True,
         verbose: bool = True,
     ) -> None:
         """
@@ -38,20 +38,22 @@ class ClearMLCheckpointUploader(ModelCheckpointUploader):
             comment: A comment / description for the model.
             framework: The framework of the model (e.g., "PyTorch", "TensorFlow").
             base_model_id: Optional ClearML model ID to use as a base for the new model
-            new_model_per_upload: Whether to create a new ClearML model
-                for every upload or update weights of the same model. When updating weights,
-                the last uploaded checkpoint will be replaced (and deleted).
+            upload_as_new_model: Whether to create a new ClearML model
+                for every upload or update weights of the same model. When True,
+                each checkpoint is uploaded as a separate model with timestamp added to the name.
+                When False, weights of the same model are updated.
             verbose: Whether to log messages during upload.
 
         """
         super().__init__()
-        if base_model_id is not None and new_model_per_upload:
+        if base_model_id is not None and upload_as_new_model:
             raise ValueError(
-                "Cannot set base_model_id when new_model_per_upload is True."
+                "Cannot set base_model_id when upload_as_new_model is True."
             )
 
         self.verbose = verbose
-        self.new_model_per_upload = new_model_per_upload
+        self.upload_as_new_model = upload_as_new_model
+        self.model_name = model_name
         self.best_model_path: str = ""
         self.config_dict = config_dict
         self._output_model: OutputModel | None = None
@@ -59,15 +61,13 @@ class ClearMLCheckpointUploader(ModelCheckpointUploader):
         self._upload_callback: Callable | None = None
 
         self._validate_tags(tags)
-        self.model_fabric = partial(
-            OutputModel,
-            name=model_name,
-            label_enumeration=label_enumeration,
-            tags=tags,
-            comment=comment,
-            framework=framework,
-            base_model_id=base_model_id,
-        )
+        self.model_fabric_kwargs = {
+            "label_enumeration": label_enumeration,
+            "tags": tags,
+            "comment": comment,
+            "framework": framework,
+            "base_model_id": base_model_id,
+        }
         return
 
     @staticmethod
@@ -78,16 +78,22 @@ class ClearMLCheckpointUploader(ModelCheckpointUploader):
             tags.append("LightningCheckpoint")
         return None
 
-    @property
-    def output_model_(self) -> OutputModel:
-        """Returns the OutputModel instance based on `new_model_per_upload` setting."""
-        if self.new_model_per_upload:
-            model = self.model_fabric()
-            self._output_model = self.model_fabric()
-        else:
-            if self._output_model is None:
-                self._output_model = self.model_fabric()
-            model = self._output_model
+    def _create_new_model(self) -> OutputModel:
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        model_name_with_timestamp = f"{self.model_name}_{timestamp}"
+        model = OutputModel(
+            name=model_name_with_timestamp,
+            **self.model_fabric_kwargs,
+        )
+        return model
+
+    def _get_output_model(self) -> OutputModel:
+        if self._output_model is None:
+            self._output_model = OutputModel(
+                name=self.model_name,
+                **self.model_fabric_kwargs,
+            )
+        model = self._output_model
         return model
 
     @override
@@ -105,12 +111,17 @@ class ClearMLCheckpointUploader(ModelCheckpointUploader):
         if self.verbose:
             logger.info(f"Uploading model from {path}")
 
-        self.output_model_.update_weights(
+        if self.upload_as_new_model:
+            output_model = self._create_new_model()
+        else:
+            output_model = self._get_output_model()
+
+        output_model.update_weights(
             path,
             auto_delete_file=False,
             async_enable=False,
         )
-        self.output_model_.update_design(config_dict=self.config_dict)
+        output_model.update_design(config_dict=self.config_dict)
 
         self._last_uploaded_model_path = path
         return

{kostyl_toolkit-0.1.38 → kostyl_toolkit-0.1.40}/kostyl/ml/integrations/clearml/config_mixin.py

@@ -10,7 +10,7 @@ from kostyl.utils.dict_manipulations import flattened_dict_to_nested
 from kostyl.utils.fs import load_config
 
 
-class BaseModelWithClearmlSyncing[TConfig: ConfigLoadingMixin]:
+class ConfigSyncingClearmlMixin[TConfig: ConfigLoadingMixin]:
     """Mixin providing ClearML task configuration syncing functionality for Pydantic models."""
 
     @classmethod

{kostyl_toolkit-0.1.38 → kostyl_toolkit-0.1.40}/kostyl/ml/integrations/clearml/loading_utils.py

@@ -31,7 +31,7 @@ except ImportError:
     LIGHTING_MIXIN_AVAILABLE = False
 
 
-def get_tokenizer_from_clearml(
+def load_tokenizer_from_clearml(
     model_id: str,
     task: Task | None = None,
     ignore_remote_overrides: bool = True,
@@ -66,7 +66,7 @@ def get_tokenizer_from_clearml(
     return tokenizer, clearml_tokenizer
 
 
-def get_model_from_clearml[
+def load_model_from_clearml[
     TModel: PreTrainedModel | LightningCheckpointLoaderMixin | AutoModel
 ](
     model_id: str,

{kostyl_toolkit-0.1.38 → kostyl_toolkit-0.1.40}/kostyl/ml/integrations/lightning/callbacks/checkpoint.py

@@ -15,7 +15,7 @@ from kostyl.ml.dist_utils import is_local_zero_rank
 from kostyl.utils import setup_logger
 
 
-logger = setup_logger("callbacks/checkpoint.py")
+logger = setup_logger()
 
 
 class ModelCheckpointWithCheckpointUploader(ModelCheckpoint):
@@ -278,6 +278,10 @@ class ModelCheckpointWithCheckpointUploader(ModelCheckpoint):
                 case "only-best":
                     if filepath == self.best_model_path:
                         self.registry_uploader.upload_checkpoint(filepath)
+                case _:
+                    logger.warning_once(
+                        "Unknown upload strategy for checkpoint uploader. Skipping upload."
+                    )
         return
 
 
@@ -286,7 +290,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.40}/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.40}/kostyl/ml/integrations/lightning/module.py

@@ -14,7 +14,7 @@ from transformers import PretrainedConfig
 from transformers import PreTrainedModel
 
 from kostyl.ml.integrations.lightning.metrics_formatting import apply_suffix
-from kostyl.ml.schedulers.base import BaseScheduler
+from kostyl.ml.optim.schedulers import BaseScheduler
 from kostyl.utils import setup_logger
 
 

{kostyl_toolkit-0.1.38 → kostyl_toolkit-0.1.40}/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.40/kostyl/ml/optim/__init__.py

@@ -0,0 +1,8 @@
+from .factory import create_optimizer
+from .factory import create_scheduler
+
+
+__all__ = [
+    "create_optimizer",
+    "create_scheduler",
+]

kostyl_toolkit-0.1.40/kostyl/ml/optim/factory.py

@@ -0,0 +1,257 @@
+from typing import Any
+
+from torch.optim import Optimizer
+
+from kostyl.ml.configs import OPTIMIZER_CONFIG
+from kostyl.ml.configs import AdamConfig
+from kostyl.ml.configs import AdamWithPrecisionConfig
+from kostyl.ml.configs import MuonConfig
+from kostyl.ml.configs import ScheduledParamConfig
+from kostyl.utils import setup_logger
+
+from .schedulers import SCHEDULER_MAPPING
+from .schedulers import CosineScheduler
+from .schedulers import LinearScheduler
+from .schedulers import PlateauWithAnnealingScheduler
+
+
+logger = setup_logger(fmt="only_message")
+
+
+def create_scheduler(
+    config: ScheduledParamConfig,
+    param_group_field: str,
+    num_iters: int,
+    optim: Optimizer,
+    multiplier_field: str | None = None,
+    skip_if_zero: bool = False,
+    apply_if_field: str | None = None,
+    ignore_if_field: str | None = None,
+) -> LinearScheduler | CosineScheduler | PlateauWithAnnealingScheduler:
+    """
+    Converts a ScheduledParamConfig to a scheduler instance.
+
+    Args:
+        config: Configuration object for the scheduler.
+        param_group_field: The field name in the optimizer's param groups to schedule.
+        num_iters: Total number of iterations.
+        optim: The optimizer instance.
+        multiplier_field: Optional per-group field name that contains a multiplier applied to the scheduled value. If None, no multiplier is applied.
+        skip_if_zero: Leave groups untouched when their target field equals zero.
+            Default is False.
+        apply_if_field: Require this key to be present in a param group before updating.
+        ignore_if_field: Skip groups that declare this key in their dictionaries.
+
+    Returns:
+        A scheduler instance based on the configuration.
+
+    """
+    if config.scheduler_type is None:
+        raise ValueError("scheduler_type must be specified in the config.")
+
+    if "plateau" in config.scheduler_type:
+        scheduler_type = "plateau"
+    else:
+        scheduler_type = config.scheduler_type
+    scheduler_cls = SCHEDULER_MAPPING[scheduler_type]  # type: ignore
+
+    if issubclass(scheduler_cls, PlateauWithAnnealingScheduler):
+        if "cosine" in config.scheduler_type:
+            annealing_type = "cosine"
+        elif "linear" in config.scheduler_type:
+            annealing_type = "linear"
+        else:
+            raise ValueError(f"Unknown annealing_type: {config.scheduler_type}")
+        scheduler = scheduler_cls(
+            optimizer=optim,
+            param_group_field=param_group_field,
+            num_iters=num_iters,
+            plateau_value=config.base_value,
+            final_value=config.final_value,  # type: ignore
+            warmup_ratio=config.warmup_ratio,
+            warmup_value=config.warmup_value,
+            freeze_ratio=config.freeze_ratio,
+            plateau_ratio=config.plateau_ratio,  # type: ignore
+            annealing_type=annealing_type,
+            multiplier_field=multiplier_field,
+            skip_if_zero=skip_if_zero,
+            apply_if_field=apply_if_field,
+            ignore_if_field=ignore_if_field,
+        )
+    elif issubclass(scheduler_cls, LinearScheduler):
+        scheduler = scheduler_cls(
+            optimizer=optim,
+            param_group_field=param_group_field,
+            num_iters=num_iters,
+            initial_value=config.base_value,
+            final_value=config.final_value,  # type: ignore
+            multiplier_field=multiplier_field,
+            skip_if_zero=skip_if_zero,
+            apply_if_field=apply_if_field,
+            ignore_if_field=ignore_if_field,
+        )
+    elif issubclass(scheduler_cls, CosineScheduler):
+        scheduler = scheduler_cls(
+            optimizer=optim,
+            param_group_field=param_group_field,
+            num_iters=num_iters,
+            base_value=config.base_value,
+            final_value=config.final_value,  # type: ignore
+            warmup_ratio=config.warmup_ratio,
+            warmup_value=config.warmup_value,
+            freeze_ratio=config.freeze_ratio,
+            multiplier_field=multiplier_field,
+            skip_if_zero=skip_if_zero,
+            apply_if_field=apply_if_field,
+            ignore_if_field=ignore_if_field,
+        )
+    else:
+        raise ValueError(f"Unsupported scheduler type: {config.scheduler_type}")
+    return scheduler
+
+
+def create_optimizer(  # noqa: C901
+    parameters_groups: dict[str, Any],
+    optimizer_config: OPTIMIZER_CONFIG,
+    lr: float,
+    weight_decay: float,
+) -> Optimizer:
+    """
+    Creates an optimizer based on the configuration.
+
+    Args:
+        parameters_groups: Dictionary containing model parameters
+            (key "params" and per-group options, i.e. "lr", "weight_decay" and etc.).
+        optimizer_config: Configuration for the optimizer.
+        lr: Learning rate.
+        weight_decay: Weight decay.
+
+    Returns:
+        An instantiated optimizer.
+
+    """
+    if isinstance(optimizer_config, AdamConfig):
+        match optimizer_config.type:
+            case "Adam":
+                from torch.optim import Adam
+
+                optimizer = Adam(
+                    params=parameters_groups["params"],
+                    lr=lr,
+                    weight_decay=weight_decay,
+                    betas=optimizer_config.betas,
+                )
+
+            case "AdamW":
+                from torch.optim import AdamW
+
+                optimizer = AdamW(
+                    params=parameters_groups["params"],
+                    lr=lr,
+                    weight_decay=weight_decay,
+                    betas=optimizer_config.betas,
+                )
+                return optimizer
+            case _:
+                raise ValueError(f"Unsupported optimizer type: {optimizer_config.type}")
+    elif isinstance(optimizer_config, MuonConfig):
+        from torch.optim import Muon
+
+        optimizer = Muon(
+            params=parameters_groups["params"],
+            lr=lr,
+            weight_decay=weight_decay,
+            momentum=optimizer_config.momentum,
+            nesterov=optimizer_config.nesterov,
+            ns_coefficients=optimizer_config.ns_coefficients,
+            ns_steps=optimizer_config.ns_steps,
+        )
+    elif isinstance(optimizer_config, AdamWithPrecisionConfig):
+        try:
+            import torchao  # noqa: F401
+        except ImportError as e:
+            raise ImportError(
+                "torchao is required for low-precision Adam optimizers. "
+                "Please install it via 'pip install torchao'."
+            ) from e
+        match optimizer_config.type:
+            case "Adam8bit":
+                from torchao.optim import Adam8bit
+
+                logger.warning(
+                    "Ignoring weight_decay for Adam8bit optimizer as it is not supported."
+                )
+
+                optimizer = Adam8bit(
+                    params=parameters_groups["params"],
+                    lr=lr,
+                    betas=optimizer_config.betas,
+                    block_size=optimizer_config.block_size,
+                    bf16_stochastic_round=optimizer_config.bf16_stochastic_round,
+                )
+            case "Adam4bit":
+                from torchao.optim import Adam4bit
+
+                logger.warning(
+                    "Ignoring weight_decay for Adam4bit optimizer as it is not supported."
+                )
+
+                optimizer = Adam4bit(
+                    params=parameters_groups["params"],
+                    lr=lr,
+                    betas=optimizer_config.betas,
+                    block_size=optimizer_config.block_size,
+                    bf16_stochastic_round=optimizer_config.bf16_stochastic_round,
+                )
+            case "AdamFp8":
+                from torchao.optim import AdamFp8
+
+                logger.warning(
+                    "Ignoring weight_decay for AdamFp8 optimizer as it is not supported."
+                )
+
+                optimizer = AdamFp8(
+                    params=parameters_groups["params"],
+                    lr=lr,
+                    betas=optimizer_config.betas,
+                    block_size=optimizer_config.block_size,
+                    bf16_stochastic_round=optimizer_config.bf16_stochastic_round,
+                )
+            case "AdamW8bit":
+                from torchao.optim import AdamW8bit
+
+                optimizer = AdamW8bit(
+                    params=parameters_groups["params"],
+                    lr=lr,
+                    weight_decay=weight_decay,
+                    betas=optimizer_config.betas,
+                    block_size=optimizer_config.block_size,
+                    bf16_stochastic_round=optimizer_config.bf16_stochastic_round,
+                )
+            case "AdamW4bit":
+                from torchao.optim import AdamW4bit
+
+                optimizer = AdamW4bit(
+                    params=parameters_groups["params"],
+                    lr=lr,
+                    weight_decay=weight_decay,
+                    betas=optimizer_config.betas,
+                    block_size=optimizer_config.block_size,
+                    bf16_stochastic_round=optimizer_config.bf16_stochastic_round,
+                )
+            case "AdamWFp8":
+                from torchao.optim import AdamWFp8
+
+                optimizer = AdamWFp8(
+                    params=parameters_groups["params"],
+                    lr=lr,
+                    weight_decay=weight_decay,
+                    betas=optimizer_config.betas,
+                    block_size=optimizer_config.block_size,
+                    bf16_stochastic_round=optimizer_config.bf16_stochastic_round,
+                )
+            case _:
+                raise ValueError(f"Unsupported optimizer type: {optimizer_config.type}")
+    else:
+        raise ValueError("Unsupported optimizer configuration type.")
+    return optimizer

kostyl_toolkit-0.1.40/kostyl/ml/optim/schedulers/__init__.py

@@ -0,0 +1,56 @@
+from typing import TypedDict
+
+from .base import BaseScheduler
+from .composite import CompositeScheduler
+from .cosine import CosineParamScheduler
+from .cosine import CosineScheduler
+from .linear import LinearParamScheduler
+from .linear import LinearScheduler
+from .plateau import PlateauWithAnnealingParamScheduler
+from .plateau import PlateauWithAnnealingScheduler
+
+
+class SchedulerMapping(TypedDict):
+    """Map names to scheduler classes."""
+
+    linear: type[LinearScheduler]
+    cosine: type[CosineScheduler]
+    plateau: type[PlateauWithAnnealingScheduler]
+    composite: type[CompositeScheduler]
+
+
+class ParamSchedulerMapping(TypedDict):
+    """Map names to scheduler classes."""
+
+    linear: type[LinearParamScheduler]
+    cosine: type[CosineParamScheduler]
+    plateau: type[PlateauWithAnnealingParamScheduler]
+
+
+SCHEDULER_MAPPING: SchedulerMapping = {
+    "linear": LinearScheduler,
+    "cosine": CosineScheduler,
+    "plateau": PlateauWithAnnealingScheduler,
+    "composite": CompositeScheduler,
+}
+
+
+PARAM_SCHEDULER_MAPPING: ParamSchedulerMapping = {
+    "linear": LinearParamScheduler,
+    "cosine": CosineParamScheduler,
+    "plateau": PlateauWithAnnealingParamScheduler,
+}
+
+
+__all__ = [
+    "PARAM_SCHEDULER_MAPPING",
+    "SCHEDULER_MAPPING",
+    "BaseScheduler",
+    "CompositeScheduler",
+    "CosineParamScheduler",
+    "CosineScheduler",
+    "LinearParamScheduler",
+    "LinearScheduler",
+    "PlateauWithAnnealingParamScheduler",
+    "PlateauWithAnnealingScheduler",
+]

{kostyl_toolkit-0.1.38/kostyl/ml → kostyl_toolkit-0.1.40/kostyl/ml/optim}/schedulers/cosine.py

@@ -145,8 +145,8 @@ class CosineScheduler(_CosineSchedulerCore):
             freeze_ratio: Optional fraction of iterations to keep the value frozen at zero at the beginning.
             multiplier_field: Optional per-group multiplier applied to the scheduled value.
             skip_if_zero: Leave groups untouched when their target field equals zero.
-            apply_if_field: Require this flag to be present in a param group before updating.
-            ignore_if_field: Skip groups that declare this flag.
+            apply_if_field: Require this key to be present in a param group before updating.
+            ignore_if_field: Skip groups that declare this key in their dictionaries.
 
         """
         self.apply_if_field = apply_if_field

{kostyl_toolkit-0.1.38/kostyl/ml → kostyl_toolkit-0.1.40/kostyl/ml/optim}/schedulers/linear.py

@@ -13,21 +13,21 @@ class _LinearScheduleBase(BaseScheduler):
         self,
         param_name: str,
         num_iters: int,
-        start_value: float,
+        initial_value: float,
         final_value: float,
     ) -> None:
         self.param_name = param_name
         self.num_iters = num_iters
-        self.start_value = start_value
+        self.initial_value = initial_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.initial_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.initial_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,
+        initial_value: float,
         final_value: float,
         multiplier_field: str | None = None,
         skip_if_zero: bool = False,
@@ -82,12 +82,12 @@ 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.
+            initial_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.
-            apply_if_field: Require this flag to be present in a param group before updating.
-            ignore_if_field: Skip groups that declare this flag.
+            apply_if_field: Require this key to be present in a param group before updating.
+            ignore_if_field: Skip groups that declare this key in their dictionaries.
 
         """
         self.apply_if_field = apply_if_field
@@ -98,7 +98,7 @@ class LinearScheduler(_LinearScheduleBase):
         super().__init__(
             param_name=param_group_field,
             num_iters=num_iters,
-            start_value=start_value,
+            initial_value=initial_value,
             final_value=final_value,
         )
         self.param_group_field = param_group_field

{kostyl_toolkit-0.1.38/kostyl/ml → kostyl_toolkit-0.1.40/kostyl/ml/optim}/schedulers/plateau.py

@@ -198,8 +198,8 @@ class PlateauWithAnnealingScheduler(_PlateauWithAnnealingCore):
             annealing_type: Type of annealing from plateau to final value ("cosine" or "linear").
             multiplier_field: Optional per-group multiplier applied to the scheduled value.
             skip_if_zero: Leave groups untouched when their target field equals zero.
-            apply_if_field: Require this flag to be present in a param group before updating.
-            ignore_if_field: Skip groups that declare this flag.
+            apply_if_field: Require this key to be present in a param group before updating.
+            ignore_if_field: Skip groups that declare this key in their dictionaries.
 
         """
         self.apply_if_field = apply_if_field

{kostyl_toolkit-0.1.38 → kostyl_toolkit-0.1.40}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "kostyl-toolkit"
-version = "0.1.38"
+version = "0.1.40"
 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,18 +28,19 @@ dev = [
     "pyarrow>=22.0.0",
 ]
 
-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",
+    "torchao>=0.15.0",
     "transformers>=4.57.1",
 ]
 
 
 [tool.uv]
-default-groups = ["dev", "ml-core"]
+default-groups = ["dev", "ml"]
 
 [build-system]
 requires = ["uv_build>=0.9.9,<0.10.0"]

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

kostyl_toolkit-0.1.38/kostyl/ml/integrations/clearml/__init__.py

@@ -1,7 +0,0 @@
-try:
-    import clearml  # noqa: F401
-except ImportError as e:
-    raise ImportError(
-        "ClearML integration requires the 'clearml' package. "
-        "Please install it via 'pip install clearml'."
-    ) from e

kostyl_toolkit-0.1.38/kostyl/ml/schedulers/__init__.py

@@ -1,18 +0,0 @@
-from .composite import CompositeScheduler
-from .cosine import CosineParamScheduler
-from .cosine import CosineScheduler
-from .linear import LinearParamScheduler
-from .linear import LinearScheduler
-from .plateau import PlateauWithAnnealingParamScheduler
-from .plateau import PlateauWithAnnealingScheduler
-
-
-__all__ = [
-    "CompositeScheduler",
-    "CosineParamScheduler",
-    "CosineScheduler",
-    "LinearParamScheduler",
-    "LinearScheduler",
-    "PlateauWithAnnealingParamScheduler",
-    "PlateauWithAnnealingScheduler",
-]