kostyl-toolkit 0.1.35__tar.gz → 0.1.37__tar.gz

{kostyl_toolkit-0.1.35 → kostyl_toolkit-0.1.37}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: kostyl-toolkit
-Version: 0.1.35
+Version: 0.1.37
 Summary: Kickass Orchestration System for Training, Yielding & Logging 
 Requires-Dist: case-converter>=1.2.0
 Requires-Dist: loguru>=0.7.3

{kostyl_toolkit-0.1.35 → kostyl_toolkit-0.1.37}/kostyl/ml/configs/hyperparams.py

@@ -1,3 +1,5 @@
+from typing import Literal
+
 from pydantic import BaseModel
 from pydantic import Field
 from pydantic import model_validator
@@ -8,11 +10,25 @@ from kostyl.utils.logging import setup_logger
 logger = setup_logger(fmt="only_message")
 
 
-class Optimizer(BaseModel):
-    """Optimizer hyperparameters configuration."""
+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
+
 
-    adamw_beta1: float = 0.9
-    adamw_beta2: float = 0.999
+Optimizer = AdamConfig | AdamWithPrecisionConfig
 
 
 class Lr(BaseModel):
@@ -73,6 +89,6 @@ class HyperparamsConfig(BaseModel):
     """Model training hyperparameters configuration."""
 
     grad_clip_val: float | None = Field(default=None, gt=0, validate_default=False)
-    optimizer: Optimizer = Optimizer()
+    optimizer: Optimizer
     lr: Lr
     weight_decay: WeightDecay

{kostyl_toolkit-0.1.35 → kostyl_toolkit-0.1.37}/kostyl/ml/configs/training_settings.py

@@ -25,21 +25,31 @@ PRECISION = Literal[
     "16",
     "bf16",
 ]
+DTYPE = Literal["float32", "float16", "bfloat16", "float64"]
+
+
+class SingleDeviceStrategyConfig(BaseModel):
+    """Single device strategy configuration."""
+
+    type: Literal["single_device"]
 
 
 class FSDP1StrategyConfig(BaseModel):
     """Fully Sharded Data Parallel (FSDP) strategy configuration."""
 
     type: Literal["fsdp1"]
-    param_dtype: Literal["float32", "float16", "bfloat16"]
-    reduce_dtype: Literal["float32", "float16", "bfloat16"]
-    buffer_dtype: Literal["float32", "float16", "bfloat16"]
+    param_dtype: DTYPE | None
+    reduce_dtype: DTYPE | None
+    buffer_dtype: DTYPE | None
 
 
-class SingleDeviceStrategyConfig(BaseModel):
-    """Single device strategy configuration."""
+class FSDP2StrategyConfig(BaseModel):
+    """Fully Sharded Data Parallel (FSDP) strategy configuration."""
 
-    type: Literal["single_device"]
+    type: Literal["fsdp2"]
+    param_dtype: DTYPE | None
+    reduce_dtype: DTYPE | None
+    buffer_dtype: DTYPE | None
 
 
 class DDPStrategyConfig(BaseModel):
@@ -82,6 +92,7 @@ class CheckpointConfig(BaseModel):
     monitor: str = "val_loss"
     mode: str = "min"
     filename: str = "{epoch:02d}-{val_loss:.2f}"
+    save_weights_only: bool = True
 
 
 class DataConfig(BaseModel):

kostyl_toolkit-0.1.37/kostyl/ml/dist_utils.py

@@ -0,0 +1,129 @@
+import math
+import os
+from typing import Literal
+
+import torch.distributed as dist
+
+from kostyl.utils.logging import KostylLogger
+from kostyl.utils.logging import setup_logger
+
+
+module_logger = setup_logger()
+
+
+def log_dist(
+    msg: str,
+    logger: KostylLogger | None = None,
+    level: Literal["info", "warning", "error", "warning_once", "debug"] = "info",
+    log_scope: Literal["only-zero-rank", "world"] = "world",
+    group: dist.ProcessGroup | None = None,
+) -> None:
+    """
+    Log a message in a distributed environment based on the specified verbosity level.
+
+    Args:
+        msg (str): The message to log.
+        log_scope (Literal["only-zero-rank", "world"]): The verbosity level for logging.
+            - "only-zero-rank": Log only from the main process (rank 0).
+            - "world": Log from all processes in the distributed environment.
+        logger (KostylLogger | None): The logger instance to use. If None, the module logger is used.
+        level (Literal["info", "warning", "error", "warning_once", "debug"]): The logging level.
+        group (dist.ProcessGroup | None): Optional process group used to determine ranks. Defaults to the global process group.
+
+    """
+    if logger is None:
+        logger = module_logger
+
+    log_attr = getattr(logger, level, None)
+    if log_attr is None:
+        raise ValueError(f"Invalid logging level: {level}")
+
+    if not dist.is_initialized():
+        module_logger.warning_once(
+            "Distributed process group is not initialized; logging from all ranks."
+        )
+        log_attr(msg)
+        return
+
+    match log_scope:
+        case "only-zero-rank":
+            if group is None:
+                module_logger.debug(
+                    "No process group provided; assuming global group for rank check."
+                )
+                group = dist.group.WORLD
+            group_rank = dist.get_rank(group=group)
+            if dist.get_global_rank(group=group, group_rank=group_rank) == 0:  # pyright: ignore[reportArgumentType]
+                log_attr(msg)
+        case "world":
+            log_attr(msg)
+        case _:
+            raise ValueError(f"Invalid logging verbosity level: {log_scope}")
+    return
+
+
+def scale_lrs_by_world_size(
+    lrs: dict[str, float],
+    group: dist.ProcessGroup | None = None,
+    config_name: str = "",
+    inv_scale: bool = False,
+    verbose_level: Literal["only-zero-rank", "world"] | None = None,
+) -> dict[str, float]:
+    """
+    Scale learning-rate configuration values to match the active distributed world size.
+
+    Note:
+        The value in the `lrs` will be modified in place.
+
+    Args:
+        lrs (dict[str, float]): A dictionary of learning rate names and their corresponding values to be scaled.
+        group (dist.ProcessGroup | None): Optional process group used to determine
+            the target world size. Defaults to the global process group.
+        config_name (str): Human-readable identifier included in log messages.
+        inv_scale (bool): If True, use the inverse square-root scale factor.
+        verbose_level (Literal["only-zero-rank", "world"] | None): Verbosity level for logging scaled values.
+            - "only-zero-rank": Log only from the main process (rank 0).
+            - "world": Log from all processes in the distributed environment.
+            -  None: No logging.
+
+    Returns:
+        dict[str, float]: The learning-rate configuration with scaled values.
+
+    """
+    world_size = dist.get_world_size(group=group)
+
+    if inv_scale:
+        scale = 1 / math.sqrt(world_size)
+    else:
+        scale = math.sqrt(world_size)
+
+    for name, value in lrs.items():
+        old_value = value
+        new_value = value * scale
+        if verbose_level is not None:
+            log_dist(
+                f"New {config_name} lr {name.upper()}: {new_value}; OLD: {old_value}",
+                log_scope=verbose_level,
+                group=group,
+            )
+        lrs[name] = new_value
+    return lrs
+
+
+def get_local_rank(group: dist.ProcessGroup | None = None) -> int:
+    """Gets the local rank of the current process in a distributed setting."""
+    if dist.is_initialized() and group is not None:
+        return dist.get_rank(group=group)
+    if "SLURM_LOCALID" in os.environ:
+        return int(os.environ["SLURM_LOCALID"])
+    if "LOCAL_RANK" in os.environ:
+        return int(os.environ["LOCAL_RANK"])
+    return 0
+
+
+def is_local_zero_rank() -> bool:
+    """Checks if the current process is the main process (rank 0) for the local node in a distributed setting."""
+    rank = get_local_rank()
+    if rank != 0:
+        return False
+    return True

{kostyl_toolkit-0.1.35 → kostyl_toolkit-0.1.37}/kostyl/ml/lightning/callbacks/checkpoint.py

@@ -10,7 +10,7 @@ from lightning.fabric.utilities.types import _PATH
 from lightning.pytorch.callbacks import ModelCheckpoint
 
 from kostyl.ml.configs import CheckpointConfig
-from kostyl.ml.dist_utils import is_main_process
+from kostyl.ml.dist_utils import is_local_zero_rank
 from kostyl.ml.lightning import KostylLightningModule
 from kostyl.ml.registry_uploader import RegistryUploaderCallback
 from kostyl.utils import setup_logger
@@ -299,9 +299,9 @@ class ModelCheckpointWithRegistryUploader(ModelCheckpoint):
 def setup_checkpoint_callback(
     dirpath: Path,
     ckpt_cfg: CheckpointConfig,
-    save_weights_only: bool = True,
     registry_uploader_callback: RegistryUploaderCallback | None = None,
     uploading_strategy: Literal["only-best", "every-checkpoint"] | None = None,
+    remove_folder_if_exists: bool = True,
 ) -> ModelCheckpointWithRegistryUploader | ModelCheckpoint:
     """
     Create and configure a checkpoint callback for model saving.
@@ -313,14 +313,13 @@ def setup_checkpoint_callback(
     Args:
         dirpath: Path to the directory for saving checkpoints.
         ckpt_cfg: Checkpoint configuration (filename, monitor, mode, save_top_k).
-        save_weights_only: If True, only model weights are saved without optimizer and lr-scheduler state.
-            Defaults to True.
         registry_uploader_callback: Optional callback for uploading checkpoints to a remote registry.
             Must be specified together with uploading_strategy.
         uploading_strategy: Checkpoint upload mode:
             - "only-best": only the best checkpoint is uploaded
             - "every-checkpoint": every saved checkpoint is uploaded
             Must be specified together with registry_uploader_callback.
+        remove_folder_if_exists: If True, removes existing checkpoint directory before creating a new one.
 
     Returns:
         ModelCheckpointWithRegistryUploader if registry_uploader_callback is provided,
@@ -331,7 +330,7 @@ def setup_checkpoint_callback(
 
     Note:
         If the dirpath directory already exists, it will be removed and recreated
-        (only on the main process in distributed training).
+        (only on the main process in distributed training) if remove_folder_if_exists is True.
 
     """
     if (registry_uploader_callback is None) != (uploading_strategy is None):
@@ -340,10 +339,11 @@ def setup_checkpoint_callback(
         )
 
     if dirpath.exists():
-        if is_main_process():
+        if is_local_zero_rank():
             logger.warning(f"Checkpoint directory {dirpath} already exists.")
-            rmtree(dirpath)
-            logger.warning(f"Removed existing checkpoint directory {dirpath}.")
+            if remove_folder_if_exists:
+                rmtree(dirpath)
+                logger.warning(f"Removed existing checkpoint directory {dirpath}.")
     else:
         logger.info(f"Creating checkpoint directory {dirpath}.")
         dirpath.mkdir(parents=True, exist_ok=True)
@@ -356,7 +356,7 @@ def setup_checkpoint_callback(
             monitor=ckpt_cfg.monitor,
             mode=ckpt_cfg.mode,
             verbose=True,
-            save_weights_only=save_weights_only,
+            save_weights_only=ckpt_cfg.save_weights_only,
             registry_uploader_callback=registry_uploader_callback,
             uploading_mode=uploading_strategy,
         )
@@ -368,6 +368,6 @@ def setup_checkpoint_callback(
             monitor=ckpt_cfg.monitor,
             mode=ckpt_cfg.mode,
             verbose=True,
-            save_weights_only=save_weights_only,
+            save_weights_only=ckpt_cfg.save_weights_only,
         )
     return checkpoint_callback

{kostyl_toolkit-0.1.35 → kostyl_toolkit-0.1.37}/kostyl/ml/lightning/extensions/custom_module.py

@@ -26,11 +26,6 @@ module_logger = setup_logger(fmt="only_message")
 class KostylLightningModule(L.LightningModule):
     """Custom PyTorch Lightning Module with logging, checkpointing, and distributed training utilities."""
 
-    @property
-    def process_group(self) -> ProcessGroup | None:
-        """Returns the data parallel process group for distributed training."""
-        return self.get_process_group()
-
     def get_process_group(self) -> ProcessGroup | None:
         """
         Retrieves the data parallel process group for distributed training.

{kostyl_toolkit-0.1.35 → kostyl_toolkit-0.1.37}/kostyl/ml/lightning/extensions/pretrained_model.py

@@ -12,12 +12,12 @@ from kostyl.utils.logging import setup_logger
 logger = setup_logger("LightningPretrainedModelMixin", fmt="only_message")
 
 
-class LightningCheckpointLoaderMixin(PreTrainedModel):
+class LightningCheckpointLoaderMixin:
     """A mixin class for loading pretrained models from PyTorch Lightning checkpoints."""
 
     @classmethod
-    def from_lightning_checkpoint[TModelInstance: LightningCheckpointLoaderMixin](  # noqa: C901
-        cls: type[TModelInstance],
+    def from_lightning_checkpoint[TModelInstance: PreTrainedModel](  # noqa: C901
+        cls: type[TModelInstance],  # pyright: ignore[reportGeneralTypeIssues]
         checkpoint_path: str | Path,
         config_key: str = "config",
         weights_prefix: str | None = "model.",
@@ -78,7 +78,7 @@ class LightningCheckpointLoaderMixin(PreTrainedModel):
             mmap=True,
         )
 
-        # 1. Восстанавливаем конфиг
+        # Load config
         config_cls = cast(type[PretrainedConfig], cls.config_class)
         config_dict = checkpoint_dict[config_key]
         config_dict.update(kwargs)
@@ -91,6 +91,7 @@ class LightningCheckpointLoaderMixin(PreTrainedModel):
 
         raw_state_dict: dict[str, torch.Tensor] = checkpoint_dict["state_dict"]
 
+        # Handle weights prefix
         if weights_prefix:
             if not weights_prefix.endswith("."):
                 weights_prefix = weights_prefix + "."
@@ -117,6 +118,7 @@ class LightningCheckpointLoaderMixin(PreTrainedModel):
         else:
             state_dict = raw_state_dict
 
+        # Instantiate model and load state dict
         model = cls.from_pretrained(
             pretrained_model_name_or_path=None,
             config=config,

{kostyl_toolkit-0.1.35 → kostyl_toolkit-0.1.37}/kostyl/ml/lightning/loggers/tb_logger.py

@@ -3,7 +3,7 @@ from shutil import rmtree
 
 from lightning.pytorch.loggers import TensorBoardLogger
 
-from kostyl.ml.dist_utils import is_main_process
+from kostyl.ml.dist_utils import is_local_zero_rank
 from kostyl.utils.logging import setup_logger
 
 
@@ -15,7 +15,7 @@ def setup_tb_logger(
 ) -> TensorBoardLogger:
     """Sets up a TensorBoardLogger for PyTorch Lightning."""
     if runs_dir.exists():
-        if is_main_process():
+        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}.")

kostyl_toolkit-0.1.37/kostyl/ml/lightning/utils.py

@@ -0,0 +1,58 @@
+from typing import cast
+
+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()
+
+
+def estimate_total_steps(
+    trainer: L.Trainer, dp_process_group: ProcessGroup | None = None
+) -> int:
+    """
+    Estimates the total number of training steps with respect to data parallelism and gradient accumulation.
+
+    Args:
+        trainer: The PyTorch Lightning Trainer instance.
+        dp_process_group: The data parallel process group. If None, the world process group will be used.
+
+    """
+    if dist.is_initialized():
+        world_size = dist.get_world_size(dp_process_group)
+    else:
+        world_size = 1
+
+    datamodule = trainer.datamodule  # type: ignore
+    if datamodule is None:
+        raise ValueError("Trainer must have a datamodule to estimate total steps.")
+    datamodule = cast(L.LightningDataModule, datamodule)
+
+    logger.info("Loading `train_dataloader` to estimate number of stepping batches.")
+    datamodule.setup("fit")
+
+    dataloader_len = len(datamodule.train_dataloader())
+    steps_per_epoch = dataloader_len // trainer.accumulate_grad_batches // world_size
+
+    if trainer.max_epochs is None:
+        raise ValueError("Trainer must have `max_epochs` set to estimate total steps.")
+    total_steps = steps_per_epoch * trainer.max_epochs
+
+    logger.info(
+        f"Total steps: {total_steps} (per-epoch: {steps_per_epoch}) "
+        f"-> Dataloader len: {dataloader_len} "
+        f"-> Accumulate grad batches: {trainer.accumulate_grad_batches} "
+        f"-> Epochs: {trainer.max_epochs} "
+        f"-> DataParallel size: {world_size}"
+    )
+    return total_steps

kostyl_toolkit-0.1.37/kostyl/ml/registry_uploader.py

@@ -0,0 +1,126 @@
+from abc import ABC
+from abc import abstractmethod
+from collections.abc import Callable
+from functools import partial
+from pathlib import Path
+from typing import override
+
+from clearml import OutputModel
+
+from kostyl.utils.logging import setup_logger
+
+
+logger = setup_logger()
+
+
+class RegistryUploaderCallback(ABC):
+    """Abstract Lightning callback responsible for tracking and uploading the best-performing model checkpoint."""
+
+    @abstractmethod
+    def upload_checkpoint(self, path: str | Path) -> None:
+        """Upload the checkpoint located at the given path to the configured registry backend."""
+        raise NotImplementedError
+
+
+class ClearMLRegistryUploaderCallback(RegistryUploaderCallback):
+    """PyTorch Lightning callback to upload the best model checkpoint to ClearML."""
+
+    def __init__(
+        self,
+        model_name: str,
+        config_dict: dict[str, str] | None = None,
+        label_enumeration: dict[str, int] | None = None,
+        tags: list[str] | None = None,
+        comment: str | None = None,
+        framework: str | None = None,
+        base_model_id: str | None = None,
+        new_model_per_upload: bool = True,
+        verbose: bool = True,
+    ) -> None:
+        """
+        Initializes the ClearMLRegistryUploaderCallback.
+
+        Args:
+            model_name: The name for the newly created model.
+            label_enumeration: The label enumeration dictionary of string (label) to integer (value) pairs.
+            config_dict: Optional configuration dictionary to associate with the model.
+            tags: A list of strings which are tags for the model.
+            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).
+            verbose: Whether to log messages during upload.
+
+        """
+        super().__init__()
+        if base_model_id is not None and new_model_per_upload:
+            raise ValueError(
+                "Cannot set base_model_id when new_model_per_upload is True."
+            )
+
+        self.verbose = verbose
+        self.new_model_per_upload = new_model_per_upload
+        self.best_model_path: str = ""
+        self.config_dict = config_dict
+        self._output_model: OutputModel | None = None
+        self._last_uploaded_model_path: str = ""
+        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,
+        )
+        return
+
+    @staticmethod
+    def _validate_tags(tags: list[str] | None) -> None:
+        if tags is None:
+            return
+        if "LightningCheckpoint" not in tags:
+            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
+        return model
+
+    @override
+    def upload_checkpoint(
+        self,
+        path: str | Path,
+    ) -> None:
+        if isinstance(path, Path):
+            path = str(path)
+        if path == self._last_uploaded_model_path:
+            if self.verbose:
+                logger.info("Model unchanged since last upload")
+            return
+
+        if self.verbose:
+            logger.info(f"Uploading model from {path}")
+
+        self.output_model_.update_weights(
+            path,
+            auto_delete_file=False,
+            async_enable=False,
+        )
+        self.output_model_.update_design(config_dict=self.config_dict)
+
+        self._last_uploaded_model_path = path
+        return

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

@@ -0,0 +1,18 @@
+from .composite import CompositeScheduler
+from .cosine import CosineParamScheduler
+from .cosine import CosineScheduler
+from .cosine_with_plateu import CosineWithPlateauParamScheduler
+from .cosine_with_plateu import CosineWithPlateuScheduler
+from .linear import LinearParamScheduler
+from .linear import LinearScheduler
+
+
+__all__ = [
+    "CompositeScheduler",
+    "CosineParamScheduler",
+    "CosineScheduler",
+    "CosineWithPlateauParamScheduler",
+    "CosineWithPlateuScheduler",
+    "LinearParamScheduler",
+    "LinearScheduler",
+]

{kostyl_toolkit-0.1.35 → kostyl_toolkit-0.1.37}/kostyl/ml/schedulers/base.py

@@ -6,18 +6,20 @@ from typing import Any
 class BaseScheduler(ABC):
     """Base class for learning rate schedulers."""
 
+    @abstractmethod
     def state_dict(self) -> dict[str, Any]:
         """Get the state as a state dictionary."""
-        return {
-            key: value
-            for key, value in self.__dict__.items()
-            if key not in ["optimizer", "scheduler_values"]
-        }
+        raise NotImplementedError
 
+    @abstractmethod
     def load_state_dict(self, state_dict: dict[str, Any]) -> None:
         """Load the state from a state dictionary."""
-        self.__dict__.update(state_dict)
-        return
+        raise NotImplementedError
+
+    @abstractmethod
+    def _verify(self) -> None:
+        """Verify the scheduler configuration."""
+        raise NotImplementedError
 
     def __getstate__(self) -> dict[str, Any]:
         """Get the state for pickling."""