Some more kron work. Figured out why some tests fail, implemented a deterministic rng state load but too slow so skipping some tests for now.

rwightman · rwightman · commit f759d1213a38 · 2025-01-26T20:44:09.000-08:00
diff --git a/tests/test_optim.py b/tests/test_optim.py
@@ -290,7 +290,7 @@ def _build_params_dict_single(weight, bias, **kwargs):
     return [dict(params=bias, **kwargs)]
 
 
-@pytest.mark.parametrize('optimizer', list_optimizers(exclude_filters=('fused*', 'bnb*')))
+@pytest.mark.parametrize('optimizer', list_optimizers(exclude_filters=('fused*', 'bnb*', 'kron*')))
 def test_optim_factory(optimizer):
     assert issubclass(get_optimizer_class(optimizer, bind_defaults=False), torch.optim.Optimizer)
 
@@ -386,6 +386,14 @@ def test_adam(optimizer):
     _test_model(optimizer, dict(lr=5e-2))
 
 
+@pytest.mark.parametrize('optimizer',  ['kron'])
+def test_kron(optimizer):
+    _test_rosenbrock(
+        lambda params: create_optimizer_v2(params, optimizer, lr=1e-3)
+    )
+    _test_model(optimizer, dict(lr=1e-3))
+
+
 @pytest.mark.parametrize('optimizer',  ['adopt', 'adoptw'])
 def test_adopt(optimizer):
     _test_rosenbrock(
diff --git a/timm/optim/_optim_factory.py b/timm/optim/_optim_factory.py
@@ -697,9 +697,16 @@ def _register_other_optimizers(registry: OptimizerRegistry) -> None:
         OptimInfo(
             name='kron',
             opt_class=Kron,
-            description='',
+            description='PSGD optimizer with Kronecker-factored preconditioner',
             has_momentum=True,
         ),
+        OptimInfo(
+            name='kronw',
+            opt_class=Kron,
+            description='PSGD optimizer with Kronecker-factored preconditioner and decoupled weight decay',
+            has_momentum=True,
+            defaults={'decoupled_decay': True}
+        ),
         OptimInfo(
             name='laprop',
             opt_class=LaProp,
diff --git a/timm/optim/kron.py b/timm/optim/kron.py

-Original file line number
+Diff line change
 * https://github.com/ClashLuke/HeavyBall (added improvements)
 """
 +import logging
 import string
 import random
 +import warnings
 +from typing import Any, Callable, Dict, Optional, Tuple, Union
 import numpy as np
 import torch
++
 try:
     # NOTE opt_einsum needed to avoid blowing up memory with einsum ops
     import opt_einsum
     opt_einsum.enabled = True
     opt_einsum.strategy = "auto-hq"
     import torch.backends.opt_einsum
 +    has_opt_einsum = True
 except ImportError:
 -    opt_einsum = None
 +    has_opt_einsum = False
 try:
     torch._dynamo.config.cache_size_limit = 1_000_000
     has_dynamo = True
 except AttributeError:
     has_dynamo = False
 +_logger = logging.getLogger(__name__)
++
 def precond_update_prob_schedule(
     n: float,
     max_prob: float = 1.0,
     min_prob: float = 0.03,
     decay: float = 0.001,
     flat_start: float = 500,
 -):
 +) -> torch.Tensor:
     """Anneal preconditioner update probability during beginning of training.
     PSGD benefits from more preconditioner updates at the beginning of training,
     """Implements PSGD Kron from https://github.com/lixilinx/psgd_torch.
     Args:
 -        params (iterable): Iterable of parameters to optimize or dicts defining parameter groups.
 -        lr (float): Learning rate.
 -        momentum (float): Momentum parameter.
 -        weight_decay (float): Weight decay (L2 penalty).
 -        preconditioner_update_probability (callable or float, optional): Probability of
 -            updating the preconditioner. If None, defaults to a schedule that anneals
 -            from 1.0 to 0.03 by 4000 steps.
 -        max_size_triangular (int): Max size for dim's preconditioner to be triangular.
 -        min_ndim_triangular (int): Minimum number of dimensions a layer needs to have triangular preconditioners.
 -        memory_save_mode: (string, optional), None, 'one_diag', or 'all_diag', None is default
 +        params: Iterable of parameters to optimize or dicts defining parameter groups.
 +        lr: Learning rate.
 +        momentum: Momentum parameter.
 +        weight_decay: Weight decay (L2 penalty).
 +        preconditioner_update_probability: Probability of updating the preconditioner.
 +            If None, defaults to a schedule that anneals from 1.0 to 0.03 by 4000 steps.
 +        max_size_triangular: Max size for dim's preconditioner to be triangular.
 +        min_ndim_triangular: Minimum number of dimensions a layer needs to have triangular preconditioners.
 +        memory_save_mode: 'one_diag', or 'all_diag', None is default
             to set all preconditioners to be triangular, 'one_diag' sets the largest
             or last dim to be diagonal per layer, and 'all_diag' sets all preconditioners to be diagonal.
 -        momentum_into_precond_update: (bool), whether to send momentum into preconditioner
 +        momentum_into_precond_update: whether to send momentum into preconditioner
             update instead of raw gradients.
 -        mu_dtype (torch.dtype, optional): Dtype of the momentum accumulator.
 -        precond_dtype (torch.dtype, optional): Dtype of the preconditioner.
 +        mu_dtype: Dtype of the momentum accumulator.
 +        precond_dtype: Dtype of the preconditioner.
 +        decoupled_decay: AdamW style decoupled-decay.
 +        deterministic: Deterministic behaviour across save / load (resume). FIXME slow, needs work
     """
     def __init__(
         self,
         params,
 -        lr=0.001,
 -        momentum=0.9,
 -        weight_decay=0.0,
 -        preconditioner_update_probability=None,
 -        max_size_triangular=2048,
 -        min_ndim_triangular=2,
 -        memory_save_mode=None,
 -        momentum_into_precond_update=True,
 -        mu_dtype=None,
 -        precond_dtype=None,
 +        lr: float = 0.001,
 +        momentum: float = 0.9,
 +        weight_decay: float = 0.0,
 +        preconditioner_update_probability: Optional[Union[Callable, float]] = None,
 +        max_size_triangular: int = 2048,
 +        min_ndim_triangular: int = 2,
 +        memory_save_mode: Optional[str] = None,
 +        momentum_into_precond_update: bool = True,
 +        mu_dtype: Optional[torch.dtype] = None,
 +        precond_dtype: Optional[torch.dtype] = None,
 +        decoupled_decay: bool = False,
 +        deterministic: bool = False,
     ):
 +        if not has_opt_einsum:
 +            warnings.warn("It is highly recommended to have 'opt_einsum' installed for this optimizer." )
         if not 0.0 <= lr:
             raise ValueError(f"Invalid learning rate: {lr}")
         if not 0.0 <= momentum < 1.0:
             precond_init_scale=1.0,  # precond init scale hardcoded to 1.0
             mu_dtype=mu_dtype,
             precond_dtype=precond_dtype,
 +            decoupled_decay=decoupled_decay,
+        )
         super(Kron, self).__init__(params, defaults)
 +        self._param_exprs = {}
         self._tiny = torch.finfo(torch.bfloat16).tiny
 -        self._prob_step = 0
 -        self._update_counter = 0
 -        self.rng = random.Random(5318008)
 +        self.rng = random.Random(1337)
 +        if deterministic:
 +            # Use a Generator to try to be more deterministic across resume (save/load)
 +            self.torch_rng = torch.Generator().manual_seed(1337)
 +        else:
 +            self.torch_rng = None
         # make compile optional (for bwd compat)
         if has_dynamo:
             self._precond_grad = _precond_grad
             self._balance_Q = _balance_Q
 +    def __getstate__(self):
 +        _dict = super().__getstate__()
 +        _dict["rng"] = self.rng
 +        _dict["torch_rng"] = self.torch_rng
 +        return _dict
++
 +    def state_dict(self) -> Dict[str, Any]:
 +        # Get the optimizer's state dict
 +        optimizer_state = super().state_dict()
++
 +        # Add the generator state
 +        optimizer_state['rng_state'] = self.rng.getstate()
 +        if self.torch_rng is not None:
 +            optimizer_state['torch_rng_state'] = self.torch_rng.get_state()
++
 +        return optimizer_state
++
 +    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
 +        # Extract and remove the RNG state from the state dict
 +        rng_state = state_dict.pop('rng_state', None)
 +        torch_rng_state = state_dict.pop('torch_rng_state', None)
++
 +        # Load the optimizer state
 +        super().load_state_dict(state_dict)
++
 +        # Restore the RNG state if it exists
 +        if rng_state is not None:
 +            self.rng.setstate(rng_state)
 +            state_dict['rng_state'] = rng_state  # put it back if caller still using state_dict
 +        if torch_rng_state is not None:
 +            if self.torch_rng is not None:
 +                self.torch_rng.set_state(torch_rng_state)
 +            state_dict['torch_rng_state'] = torch_rng_state  # put it back if caller still using state_dict
++
 +    def __setstate__(self, state):
 +        super().__setstate__(state)
 +        self._param_exprs = {}
++
     @torch.no_grad()
     def step(self, closure=None):
         loss = None
         total_precond_size = 0
         total_precond_mb = 0
 -        # update preconditioners all together deterministically
 -        update_prob = self.param_groups[0]["preconditioner_update_probability"]
 -        if update_prob is None:
 -            update_prob = precond_update_prob_schedule
 -        if callable(update_prob):
 -            update_prob = update_prob(self._prob_step)
 -        self._update_counter += 1
 -        do_update = self._update_counter >= 1 / update_prob
 -        if do_update:
 -            self._update_counter = 0
 -        self._prob_step += 1
+-
 -        # balance preconditioners roughly every 100 updates
 -        balance = self.rng.random() < 0.01 and do_update
+-
         for group in self.param_groups:
             mu_dtype = group.get("mu_dtype")
             precond_dtype = group.get("precond_dtype", torch.float32)
             momentum_into_precond_update = group.get("momentum_into_precond_update", True)
 +            update_prob = group.get("preconditioner_update_probability", None)
             for p in group["params"]:
                 if p.grad is None:
                 if len(state) == 0:
                     state["step"] = 0
 +                    state["update_counter"] = 0
                     state["momentum_buffer"] = torch.zeros_like(p, dtype=mu_dtype or p.dtype)
 -                    state["Q"], state["exprs"] = _init_Q_exprs(
 +                    state["Q"], exprs = _init_Q_exprs(
                         p,
                         group["precond_init_scale"],
                         group["max_size_triangular"],
                         group["min_ndim_triangular"],
                         group["memory_save_mode"],
                         dtype=precond_dtype,
+                    )
 +                    self._param_exprs[p] = exprs
 -                    # Print sizes
 +                    # Accumulate sizes for log
                     momentum_size = state["momentum_buffer"].numel()
                     momentum_mb = momentum_size * state["momentum_buffer"].element_size() / 2**20
                     total_momentum_size += momentum_size
                     precond_mb = sum(q.numel() * q.element_size() for q in state["Q"]) / 2**20
                     total_precond_size += precond_size
                     total_precond_mb += precond_mb
 +                elif p not in self._param_exprs:
 +                    exprs = _init_Q_exprs(
 +                        p,
 +                        group["precond_init_scale"],
 +                        group["max_size_triangular"],
 +                        group["min_ndim_triangular"],
 +                        group["memory_save_mode"],
 +                        dtype=precond_dtype,
 +                        init_q=False,
 +                    )
 +                    self._param_exprs[p] = exprs
 +                else:
 +                    exprs = self._param_exprs[p]
++
 +                # update preconditioners all together deterministically
 +                if update_prob is None:
 +                    update_prob = precond_update_prob_schedule
 +                if callable(update_prob):
 +                    update_prob = update_prob(state["step"])
 +                state["update_counter"] += 1
 +                do_update = state["update_counter"] >= 1 / update_prob
 +                if do_update:
 +                    state["update_counter"] = 0
                 state["step"] += 1
                 bias_correction = 1 - beta ** state["step"]
                 momentum_buffer = state["momentum_buffer"]
                 momentum_buffer.mul_(group["momentum"]).add_(grad, alpha=1 - group["momentum"])
++
                 # Restore momentum dtype
                 if mu_dtype is not None:
 -                    momentum_buffer.copy_(momentum_buffer.to(dtype=mu_dtype, non_blocking=True))
 -                debiased_momentum = momentum_buffer / bias_correction
 -                debiased_momentum = debiased_momentum.to(dtype=precond_dtype, non_blocking=True)
 +                    momentum_buffer.copy_(momentum_buffer.to(dtype=mu_dtype))
 +                debiased_momentum = (momentum_buffer / bias_correction).to(dtype=precond_dtype)
 -                # balance preconditioners about every 100 updates
 +                # Balance preconditioners roughly every 100 updates
 +                balance = self.rng.random() < 0.01 and do_update
                 if grad.dim() > 1 and balance:
                     self._balance_Q(state["Q"])
                 # Update preconditioner
                 if do_update:
 -                    exprA, exprGs, _ = state["exprs"]
 +                    exprA, exprGs, _ = exprs
                     Q = state["Q"]
 -                    V = torch.randn_like(debiased_momentum, dtype=precond_dtype)
 +                    if self.torch_rng is None:
 +                        V = torch.randn_like(debiased_momentum, dtype=precond_dtype)
 +                    else:
 +                        # Restoring generator state to device is messy. For now,
 +                        # we keep RNG on CPU, but this slows the optimizer down quite a bit.
 +                        # FIXME Need a better approach
 +                        V = torch.randn(
 +                            debiased_momentum.shape, generator=self.torch_rng, dtype=precond_dtype, device='cpu')
 +                        V = V.to(debiased_momentum.device)
                     G = debiased_momentum if momentum_into_precond_update else grad
                     A, conjB = self._calc_A_and_conjB(exprA, G, Q, V)
                         if q.dim() < 2:
                             tmp *= q
                             tmp /= (term1 + term2).norm(float("inf")) + self._tiny
 -                            q.sub_(tmp)
                         else:
                             tmp = torch.triu(tmp)
                             tmp /= _norm_lower_bound(term1 + term2) + self._tiny
                             tmp @= q
 -                            q.sub_(tmp)
+-
 -                    # _update_precond(
 -                    #     state["Q"],
 -                    #     state["exprs"],
 -                    #     torch.randn_like(debiased_momentum, dtype=precond_dtype),
 -                    #     debiased_momentum if momentum_into_precond_update else grad,
 -                    #     group["precond_lr"],
 -                    #     self._tiny,
 -                    # )
 +                        q.sub_(tmp)
                 # Precondition gradients
                 pre_grad = self._precond_grad(
                     state["Q"],
 -                    state["exprs"],
 +                    exprs,
                     debiased_momentum,
 -                ).to(dtype=p.dtype, non_blocking=True)
 +                ).to(dtype=p.dtype)
                 # RMS of pre_grad should be 1.0, so let's cap at 1.1
 -                pre_grad.mul_(torch.clamp(1.1 / (pre_grad.square().mean().sqrt() + 1e-6), max=1.0))
 +                pre_grad.mul_(torch.clamp(1.1 / (pre_grad.square().mean().sqrt_() + 1e-8), max=1.0))
 -                # Apply weight decay and update parameters
 -                if group["weight_decay"] != 0 and p.dim() >= 2:
 -                    pre_grad.add_(p, alpha=group["weight_decay"])
 +                # Apply weight decay
 +                if group["weight_decay"] != 0:
 +                    if group["decoupled_decay"]:
 +                        p.mul_(1. - group["lr"] * group["weight_decay"])
 +                    else:
 +                        pre_grad.add_(p, alpha=group["weight_decay"])
++
 +                # Update parameters
                 p.add_(pre_grad, alpha=-group["lr"])
         if total_momentum_size > 0:
 -            print(f"PSGD Momentum buffer size: {total_momentum_size} elements, {total_momentum_mb:.2f} MB")
 -            print(f"PSGD Preconditioners size: {total_precond_size} elements, {total_precond_mb:.2f} MB")
 +            _logger.info(f"PSGD Momentum buffer size: {total_momentum_size} elements, {total_momentum_mb:.2f} MB")
 +            _logger.info(f"PSGD Preconditioners size: {total_precond_size} elements, {total_precond_mb:.2f} MB")
         return loss
 -def _init_Q_exprs(t, scale, max_size, min_ndim_triangular, memory_save_mode, dtype=None):
 +def _init_Q_exprs(
 +        t,
 +        scale,
 +        max_size,
 +        min_ndim_triangular,
 +        memory_save_mode,
 +        dtype=None,
 +        init_q=True,
 +):
     """For a scalar or tensor t, we initialize its preconditioner Q and
     reusable einsum expressions for updating Q and preconditioning gradient.
     """
     letters = string.ascii_lowercase + string.ascii_uppercase
     dtype = dtype if dtype is not None else t.dtype
     shape = t.shape
 +    Q = []
     if len(shape) == 0:  # scalar
 -        Q = [scale * torch.ones_like(t, dtype=dtype)]
 +        if init_q:
 +            Q.append(scale * torch.ones_like(t, dtype=dtype))
         exprA = ",->"
         exprGs = [",->"]
         exprP = ",,->"
             rev_sorted_dims = np.argsort(shape)[::-1]
             dim_diag = [False for _ in shape]
             dim_diag[rev_sorted_dims[0]] = True
 +        elif memory_save_mode == "smart_one_diag":
 +            dim_diag = [False for _ in shape]
 +            rev_sorted_dims = np.argsort(shape)[::-1]
 +            sorted_shape = sorted(shape)
 +            if len(shape) >= 2 and sorted_shape[-1] > sorted_shape[-2]:
 +                dim_diag[rev_sorted_dims[0]] = True
         elif memory_save_mode == "all_diag":
             dim_diag = [True for _ in shape]
         else:
             raise ValueError(
                 f"Invalid memory_save_mode: {memory_save_mode}, must be one of [None, 'one_diag', 'all_diag']")
 -        Q = []
         piece1A, piece2A, piece3A = ([], "", "")
         exprGs = []
         piece1P, piece2P, piece3P, piece4P = ([], [], "", "")
                 or dim_d
             ):
                 # use diagonal matrix as preconditioner for this dim
 -                Q.append(scale * torch.ones(size, dtype=dtype, device=t.device))
 +                if init_q:
 +                    Q.append(scale * torch.ones(size, dtype=dtype, device=t.device))
                 piece1A.append(letters[i])
                 piece2A = piece2A + letters[i]
                 piece4P = piece4P + letters[i + 13]
             else:
                 # use triangular matrix as preconditioner for this dim
 -                Q.append(scale * torch.eye(size, dtype=dtype, device=t.device))
 +                if init_q:
 +                    Q.append(scale * torch.eye(size, dtype=dtype, device=t.device))
                 piece1A.append(letters[i] + letters[i + 13])
                 piece2A = piece2A + letters[i + 13]
         exprP = ",".join(piece1P) + "," + ",".join(piece2P) + "," + piece3P + "->" + piece4P
     exprGs = tuple(exprGs)
 -    return [Q, (exprA, exprGs, exprP)]
 +    if init_q:
 +        return [Q, (exprA, exprGs, exprP)]
 +    else:
 +        return exprA, exprGs, exprP
 def _lb(A, max_abs):
 def _solve_triangular_right(X, A):
     """X @ inv(A)"""
     orig_dtype = X.dtype
 -    X = X.to(dtype=torch.float32, non_blocking=True)
 -    A = A.to(dtype=torch.float32, non_blocking=True)
 +    X = X.to(dtype=torch.float32)
 +    A = A.to(dtype=torch.float32)
     out = torch.linalg.solve_triangular(A, X.reshape(-1, X.size(-1)), upper=True, left=False).reshape_as(X)
 -    return out.to(dtype=orig_dtype, non_blocking=True)
 +    return out.to(dtype=orig_dtype)
 def _balance_Q(Q_in):