nshtrainer 0.26.2__tar.gz → 0.28.0__tar.gz

{nshtrainer-0.26.2 → nshtrainer-0.28.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.26.2
+Version: 0.28.0
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-0.26.2 → nshtrainer-0.28.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshtrainer"
-version = "0.26.2"
+version = "0.28.0"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"

{nshtrainer-0.26.2 → nshtrainer-0.28.0}/src/nshtrainer/_hf_hub.py

@@ -359,19 +359,3 @@ class HFHubCallback(NTCallbackBase):
         # NOTE: This file is fairly small, so we can just upload it directly.
         # No need to copy.
         self._save_file(metadata_path)
-
-    @override
-    def state_dict(self):
-        return {
-            "repo_id": self._repo_id,
-            "checksum_to_path_in_repo": {
-                k: str(v) for k, v in self._checksum_to_path_in_repo.items()
-            },
-        }
-
-    @override
-    def load_state_dict(self, state_dict):
-        self._repo_id = state_dict["repo_id"]
-        self._checksum_to_path_in_repo = {
-            k: Path(v) for k, v in state_dict["checksum_to_path_in_repo"].items()
-        }

{nshtrainer-0.26.2 → nshtrainer-0.28.0}/src/nshtrainer/callbacks/checkpoint/_base.py

@@ -152,12 +152,7 @@ class CheckpointBase(Checkpoint, ABC, Generic[TConfig]):
 
         # Save the new checkpoint
         filepath = self.resolve_checkpoint_path(self.current_metrics(trainer))
-        trainer._nshtrainer_save_checkpoint(
-            filepath,
-            self.config.save_weights_only,
-            use_checkpoint_cache=None,
-            ckpt_cache_use_symlink=False,
-        )
+        trainer.save_checkpoint(filepath, self.config.save_weights_only)
 
         if trainer.is_global_zero:
             # Create the latest symlink

{nshtrainer-0.26.2 → nshtrainer-0.28.0}/src/nshtrainer/model/config.py

@@ -1012,8 +1012,6 @@ class TrainerConfig(C.Config):
     """If enabled, the model supports scaling the gradients of shared parameters that are registered using `LightningModuleBase.register_shared_parameters(...)`"""
     save_checkpoint_metadata: bool = True
     """If enabled, will save additional metadata whenever a checkpoint is saved."""
-    use_checkpoint_cache: bool = True
-    """If enabled, will optimize the saving of duplicate checkpoints by creating symlinks instead of copying the file."""
 
     lightning_kwargs: LightningTrainerKwargs = LightningTrainerKwargs()
     """

{nshtrainer-0.26.2 → nshtrainer-0.28.0}/src/nshtrainer/trainer/trainer.py

@@ -1,6 +1,5 @@
 import logging
 import os
-import shutil
 from collections.abc import Sequence
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, cast
@@ -18,7 +17,6 @@ from lightning.pytorch.utilities.types import _EVALUATE_OUTPUT, _PREDICT_OUTPUT
 from typing_extensions import Unpack, assert_never, override
 
 from .._checkpoint.metadata import _write_checkpoint_metadata
-from .._checkpoint.saver import _link_checkpoint
 from ..callbacks.base import resolve_all_callbacks
 from ..model.config import (
     AcceleratorConfigProtocol,
@@ -287,8 +285,6 @@ class Trainer(LightningTrainer):
         /,
         **kwargs: Unpack[LightningTrainerKwargs],
     ):
-        self._nshtrainer_checkpoint_cache: dict[tuple[int, int], Path] = {}
-
         self._pre_init(config)
 
         kwargs = self._update_kwargs(config, kwargs)
@@ -411,54 +407,24 @@ class Trainer(LightningTrainer):
 
         return super()._run(model, ckpt_path)
 
-    def _nshtrainer_save_checkpoint(
+    @override
+    def save_checkpoint(
         self,
         filepath: str | Path,
         weights_only: bool = False,
         storage_options: Any | None = None,
-        use_checkpoint_cache: bool | None = None,
-        ckpt_cache_use_symlink: bool = False,
     ):
-        lm = self._base_module
-        root_config = cast(BaseConfig, lm.hparams)
-        if use_checkpoint_cache is None:
-            use_checkpoint_cache = root_config.trainer.use_checkpoint_cache
-
         filepath = Path(filepath)
 
         # List of files that we should upload to HF
         written_files: list[Path] = [filepath]
 
-        cached_path = None
-        if (
-            use_checkpoint_cache
-            and (
-                cached_path := self._nshtrainer_checkpoint_cache.get(
-                    (self.current_epoch, self.global_step)
-                )
-            )
-            is not None
-        ):
-            # If we have a cached path, then we symlink it to the new path.
-            log.info(f"Re-using cached checkpoint {cached_path} for {filepath}.")
-            if self.is_global_zero:
-                if ckpt_cache_use_symlink:
-                    _link_checkpoint(cached_path, filepath, metadata=False)
-                else:
-                    shutil.copy(cached_path, filepath)
-            self.strategy.barrier("Trainer.save_checkpoint")
-        else:
-            super().save_checkpoint(filepath, weights_only, storage_options)
-
-        # If we are using the cache but we don't have a cached path, then we save the checkpoint to the cache.
-        if use_checkpoint_cache and cached_path is None:
-            self._nshtrainer_checkpoint_cache[
-                (self.current_epoch, self.global_step)
-            ] = filepath
-            log.debug(f"Checkpoint saved to cache: {filepath}")
+        super().save_checkpoint(filepath, weights_only, storage_options)
 
         # Save the checkpoint metadata
         metadata_path = None
+        lm = self._base_module
+        root_config = cast(BaseConfig, lm.hparams)
         if root_config.trainer.save_checkpoint_metadata and self.is_global_zero:
             # Generate the metadata and write to disk
             if (
@@ -470,17 +436,3 @@ class Trainer(LightningTrainer):
         from .. import _callback
 
         _callback._call_on_checkpoint_saved(self, filepath, metadata_path)
-
-    @override
-    def save_checkpoint(
-        self,
-        filepath: str | Path,
-        weights_only: bool = False,
-        storage_options: Any | None = None,
-    ):
-        return self._nshtrainer_save_checkpoint(
-            filepath=filepath,
-            weights_only=weights_only,
-            storage_options=storage_options,
-            use_checkpoint_cache=False,
-        )