Source code for tensorplay.amp.grad_scaler

# mypy: allow-untyped-defs
from __future__ import annotations

import inspect
import warnings
from collections import abc, defaultdict
from enum import Enum
from typing import Any, cast, overload, TYPE_CHECKING

import tensorplay


if TYPE_CHECKING:
    from collections.abc import Iterable


__all__ = ["OptState", "GradScaler"]


def _cuda_amp_definitely_not_available() -> bool:
    try:
        import tensorplay.cuda as cuda

        return not cuda.is_available()
    except Exception:
        return True


def _device_key(device: Any) -> tuple[str, int]:
    """

    ``tensorplay.Device`` is not hashable, so per-device maps are keyed by
    ``(type, index)`` tuples instead.
    """
    return (device.type, device.index)


class _MultiDeviceReplicator:
    """Lazily serves copies of a tensor to requested devices.

    Copies are cached per-device.  ``device`` may be a
    :class:`tensorplay.Device` or a ``(type, index)`` key from
    :func:`_device_key`.
    """

    def __init__(self, master_tensor: tensorplay.Tensor) -> None:
        self.master = master_tensor
        self._per_device_tensors: dict[Any, tensorplay.Tensor] = {}

    def get(self, device: Any) -> tensorplay.Tensor:
        retval = self._per_device_tensors.get(device, None)
        if retval is None:
            target = (
                device
                if isinstance(device, tensorplay.Device)
                else tensorplay.Device(*device)
            )
            retval = self.master.to(device=target, non_blocking=True, copy=True)
            self._per_device_tensors[device] = retval
        return retval


# Defines default_factory for GradScaler's _per_optimizer_states defaultdict,
# as well as associated "enum" values.  Prefers defining these at top level because
# - Lambdas can't be pickled, so we don't want to supply a lambda as the factory.
# - Defining READY, UNSCALED, STEPPED and _refresh_per_optimizer_state within GradScaler
#   causes a circular reference, which we'd rather avoid.
class OptState(Enum):
    READY = 0
    UNSCALED = 1
    STEPPED = 2


def _refresh_per_optimizer_state() -> dict[str, Any]:
    return {"stage": OptState.READY, "found_inf_per_device": {}}



[docs]
class GradScaler:
    """An instance ``scaler`` of :class:`GradScaler`.

    Helps perform the steps of gradient scaling
    conveniently.

    * ``scaler.scale(loss)`` multiplies a given loss by ``scaler``'s current scale factor.
    * ``scaler.step(optimizer)`` safely unscales gradients and calls ``optimizer.step()``.
    * ``scaler.update()`` updates ``scaler``'s scale factor.

    Example::

        # Creates a GradScaler once at the beginning of training.
        scaler = GradScaler()

        for epoch in epochs:
            for input, target in data:
                optimizer.zero_grad()
                output = model(input)
                loss = loss_fn(output, target)

                # Scales loss.  Calls backward() on scaled loss to create scaled gradients.
                scaler.scale(loss).backward()

                # scaler.step() first unscales gradients of the optimizer's params.
                # If gradients don't contain infs/NaNs, optimizer.step() is then called,
                # otherwise, optimizer.step() is skipped.
                scaler.step(optimizer)

                # Updates the scale for next iteration.
                scaler.update()

    ``scaler`` dynamically estimates the scale factor each iteration.  To minimize gradient underflow,
    a large scale factor should be used.  However, ``float16`` values can "overflow" (become inf or NaN) if
    the scale factor is too large.  Therefore, the optimal scale factor is the largest factor that can be used
    without incurring inf or NaN gradient values.
    ``scaler`` approximates the optimal scale factor over time by checking the gradients for infs and NaNs during every
    ``scaler.step(optimizer)`` (or optional separate ``scaler.unscale_(optimizer)``, see :meth:`unscale_`).

    * If infs/NaNs are found, ``scaler.step(optimizer)`` skips the underlying ``optimizer.step()`` (so the params
      themselves remain uncorrupted) and ``update()`` multiplies the scale by ``backoff_factor``.

    * If no infs/NaNs are found, ``scaler.step(optimizer)`` runs the underlying ``optimizer.step()`` as usual.
      If ``growth_interval`` unskipped iterations occur consecutively, ``update()`` multiplies the scale by
      ``growth_factor``.

    Args:
        device (str, optional, default="cuda"): Device type to use. Possible values are: 'cuda' and 'cpu'.
            The type is the same as the `type` attribute of a :class:`tensorplay.device`.
            Thus, you may obtain the device type of a tensor using `Tensor.device.type`.
        init_scale (float, optional, default=2.**16):  Initial scale factor.
        growth_factor (float, optional, default=2.0):  Factor by which the scale is multiplied during
            :meth:`update` if no inf/NaN gradients occur for ``growth_interval`` consecutive iterations.
        backoff_factor (float, optional, default=0.5):  Factor by which the scale is multiplied during
            :meth:`update` if inf/NaN gradients occur in an iteration.
        growth_interval (int, optional, default=2000):  Number of consecutive iterations without inf/NaN gradients
            that must occur for the scale to be multiplied by ``growth_factor``.
        enabled (bool, optional):  If ``False``, disables gradient scaling. :meth:`step` simply
            invokes the underlying ``optimizer.step()``, and other methods become no-ops.
            Default: ``True``
    """

    def __init__(
        self,
        device: str = "cuda",
        init_scale: float = 2.0**16,
        growth_factor: float = 2.0,
        backoff_factor: float = 0.5,
        growth_interval: int = 2000,
        enabled: bool = True,
    ) -> None:
        self._device = device
        self._enabled = enabled
        if self._device == "cuda":
            if enabled and _cuda_amp_definitely_not_available():
                warnings.warn(
                    "tensorplay.cuda.amp.GradScaler is enabled, but CUDA is not available.  Disabling.",
                    stacklevel=2,
                )
                self._enabled = False

        if self._enabled:
            if growth_factor <= 1.0:
                raise AssertionError("The growth factor must be > 1.0.")
            if backoff_factor >= 1.0:
                raise AssertionError("The backoff factor must be < 1.0.")

            self._init_scale = init_scale
            # self._scale will be lazily initialized during the first call to scale()
            self._scale: tensorplay.Tensor | None = None
            self._growth_factor = growth_factor
            self._backoff_factor = backoff_factor
            self._growth_interval = growth_interval
            self._init_growth_tracker = 0
            # self._growth_tracker will be lazily initialized during the first call to scale()
            self._growth_tracker: tensorplay.Tensor | None = None
            self._per_optimizer_states: dict[int, dict[str, Any]] = defaultdict(
                _refresh_per_optimizer_state
            )

    def _check_scale_growth_tracker(
        self, funcname: str
    ) -> tuple[tensorplay.Tensor, tensorplay.Tensor]:
        fix = "This may indicate your script did not use scaler.scale(loss or outputs) earlier in the iteration."
        if self._scale is None:
            raise AssertionError(f"Attempted {funcname} but _scale is None.  " + fix)
        if self._growth_tracker is None:
            raise AssertionError(
                f"Attempted {funcname} but _growth_tracker is None.  " + fix
            )
        return (self._scale, self._growth_tracker)

    def _lazy_init_scale_growth_tracker(self, dev: Any) -> None:
        if self._growth_tracker is not None:
            raise AssertionError("_growth_tracker initialized before _scale")
        self._scale = tensorplay.full(
            (), self._init_scale, dtype=tensorplay.float32, device=dev
        )
        self._growth_tracker = tensorplay.full(
            (), self._init_growth_tracker, dtype=tensorplay.int32, device=dev
        )

    @overload
    def scale(self, outputs: tensorplay.Tensor) -> tensorplay.Tensor: ...

    @overload
    def scale(self, outputs: list[tensorplay.Tensor]) -> list[tensorplay.Tensor]: ...

    @overload
    def scale(
        self, outputs: tuple[tensorplay.Tensor, ...]
    ) -> tuple[tensorplay.Tensor, ...]: ...

    @overload
    def scale(
        self, outputs: Iterable[tensorplay.Tensor]
    ) -> Iterable[tensorplay.Tensor]: ...


[docs]
    def scale(
        self,
        outputs: tensorplay.Tensor | Iterable[tensorplay.Tensor],
    ) -> tensorplay.Tensor | Iterable[tensorplay.Tensor]:
        """
        Multiplies ('scales') a tensor or list of tensors by the scale factor.

        Returns scaled outputs.  If this instance of :class:`GradScaler` is not enabled, outputs are returned
        unmodified.

        Args:
            outputs (Tensor or iterable of Tensors):  Outputs to scale.
        """
        if not self._enabled:
            return outputs

        # Short-circuit for the common case.
        if isinstance(outputs, tensorplay.Tensor):
            if self._scale is None:
                self._lazy_init_scale_growth_tracker(outputs.device)
            # _scale is now guaranteed to be set
            scale = self._scale
            if scale is None:
                raise AssertionError("_scale should not be None after lazy init")
            return outputs * scale.to(device=outputs.device, non_blocking=True)

        # Invoke the more complex machinery only if we're treating multiple outputs.
        stash: list[
            _MultiDeviceReplicator
        ] = []  # holds a reference that can be overwritten by apply_scale

        def apply_scale(val: tensorplay.Tensor | Iterable[tensorplay.Tensor]):
            if isinstance(val, tensorplay.Tensor):
                if len(stash) == 0:
                    if self._scale is None:
                        self._lazy_init_scale_growth_tracker(val.device)
                    # _scale is now guaranteed to be set
                    scale = self._scale
                    if scale is None:
                        raise AssertionError(
                            "_scale should not be None after lazy init"
                        )
                    stash.append(_MultiDeviceReplicator(scale))
                return val * stash[0].get(_device_key(val.device))
            if isinstance(val, abc.Iterable):
                iterable = map(apply_scale, val)
                if isinstance(val, (list, tuple)):
                    return type(val)(iterable)
                return iterable
            raise ValueError("outputs must be a Tensor or an iterable of Tensors")

        return apply_scale(outputs)


    def _unscale_grads_(
        self,
        optimizer: tensorplay.optim.Optimizer,
        inv_scale: tensorplay.Tensor,
        found_inf: tensorplay.Tensor,
        allow_fp16: bool,
    ) -> dict[Any, tensorplay.Tensor]:
        per_device_inv_scale = _MultiDeviceReplicator(inv_scale)
        per_device_found_inf = _MultiDeviceReplicator(found_inf)

        # To set up _amp_foreach_non_finite_check_and_unscale_, split grads by device and dtype.
        # There could be hundreds of grads, so we'd like to iterate through them just once.
        # However, we don't know their devices or dtypes in advance.
        per_device_and_dtype_grads: dict[
            Any, dict[Any, list[tensorplay.Tensor]]
        ] = defaultdict(lambda: defaultdict(list))
        with tensorplay.no_grad():
            for group in optimizer.param_groups:
                for param in group["params"]:
                    if not isinstance(param, tensorplay.Tensor):
                        raise AssertionError(
                            f"expected param to be tensorplay.Tensor, got {type(param).__name__}"
                        )
                    if param.grad is None:
                        continue
                    if (not allow_fp16) and param.grad.dtype == tensorplay.float16:
                        raise ValueError("Attempting to unscale FP16 gradients.")
                    if param.grad.is_sparse:
                        raise NotImplementedError(
                            "GradScaler does not support sparse gradients yet."
                        )
                    to_unscale = param.grad

                    per_device_and_dtype_grads[_device_key(to_unscale.device)][
                        to_unscale.dtype
                    ].append(to_unscale)

            for device, per_dtype_grads in per_device_and_dtype_grads.items():
                for grads in per_dtype_grads.values():
                    tensorplay._amp_foreach_non_finite_check_and_unscale_(
                        grads,
                        per_device_found_inf.get(device),
                        per_device_inv_scale.get(device),
                    )

        return per_device_found_inf._per_device_tensors


[docs]
    def unscale_(self, optimizer: tensorplay.optim.Optimizer) -> None:
        """
        Divides ("unscales") the optimizer's gradient tensors by the scale factor.

        :meth:`unscale_` is optional, serving cases where you need to
        modify or inspect gradients between the backward pass(es) and :meth:`step`.
        If :meth:`unscale_` is not called explicitly,  gradients will be unscaled  automatically during :meth:`step`.

        Simple example, using :meth:`unscale_` to enable clipping of unscaled gradients::

            ...
            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            tensorplay.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
            scaler.step(optimizer)
            scaler.update()

        Args:
            optimizer (tensorplay.optim.Optimizer):  Optimizer that owns the gradients to be unscaled.

        .. warning::
            :meth:`unscale_` should only be called once per optimizer per :meth:`step` call,
            and only after all gradients for that optimizer's assigned parameters have been accumulated.
            Calling :meth:`unscale_` twice for a given optimizer between each :meth:`step` triggers a RuntimeError.
        """
        if not self._enabled:
            return

        self._check_scale_growth_tracker("unscale_")

        optimizer_state = self._per_optimizer_states[id(optimizer)]

        if optimizer_state["stage"] is OptState.UNSCALED:
            raise RuntimeError(
                "unscale_() has already been called on this optimizer since the last update()."
            )
        elif optimizer_state["stage"] is OptState.STEPPED:
            raise RuntimeError("unscale_() is being called after step().")

        # FP32 division can be imprecise for certain compile options, so we carry out the reciprocal in FP64.
        if self._scale is None:
            raise AssertionError("_scale is None in unscale_")
        inv_scale = (
            tensorplay.full((), 1.0, dtype=tensorplay.float64, device=self._scale.device)
            .div(self._scale.to(tensorplay.float64))
            .to(tensorplay.float32)
        )
        found_inf = tensorplay.full(
            (), 0.0, dtype=tensorplay.float32, device=self._scale.device
        )

        optimizer_state["found_inf_per_device"] = self._unscale_grads_(
            optimizer, inv_scale, found_inf, False
        )
        optimizer_state["stage"] = OptState.UNSCALED


    def _maybe_opt_step(
        self,
        optimizer: tensorplay.optim.Optimizer,
        optimizer_state: dict[str, Any],
        *args: Any,
        **kwargs: Any,
    ) -> Any:
        retval: Any = None
        if not sum(v.item() for v in optimizer_state["found_inf_per_device"].values()):
            retval = optimizer.step(*args, **kwargs)
        return retval


[docs]
    def step(
        self, optimizer: tensorplay.optim.Optimizer, *args: Any, **kwargs: Any
    ) -> Any:
        """Invoke ``unscale_(optimizer)`` followed by parameter update, if gradients are not infs/NaN.

        :meth:`step` carries out the following two operations:

        1.  Internally invokes ``unscale_(optimizer)`` (unless :meth:`unscale_` was explicitly called for ``optimizer``
            earlier in the iteration).  As part of the :meth:`unscale_`, gradients are checked for infs/NaNs.
        2.  If no inf/NaN gradients are found, invokes ``optimizer.step()`` using the unscaled
            gradients.  Otherwise, ``optimizer.step()`` is skipped to avoid corrupting the params.

        ``*args`` and ``**kwargs`` are forwarded to ``optimizer.step()``.

        Returns the return value of ``optimizer.step(*args, **kwargs)``.

        Args:
            optimizer (tensorplay.optim.Optimizer):  Optimizer that applies the gradients.
            args:  Any arguments.
            kwargs:  Any keyword arguments.

        .. warning::
            Closure use is not currently supported.
        """
        if not self._enabled:
            return optimizer.step(*args, **kwargs)

        if "closure" in kwargs:
            raise RuntimeError(
                "Closure use is not currently supported if GradScaler is enabled."
            )

        self._check_scale_growth_tracker("step")

        optimizer_state = self._per_optimizer_states[id(optimizer)]

        if optimizer_state["stage"] is OptState.STEPPED:
            raise RuntimeError(
                "step() has already been called since the last update()."
            )

        retval: Any = None

        if getattr(optimizer, "_step_supports_amp_scaling", False):
            # This optimizer has customized scale-handling logic, so we can call optimizer.step() directly.
            # The contract with custom optimizers is that their step() should accept an additional,
            # optional grad_scaler kwarg.  We append self to the kwargs so the custom optimizer has full information:
            # it can query its own state, invoke unscale_ on itself, etc.
            # Otherwise, two Tensor attributes of `grad_scale` and `found_inf` are registered on the passed
            # optimizer so that the optimizer can utilize those to skip the parameter updates or unscale
            # gradients before updating parameters.
            kwargs_ = kwargs
            has_grad_scaler_kwarg = (
                "grad_scaler" in inspect.signature(optimizer.step).parameters
            )
            if has_grad_scaler_kwarg:
                warnings.warn(
                    "GradScaler is going to stop passing itself as a keyword argument to the passed "
                    "optimizer. In the near future GradScaler registers `grad_scale: Tensor` and "
                    "`found_inf: Tensor` to the passed optimizer and let the optimizer use them directly.",
                    FutureWarning,
                    stacklevel=2,
                )
                kwargs_.update({"grad_scaler": self})
            else:
                if optimizer_state["stage"] is OptState.READY:
                    self._check_inf_per_device(optimizer)
                scaler = self._get_scale_async()
                if scaler is None:
                    raise AssertionError("_get_scale_async returned None")
                found_sum: Any = None
                for t in optimizer_state["found_inf_per_device"].values():
                    moved = t.to(scaler.device, non_blocking=True)
                    found_sum = (
                        moved if found_sum is None else found_sum + moved
                    )
                found_inf_combined = 0 if found_sum is None else found_sum
                # Take the product of the scales, if the user has already set `optimizer.grad_scale`.
                optimizer.grad_scale = (  # type: ignore[attr-defined]
                    getattr(optimizer, "grad_scale", None)
                    if optimizer_state["stage"] == OptState.UNSCALED
                    else scaler * getattr(optimizer, "grad_scale", 1)
                )
                optimizer.found_inf = found_inf_combined  # type: ignore[attr-defined]
            retval = optimizer.step(*args, **kwargs_)
            optimizer_state["stage"] = OptState.STEPPED
            if not has_grad_scaler_kwarg:
                del optimizer.grad_scale  # type: ignore[attr-defined]
                del optimizer.found_inf  # type: ignore[attr-defined]
            return retval

        if optimizer_state["stage"] is OptState.READY:
            self.unscale_(optimizer)

        if len(optimizer_state["found_inf_per_device"]) == 0:
            raise AssertionError("No inf checks were recorded for this optimizer.")

        retval = self._maybe_opt_step(optimizer, optimizer_state, *args, **kwargs)

        optimizer_state["stage"] = OptState.STEPPED

        return retval



[docs]
    def update(self, new_scale: float | tensorplay.Tensor | None = None) -> None:
        """Update the scale factor.

        If any optimizer steps were skipped the scale is multiplied by ``backoff_factor``
        to reduce it. If ``growth_interval`` unskipped iterations occurred consecutively,
        the scale is multiplied by ``growth_factor`` to increase it.

        Passing ``new_scale`` sets the new scale value manually. (``new_scale`` is not
        used directly, it's used to fill GradScaler's internal scale tensor. So if
        ``new_scale`` was a tensor, later in-place changes to that tensor will not further
        affect the scale GradScaler uses internally.)

        Args:
            new_scale (float or :class:`tensorplay.Tensor`, optional, default=None):  New scale factor.

        .. warning::
            :meth:`update` should only be called at the end of the iteration, after ``scaler.step(optimizer)`` has
            been invoked for all optimizers used this iteration.
        """
        if not self._enabled:
            return

        _scale, _growth_tracker = self._check_scale_growth_tracker("update")

        if new_scale is not None:
            if self._scale is None:
                raise AssertionError("_scale is None in update")
            # Accept a new user-defined scale.
            if isinstance(new_scale, float):
                self._scale.fill_(new_scale)
            else:
                reason = (
                    "new_scale should be a float or a 1-element tensor with requires_grad=False."
                )
                if new_scale.device.type != self._device:
                    raise AssertionError(reason)
                if new_scale.numel() != 1:
                    raise AssertionError(reason)
                if new_scale.requires_grad is True:
                    raise AssertionError(reason)
                self._scale.copy_(new_scale)
        else:
            # Consume shared inf/nan data collected from optimizers to update the scale.
            found_infs = [
                found_inf.to(device=_scale.device, non_blocking=True)
                for state in self._per_optimizer_states.values()
                for found_inf in state["found_inf_per_device"].values()
            ]

            if len(found_infs) == 0:
                raise AssertionError("No inf checks were recorded prior to update.")

            found_inf_combined = found_infs[0]
            if len(found_infs) > 1:
                for i in range(1, len(found_infs)):
                    found_inf_combined += found_infs[i]

            tensorplay._amp_update_scale_(
                _scale,
                _growth_tracker,
                found_inf_combined,
                self._growth_factor,
                self._backoff_factor,
                self._growth_interval,
            )

        # To prepare for next iteration, clear the data collected from optimizers this iteration.
        self._per_optimizer_states = defaultdict(_refresh_per_optimizer_state)


    def _get_scale_async(self) -> tensorplay.Tensor | None:
        return self._scale


[docs]
    def get_scale(self) -> float:
        """Return a Python float containing the current scale, or 1.0 if scaling is disabled."""
        if self._enabled:
            return (
                self._init_scale
                if (scale := self._get_scale_async()) is None
                else cast(float, scale.item())
            )
        return 1.0



[docs]
    def get_growth_factor(self) -> float:
        r"""Return a Python float containing the scale growth factor."""
        return self._growth_factor



[docs]
    def set_growth_factor(self, new_factor: float) -> None:
        r"""Set a new scale growth factor.

        Args:
            new_scale (float):  Value to use as the new scale growth factor.
        """
        self._growth_factor = new_factor



[docs]
    def get_backoff_factor(self) -> float:
        r"""Return a Python float containing the scale backoff factor."""
        return self._backoff_factor



[docs]
    def set_backoff_factor(self, new_factor: float) -> None:
        r"""Set a new scale backoff factor.

        Args:
            new_scale (float):  Value to use as the new scale backoff factor.
        """
        self._backoff_factor = new_factor



[docs]
    def get_growth_interval(self) -> int:
        r"""Return a Python int containing the growth interval."""
        return self._growth_interval



[docs]
    def set_growth_interval(self, new_interval: int) -> None:
        r"""Set a new growth interval.

        Args:
            new_interval (int):  Value to use as the new growth interval.
        """
        self._growth_interval = new_interval


    def _get_growth_tracker(self) -> int:
        if self._enabled:
            return (
                self._init_growth_tracker
                if self._growth_tracker is None
                else cast(int, self._growth_tracker.item())
            )
        return 0


[docs]
    def is_enabled(self) -> bool:
        r"""Return a bool indicating whether this instance is enabled."""
        return self._enabled



[docs]
    def state_dict(self) -> dict[str, Any]:
        r"""Return the state of the scaler as a :class:`dict`.

        It contains five entries:

        * ``"scale"`` - a Python float containing the current scale
        * ``"growth_factor"`` - a Python float containing the current growth factor
        * ``"backoff_factor"`` - a Python float containing the current backoff factor
        * ``"growth_interval"`` - a Python int containing the current growth interval
        * ``"_growth_tracker"`` - a Python int containing the number of recent consecutive unskipped steps.

        If this instance is not enabled, returns an empty dict.

        .. note::
           If you wish to checkpoint the scaler's state after a particular iteration, :meth:`state_dict`
           should be called after :meth:`update`.
        """
        if self._enabled:
            return {
                "scale": self.get_scale(),
                "growth_factor": self._growth_factor,
                "backoff_factor": self._backoff_factor,
                "growth_interval": self._growth_interval,
                "_growth_tracker": self._get_growth_tracker(),
            }
        return {}



[docs]
    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
        r"""Load the scaler state.

        If this instance is disabled, :meth:`load_state_dict` is a no-op.

        Args:
           state_dict(dict): scaler state.  Should be an object returned from a call to :meth:`state_dict`.
        """
        if not self._enabled:
            return

        if len(state_dict) == 0:
            raise RuntimeError(
                "The source state dict is empty, possibly because it was saved "
                "from a disabled instance of GradScaler."
            )

        self._init_scale = cast(float, state_dict["scale"])
        if self._scale is not None:
            self._scale.fill_(state_dict["scale"])
        self._growth_factor = cast(float, state_dict["growth_factor"])
        self._backoff_factor = cast(float, state_dict["backoff_factor"])
        self._growth_interval = cast(int, state_dict["growth_interval"])
        self._init_growth_tracker = cast(int, state_dict["_growth_tracker"])
        if self._growth_tracker is not None:
            self._growth_tracker.fill_(state_dict["_growth_tracker"])


    def __getstate__(self) -> dict[str, Any]:
        state = self.__dict__.copy()
        if self._enabled:
            if len(self._per_optimizer_states) != 0:
                raise AssertionError(
                    "A GradScaler instance may only be pickled at the beginning "
                    "of an iteration, or at the end after scaler.update()."
                )
            # Pickling _scale and _growth_tracker Tensors directly triggers storage
            # pickling, so instead, we set the unpickled instance up to reinitialize
            # them lazily.
            state["_init_scale"] = self.get_scale()
            state["_init_growth_tracker"] = self._get_growth_tracker()
            state["_scale"] = None
            state["_growth_tracker"] = None
        return state

    def __setstate__(self, state: dict[str, Any]) -> None:
        self.__dict__.update(state)

    def _check_inf_per_device(
        self, optimizer: tensorplay.optim.Optimizer
    ) -> dict[str, Any]:
        _scale, _ = self._check_scale_growth_tracker("_check_inf_per_device")

        dummy_inv_scale = tensorplay.full(
            (), 1.0, dtype=tensorplay.float32, device=_scale.device
        )
        found_inf = tensorplay.full(
            (), 0.0, dtype=tensorplay.float32, device=_scale.device
        )

        self._per_optimizer_states[id(optimizer)]["found_inf_per_device"] = (
            self._unscale_grads_(optimizer, dummy_inv_scale, found_inf, True)
        )

        return self._per_optimizer_states[id(optimizer)]["found_inf_per_device"]

    def _found_inf_per_device(
        self, optimizer: tensorplay.optim.Optimizer
    ) -> dict[str, Any]:
        return self._per_optimizer_states[id(optimizer)]["found_inf_per_device"]