kostyl-toolkit 0.1.34__py3-none-any.whl → 0.1.36__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/lightning/callbacks/checkpoint.py

@@ -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):
@@ -342,8 +341,9 @@ def setup_checkpoint_callback(
     if dirpath.exists():
         if is_main_process():
             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/pretrained_model.py

@@ -20,7 +20,8 @@ class LightningCheckpointLoaderMixin(PreTrainedModel):
         cls: type[TModelInstance],
         checkpoint_path: str | Path,
         config_key: str = "config",
-        weights_prefix: str = "model.",
+        weights_prefix: str | None = "model.",
+        strict_prefix: bool = False,
         **kwargs: Any,
     ) -> TModelInstance:
         """
@@ -39,8 +40,10 @@ class LightningCheckpointLoaderMixin(PreTrainedModel):
             checkpoint_path (str | Path): Path to the checkpoint file. Must be a .ckpt file.
             config_key (str, optional): Key in the checkpoint dictionary where the config is stored.
                 Defaults to "config".
-            weights_prefix (str, optional): Prefix to strip from state dict keys. Defaults to "model.".
-                If not empty and doesn't end with ".", a "." is appended.
+            weights_prefix (str | None, optional): Prefix to strip from state dict keys. Defaults to "model.".
+                If not empty and doesn't end with ".", a "." is appended. If empty or None, no prefix stripping will be skipped.
+            strict_prefix (bool, optional): If True, drop tensors those keys that do not start with the
+                specified prefix. Defaults to False.
             kwargs: Additional keyword arguments to pass to the model's `from_pretrained` method.
 
         Returns:
@@ -53,6 +56,13 @@ class LightningCheckpointLoaderMixin(PreTrainedModel):
         """
         if isinstance(checkpoint_path, str):
             checkpoint_path = Path(checkpoint_path)
+        if weights_prefix is None:
+            weights_prefix = ""
+        weights_prefix = cast(str, weights_prefix)
+        if weights_prefix == "" and strict_prefix:
+            logger.warning(
+                "strict_prefix=True has no effect when weights_prefix is empty or None."
+            )
 
         if checkpoint_path.is_dir():
             raise ValueError(f"{checkpoint_path} is a directory")
@@ -85,13 +95,25 @@ class LightningCheckpointLoaderMixin(PreTrainedModel):
             if not weights_prefix.endswith("."):
                 weights_prefix = weights_prefix + "."
             state_dict: dict[str, torch.Tensor] = {}
-
+            matched_keys_counter = 0
             for key, value in raw_state_dict.items():
                 if key.startswith(weights_prefix):
                     new_key = key[len(weights_prefix) :]
                     state_dict[new_key] = value
-                else:
+                    matched_keys_counter += 1
+                elif not strict_prefix:
                     state_dict[key] = value
+
+            if matched_keys_counter == 0:
+                if strict_prefix:
+                    raise ValueError(
+                        f"No keys in the checkpoint start with the specified prefix '{weights_prefix}'. "
+                        "Try to load with `strict_prefix=False` or verify the prefix."
+                    )
+                else:
+                    logger.warning(
+                        f"No keys in the checkpoint start with the specified prefix '{weights_prefix}'. "
+                    )
         else:
             state_dict = raw_state_dict
 

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(add_rank=True)
+
+
+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."""

kostyl/ml/schedulers/cosine.py

@@ -2,7 +2,6 @@ from typing import Any
 from typing import override
 
 import numpy as np
-import numpy.typing as npt
 import torch
 
 from .base import BaseScheduler
@@ -29,18 +28,24 @@ class _CosineSchedulerCore(BaseScheduler):
         if freeze_ratio is not None:
             if not (0 < freeze_ratio < 1):
                 raise ValueError(f"Freeze ratio must be in (0, 1), got {freeze_ratio}.")
+        pre_annealing_ratio = (warmup_ratio if warmup_ratio is not None else 0) + (
+            freeze_ratio if freeze_ratio is not None else 0
+        )
+        if pre_annealing_ratio > 1:
+            raise ValueError(
+                "The sum of warmup_ratio and freeze_ratio must <= 1, got "
+                f"{pre_annealing_ratio}."
+            )
 
         self.param_name = param_name
         self.num_iters = num_iters
         self.base_value = base_value
         self.final_value = final_value
-
         self.warmup_ratio = warmup_ratio
         self.warmup_value = warmup_value
-
         self.freeze_ratio = freeze_ratio
 
-        self.scheduler_values: npt.NDArray[np.float64] = np.array([], dtype=np.float64)
+        self.scheduled_values: np.ndarray = np.array([], dtype=np.float64)
         self.current_value_ = self.base_value
         return
 
@@ -63,31 +68,29 @@ class _CosineSchedulerCore(BaseScheduler):
             warmup_iters = 0
             warmup_schedule = np.array([], dtype=np.float64)
 
+        # Create cosine annealing schedule
         cosine_annealing_iters = self.num_iters - warmup_iters - freeze_iters
-        if cosine_annealing_iters <= 0:
-            raise ValueError("Cosine annealing iters must be > 0.")
-
-        # Create cosine schedule
-        iters = np.arange(cosine_annealing_iters)
-        schedule = self.final_value + 0.5 * (self.base_value - self.final_value) * (
-            1 + np.cos(np.pi * iters / len(iters))
-        )
+        if cosine_annealing_iters > 0:
+            iters = np.arange(cosine_annealing_iters)
+            cosine_annealing_schedule = self.final_value + 0.5 * (
+                self.base_value - self.final_value
+            ) * (1 + np.cos(np.pi * iters / len(iters)))
+        else:
+            cosine_annealing_schedule = np.array([], dtype=np.float64)
 
         # Concatenate all parts of the schedule
-        self.scheduler_values = np.concatenate(
-            (freeze_schedule, warmup_schedule, schedule)
+        self.scheduled_values = np.concatenate(
+            (freeze_schedule, warmup_schedule, cosine_annealing_schedule)
         )
-
-        if len(self.scheduler_values) != self.num_iters:
-            raise ValueError(
-                f"Scheduler length ({len(self.scheduler_values)}) does not match num_iters ({self.num_iters})."
-            )
+        self._verify()
         return
 
     @override
-    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
-        super().load_state_dict(state_dict)
-        self.scheduler_values = np.array([], dtype=np.float64)
+    def _verify(self) -> None:
+        if len(self.scheduled_values) != self.num_iters:
+            raise ValueError(
+                f"Scheduler length ({len(self.scheduled_values)}) does not match num_iters ({self.num_iters})."
+            )
         return
 
     @override
@@ -95,13 +98,13 @@ class _CosineSchedulerCore(BaseScheduler):
         raise NotImplementedError
 
     def _get_value(self, it: int) -> float:
-        if len(self.scheduler_values) == 0:
+        if len(self.scheduled_values) == 0:
             self._create_scheduler()
 
         if it >= self.num_iters:
             value: float = self.final_value
         else:
-            value: float = self.scheduler_values[it]
+            value: float = self.scheduled_values[it]
         self.current_value_ = value
         return value
 
@@ -163,6 +166,21 @@ class CosineScheduler(_CosineSchedulerCore):
         self.param_group_field = param_group_field
         return
 
+    @override
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        self.__dict__.update(state_dict)
+        self.scheduled_values = np.array([], dtype=np.float64)
+        return
+
+    @override
+    def state_dict(self) -> dict[str, Any]:
+        state = {
+            k: v
+            for k, v in self.__dict__.items()
+            if k not in ["scheduled_values", "optimizer"]
+        }
+        return state
+
     @override
     def step(self, it: int) -> None:
         value = self._get_value(it)
@@ -209,3 +227,14 @@ class CosineParamScheduler(_CosineSchedulerCore):
         """
         value = self._get_value(it)
         return value
+
+    @override
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        self.__dict__.update(state_dict)
+        self.scheduled_values = np.array([], dtype=np.float64)
+        return
+
+    @override
+    def state_dict(self) -> dict[str, Any]:
+        state = {k: v for k, v in self.__dict__.items() if k != "scheduled_values"}
+        return state

kostyl/ml/schedulers/cosine_with_plateu.py

@@ -0,0 +1,277 @@
+from typing import Any
+from typing import override
+
+import numpy as np
+import torch
+
+from .base import BaseScheduler
+
+
+class _CosineWithPlateauSchedulerCore(BaseScheduler):
+    """Core cosine with plateau scheduler logic."""
+
+    def __init__(
+        self,
+        param_name: str,
+        num_iters: int,
+        base_value: float,
+        final_value: float,
+        plateau_ratio: float,
+        warmup_value: float | None = None,
+        warmup_ratio: float | None = None,
+        freeze_ratio: float | None = None,
+    ) -> None:
+        if warmup_ratio is not None:
+            if not (0 < warmup_ratio < 1):
+                raise ValueError(f"Warmup ratio must be in (0, 1), got {warmup_ratio}.")
+        if (warmup_value is None) != (warmup_ratio is None):
+            raise ValueError(
+                "Both warmup_ratio and warmup_value must be provided or neither."
+            )
+        if freeze_ratio is not None:
+            if not (0 < freeze_ratio < 1):
+                raise ValueError(f"Freeze ratio must be in (0, 1), got {freeze_ratio}.")
+        if not (0 < plateau_ratio < 1):
+            raise ValueError(f"Plateau ratio must be in (0, 1), got {plateau_ratio}.")
+
+        pre_annealing_ratio = (
+            plateau_ratio
+            + (warmup_ratio if warmup_ratio is not None else 0)
+            + (freeze_ratio if freeze_ratio is not None else 0)
+        )
+        if pre_annealing_ratio > 1:
+            raise ValueError(
+                "The sum of plateau_ratio, warmup_ratio, and freeze_ratio must <= 1, got "
+                f"{pre_annealing_ratio}."
+            )
+
+        self.param_name = param_name
+        self.num_iters = num_iters
+        self.base_value = base_value
+        self.final_value = final_value
+        self.cosine_annealing_ratio = 1 - pre_annealing_ratio
+        self.plateau_ratio = plateau_ratio
+        self.warmup_ratio = warmup_ratio
+        self.warmup_value = warmup_value
+        self.freeze_ratio = freeze_ratio
+
+        self.scheduled_values: np.ndarray = np.array([], dtype=np.float64)
+        self.current_value_ = self.base_value
+        return
+
+    def _create_scheduler(self) -> None:
+        # Create freeze schedule
+        if self.freeze_ratio is not None:
+            freeze_iters = int(self.num_iters * self.freeze_ratio)
+            freeze_schedule = np.zeros(freeze_iters, dtype=np.float64)
+        else:
+            freeze_iters = 0
+            freeze_schedule = np.array([], dtype=np.float64)
+
+        # Create linear warmup schedule
+        if self.warmup_ratio is not None and self.warmup_value is not None:
+            warmup_iters = int(self.num_iters * self.warmup_ratio)
+            warmup_schedule = np.linspace(
+                self.warmup_value, self.base_value, warmup_iters, dtype=np.float64
+            )
+        else:
+            warmup_iters = 0
+            warmup_schedule = np.array([], dtype=np.float64)
+
+        # Create cosine annealing schedule
+        if self.cosine_annealing_ratio > 0:
+            cosine_annealing_iters = int(self.num_iters * self.cosine_annealing_ratio)
+            iters = np.arange(cosine_annealing_iters)
+            cosine_annealing_schedule = self.final_value + 0.5 * (
+                self.base_value - self.final_value
+            ) * (1 + np.cos(np.pi * iters / len(iters)))
+        else:
+            cosine_annealing_iters = 0
+            cosine_annealing_schedule = np.array([], dtype=np.float64)
+
+        plateau_iters = (
+            self.num_iters - warmup_iters - freeze_iters - cosine_annealing_iters
+        )
+        if plateau_iters > 0:
+            plateau_schedule = np.full(plateau_iters, self.base_value, dtype=np.float64)
+        else:
+            plateau_schedule = np.array([], dtype=np.float64)
+
+        # Concatenate all parts of the schedule
+        self.scheduled_values = np.concatenate(
+            (
+                freeze_schedule,
+                warmup_schedule,
+                plateau_schedule,
+                cosine_annealing_schedule,
+            )
+        )
+        self._verify()
+        return
+
+    @override
+    def _verify(self) -> None:
+        if len(self.scheduled_values) != self.num_iters:
+            raise ValueError(
+                f"Scheduler length ({len(self.scheduled_values)}) does not match num_iters ({self.num_iters})."
+            )
+        return
+
+    @override
+    def step(self, it: int) -> None | float:
+        raise NotImplementedError
+
+    def _get_value(self, it: int) -> float:
+        if len(self.scheduled_values) == 0:
+            self._create_scheduler()
+
+        if it >= self.num_iters:
+            value: float = self.final_value
+        else:
+            value: float = self.scheduled_values[it]
+        self.current_value_ = value
+        return value
+
+    @override
+    def current_value(self) -> dict[str, float]:
+        return {self.param_name: self.current_value_}
+
+
+class CosineWithPlateuScheduler(_CosineWithPlateauSchedulerCore):
+    """
+    Applies a cosine schedule with plateau to an optimizer param-group field.
+
+    Schedule phases: freeze (0) → warmup → plateau (base_value) → cosine annealing to final_value.
+    The plateau phase maintains the base_value before cosine annealing begins.
+    """
+
+    def __init__(
+        self,
+        optimizer: torch.optim.Optimizer,
+        param_group_field: str,
+        num_iters: int,
+        base_value: float,
+        final_value: float,
+        plateau_ratio: float,
+        warmup_value: float | None = None,
+        warmup_ratio: float | None = None,
+        freeze_ratio: float | None = None,
+        multiplier_field: str | None = None,
+        skip_if_zero: bool = False,
+        apply_if_field: str | None = None,
+        ignore_if_field: str | None = None,
+    ) -> None:
+        """
+        Configure cosine scheduling for matching optimizer groups.
+
+        Args:
+            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``.
+            base_value: Value maintained during plateau phase and used as cosine start.
+            final_value: Value approached as iterations progress during cosine annealing.
+            plateau_ratio: Fraction of iterations to maintain ``base_value`` before cosine annealing.
+            warmup_ratio: Optional fraction of iterations to linearly ramp from ``warmup_value`` to ``base_value``.
+            warmup_value: Starting value for the warmup ramp.
+            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.
+
+        """
+        self.apply_if_field = apply_if_field
+        self.ignore_if_field = ignore_if_field
+        self.optimizer = optimizer
+        self.multiplier_field = multiplier_field
+        self.skip_if_zero = skip_if_zero
+        super().__init__(
+            param_name=param_group_field,
+            num_iters=num_iters,
+            base_value=base_value,
+            final_value=final_value,
+            plateau_ratio=plateau_ratio,
+            warmup_ratio=warmup_ratio,
+            warmup_value=warmup_value,
+            freeze_ratio=freeze_ratio,
+        )
+        self.param_group_field = param_group_field
+        return
+
+    @override
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        self.__dict__.update(state_dict)
+        self.scheduled_values = np.array([], dtype=np.float64)
+        return
+
+    @override
+    def state_dict(self) -> dict[str, Any]:
+        state = {
+            k: v
+            for k, v in self.__dict__.items()
+            if k not in ["scheduled_values", "optimizer"]
+        }
+        return state
+
+    @override
+    def step(self, it: int) -> None:
+        value = self._get_value(it)
+        for pg in self.optimizer.param_groups:
+            if self.param_group_field not in pg:
+                raise ValueError(
+                    f"Parameter group field '{self.param_group_field}' not found in optimizer parameter groups."
+                )
+
+            if (self.apply_if_field is not None) and (self.apply_if_field not in pg):
+                continue
+
+            if (self.ignore_if_field is not None) and (self.ignore_if_field in pg):
+                continue
+
+            if self.skip_if_zero and pg[self.param_group_field] == 0:
+                continue
+
+            if self.multiplier_field is not None:
+                if self.multiplier_field not in pg:
+                    multiplier = 1.0
+                else:
+                    multiplier = pg[self.multiplier_field]
+                pg[self.param_group_field] = value * multiplier
+            else:
+                pg[self.param_group_field] = value
+        return
+
+
+class CosineWithPlateauParamScheduler(_CosineWithPlateauSchedulerCore):
+    """
+    Standalone cosine scheduler with plateau for non-optimizer parameters.
+
+    Schedule phases: freeze (0) → warmup → plateau (base_value) → cosine annealing to final_value.
+    The plateau phase maintains the base_value before cosine annealing begins.
+    """
+
+    @override
+    def step(self, it: int) -> float:
+        """
+        Computes the value corresponding to the given iteration step.
+
+        Args:
+            it: The current iteration index used for value computation.
+
+        Returns:
+            The computed value for the provided iteration step as a float.
+
+        """
+        value = self._get_value(it)
+        return value
+
+    @override
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        self.__dict__.update(state_dict)
+        self.scheduled_values = np.array([], dtype=np.float64)
+        return
+
+    @override
+    def state_dict(self) -> dict[str, Any]:
+        state = {k: v for k, v in self.__dict__.items() if k != "scheduled_values"}
+        return state

kostyl/ml/schedulers/linear.py

@@ -21,24 +21,23 @@ class _LinearScheduleBase(BaseScheduler):
         self.start_value = start_value
         self.final_value = final_value
 
-        self.scheduler_values: npt.NDArray[np.float64] = np.array([], dtype=np.float64)
+        self.scheduled_values: npt.NDArray[np.float64] = np.array([], dtype=np.float64)
         self.current_value_ = self.start_value
         return
 
     def _create_scheduler(self) -> None:
-        self.scheduler_values = np.linspace(
+        self.scheduled_values = np.linspace(
             self.start_value, self.final_value, num=self.num_iters, dtype=np.float64
         )
-        if len(self.scheduler_values) != self.num_iters:
-            raise ValueError(
-                f"Scheduler length ({len(self.scheduler_values)}) does not match total_iters ({self.num_iters})."
-            )
+        self._verify()
         return
 
     @override
-    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
-        super().load_state_dict(state_dict)
-        self.scheduler_values = np.array([], dtype=np.float64)
+    def _verify(self) -> None:
+        if len(self.scheduled_values) != self.num_iters:
+            raise ValueError(
+                f"Scheduler length ({len(self.scheduled_values)}) does not match total_iters ({self.num_iters})."
+            )
         return
 
     @override
@@ -46,13 +45,13 @@ class _LinearScheduleBase(BaseScheduler):
         raise NotImplementedError
 
     def _get_value(self, it: int) -> float:
-        if len(self.scheduler_values) == 0:
+        if len(self.scheduled_values) == 0:
             self._create_scheduler()
 
         if it >= self.num_iters:
             value: float = self.final_value
         else:
-            value: float = self.scheduler_values[it]
+            value: float = self.scheduled_values[it]
         self.current_value_ = value
         return value
 
@@ -105,6 +104,21 @@ class LinearScheduler(_LinearScheduleBase):
         self.param_group_field = param_group_field
         return
 
+    @override
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        self.__dict__.update(state_dict)
+        self.scheduled_values = np.array([], dtype=np.float64)
+        return
+
+    @override
+    def state_dict(self) -> dict[str, Any]:
+        state = {
+            k: v
+            for k, v in self.__dict__.items()
+            if k not in ["scheduled_values", "optimizer"]
+        }
+        return state
+
     @override
     def step(self, it: int) -> None:
         value = self._get_value(it)
@@ -137,6 +151,17 @@ class LinearScheduler(_LinearScheduleBase):
 class LinearParamScheduler(_LinearScheduleBase):
     """LinearParamScheduler adjusts a parameter value using a linear scheduler."""
 
+    @override
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        self.__dict__.update(state_dict)
+        self.scheduled_values = np.array([], dtype=np.float64)
+        return
+
+    @override
+    def state_dict(self) -> dict[str, Any]:
+        state = {k: v for k, v in self.__dict__.items() if k != "scheduled_values"}
+        return state
+
     @override
     def step(self, it: int) -> float:
         """

kostyl/utils/logging.py

@@ -94,7 +94,7 @@ _PRESETS = {"default": _DEFAULT_FMT, "only_message": _ONLY_MESSAGE_FMT}
 
 def setup_logger(
     name: str | None = None,
-    fmt: Literal["default", "only_message"] | str = "default",
+    fmt: Literal["default", "only_message"] | str = "only_message",
     level: str = "INFO",
     add_rank: bool | None = None,
     sink=sys.stdout,

{kostyl_toolkit-0.1.34.dist-info → kostyl_toolkit-0.1.36.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: kostyl-toolkit
-Version: 0.1.34
+Version: 0.1.36
 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.34.dist-info → kostyl_toolkit-0.1.36.dist-info}/RECORD

@@ -6,32 +6,33 @@ kostyl/ml/clearml/logging_utils.py,sha256=GBjIIZbH_itd5sj7XpvxjkyZwxxGOpEcQ3BiWa
 kostyl/ml/clearml/pulling_utils.py,sha256=jMlVXcYRumwWnPlELRlgEdfq5L6Wir_EcfTmOoWBLTA,4077
 kostyl/ml/configs/__init__.py,sha256=IetcivbqYGutowLqxdKp7QR4tkXKBr4m8t4Zkk9jHZU,911
 kostyl/ml/configs/base_model.py,sha256=Eofn14J9RsjpVx_J4rp6C19pDDCANU4hr3JtX-d0FpQ,4820
-kostyl/ml/configs/hyperparams.py,sha256=2S_VEZ07RWquNFSWjHBb3OUpBlTznbUpFSchzMpSBOc,2879
-kostyl/ml/configs/training_settings.py,sha256=Sq2tiRuwkbmi9zKDG2JghZLXo5DDt_eQqN_KYJSdcTY,2509
+kostyl/ml/configs/hyperparams.py,sha256=lvtbvOFEoTBAJug7FR35xMQdPLgDQjRoP2fyDP-jD7E,3305
+kostyl/ml/configs/training_settings.py,sha256=wT9CHuLaKrLwonsc87Ee421EyFis_c9fqOgn9bSClm8,2747
 kostyl/ml/data_processing_utils.py,sha256=jjEjV0S0wREgZkzg27ip0LpI8cQqkwe2QwATmAqm9-g,3832
 kostyl/ml/dist_utils.py,sha256=Onf0KHVLA8oeUgZTcTdmR9qiM22f2uYLoNwgLbMGJWk,3495
 kostyl/ml/lightning/__init__.py,sha256=R36PImjVvzBF9t_z9u6RYVnUFJJ-sNDUOdboWUojHmM,173
 kostyl/ml/lightning/callbacks/__init__.py,sha256=EnKkNwwNDZnEqKRlpY4FVrqP88ECPF6nlT2bSLUIKRk,194
-kostyl/ml/lightning/callbacks/checkpoint.py,sha256=sZ9OqudO-gXp7FqtWaOH46TXVpeCJxV-EowyRPN836k,18983
+kostyl/ml/lightning/callbacks/checkpoint.py,sha256=COW7WErj4EMxJNMn97WQO-G2A3LbI6GQOCpIZu3Cblk,19060
 kostyl/ml/lightning/callbacks/early_stopping.py,sha256=D5nyjktCJ9XYAf28-kgXG8jORvXLl1N3nbDQnvValPM,615
 kostyl/ml/lightning/extensions/__init__.py,sha256=OY6QGv1agYgqqKf1xJBrxgp_i8FunVfPzYezfaRrGXU,182
 kostyl/ml/lightning/extensions/custom_module.py,sha256=iQrnPz-WTmRfvLo94C5fQc2Qwa1IpHtUh1sCpVwTSFM,6602
-kostyl/ml/lightning/extensions/pretrained_model.py,sha256=QJGr2UvYJcU2Gy2w8z_cEvTodjv7hGdd2PPPfdOI-Mw,4017
+kostyl/ml/lightning/extensions/pretrained_model.py,sha256=eRfQBzAjVernHl9A4PP5uTLvjjmcNKPdTu7ABFLq7HI,5196
 kostyl/ml/lightning/loggers/__init__.py,sha256=e51dszaoJbuzwBkbdugmuDsPldoSO4yaRgmZUg1Bdy0,71
 kostyl/ml/lightning/loggers/tb_logger.py,sha256=j02HK5ue8yzXXV8FWKmmXyHkFlIxgHx-ahHWk_rFCZs,893
-kostyl/ml/lightning/training_utils.py,sha256=u7X9ysF9Gqy8CdwacdcDlNQNsbagYAhslbv-1WLJ45k,9052
+kostyl/ml/lightning/utils.py,sha256=imvMbgOKRtCUiiRGEcVtN-hxw-aEFKHdCWc0J_CIZp4,1980
 kostyl/ml/metrics_formatting.py,sha256=U6vdNENZLvp2dT1L3HqFKtXrHwGKoDXN93hvamPGHjM,1341
 kostyl/ml/params_groups.py,sha256=nUyw5d06Pvy9QPiYtZzLYR87xwXqJLxbHthgQH8oSCM,3583
-kostyl/ml/registry_uploader.py,sha256=W90TYo_WKv2oBE6nqEJl4hecYmJyyuKwQJ9_uUPGnJQ,3346
-kostyl/ml/schedulers/__init__.py,sha256=bxXbsU_WYnVbhvNNnuI7cOAh2Axz7D25TaleBTZhYfc,197
-kostyl/ml/schedulers/base.py,sha256=9M2iOoOVSRojR_liPX1qo3Nn4iMXSM5ZJuAFWZTulUk,1327
+kostyl/ml/registry_uploader.py,sha256=BbyLXvF8AL145k7g6MRkJ7gf_3Um53p3Pn5280vVD9U,4384
+kostyl/ml/schedulers/__init__.py,sha256=_EtZu8DwTCSv4-eR84kRstEZblHylVqda7WQUOXIKfw,534
+kostyl/ml/schedulers/base.py,sha256=bjmwgdZpnSqpCnHPnKC6MEiRO79cwxMJpZq-eQVNs2M,1353
 kostyl/ml/schedulers/composite.py,sha256=ee4xlMDMMtjKPkbTF2ue9GTr9DuGCGjZWf11mHbi6aE,2387
-kostyl/ml/schedulers/cosine.py,sha256=t74_ByT22L5NQKpnBVU9UGzBVx1ZM2GTylb9ct3_PVg,7627
-kostyl/ml/schedulers/linear.py,sha256=7HPkVWcPa0lbaZywutXSDdVLLSihAyWk5XIE2Dzj_5Q,5168
+kostyl/ml/schedulers/cosine.py,sha256=y8ylrgVOkVcr2-ExoqqNW--tdDX88TBYPQCOppIf2_M,8685
+kostyl/ml/schedulers/cosine_with_plateu.py,sha256=0-X6wl3HgsTiLIbISb9lOxIVWXHDEND7rILitMWtIiM,10195
+kostyl/ml/schedulers/linear.py,sha256=RnnnblRuRXP3LT03QVIHUaK2kNsiMP1AedrMoeyh3qk,5843
 kostyl/utils/__init__.py,sha256=hkpmB6c5pr4Ti5BshOROebb7cvjDZfNCw83qZ_FFKMM,240
 kostyl/utils/dict_manipulations.py,sha256=e3vBicID74nYP8lHkVTQc4-IQwoJimrbFELy5uSF6Gk,1073
 kostyl/utils/fs.py,sha256=gAQNIU4R_2DhwjgzOS8BOMe0gZymtY1eZwmdgOdDgqo,510
-kostyl/utils/logging.py,sha256=Vye0u4-yeOSUc-f03gpQbxSktTbFiilTWLEVr00ZHvc,5796
-kostyl_toolkit-0.1.34.dist-info/WHEEL,sha256=ZyFSCYkV2BrxH6-HRVRg3R9Fo7MALzer9KiPYqNxSbo,79
-kostyl_toolkit-0.1.34.dist-info/METADATA,sha256=AduVYS1jJViZoxxpO7KYHO0_YNZ23f3P9eKgZAfyteE,4269
-kostyl_toolkit-0.1.34.dist-info/RECORD,,
+kostyl/utils/logging.py,sha256=LSbyQFLAIa89xPb4tcobE2BwVIHHUSaDXqOIKVzLoWs,5801
+kostyl_toolkit-0.1.36.dist-info/WHEEL,sha256=eycQt0QpYmJMLKpE3X9iDk8R04v2ZF0x82ogq-zP6bQ,79
+kostyl_toolkit-0.1.36.dist-info/METADATA,sha256=Lfyx6u3LKZ6co4s7GZgJp31zoy-NViriSGqwjIzOQFA,4269
+kostyl_toolkit-0.1.36.dist-info/RECORD,,

{kostyl_toolkit-0.1.34.dist-info → kostyl_toolkit-0.1.36.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: uv 0.9.18
+Generator: uv 0.9.24
 Root-Is-Purelib: true
 Tag: py3-none-any

kostyl/ml/lightning/training_utils.py

@@ -1,241 +0,0 @@
-from dataclasses import dataclass
-from dataclasses import fields
-from pathlib import Path
-from typing import Literal
-from typing import cast
-
-import lightning as L
-import torch
-import torch.distributed as dist
-from clearml import OutputModel
-from clearml import Task
-from lightning.pytorch.callbacks import Callback
-from lightning.pytorch.callbacks import EarlyStopping
-from lightning.pytorch.callbacks import LearningRateMonitor
-from lightning.pytorch.callbacks import ModelCheckpoint
-from lightning.pytorch.loggers import TensorBoardLogger
-from lightning.pytorch.strategies import DDPStrategy
-from lightning.pytorch.strategies import FSDPStrategy
-from torch.distributed import ProcessGroup
-from torch.distributed.fsdp import MixedPrecision
-from torch.nn import Module
-
-from kostyl.ml.configs import CheckpointConfig
-from kostyl.ml.configs import DDPStrategyConfig
-from kostyl.ml.configs import EarlyStoppingConfig
-from kostyl.ml.configs import FSDP1StrategyConfig
-from kostyl.ml.configs import SingleDeviceStrategyConfig
-from kostyl.ml.lightning.callbacks import setup_checkpoint_callback
-from kostyl.ml.lightning.callbacks import setup_early_stopping_callback
-from kostyl.ml.lightning.loggers import setup_tb_logger
-from kostyl.ml.registry_uploader import ClearMLRegistryUploaderCallback
-from kostyl.utils.logging import setup_logger
-
-
-TRAINING_STRATEGIES = (
-    FSDP1StrategyConfig | DDPStrategyConfig | SingleDeviceStrategyConfig
-)
-
-logger = setup_logger(add_rank=True)
-
-
-def estimate_total_steps(
-    trainer: L.Trainer, process_group: ProcessGroup | None = None
-) -> int:
-    """
-    Estimates the total number of training steps based on the
-    dataloader length, accumulation steps, and distributed world size.
-    """  # noqa: D205
-    if dist.is_initialized():
-        world_size = dist.get_world_size(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})\n"
-        f"-> Dataloader len: {dataloader_len}\n"
-        f"-> Accumulate grad batches: {trainer.accumulate_grad_batches}\n"
-        f"-> Epochs: {trainer.max_epochs}\n "
-        f"-> World size: {world_size}"
-    )
-    return total_steps
-
-
-@dataclass
-class Callbacks:
-    """Dataclass to hold PyTorch Lightning callbacks."""
-
-    checkpoint: ModelCheckpoint
-    lr_monitor: LearningRateMonitor
-    early_stopping: EarlyStopping | None = None
-
-    def to_list(self) -> list[Callback]:
-        """Convert dataclass fields to a list of Callbacks. None values are omitted."""
-        callbacks: list[Callback] = [
-            getattr(self, field.name)
-            for field in fields(self)
-            if getattr(self, field.name) is not None
-        ]
-        return callbacks
-
-
-def setup_callbacks(
-    task: Task,
-    root_path: Path,
-    checkpoint_cfg: CheckpointConfig,
-    early_stopping_cfg: EarlyStoppingConfig | None,
-    output_model: OutputModel,
-    checkpoint_upload_strategy: Literal["only-best", "every-checkpoint"],
-    config_dict: dict[str, str] | None = None,
-    enable_tag_versioning: bool = False,
-) -> Callbacks:
-    """
-    Set up PyTorch Lightning callbacks for training.
-
-    Creates and configures a set of callbacks including checkpoint saving,
-    learning rate monitoring, model registry uploading, and optional early stopping.
-
-    Args:
-        task: ClearML task for organizing checkpoints by task name and ID.
-        root_path: Root directory for saving checkpoints.
-        checkpoint_cfg: Configuration for checkpoint saving behavior.
-        checkpoint_upload_strategy: Model upload strategy:
-            - `"only-best"`: Upload only the best checkpoint based on monitored metric.
-            - `"every-checkpoint"`: Upload every saved checkpoint.
-        output_model: ClearML OutputModel instance for model registry integration.
-        early_stopping_cfg: Configuration for early stopping. If None, early stopping
-            is disabled.
-        config_dict: Optional configuration dictionary to store with the model
-            in the registry.
-        enable_tag_versioning: Whether to auto-increment version tags (e.g., "v1.0")
-            on the uploaded model.
-
-    Returns:
-        Callbacks dataclass containing configured checkpoint, lr_monitor,
-        and optionally early_stopping callbacks.
-
-    """
-    lr_monitor = LearningRateMonitor(
-        logging_interval="step", log_weight_decay=True, log_momentum=False
-    )
-    model_uploader = ClearMLRegistryUploaderCallback(
-        output_model=output_model,
-        config_dict=config_dict,
-        verbose=True,
-        enable_tag_versioning=enable_tag_versioning,
-    )
-    checkpoint_callback = setup_checkpoint_callback(
-        root_path / "checkpoints" / task.name / task.id,
-        checkpoint_cfg,
-        registry_uploader_callback=model_uploader,
-        uploading_strategy=checkpoint_upload_strategy,
-    )
-    if early_stopping_cfg is not None:
-        early_stopping_callback = setup_early_stopping_callback(early_stopping_cfg)
-    else:
-        early_stopping_callback = None
-
-    callbacks = Callbacks(
-        checkpoint=checkpoint_callback,
-        lr_monitor=lr_monitor,
-        early_stopping=early_stopping_callback,
-    )
-    return callbacks
-
-
-def setup_loggers(task: Task, root_path: Path) -> list[TensorBoardLogger]:
-    """
-    Set up PyTorch Lightning loggers for training.
-
-    Args:
-        task: ClearML task used to organize log directories by task name and ID.
-        root_path: Root directory for storing TensorBoard logs.
-
-    Returns:
-        List of configured TensorBoard loggers.
-
-    """
-    loggers = [
-        setup_tb_logger(root_path / "runs" / task.name / task.id),
-    ]
-    return loggers
-
-
-def setup_strategy(
-    strategy_settings: TRAINING_STRATEGIES,
-    devices: list[int] | int,
-    auto_wrap_policy: set[type[Module]] | None = None,
-) -> Literal["auto"] | FSDPStrategy | DDPStrategy:
-    """
-    Configure and return a PyTorch Lightning training strategy.
-
-    Args:
-        strategy_settings: Strategy configuration object. Must be one of:
-            - `FSDP1StrategyConfig`: Fully Sharded Data Parallel strategy (requires 2+ devices).
-            - `DDPStrategyConfig`: Distributed Data Parallel strategy (requires 2+ devices).
-            - `SingleDeviceStrategyConfig`: Single device training (requires exactly 1 device).
-        devices: Device(s) to use for training. Either a list of device IDs or
-            a single integer representing the number of devices.
-        auto_wrap_policy: Set of module types that should be wrapped for FSDP.
-            Required when using `FSDP1StrategyConfig`, ignored otherwise.
-
-    Returns:
-        Configured strategy: `FSDPStrategy`, `DDPStrategy`, or `"auto"` for single device.
-
-    Raises:
-        ValueError: If device count doesn't match strategy requirements or
-            if `auto_wrap_policy` is missing for FSDP.
-
-    """
-    if isinstance(devices, list):
-        num_devices = len(devices)
-    else:
-        num_devices = devices
-
-    match strategy_settings:
-        case FSDP1StrategyConfig():
-            if num_devices < 2:
-                raise ValueError("FSDP strategy requires multiple devices.")
-
-            if auto_wrap_policy is None:
-                raise ValueError("auto_wrap_policy must be provided for FSDP strategy.")
-
-            mixed_precision_config = MixedPrecision(
-                param_dtype=getattr(torch, strategy_settings.param_dtype),
-                reduce_dtype=getattr(torch, strategy_settings.reduce_dtype),
-                buffer_dtype=getattr(torch, strategy_settings.buffer_dtype),
-            )
-            strategy = FSDPStrategy(
-                auto_wrap_policy=auto_wrap_policy,
-                mixed_precision=mixed_precision_config,
-            )
-        case DDPStrategyConfig():
-            if num_devices < 2:
-                raise ValueError("DDP strategy requires at least two devices.")
-            strategy = DDPStrategy(
-                find_unused_parameters=strategy_settings.find_unused_parameters
-            )
-        case SingleDeviceStrategyConfig():
-            if num_devices != 1:
-                raise ValueError("SingleDevice strategy requires exactly one device.")
-            strategy = "auto"
-        case _:
-            raise ValueError(
-                f"Unsupported strategy type: {type(strategy_settings.trainer.strategy)}"
-            )
-    return strategy