kostyl-toolkit 0.1.35__py3-none-any.whl → 0.1.37__py3-none-any.whl

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/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/ml/dist_utils.py

@@ -4,38 +4,61 @@ from typing import Literal
 
 import torch.distributed as dist
 
+from kostyl.utils.logging import KostylLogger
 from kostyl.utils.logging import setup_logger
 
 
-logger = setup_logger(add_rank=True)
+module_logger = setup_logger()
 
 
-def log_dist(msg: str, how: Literal["only-zero-rank", "world"]) -> None:
+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.
-        how (Literal["only-zero-rank", "world"]): The verbosity level for logging.
+        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.
 
     """
-    match how:
-        case _ if not dist.is_initialized():
-            logger.warning_once(
-                "Distributed logging requested but torch.distributed is not initialized."
-            )
-            logger.info(msg)
+    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 is_main_process():
-                logger.info(msg)
+            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":
-            logger.info(msg)
+            log_attr(msg)
         case _:
-            logger.warning_once(
-                f"Invalid logging verbosity level requested: {how}. Message not logged."
-            )
+            raise ValueError(f"Invalid logging verbosity level: {log_scope}")
     return
 
 
@@ -44,7 +67,7 @@ def scale_lrs_by_world_size(
     group: dist.ProcessGroup | None = None,
     config_name: str = "",
     inv_scale: bool = False,
-    verbose: Literal["only-zero-rank", "world"] | None = None,
+    verbose_level: Literal["only-zero-rank", "world"] | None = None,
 ) -> dict[str, float]:
     """
     Scale learning-rate configuration values to match the active distributed world size.
@@ -58,7 +81,7 @@ def scale_lrs_by_world_size(
             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 (Literal["only-zero-rank", "world"] | None): Verbosity level for logging scaled values.
+        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.
@@ -77,31 +100,30 @@ def scale_lrs_by_world_size(
     for name, value in lrs.items():
         old_value = value
         new_value = value * scale
-        if verbose is not None:
+        if verbose_level is not None:
             log_dist(
                 f"New {config_name} lr {name.upper()}: {new_value}; OLD: {old_value}",
-                verbose,
+                log_scope=verbose_level,
+                group=group,
             )
         lrs[name] = new_value
     return lrs
 
 
-def get_rank() -> int:
-    """Gets the rank of the current process in a distributed setting."""
-    if dist.is_initialized():
-        return dist.get_rank()
-    if "RANK" in os.environ:
-        return int(os.environ["RANK"])
-    if "SLURM_PROCID" in os.environ:
-        return int(os.environ["SLURM_PROCID"])
+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_main_process() -> bool:
-    """Checks if the current process is the main process (rank 0) in a distributed setting."""
-    rank = get_rank()
+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/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/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/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/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/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/ml/registry_uploader.py

@@ -1,13 +1,12 @@
 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.ml.clearml.logging_utils import find_version_in_tags
-from kostyl.ml.clearml.logging_utils import increment_version
 from kostyl.utils.logging import setup_logger
 
 
@@ -28,51 +27,79 @@ class ClearMLRegistryUploaderCallback(RegistryUploaderCallback):
 
     def __init__(
         self,
-        output_model: OutputModel,
+        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,
-        enable_tag_versioning: bool = False,
     ) -> None:
         """
         Initializes the ClearMLRegistryUploaderCallback.
 
         Args:
-            output_model: ClearML OutputModel instance representing the model to upload.
-            verbose: Whether to log messages during upload.
+            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.
-            enable_tag_versioning: Whether to enable versioning in tags. If True,
-                the version tag (e.g., "v1.0") will be automatically incremented or if not present, added as "v1.0".
+            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__()
-        self.output_model = output_model
-        self.config_dict = config_dict
-        self.verbose = verbose
-        self.enable_tag_versioning = enable_tag_versioning
+        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()
+        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
 
-    def _validate_tags(self) -> None:
-        output_model_tags = self.output_model.tags or []
-        if self.enable_tag_versioning:
-            version = find_version_in_tags(output_model_tags)
-            if version is None:
-                output_model_tags.append("v1.0")
-            else:
-                new_version = increment_version(version)
-                output_model_tags.remove(version)
-                output_model_tags.append(new_version)
-        if "LightningCheckpoint" not in output_model_tags:
-            output_model_tags.append("LightningCheckpoint")
-        self.output_model.tags = output_model_tags
+    @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,
@@ -88,12 +115,12 @@ class ClearMLRegistryUploaderCallback(RegistryUploaderCallback):
         if self.verbose:
             logger.info(f"Uploading model from {path}")
 
-        self.output_model.update_weights(
+        self.output_model_.update_weights(
             path,
             auto_delete_file=False,
             async_enable=False,
         )
-        self.output_model.update_design(config_dict=self.config_dict)
+        self.output_model_.update_design(config_dict=self.config_dict)
 
         self._last_uploaded_model_path = path
         return

kostyl/ml/schedulers/__init__.py

@@ -1,6 +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"]
+__all__ = [
+    "CompositeScheduler",
+    "CosineParamScheduler",
+    "CosineScheduler",
+    "CosineWithPlateauParamScheduler",
+    "CosineWithPlateuScheduler",
+    "LinearParamScheduler",
+    "LinearScheduler",
+]

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."""