optim.py

#
# Copyright (C) 2025 Intel Corporation
#
# SPDX-License-Identifier: Apache-2.0
#

import torch
import bisect
import inspect
import warnings
from random import shuffle
from operator import itemgetter
from collections import UserDict
from collections.abc import Callable
from accelerate.state import PartialState
from accelerate.utils import gather, reduce, pad_across_processes
from accelerate.utils import broadcast_object_list, tqdm
from accelerate.utils.memory import should_reduce_batch_size
from torch.optim.optimizer import Optimizer, ParamsT  # type: ignore[reportPrivateImportUsage]
from torch.optim.sgd import SGD  # type: ignore[reportPrivateImportUsage]
from torch.optim.adam import Adam  # type: ignore[reportPrivateImportUsage]

from llmart.config import OptimConf
from llmart.pickers import pick_coord_topk, pick_global_topk, pick_unique_swaps
from llmart.samplers import sample_coord_randk, sample_swaps


class Coordinate(UserDict[int, int]):
    """A dictionary-like class for storing coordinate mappings.

    Stores index-to-value mappings for coordinates and provides conversion
    utilities to/from tensors.
    """

    @property
    def index(self) -> list[int]:
        return list(self.keys())

    @property
    def value(self) -> list[int]:
        return list(self.values())

    def to_tensor(self, device: str) -> torch.Tensor:
        return torch.tensor(
            [self.index, self.value], device=device
        ).T  # Index coordinates on row axis

    @classmethod
    def from_tensor(cls, tensor: torch.Tensor) -> "Coordinate":
        if len(tensor.shape) == 1:
            return cls(enumerate(tensor.tolist()))
        else:
            return cls(zip(tensor[..., 0].tolist(), tensor[..., 1].tolist()))


class GreedyCoordinateGradient(Optimizer):
    """An optimizer implementing greedy coordinate gradient descent.

    Args:
        params: Parameters to optimize.
        negative_only: If True, only consider negative gradients.
        coord_randk: Number of random dictionary entries to sample per token.
        coord_topk: Number of top dictionary entries to select per token.
        global_topk: Number of top dictionary entries to select globally.
        n_tokens: Number of tokens to swap at once.
        n_swaps: Maximum number of swaps to attempt.
        n_buffers: Size of the buffer for storing best coordinates.
        ignore_curr_marginals: If True, ignore current token positions.
        ignored_values: Tensor of dictionary entries to ignore.
        embedding: Token embedding matrix.
        world_size: Number of distributed processes.

    Raises:
        ValueError: If params or embedding configuration is invalid.
    """

    def __init__(
        self,
        params: ParamsT,
        negative_only: bool = False,
        coord_randk: int = 0,
        coord_topk: int = 256,
        global_topk: int = 0,
        n_tokens: int = 20,
        n_swaps: int = 1024,
        n_buffers: int = 1,
        ignore_curr_marginals: bool = False,
        ignored_values: torch.Tensor | None = None,
        embedding: torch.nn.Module | None = None,
        world_size: int = 1,
    ):
        if ignored_values is None:
            ignored_values = torch.LongTensor()

        defaults = dict(
            negative_only=negative_only,
            coord_randk=coord_randk,
            coord_topk=coord_topk,
            global_topk=global_topk,
            n_tokens=n_tokens,
            n_swaps=n_swaps,
            n_buffers=max(n_buffers, 1),
            ignore_curr_marginals=ignore_curr_marginals,
            world_size=world_size,
            local_swap_count=torch.tensor(0),
        )
        super().__init__(params, defaults)

        if len(self.param_groups) != 1:
            raise ValueError(
                f"{self.__class__.__name__} doesn't support per-parameter options (parameter groups)"
            )
        if len(self.param_groups[0]["params"]) != 1:
            raise ValueError(
                f"{self.__class__.__name__} doesn't support more than 1 parameter"
            )
        if len(self.param_groups[0]["params"][0].shape) != 2:
            raise ValueError(
                f"{self.__class__.__name__} requires 2-dimensional parameters"
            )
        if embedding is not None and not isinstance(embedding, torch.nn.Embedding):
            raise ValueError("embedding must be an nn.Embedding")

        if (
            embedding is not None
            and self.param_groups[0]["params"][0].shape[-1] != embedding.embedding_dim
        ):
            raise ValueError(
                f"Parameter must have same dimension as embedding ({embedding.embedding_dim})!"
            )

        if (
            embedding is None
            and (self.param_groups[0]["params"][0].sum(dim=-1) != 1.0).any()
        ):
            raise ValueError(
                f"{self.__class__.__name__} only works with 1-hot encoded parameters when embedding is None"
            )

        self._update_hyperparams()

        self._ignored_values = ignored_values
        self._embedding = embedding
        self._replacements: list[Coordinate] | None = None

    def _update_hyperparams(self):
        self._coord_topk = self.param_groups[0]["coord_topk"]
        self._n_tokens = self.param_groups[0]["n_tokens"]
        self._n_swaps = self.param_groups[0]["n_swaps"]
        self._local_swap_count = self.param_groups[0]["local_swap_count"]
        self._param = self.param_groups[0]["params"][0]
        self._negative_only = self.param_groups[0]["negative_only"]
        self._coord_randk = self.param_groups[0]["coord_randk"]
        self._global_topk = self.param_groups[0]["global_topk"]
        self._n_buffers = self.param_groups[0]["n_buffers"]
        self._ignore_curr_marginals = self.param_groups[0]["ignore_curr_marginals"]

    @property
    def coordinate(self) -> Coordinate:
        param = self._param
        if self._embedding:
            param = torch.matmul(param, self._embedding.weight.T)
        return Coordinate.from_tensor(param.argmax(-1))

    @property
    def coordinate_replacements(self) -> list[Coordinate]:
        if self._replacements is None:
            if self._param.grad is None:
                raise ValueError("Params must have gradients!")
            param = self._param
            if self._embedding:
                param = torch.matmul(param, self._embedding.weight.T)
                param.grad = torch.matmul(self._param.grad, self._embedding.weight.T)

            replacements = self._select_replacements(param)
            # _select_replacements is random so synchronize with all other processes
            self._replacements = broadcast_object_list(replacements, from_process=0)

        return self._replacements

    @coordinate_replacements.setter
    def coordinate_replacements(self, replacements: list[Coordinate] | None):
        self._replacements = replacements

    def _select_replacements(self, param) -> list[Coordinate]:
        # Explicitly ignore universally banned tokens by making gradients very large
        param.grad[..., self._ignored_values] = torch.inf

        # Explicitly ignore each of the current tokens and dictionary for all possible picks
        # FIXME: This is not right at all, it ignores everything?
        if self._ignore_curr_marginals:
            param.grad[torch.where(param)] = torch.inf

        assert param.grad.ndim == 2, (
            "Need 'grad' to be 2D tensor of shape (n_tokens, n_dictionary)!"
        )
        coords = torch.where(torch.isfinite(param.grad))
        coords = torch.stack(coords, dim=0).T

        # Select only negative dictionary gradients for each token position
        if self._negative_only:
            coords = torch.where(param.grad < 0)
            coords = torch.stack(coords, dim=0).T
            # Check if no tokens survive after self._negative_only
            if len(coords) == 0:
                # The entire optimization step will be skipped
                warnings.warn(
                    f"No tokens can be swapped at all! Check {self._negative_only = }!"
                )
                return []

        # Sample coord_randk random gradients for each token position separately
        if self._coord_randk > 0:
            coords = sample_coord_randk(coords, self._coord_randk)

        # Select coord_topk lowest gradients for each token position separately
        if self._coord_topk > 0:
            coords = pick_coord_topk(coords, param.grad, self._coord_topk)

        # Select global_topk lowest gradients with survival guarantees for all token positions
        if self._global_topk > 0:
            coords = pick_global_topk(coords, param.grad, self._global_topk)

        # Swap as many tokens as we can up to the requested tuple size
        _n_swap_tokens = self._n_tokens
        if len(torch.unique(coords[:, 0])) < self._n_tokens:
            _n_swap_tokens = len(torch.unique(coords[:, 0]))
            warnings.warn(
                f"Fewer than {self._n_tokens} token(s) survived, will instead swap tuples of {_n_swap_tokens} token(s)!"
            )

        # Sample _n_swap_tokens-tuples
        swaps = sample_swaps(coords, _n_swap_tokens, self._n_swaps)

        # Remove duplicate swaps
        swaps = pick_unique_swaps(swaps)

        # Verify we're not swapping the same token position twice in the same swap
        assert (
            torch.count_nonzero(swaps[:, :, None, 0] == swaps[:, None, :, 0], dim=2)
            == 1
        ).all(), "Attempting to swap the same token more than once in the same tuple!"

        return [Coordinate.from_tensor(swap) for swap in swaps]

    def update_parameter(self, coord: Coordinate):
        self._param.data[coord.index] = 0.0
        if self._embedding is None:
            self._param.data[coord.index, coord.value] = 1.0
        else:
            # Convert values to embeddings and use basic indexing to set embeds
            embeds = self._embedding(coord.to_tensor(self._param.device)[..., 1])
            for idx, embed in zip(coord.index, embeds):
                self._param.data[idx].copy_(embed)

    @torch.inference_mode()
    def step(self, closure: Callable):  # type: ignore
        # Stay up to date with the latest hyper-parameters
        self._update_hyperparams()

        # Lazily initialize state
        p_state = self.state[self._param]
        p_state["swap_count"] = p_state.get(
            "swap_count", torch.tensor(0, device=self._param.device)
        )
        p_state["buffers"] = p_state.get("buffers", [])

        # Split indexed replacements across devices and compute losses
        replacements = self.coordinate_replacements
        idxes = torch.arange(len(replacements), device=self._param.device)
        idxes_replacements = list(zip(idxes, replacements))
        shuffle(idxes_replacements)
        with PartialState().split_between_processes(
            idxes_replacements
        ) as local_replacements:
            # Protect against OOMs happening in one process but not another.
            had_oom = torch.zeros((1,), device=self._param.device)
            try:
                local_outputs = self._step(closure, local_replacements)  # type: ignore
            except Exception as e:
                if should_reduce_batch_size(e):
                    had_oom.add_(1)
                else:
                    raise
            # If any process had an OOM then raise an OOM in all processes using a generic "CUDA" message
            had_oom: torch.Tensor = reduce(had_oom, "sum")  # type: ignore
            if had_oom > 0:
                raise torch.OutOfMemoryError("CUDA out of memory.")
            local_outputs = pad_across_processes(local_outputs, pad_index=-1)  # type: ignore

        # Gather global indices and losses and select valid ones
        idxes, losses = gather(local_outputs)
        is_valid = torch.where((idxes != -1) & torch.isfinite(losses))[0]  # type: ignore
        idxes, losses = idxes[is_valid], losses[is_valid]

        if len(losses) == 0:
            warnings.warn(f"All {self._n_swaps} attempted swaps do not re-encode!")

        # Add losses to buffer popping max loss if we exceed length
        old_coord = self.coordinate
        for idx, loss in zip(idxes, losses):
            coord = old_coord | replacements[idx]
            bisect.insort(p_state["buffers"], (loss, coord), key=itemgetter(0))
            if len(p_state["buffers"]) > self._n_buffers:
                p_state["buffers"].pop()

        # Update parameter with smallest loss coordinate
        if len(p_state["buffers"]) > 0:
            _, best_coord = p_state["buffers"].pop(0)
            self.update_parameter(best_coord)

        # Save information about replaced tokens
        p_state["swap_count"] += reduce(self._local_swap_count, "sum")
        self._local_swap_count.zero_()

        # Force next call to step to update replacements
        self.coordinate_replacements = None

        # Safety check that param is still 1-hot or a direct embedding
        if self._embedding is None:
            assert (self._param.sum(dim=-1) == 1.0).all()
        else:
            assert torch.allclose(
                self._param,
                self._embedding(
                    torch.cdist(self._param, self._embedding.weight).argmin(-1)
                ),
            )

    def _step(
        self, closure: Callable, replacements: list[tuple[torch.Tensor, Coordinate]]
    ) -> tuple[torch.Tensor, torch.Tensor]:
        generator = None
        if inspect.isgeneratorfunction(closure):
            generator = closure()
            generator.send(None)

        old_coord = self.coordinate

        idx_losses = []
        for idx, new_coord in tqdm(replacements, desc="replacements", leave=False):
            # Vector replace coordinates with new value
            self.update_parameter(new_coord)
            idx_losses.extend(generator.send(idx) if generator else [(idx, closure())])
            self.update_parameter(old_coord)
            self._local_swap_count += 1

        # Get remaining losses, if any
        if generator:
            idx_losses.extend(generator.send(None))
            generator.close()

        # Return dummy index and loss when empty otherwise stack them
        if len(idx_losses) == 0:
            return (
                torch.tensor([-1], device=self._param.device),
                torch.tensor([torch.inf], device=self._param.device),
            )
        idxes, losses = list(zip(*idx_losses))
        return torch.stack(idxes), torch.stack(losses)


def from_config(cfg: OptimConf, params: ParamsT, **gcg_kwargs) -> Optimizer:
    """Creates an optimizer instance from configuration.

    Args:
        cfg: Optimizer configuration object.
        params: Model parameters to optimize.
        **gcg_kwargs: Additional keyword arguments for GCG optimizer.

    Returns:
        Configured optimizer instance.

    Raises:
        ValueError: If optimizer name is unknown.
    """

    if cfg.name == "gcg":
        return GreedyCoordinateGradient(
            params,
            negative_only=cfg.negative_only,
            coord_randk=cfg.coord_randk,
            coord_topk=cfg.coord_topk,
            global_topk=cfg.global_topk,
            n_tokens=cfg.n_tokens,
            n_swaps=cfg.n_swaps,
            n_buffers=cfg.n_buffers,
            **gcg_kwargs,
        )

    if cfg.name == "sgd":
        return SGD(
            params,
            lr=cfg.lr,
            momentum=cfg.momentum,
            dampening=cfg.dampening,
            weight_decay=cfg.weight_decay,
            nesterov=cfg.nesterov,
        )

    if cfg.name == "adam":
        return Adam(
            params,
            lr=cfg.lr,
            betas=cfg.betas,
            eps=cfg.eps,
            weight_decay=cfg.weight_decay,
            amsgrad=cfg.amsgrad,
        )

    raise ValueError(f"Unknown optimizer: {cfg.name}")