nshtrainer 0.1.0__py3-none-any.whl

nshtrainer/callbacks/early_stopping.py

@@ -0,0 +1,112 @@
+import math
+from logging import getLogger
+
+from lightning.fabric.utilities.rank_zero import _get_rank
+from lightning.pytorch import Trainer
+from lightning.pytorch.callbacks import EarlyStopping as _EarlyStopping
+from lightning.pytorch.utilities.rank_zero import rank_prefixed_message
+from typing_extensions import override
+
+log = getLogger(__name__)
+
+
+class EarlyStopping(_EarlyStopping):
+    def __init__(
+        self,
+        monitor: str,
+        min_delta: float = 0,
+        min_lr: float | None = None,
+        patience: int = 3,
+        verbose: bool = True,
+        mode: str = "min",
+        strict: bool = True,
+        check_finite: bool = True,
+        stopping_threshold: float | None = None,
+        divergence_threshold: float | None = None,
+        check_on_train_epoch_end: bool | None = None,
+        log_rank_zero_only: bool = False,
+    ):
+        super().__init__(
+            monitor,
+            min_delta,
+            patience,
+            verbose,
+            mode,
+            strict,
+            check_finite,
+            stopping_threshold,
+            divergence_threshold,
+            check_on_train_epoch_end,
+            log_rank_zero_only,
+        )
+
+        self.min_lr = min_lr
+
+    @override
+    @staticmethod
+    def _log_info(
+        trainer: Trainer | None, message: str, log_rank_zero_only: bool
+    ) -> None:
+        rank = _get_rank()
+        if trainer is not None and trainer.world_size <= 1:
+            rank = None
+        message = rank_prefixed_message(message, rank)
+        if rank is None or not log_rank_zero_only or rank == 0:
+            log.critical(message)
+
+    @override
+    def _run_early_stopping_check(self, trainer: Trainer):
+        """Checks whether the early stopping condition is met and if so tells the trainer to stop the training."""
+        logs = trainer.callback_metrics
+
+        # Disable early_stopping with fast_dev_run
+        if getattr(trainer, "fast_dev_run", False):
+            return
+
+        should_stop, reason = False, None
+
+        if not should_stop:
+            should_stop, reason = self._evaluate_stopping_criteria_min_lr(trainer)
+
+        # If metric present
+        if not should_stop and self._validate_condition_metric(logs):
+            current = logs[self.monitor].squeeze()
+            should_stop, reason = self._evaluate_stopping_criteria(current)
+
+        # stop every ddp process if any world process decides to stop
+        should_stop = trainer.strategy.reduce_boolean_decision(should_stop, all=False)
+        trainer.should_stop = trainer.should_stop or should_stop
+        if should_stop:
+            self.stopped_epoch = trainer.current_epoch
+        if reason and self.verbose:
+            self._log_info(trainer, reason, self.log_rank_zero_only)
+
+    def _evaluate_stopping_criteria_min_lr(
+        self, trainer: Trainer
+    ) -> tuple[bool, str | None]:
+        if self.min_lr is None:
+            return False, None
+
+        # Get the maximum LR across all param groups in all optimizers
+        model_max_lr = max(
+            [
+                param_group["lr"]
+                for optimizer in trainer.optimizers
+                for param_group in optimizer.param_groups
+            ]
+        )
+        if not isinstance(model_max_lr, float) or not math.isfinite(model_max_lr):
+            return False, None
+
+        # If the maximum LR is less than the minimum LR, stop training
+        if model_max_lr >= self.min_lr:
+            return False, None
+
+        return True, (
+            "Stopping threshold reached: "
+            f"The maximum LR of the model across all param groups is {model_max_lr:.2e} "
+            f"which is less than the minimum LR {self.min_lr:.2e}"
+        )
+
+    def on_early_stopping(self, trainer: Trainer):
+        pass

nshtrainer/callbacks/ema.py

@@ -0,0 +1,383 @@
+import contextlib
+import copy
+import threading
+from collections.abc import Callable, Iterable
+from typing import Any, Literal, overload
+
+import lightning.pytorch as pl
+import torch
+from lightning.pytorch import Callback
+from lightning.pytorch.utilities.exceptions import MisconfigurationException
+from typing_extensions import override
+
+from .base import CallbackConfigBase
+
+
+class EMA(Callback):
+    """
+    Implements Exponential Moving Averaging (EMA).
+
+    When training a model, this callback will maintain moving averages of the trained parameters.
+    When evaluating, we use the moving averages copy of the trained parameters.
+    When saving, we save an additional set of parameters with the prefix `ema`.
+
+    Args:
+        decay: The exponential decay used when calculating the moving average. Has to be between 0-1.
+        validate_original_weights: Validate the original weights, as apposed to the EMA weights.
+        every_n_steps: Apply EMA every N steps.
+        cpu_offload: Offload weights to CPU.
+    """
+
+    @override
+    def __init__(
+        self,
+        decay: float,
+        validate_original_weights: bool = False,
+        every_n_steps: int = 1,
+        cpu_offload: bool = False,
+    ):
+        if not (0 <= decay <= 1):
+            raise MisconfigurationException("EMA decay value must be between 0 and 1")
+        self.decay = decay
+        self.validate_original_weights = validate_original_weights
+        self.every_n_steps = every_n_steps
+        self.cpu_offload = cpu_offload
+
+    @override
+    def on_fit_start(
+        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule"
+    ) -> None:
+        device = pl_module.device if not self.cpu_offload else torch.device("cpu")
+        trainer.optimizers = [
+            EMAOptimizer(
+                optim,
+                device=device,
+                decay=self.decay,
+                every_n_steps=self.every_n_steps,
+                current_step=trainer.global_step,
+            )
+            for optim in trainer.optimizers
+            if not isinstance(optim, EMAOptimizer)
+        ]
+
+    @override
+    def on_validation_start(
+        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule"
+    ) -> None:
+        if self._should_validate_ema_weights(trainer):
+            self.swap_model_weights(trainer)
+
+    @override
+    def on_validation_end(
+        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule"
+    ) -> None:
+        if self._should_validate_ema_weights(trainer):
+            self.swap_model_weights(trainer)
+
+    @override
+    def on_test_start(
+        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule"
+    ) -> None:
+        if self._should_validate_ema_weights(trainer):
+            self.swap_model_weights(trainer)
+
+    @override
+    def on_test_end(
+        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule"
+    ) -> None:
+        if self._should_validate_ema_weights(trainer):
+            self.swap_model_weights(trainer)
+
+    def _should_validate_ema_weights(self, trainer: "pl.Trainer") -> bool:
+        return not self.validate_original_weights and self._ema_initialized(trainer)
+
+    def _ema_initialized(self, trainer: "pl.Trainer") -> bool:
+        return any(
+            isinstance(optimizer, EMAOptimizer) for optimizer in trainer.optimizers
+        )
+
+    def swap_model_weights(self, trainer: "pl.Trainer", saving_ema_model: bool = False):
+        for optimizer in trainer.optimizers:
+            assert isinstance(optimizer, EMAOptimizer)
+            optimizer.switch_main_parameter_weights(saving_ema_model)
+
+    @contextlib.contextmanager
+    def save_ema_model(self, trainer: "pl.Trainer"):
+        """
+        Saves an EMA copy of the model + EMA optimizer states for resume.
+        """
+        self.swap_model_weights(trainer, saving_ema_model=True)
+        try:
+            yield
+        finally:
+            self.swap_model_weights(trainer, saving_ema_model=False)
+
+    @contextlib.contextmanager
+    def save_original_optimizer_state(self, trainer: "pl.Trainer"):
+        for optimizer in trainer.optimizers:
+            assert isinstance(optimizer, EMAOptimizer)
+            optimizer.save_original_optimizer_state = True
+        try:
+            yield
+        finally:
+            for optimizer in trainer.optimizers:
+                assert isinstance(optimizer, EMAOptimizer)
+                optimizer.save_original_optimizer_state = False
+
+
+@torch.no_grad()
+def ema_update(ema_model_tuple, current_model_tuple, decay):
+    torch._foreach_mul_(ema_model_tuple, decay)
+    torch._foreach_add_(
+        ema_model_tuple,
+        current_model_tuple,
+        alpha=(1.0 - decay),
+    )
+
+
+def run_ema_update_cpu(
+    ema_model_tuple, current_model_tuple, decay, pre_sync_stream=None
+):
+    if pre_sync_stream is not None:
+        pre_sync_stream.synchronize()
+
+    ema_update(ema_model_tuple, current_model_tuple, decay)
+
+
+class EMAOptimizer(torch.optim.Optimizer):
+    r"""
+    EMAOptimizer is a wrapper for torch.optim.Optimizer that computes
+    Exponential Moving Average of parameters registered in the optimizer.
+
+    EMA parameters are automatically updated after every step of the optimizer
+    with the following formula:
+
+        ema_weight = decay * ema_weight + (1 - decay) * training_weight
+
+    To access EMA parameters, use ``swap_ema_weights()`` context manager to
+    perform a temporary in-place swap of regular parameters with EMA
+    parameters.
+
+    Notes:
+        - EMAOptimizer is not compatible with APEX AMP O2.
+
+    Args:
+        optimizer (torch.optim.Optimizer): optimizer to wrap
+        device (torch.device): device for EMA parameters
+        decay (float): decay factor
+
+    Returns:
+        returns an instance of torch.optim.Optimizer that computes EMA of
+        parameters
+
+    Example:
+        model = Model().to(device)
+        opt = torch.optim.Adam(model.parameters())
+
+        opt = EMAOptimizer(opt, device, 0.9999)
+
+        for epoch in range(epochs):
+            training_loop(model, opt)
+
+            regular_eval_accuracy = evaluate(model)
+
+            with opt.swap_ema_weights():
+                ema_eval_accuracy = evaluate(model)
+    """
+
+    stream: Any | None
+
+    @override
+    def __init__(
+        self,
+        optimizer: torch.optim.Optimizer,
+        device: torch.device,
+        decay: float = 0.9999,
+        every_n_steps: int = 1,
+        current_step: int = 0,
+    ):
+        self.optimizer = optimizer
+        self.decay = decay
+        self.device = device
+        self.current_step = current_step
+        self.every_n_steps = every_n_steps
+        self.save_original_optimizer_state = False
+
+        self.first_iteration = True
+        self.rebuild_ema_params = True
+        self.stream = None
+        self.thread = None
+
+        self.ema_params = ()
+        self.in_saving_ema_model_context = False
+
+    def all_parameters(self) -> Iterable[torch.Tensor]:
+        return (param for group in self.param_groups for param in group["params"])
+
+    @overload
+    def step(self, closure: None = ...) -> None: ...
+
+    @overload
+    def step(self, closure: Callable[[], float]) -> float: ...
+
+    @override
+    def step(self, closure: Callable[[], float] | None = None) -> float | None:
+        self.join()
+
+        if self.first_iteration:
+            if any(p.is_cuda for p in self.all_parameters()):
+                self.stream = torch.cuda.Stream()
+
+            self.first_iteration = False
+
+        if self.rebuild_ema_params:
+            opt_params = list(self.all_parameters())
+
+            self.ema_params += tuple(
+                copy.deepcopy(param.data.detach()).to(self.device)
+                for param in opt_params[len(self.ema_params) :]
+            )
+            self.rebuild_ema_params = False
+
+        loss = self.optimizer.step(closure)
+
+        if self._should_update_at_step():
+            self.update()
+        self.current_step += 1
+        return loss
+
+    def _should_update_at_step(self) -> bool:
+        return self.current_step % self.every_n_steps == 0
+
+    @torch.no_grad()
+    def update(self):
+        if self.stream is not None:
+            self.stream.wait_stream(torch.cuda.current_stream())
+
+        with torch.cuda.stream(self.stream):
+            current_model_state = tuple(
+                param.data.to(self.device, non_blocking=True)
+                for param in self.all_parameters()
+            )
+
+            if self.device.type == "cuda":
+                ema_update(self.ema_params, current_model_state, self.decay)
+
+        if self.device.type == "cpu":
+            self.thread = threading.Thread(
+                target=run_ema_update_cpu,
+                args=(
+                    self.ema_params,
+                    current_model_state,
+                    self.decay,
+                    self.stream,
+                ),
+            )
+            self.thread.start()
+
+    def swap_tensors(self, tensor1, tensor2):
+        tmp = torch.empty_like(tensor1)
+        tmp.copy_(tensor1)
+        tensor1.copy_(tensor2)
+        tensor2.copy_(tmp)
+
+    def switch_main_parameter_weights(self, saving_ema_model: bool = False):
+        self.join()
+        self.in_saving_ema_model_context = saving_ema_model
+        for param, ema_param in zip(self.all_parameters(), self.ema_params):
+            self.swap_tensors(param.data, ema_param)
+
+    @contextlib.contextmanager
+    def swap_ema_weights(self, enabled: bool = True):
+        r"""
+        A context manager to in-place swap regular parameters with EMA
+        parameters.
+        It swaps back to the original regular parameters on context manager
+        exit.
+
+        Args:
+            enabled (bool): whether the swap should be performed
+        """
+
+        if enabled:
+            self.switch_main_parameter_weights()
+        try:
+            yield
+        finally:
+            if enabled:
+                self.switch_main_parameter_weights()
+
+    def __getattr__(self, name):
+        return getattr(self.optimizer, name)
+
+    def join(self):
+        if self.stream is not None:
+            self.stream.synchronize()
+
+        if self.thread is not None:
+            self.thread.join()
+
+    @override
+    def state_dict(self):
+        self.join()
+
+        if self.save_original_optimizer_state:
+            return self.optimizer.state_dict()
+
+        # if we are in the context of saving an EMA model, the EMA weights are in the modules' actual weights
+        ema_params = (
+            self.ema_params
+            if not self.in_saving_ema_model_context
+            else list(self.all_parameters())
+        )
+        state_dict = {
+            "opt": self.optimizer.state_dict(),
+            "ema": ema_params,
+            "current_step": self.current_step,
+            "decay": self.decay,
+            "every_n_steps": self.every_n_steps,
+        }
+        return state_dict
+
+    @override
+    def load_state_dict(self, state_dict):
+        self.join()
+
+        self.optimizer.load_state_dict(state_dict["opt"])
+        self.ema_params = tuple(
+            param.to(self.device) for param in copy.deepcopy(state_dict["ema"])
+        )
+        self.current_step = state_dict["current_step"]
+        self.decay = state_dict["decay"]
+        self.every_n_steps = state_dict["every_n_steps"]
+        self.rebuild_ema_params = False
+
+    @override
+    def add_param_group(self, param_group):
+        self.optimizer.add_param_group(param_group)
+        self.rebuild_ema_params = True
+
+
+class EMAConfig(CallbackConfigBase):
+    name: Literal["ema"] = "ema"
+
+    decay: float
+    """The exponential decay used when calculating the moving average. Has to be between 0-1."""
+
+    validate_original_weights: bool = False
+    """Validate the original weights, as apposed to the EMA weights."""
+
+    every_n_steps: int = 1
+    """Apply EMA every N steps."""
+
+    cpu_offload: bool = False
+    """Offload weights to CPU."""
+
+    @override
+    def construct_callbacks(self, root_config):
+        yield EMA(
+            decay=self.decay,
+            validate_original_weights=self.validate_original_weights,
+            every_n_steps=self.every_n_steps,
+            cpu_offload=self.cpu_offload,
+        )

nshtrainer/callbacks/finite_checks.py

@@ -0,0 +1,75 @@
+from logging import getLogger
+from typing import Literal
+
+import torch
+from lightning.pytorch import Callback, LightningModule, Trainer
+from typing_extensions import override
+
+from .base import CallbackConfigBase
+
+log = getLogger(__name__)
+
+
+def finite_checks(
+    module: LightningModule,
+    nonfinite_grads: bool = True,
+    none_grads: bool = False,
+):
+    for name, param in module.named_parameters():
+        if not param.requires_grad:
+            continue
+
+        if param.grad is None:
+            if none_grads:
+                log.critical(f"Parameter {name} ({param.shape}) has None gradients")
+            continue
+
+        if not nonfinite_grads or torch.isfinite(param.grad.float()).all():
+            continue
+
+        has_nan = torch.isnan(param.grad.float()).any()
+        has_inf = torch.isinf(param.grad.float()).any()
+        kinds = [
+            "NaN" if has_nan else None,
+            "Inf" if has_inf else None,
+        ]
+        kinds = " and ".join(prop for prop in kinds if prop is not None)
+        log.critical(f"{name} ({param.shape}) has {kinds} gradients")
+
+
+class FiniteChecksCallback(Callback):
+    def __init__(
+        self,
+        *,
+        nonfinite_grads: bool = True,
+        none_grads: bool = True,
+    ):
+        super().__init__()
+
+        self._nonfinite_grads = nonfinite_grads
+        self._none_grads = none_grads
+
+    @override
+    def on_after_backward(self, trainer: Trainer, pl_module: LightningModule):
+        finite_checks(
+            pl_module,
+            nonfinite_grads=self._nonfinite_grads,
+            none_grads=self._none_grads,
+        )
+
+
+class FiniteChecksConfig(CallbackConfigBase):
+    name: Literal["finite_checks"] = "finite_checks"
+
+    nonfinite_grads: bool = True
+    """Whether to check for non-finite (i.e. NaN or Inf) gradients"""
+
+    none_grads: bool = True
+    """Whether to check for None gradients"""
+
+    @override
+    def construct_callbacks(self, root_config):
+        yield FiniteChecksCallback(
+            nonfinite_grads=self.nonfinite_grads,
+            none_grads=self.none_grads,
+        )

nshtrainer/callbacks/gradient_skipping.py

@@ -0,0 +1,103 @@
+from logging import getLogger
+from typing import Literal, Protocol, runtime_checkable
+
+import torch
+import torchmetrics
+from lightning.pytorch import Callback, LightningModule, Trainer
+from torch.optim import Optimizer
+from typing_extensions import override
+
+from .base import CallbackConfigBase
+from .norm_logging import compute_norm
+
+log = getLogger(__name__)
+
+
+@runtime_checkable
+class HasGradSkippedSteps(Protocol):
+    grad_skipped_steps: torchmetrics.SumMetric
+
+
+class GradientSkipping(Callback):
+    def __init__(self, config: "GradientSkippingConfig"):
+        super().__init__()
+
+        self.config = config
+
+    @override
+    def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str) -> None:
+        if not isinstance(pl_module, HasGradSkippedSteps):
+            pl_module.grad_skipped_steps = torchmetrics.SumMetric()
+
+    @override
+    def on_before_optimizer_step(
+        self, trainer: Trainer, pl_module: LightningModule, optimizer: Optimizer
+    ):
+        # This should never happen, but just in case
+        if not isinstance(pl_module, HasGradSkippedSteps):
+            raise TypeError(
+                f"Expected HasGradSkippedSteps, got {type(pl_module)} instead"
+            )
+
+        # Skip the step if the global step is less than the start_after_n_steps
+        # This is because we want to let AMP adjust the loss scale before we start
+        if (
+            self.config.start_after_n_steps is not None
+            and pl_module.global_step < self.config.start_after_n_steps
+        ):
+            return
+
+        norm = compute_norm(
+            pl_module,
+            optimizer,
+            self.config.norm_type,
+            grad=True,
+        )
+
+        # If the norm is NaN/Inf, we don't want to skip the step
+        # beacuse AMP checks for NaN/Inf grads to adjust the loss scale.
+        if self.config.skip_non_finite and not torch.isfinite(norm).all():
+            optimizer.zero_grad()
+            pl_module.grad_skipped_steps(1)
+            log.warning(
+                f"Skipping step at global step {pl_module.global_step} with non-finite norm {norm:.2f}"
+            )
+        elif (norm > self.config.threshold).any():
+            optimizer.zero_grad()
+            pl_module.grad_skipped_steps(1)
+            log.warning(
+                f"Skipping step at global step {pl_module.global_step} with norm {norm:.2f} > {self.config.threshold:.2f}"
+            )
+        else:
+            pl_module.grad_skipped_steps(0)
+
+        pl_module.log(
+            "train/grad_skipped_steps",
+            pl_module.grad_skipped_steps,
+            on_step=True,
+            on_epoch=False,
+        )
+
+
+class GradientSkippingConfig(CallbackConfigBase):
+    name: Literal["gradient_skipping"] = "gradient_skipping"
+
+    threshold: float
+    """Threshold to use for gradient skipping."""
+
+    norm_type: str | float = 2.0
+    """Norm type to use for gradient skipping."""
+
+    start_after_n_steps: int | None = 100
+    """Number of steps to wait before starting gradient skipping."""
+
+    skip_non_finite: bool = False
+    """
+    If False, it doesn't skip steps with non-finite norms. This is useful when using AMP, as AMP checks for NaN/Inf grads to adjust the loss scale. Otherwise, skips steps with non-finite norms.
+
+    Should almost always be False, especially when using AMP (unless you know what you're doing!).
+    """
+
+    @override
+    def construct_callbacks(self, root_config):
+        yield GradientSkipping(self)