Pass all tests

Author: ferris

pytorch_optimizer/base/optimizer.py

@@ -4,7 +4,7 @@
 
 import torch
 
-from pytorch_optimizer.base.exception import NegativeLRError, NegativeStepError
+from pytorch_optimizer.base.exception import NegativeLRError, NegativeStepError, NoSparseGradientError
 from pytorch_optimizer.base.types import BETAS, HUTCHINSON_G
 
 
@@ -48,7 +48,9 @@ def compute_hutchinson_hessian(self, nsamples: int = 1, pre_zero=True, alpha=1.0
         params = []
         for group in self.param_groups:
             for p in group['params']:
-                if p.grad is not None:
+                if p.requires_grad and p.grad is not None:
+                    if p.grad.is_sparse:
+                        raise NoSparseGradientError(str(self))
                     # Initialize Hessian state
                     if 'hessian' in self.state[p]:
                         if pre_zero:

pytorch_optimizer/optimizer/adahessian.py

@@ -69,13 +69,15 @@ def reset(self):
                 state['exp_hessian_diag_sq'] = torch.zero_like(p)
 
     @torch.no_grad()
-    def step(self, closure: CLOSURE = None) -> LOSS:
+    def step(self, closure: CLOSURE = None, hessian: tuple[torch.Tensor] = None) -> LOSS:
         loss: LOSS = None
         if closure is not None:
             with torch.enable_grad():
                 loss = closure()
 
-        if self._step % self.update_period == 0:
+        if hessian is not None:
+            self.set_hessian(hessian)
+        elif self._step % self.update_period == 0:
             self.compute_hutchinson_hessian(self.n_samples, distribution=self.distribution)
 
         for group in self.param_groups:
@@ -107,7 +109,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
 
                 # Decay the first and second moment running average coefficient
                 exp_avg.mul_(beta1).add_(p.grad, alpha=1 - beta1)
-                if self._step % self.update_period == 0:
+                if (self._step % self.update_period == 0 or hessian is not None) and 'hessian' in state:
                     # if self.average_conv_kernel and p.dim() == 4:
                     #     state['hessian'] = torch.abs(state['hessian']).mean(dim=[2, 3], keepdim=True).expand_as(state['hessian']).clone()
                     exp_hessian_diag_sq.mul_(beta2).addcmul_(state['hessian'], state['hessian'], value=1 - beta2)

pytorch_optimizer/optimizer/sophiah.py

@@ -83,12 +83,12 @@ def step(self, closure: CLOSURE = None, hessian: tuple[torch.Tensor] = None) ->
                 if p.grad is None:
                     continue
 
+                state = self.state[p]
                 grad = p.grad
                 if grad.is_sparse:
                     raise NoSparseGradientError(str(self))
 
                 # State initialization
-                state = self.state[p]
                 if 'momentum' not in state:
                     state['momentum'] = torch.zeros_like(p)
                     state['hessian_moment'] = torch.zeros_like(p)
@@ -106,7 +106,7 @@ def step(self, closure: CLOSURE = None, hessian: tuple[torch.Tensor] = None) ->
                 momentum, hessian_moment = state['momentum'], state['hessian_moment']
 
                 momentum.mul_(beta1).add_(p.grad, alpha=1.0-beta1)
-                if self._step % self.update_period == 0 or hessian is not None:
+                if (self._step % self.update_period == 0 or hessian is not None) and 'hessian' in state:
                     hessian_moment.mul_(beta2).add_(state['hessian'], alpha=1.0-beta2)
 
                 # See https://shreyansh26.github.io/post/2023-05-28_sophia_scalable_second_order_optimizer_llms/#per-coordinate-clipping

tests/test_general_optimizer_parameters.py

@@ -41,6 +41,7 @@ def test_epsilon(optimizer_name):
         'alig',
         'gravity',
         'srmm',
+        'signsgd'
     ):
         pytest.skip(f'skip {optimizer_name} optimizer')
 

tests/test_gradients.py

@@ -24,12 +24,18 @@ def test_no_gradients(optimizer_name):
     else:
         optimizer = load_optimizer(optimizer_name)(params)
 
+    def sphere_loss(x) -> torch.Tensor:
+        return (x ** 2).sum()
+
     optimizer.zero_grad()
-    p1.grad = torch.zeros(1, 1)
-    p2.grad = None
-    p3.grad = torch.zeros(1, 1)
-    p4.grad = None
+    sphere_loss(p1 + p3).backward(create_graph=True)
+    # p1.grad = torch.zeros(1, 1)
+    # p2.grad = None
+    # p3.grad = torch.zeros(1, 1)
+    # p4.grad = None
     optimizer.step(lambda: 0.1)  # for AliG optimizer
+    if optimizer_name != 'lookahead':
+        optimizer.zero_grad(set_to_none=True)
 
 
 @pytest.mark.parametrize('no_sparse_optimizer', NO_SPARSE_OPTIMIZERS)
@@ -109,12 +115,17 @@ def test_bf16_gradient(optimizer_name):
     if optimizer_name == 'shampoo':
         pytest.skip(f'skip {optimizer_name}')
 
+    def sphere_loss(x) -> torch.Tensor:
+        return (x ** 2).sum()
+
     param = torch.randn(1, 1).bfloat16().requires_grad_(True)
-    param.grad = torch.randn(1, 1).bfloat16()
 
     opt = load_optimizer(optimizer=optimizer_name)
     optimizer = opt([param], num_iterations=1) if optimizer_name == 'ranger21' else opt([param])
+
+    sphere_loss(param).backward(create_graph=True)
     optimizer.step(lambda: 0.1)
+    optimizer.zero_grad(True)
 
 
 def test_sam_no_gradient():

tests/test_load_optimizers.py

@@ -16,4 +16,4 @@ def test_load_optimizers_invalid(invalid_optimizer_names):
 
 
 def test_get_supported_optimizers():
-    assert len(get_supported_optimizers()) == 51
+    assert len(get_supported_optimizers()) == 54

tests/test_optimizers.py

@@ -322,6 +322,27 @@ def test_rectified_optimizer(optimizer_name):
     optimizer.step()
 
 
+@pytest.mark.parametrize('optimizer_name', ['sophiah', 'adahessian'])
+def test_hessian_optimizer(optimizer_name):
+    param = simple_parameter()
+
+    def sphere_loss(x) -> torch.Tensor:
+        return (x ** 2).sum()
+
+    parameters = {'hessian_distribution': 'gaussian', 'n_samples': 2}
+    optimizer = load_optimizer(optimizer_name)([param], **parameters)
+    optimizer.zero_grad(set_to_none=True)
+
+    # Hutchinson (internal) estimator
+    sphere_loss(param).backward(create_graph=True)
+    optimizer.step()
+    optimizer.zero_grad(set_to_none=True)
+
+    # External estimator
+    sphere_loss(param).backward()
+    optimizer.step(hessian=torch.zeros_like(param).unsqueeze(0))
+
+
 @pytest.mark.parametrize('optimizer_config', OPTIMIZERS + ADANORM_SUPPORTED_OPTIMIZERS, ids=ids)
 def test_reset(optimizer_config):
     optimizer_class, config, _ = optimizer_config