chore: import upstream snapshot with attribution

modules/optimizer/chained_optimizer.py

@@ -0,0 +1,132 @@
+import torch
+from torch import Tensor
+from torch.optim import Optimizer
+from torch.optim.optimizer import ParamsT
+from dataclasses import dataclass
+from typing import Any, Dict, List, Type, Callable, Optional
+
+
+@dataclass
+class OptimizerSpec:
+    """Spec for creating an optimizer that is part of a `ChainedOptimizer`."""
+
+    class_type: Type[Optimizer]
+    init_args: Dict[str, Any]
+    param_filter: Optional[Callable[[Tensor], bool]]
+
+
+class ChainedOptimizer(Optimizer):
+    """
+    A wrapper around multiple optimizers that allows for chaining them together.
+    The optimizers are applied in the order they are passed in the constructor.
+    Each optimizer is responsible for updating a subset of the parameters, which
+    is determined by the `param_filter` function. If no optimizer is found for a
+    parameter group, an exception is raised.
+    """
+
+    def __init__(
+        self,
+        params: ParamsT,
+        optimizer_specs: List[OptimizerSpec],
+        lr: float,
+        weight_decay: float = 0.0,
+        optimizer_selection_callback: Optional[Callable[[Tensor, int], None]] = None,
+        **common_kwargs,
+    ):
+        self.optimizer_specs = optimizer_specs
+        self.optimizer_selection_callback = optimizer_selection_callback
+        self.optimizers: List[Optimizer] = []
+        defaults = dict(lr=lr, weight_decay=weight_decay)
+        super().__init__(params, defaults)
+
+        # Split the params for each optimizer
+        params_for_optimizers = [[] for _ in optimizer_specs]
+        for param_group in self.param_groups:
+            params = param_group["params"]
+            indices = param_group["optimizer_and_param_group_indices"] = set()
+            for param in params:
+                assert isinstance(param, Tensor), f"Expected a Tensor, got {type(param)}"
+                found_optimizer = False
+                for index, spec in enumerate(optimizer_specs):
+                    if spec.param_filter is None or spec.param_filter(param):
+                        if self.optimizer_selection_callback is not None:
+                            self.optimizer_selection_callback(param, index)
+                        params_for_optimizers[index].append(param)
+                        indices.add((index, 0))
+                        found_optimizer = True
+                        break
+                if not found_optimizer:
+                    raise ValueError("No valid optimizer found for the given parameter")
+
+        # Initialize the optimizers
+        for spec, selected_params in zip(optimizer_specs, params_for_optimizers):
+            optimizer_args = {
+                'lr': lr,
+                'weight_decay': weight_decay,
+            }
+            optimizer_args.update(common_kwargs)
+            optimizer_args.update(spec.init_args)
+            optimizer = spec.class_type(selected_params, **optimizer_args)
+            self.optimizers.append(optimizer)
+
+    def state_dict(self) -> Dict[str, Any]:
+        return {
+            "optimizers": [opt.state_dict() for opt in self.optimizers],
+            **super().state_dict(),
+        }
+
+    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+        optimizers = state_dict.pop("optimizers")
+        super().load_state_dict(state_dict)
+        for i in range(len(self.optimizers)):
+            self.optimizers[i].load_state_dict(optimizers[i])
+
+    def zero_grad(self, set_to_none: bool = True) -> None:
+        for opt in self.optimizers:
+            opt.zero_grad(set_to_none=set_to_none)
+
+    def _copy_lr_to_optimizers(self) -> None:
+        for param_group in self.param_groups:
+            indices = param_group["optimizer_and_param_group_indices"]
+            for optimizer_idx, param_group_idx in indices:
+                self.optimizers[optimizer_idx].param_groups[param_group_idx]["lr"] = param_group["lr"]
+
+    def step(self, closure=None) -> None:
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+        self._copy_lr_to_optimizers()
+        for opt in self.optimizers:
+            opt.step(closure=None)
+        return loss
+
+    def add_param_group(self, param_group: Dict[str, Any]) -> None:
+        super().add_param_group(param_group)
+
+        # If optimizer has not been initialized, skip adding the param groups
+        if not self.optimizers:
+            return
+
+        # Split the params for each optimizer
+        params_for_optimizers = [[] for _ in self.optimizer_specs]
+        params = param_group["params"]
+        indices = param_group["optimizer_and_param_group_indices"] = set()
+        for param in params:
+            assert isinstance(param, Tensor), f"Expected a Tensor, got {type(param)}"
+            found_optimizer = False
+            for index, spec in enumerate(self.optimizer_specs):
+                if spec.param_filter is None or spec.param_filter(param):
+                    if self.optimizer_selection_callback is not None:
+                        self.optimizer_selection_callback(param, index)
+                    params_for_optimizers[index].append(param)
+                    indices.add((index, len(self.optimizers[index].param_groups)))
+                    found_optimizer = True
+                    break
+            if not found_optimizer:
+                raise ValueError("No valid optimizer found for the given parameter group")
+
+        # Add the selected param group to the optimizers
+        for optimizer, selected_params in zip(self.optimizers, params_for_optimizers):
+            if selected_params:
+                optimizer.add_param_group({"params": selected_params})