kostyl-toolkit 0.1.41__py3-none-any.whl → 0.1.43__py3-none-any.whl

kostyl/ml/configs/training_settings.py

@@ -38,18 +38,18 @@ class FSDP1StrategyConfig(BaseModel):
     """Fully Sharded Data Parallel (FSDP) strategy configuration."""
 
     type: Literal["fsdp1"]
-    param_dtype: DTYPE | None
-    reduce_dtype: DTYPE | None
-    buffer_dtype: DTYPE | None
+    param_dtype: DTYPE | None = None
+    reduce_dtype: DTYPE | None = None
+    buffer_dtype: DTYPE | None = None
 
 
 class FSDP2StrategyConfig(BaseModel):
     """Fully Sharded Data Parallel (FSDP) strategy configuration."""
 
     type: Literal["fsdp2"]
-    param_dtype: DTYPE | None
-    reduce_dtype: DTYPE | None
-    buffer_dtype: DTYPE | None
+    param_dtype: DTYPE | None = None
+    reduce_dtype: DTYPE | None = None
+    buffer_dtype: DTYPE | None = None
 
 
 class DDPStrategyConfig(BaseModel):
@@ -59,12 +59,20 @@ class DDPStrategyConfig(BaseModel):
     find_unused_parameters: bool = False
 
 
+SUPPORTED_STRATEGIES = (
+    FSDP1StrategyConfig
+    | FSDP2StrategyConfig
+    | SingleDeviceStrategyConfig
+    | DDPStrategyConfig
+)
+
+
 class LightningTrainerParameters(BaseModel):
     """Lightning Trainer parameters configuration."""
 
     accelerator: str
     max_epochs: int
-    strategy: FSDP1StrategyConfig | SingleDeviceStrategyConfig | DDPStrategyConfig
+    strategy: SUPPORTED_STRATEGIES
     val_check_interval: int | float
     devices: list[int] | int
     precision: PRECISION

kostyl/ml/data_collator.py

@@ -93,10 +93,10 @@ class BatchCollatorWithKeyAlignment:
                 if new_key is None:
                     continue
                 value = item[k]
-                if self.max_length is not None and new_key in (
+                if self.max_length is not None and new_key in {
                     "input_ids",
                     "attention_mask",
-                ):
+                }:
                     value = self._truncate_data(new_key, value)
                 new_item[new_key] = value
             aligned_batch.append(new_item)

kostyl/ml/optim/factory.py

@@ -1,8 +1,11 @@
-from typing import Any
+from typing import TypedDict
+from typing import Unpack
 
 from torch.optim import Optimizer
+from torch.optim.optimizer import ParamsT
 
 from kostyl.ml.configs import OPTIMIZER_CONFIG
+from kostyl.ml.configs import SCHEDULER
 from kostyl.ml.configs import AdamConfig
 from kostyl.ml.configs import AdamWithPrecisionConfig
 from kostyl.ml.configs import MuonConfig
@@ -18,6 +21,17 @@ from .schedulers import PlateauWithAnnealingScheduler
 logger = setup_logger(fmt="only_message")
 
 
+class OVERRIDABLE_CONFIG_KWARGS(TypedDict, total=False):  # noqa: D101, N801
+    scheduler_type: SCHEDULER | None
+
+    freeze_ratio: float | None
+    warmup_ratio: float | None
+    warmup_value: float | None
+    base_value: float | None
+    final_value: float | None
+    plateau_ratio: float | None
+
+
 def create_scheduler(
     config: ScheduledParamConfig,
     param_group_field: str,
@@ -27,6 +41,7 @@ def create_scheduler(
     skip_if_zero: bool = False,
     apply_if_field: str | None = None,
     ignore_if_field: str | None = None,
+    **kwargs: Unpack[OVERRIDABLE_CONFIG_KWARGS],
 ) -> LinearScheduler | CosineScheduler | PlateauWithAnnealingScheduler:
     """
     Converts a ScheduledParamConfig to a scheduler instance.
@@ -41,37 +56,48 @@ def create_scheduler(
             Default is False.
         apply_if_field: Require this key to be present in a param group before updating.
         ignore_if_field: Skip groups that declare this key in their dictionaries.
+        **kwargs: Optional overrides for scheduler configuration parameters (e.g., base_value,
+            final_value, warmup_ratio, etc.). These overrides take precedence over the values
+            provided in the `config` object.
 
     Returns:
         A scheduler instance based on the configuration.
 
     """
-    if config.scheduler_type is None:
+    scheduler_type = kwargs.get("scheduler_type", config.scheduler_type)
+    base_value = kwargs.get("base_value", config.base_value)
+    final_value = kwargs.get("final_value", config.final_value)
+    warmup_ratio = kwargs.get("warmup_ratio", config.warmup_ratio)
+    warmup_value = kwargs.get("warmup_value", config.warmup_value)
+    freeze_ratio = kwargs.get("freeze_ratio", config.freeze_ratio)
+    plateau_ratio = kwargs.get("plateau_ratio", config.plateau_ratio)
+
+    if scheduler_type is None:
         raise ValueError("scheduler_type must be specified in the config.")
 
-    if "plateau" in config.scheduler_type:
-        scheduler_type = "plateau"
+    if "plateau" in scheduler_type:
+        lookup_scheduler_type = "plateau"
     else:
-        scheduler_type = config.scheduler_type
-    scheduler_cls = SCHEDULER_MAPPING[scheduler_type]  # type: ignore
+        lookup_scheduler_type = scheduler_type
+    scheduler_cls = SCHEDULER_MAPPING[lookup_scheduler_type]  # type: ignore
 
     if issubclass(scheduler_cls, PlateauWithAnnealingScheduler):
-        if "cosine" in config.scheduler_type:
+        if "cosine" in scheduler_type:
             annealing_type = "cosine"
-        elif "linear" in config.scheduler_type:
+        elif "linear" in scheduler_type:
             annealing_type = "linear"
         else:
-            raise ValueError(f"Unknown annealing_type: {config.scheduler_type}")
+            raise ValueError(f"Unknown annealing_type: {scheduler_type}")
         scheduler = scheduler_cls(
             optimizer=optim,
             param_group_field=param_group_field,
             num_iters=num_iters,
-            plateau_value=config.base_value,
-            final_value=config.final_value,  # type: ignore
-            warmup_ratio=config.warmup_ratio,
-            warmup_value=config.warmup_value,
-            freeze_ratio=config.freeze_ratio,
-            plateau_ratio=config.plateau_ratio,  # type: ignore
+            plateau_value=base_value,
+            final_value=final_value,  # type: ignore
+            warmup_ratio=warmup_ratio,
+            warmup_value=warmup_value,
+            freeze_ratio=freeze_ratio,
+            plateau_ratio=plateau_ratio,  # type: ignore
             annealing_type=annealing_type,
             multiplier_field=multiplier_field,
             skip_if_zero=skip_if_zero,
@@ -83,8 +109,8 @@ def create_scheduler(
             optimizer=optim,
             param_group_field=param_group_field,
             num_iters=num_iters,
-            initial_value=config.base_value,
-            final_value=config.final_value,  # type: ignore
+            initial_value=base_value,
+            final_value=final_value,  # type: ignore
             multiplier_field=multiplier_field,
             skip_if_zero=skip_if_zero,
             apply_if_field=apply_if_field,
@@ -95,23 +121,23 @@ def create_scheduler(
             optimizer=optim,
             param_group_field=param_group_field,
             num_iters=num_iters,
-            base_value=config.base_value,
-            final_value=config.final_value,  # type: ignore
-            warmup_ratio=config.warmup_ratio,
-            warmup_value=config.warmup_value,
-            freeze_ratio=config.freeze_ratio,
+            base_value=base_value,
+            final_value=final_value,  # type: ignore
+            warmup_ratio=warmup_ratio,
+            warmup_value=warmup_value,
+            freeze_ratio=freeze_ratio,
             multiplier_field=multiplier_field,
             skip_if_zero=skip_if_zero,
             apply_if_field=apply_if_field,
             ignore_if_field=ignore_if_field,
         )
     else:
-        raise ValueError(f"Unsupported scheduler type: {config.scheduler_type}")
+        raise ValueError(f"Unsupported scheduler type: {scheduler_type}")
     return scheduler
 
 
 def create_optimizer(  # noqa: C901
-    parameters_groups: dict[str, Any],
+    parameters_groups: ParamsT,
     optimizer_config: OPTIMIZER_CONFIG,
     lr: float,
     weight_decay: float,
@@ -120,8 +146,7 @@ def create_optimizer(  # noqa: C901
     Creates an optimizer based on the configuration.
 
     Args:
-        parameters_groups: Dictionary containing model parameters
-            (key "params" and per-group options, i.e. "lr", "weight_decay" and etc.).
+        parameters_groups: Parameter groups for the optimizer.
         optimizer_config: Configuration for the optimizer.
         lr: Learning rate.
         weight_decay: Weight decay.

kostyl/ml/params_groups.py

@@ -28,8 +28,8 @@ def create_params_groups(
             Defaults to None, which uses an empty set.
         no_decay_keywords (set[str] | None, optional): A set of string keywords. If a parameter's
             name contains any of these keywords, its weight decay is set to 0.0.
-            If additional keywords are provided, they will be added to the default set.
-            Defaults to None, which uses a standard set of exclusion keywords:
+            If keywords are provided, they will be added to the default set, otherwise the default set is used.
+            Default set of keywords:
             {"norm", "bias", "embedding", "tokenizer", "ln", "scale"}.
 
     Returns:

{kostyl_toolkit-0.1.41.dist-info → kostyl_toolkit-0.1.43.dist-info}/METADATA

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

@@ -4,8 +4,8 @@ kostyl/ml/base_uploader.py,sha256=KxHuohCcNK18kTVFBBqDu_IOQefluhSXOzwC56O66wc,48
 kostyl/ml/configs/__init__.py,sha256=ytr2HFtsqMIs1pqQ-Ma0jv8c6Ni1-UxPnDTLu16cL68,1189
 kostyl/ml/configs/hyperparams.py,sha256=poOsR6xbE6jHO02GzxXfbz-6PA0HqULgcTgiuT_4oOM,4401
 kostyl/ml/configs/mixins.py,sha256=xHHAoRoPbzP9ECFP9duzg6SzegHcoLI8Pr9NrLoWNHs,1411
-kostyl/ml/configs/training_settings.py,sha256=wT9CHuLaKrLwonsc87Ee421EyFis_c9fqOgn9bSClm8,2747
-kostyl/ml/data_collator.py,sha256=kxiaMDKwSKXGBtrF8yXxHcypf7t_6syU-NwO1LcX50k,4062
+kostyl/ml/configs/training_settings.py,sha256=U3K45eLVrHTOfZ7p9TqTKjGWd4sdjKmJOe3a6U9kUWw,2877
+kostyl/ml/data_collator.py,sha256=NHfV5HcebHzg6iSfXj_cSUcBx9l8qP1HTbIPvqCDlTs,4062
 kostyl/ml/dist_utils.py,sha256=UFNMLEHc0A5F6KvTRG8GQPpRDwG4m5dvM__UvXNc2aQ,4526
 kostyl/ml/integrations/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 kostyl/ml/integrations/clearml/__init__.py,sha256=cIJMmiTgRrBNlCwejs6x4DrJydv7zYRPqUCWpVgQJYo,970
@@ -25,18 +25,18 @@ kostyl/ml/integrations/lightning/mixins.py,sha256=hVIsIUu6Iryrz6S7GQTqog9vNq8LQy
 kostyl/ml/integrations/lightning/module.py,sha256=APodFpspSoBYmrtTnCpUFq3WFEvpmHp2hlWUxwwAChM,5219
 kostyl/ml/integrations/lightning/utils.py,sha256=QhbK5iTv07xkRKuyXiK05aoY_ObSGorzGu1WcFvvFtI,1712
 kostyl/ml/optim/__init__.py,sha256=GgCr94WbiM-7MVDf6fPlbeUSxS33ufn_UAmqfKerElc,140
-kostyl/ml/optim/factory.py,sha256=1WYbAOOGepVL55gSVCJs5WBsiiE4k4ayrS5svFh9kbE,9760
+kostyl/ml/optim/factory.py,sha256=sPFIQkzIA4DKpny1hrcxHTKjEOXcbwLVy2CI9ZGZ938,10713
 kostyl/ml/optim/schedulers/__init__.py,sha256=lVwWj_mtUDo5OLtr8NNljLQgWGJUKPJWV7bm6-997ag,1486
 kostyl/ml/optim/schedulers/base.py,sha256=bjmwgdZpnSqpCnHPnKC6MEiRO79cwxMJpZq-eQVNs2M,1353
 kostyl/ml/optim/schedulers/composite.py,sha256=ee4xlMDMMtjKPkbTF2ue9GTr9DuGCGjZWf11mHbi6aE,2387
 kostyl/ml/optim/schedulers/cosine.py,sha256=LnFujAKaM8v5soPwiWG899nUMeaeAe8oTiaJqvoQxg8,8705
 kostyl/ml/optim/schedulers/linear.py,sha256=hlHfU88WLwfke8XGmRUzRKmEpqyyaLb4RRqq6EIzIxU,5881
 kostyl/ml/optim/schedulers/plateau.py,sha256=kjHivSp2FIMXExgS0q_Wbxk9pEHnuzhrIu4TQ59kYNw,11255
-kostyl/ml/params_groups.py,sha256=nUyw5d06Pvy9QPiYtZzLYR87xwXqJLxbHthgQH8oSCM,3583
+kostyl/ml/params_groups.py,sha256=4jMoNKvWTfqyQkfFCvU75VRk_dJaVTgShqbuTOnhOYU,3565
 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=CgNFNogcK0hoZmygvBWlTcq5A3m2Pfv9eOAP_gwx0pM,6633
-kostyl_toolkit-0.1.41.dist-info/WHEEL,sha256=e_m4S054HL0hyR3CpOk-b7Q7fDX6BuFkgL5OjAExXas,80
-kostyl_toolkit-0.1.41.dist-info/METADATA,sha256=RoZ8dXiU-CwUldPkhjY2achhvcQ2EE4RzhsJKTXe-AM,4121
-kostyl_toolkit-0.1.41.dist-info/RECORD,,
+kostyl_toolkit-0.1.43.dist-info/WHEEL,sha256=fAguSjoiATBe7TNBkJwOjyL1Tt4wwiaQGtNtjRPNMQA,80
+kostyl_toolkit-0.1.43.dist-info/METADATA,sha256=0Qt13h-ngr5NzIJd9m0NfELBQug8M5-YzzW0YBqjYXs,4121
+kostyl_toolkit-0.1.43.dist-info/RECORD,,

{kostyl_toolkit-0.1.41.dist-info → kostyl_toolkit-0.1.43.dist-info}/WHEEL

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