heavyball 0.24.3__tar.gz → 0.25.0__tar.gz

{heavyball-0.24.3 → heavyball-0.25.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: heavyball
-Version: 0.24.3
+Version: 0.25.0
 Summary: Efficient optimizers
 Home-page: https://github.com/clashluke/heavyball
 Author: Lucas Nestler

{heavyball-0.24.3 → heavyball-0.25.0}/heavyball/foreach_soap.py

@@ -71,12 +71,11 @@ class ForeachSOAP(StatefulOptimizer):
         # Decay the first and second moment running average coefficient
         # In-place operations to update the averages at the same time
         step_tensor = torch.empty((), dtype=torch.int32, device=p_list[0].device).fill_(step)
-        denom = exp_avg_(exp_avg, exp_avg_sq, grad, grad_projected, beta1, beta2, step_tensor)
 
         step_size = -group["lr"] * min(step / group['warmup_steps'], 1)
 
-        for p, g, ea, d in zip(p_list, grad, exp_avg, denom):
-            state = self.state_(p)
+        for p, g, gp, ea, eas in zip(p_list, grad, grad_projected, exp_avg, exp_avg_sq):
+            d = exp_avg_(ea, eas, g, gp, beta1, beta2, step_tensor)[0]
             # Projecting the exponential moving average of gradients to the eigenbases of Shampoo's preconditioner
             # i.e. projecting to the eigenbases of matrices in state['GG']
             exp_avg_projected = project(ea, state['Q'], False)

{heavyball-0.24.3 → heavyball-0.25.0}/heavyball/palm_foreach_soap.py

@@ -81,11 +81,10 @@ class PaLMForeachSOAP(StatefulOptimizer):
         # In-place operations to update the averages at the same time
         beta2 = torch.empty((), dtype=torch.float32, device=p_list[0].device).fill_(beta2)
         step_tensor = torch.empty((), dtype=torch.int32, device=p_list[0].device).fill_(step)
-        denom = exp_avg_(exp_avg, exp_avg_sq, grad, grad_projected, beta1, beta2, step_tensor)
         step_size = -group["lr"] * min(step / group['warmup_steps'], 1)
 
-        for p, g, ea, d in zip(p_list, grad, exp_avg, denom):
-            state = self.state_(p)
+        for p, g, gp, ea, eas in zip(p_list, grad, grad_projected, exp_avg, exp_avg_sq):
+            d = exp_avg_(ea, eas, g, gp, beta1, beta2, step_tensor)[0]
             # Projecting the exponential moving average of gradients to the eigenbases of Shampoo's preconditioner
             # i.e. projecting to the eigenbases of matrices in state['GG']
             exp_avg_projected = project(ea, state['Q'], False)

{heavyball-0.24.3 → heavyball-0.25.0}/heavyball/precond_schedule_foreach_soap.py

@@ -73,12 +73,12 @@ class PrecondScheduleForeachSOAP(StatefulOptimizer):
         # Decay the first and second moment running average coefficient
         # In-place operations to update the averages at the same time
         step_tensor = torch.empty((), dtype=torch.int32, device=p_list[0].device).fill_(step)
-        denom = exp_avg_(exp_avg, exp_avg_sq, grad, grad_projected, beta1, beta2, step_tensor)
 
         update_precond = precond_schedule(step, group['precond_scheduler'], self.rng)
         step_size = -group["lr"] * min(step / group['warmup_steps'], 1)
 
-        for p, g, ea, d in zip(p_list, grad, exp_avg, denom):
+        for p, g, gp, ea, eas in zip(p_list, grad, grad_projected, exp_avg, exp_avg_sq):
+            d = exp_avg_(ea, eas, g, gp, beta1, beta2, step_tensor)[0]
             state = self.state_(p)
             # Projecting the exponential moving average of gradients to the eigenbases of Shampoo's preconditioner
             # i.e. projecting to the eigenbases of matrices in state['GG']

{heavyball-0.24.3 → heavyball-0.25.0}/heavyball/precond_schedule_palm_foreach_soap.py

@@ -84,12 +84,12 @@ class PrecondSchedulePaLMForeachSOAP(StatefulOptimizer):
         # In-place operations to update the averages at the same time
         beta2 = torch.empty((), dtype=torch.float32, device=p_list[0].device).fill_(beta2)
         step_tensor = torch.empty((), dtype=torch.int32, device=p_list[0].device).fill_(step)
-        denom = exp_avg_(exp_avg, exp_avg_sq, grad, grad_projected, beta1, beta2, step_tensor)
 
         update_precond = precond_schedule(step, group['precond_scheduler'], self.rng)
         step_size = -group["lr"] * min(step / group['warmup_steps'], 1)
 
-        for p, g, ea, d in zip(p_list, grad, exp_avg, denom):
+        for p, g, gp, ea, eas in zip(p_list, grad, grad_projected, exp_avg, exp_avg_sq):
+            d = exp_avg_(ea, eas, g, gp, beta1, beta2, step_tensor)[0]
             state = self.state_(p)
             # Projecting the exponential moving average of gradients to the eigenbases of Shampoo's preconditioner
             # i.e. projecting to the eigenbases of matrices in state['GG']

{heavyball-0.24.3 → heavyball-0.25.0}/heavyball/utils.py

@@ -492,22 +492,20 @@ class StatefulOptimizer(torch.optim.Optimizer):
         super().__init__(params, {**defaults, 'foreach': foreach})
         self.fake_groups = {}
         self.use_ema = use_ema
-
-    def key(self, param: Tensor):
-        return (param.data_ptr(), tuple(param.shape))
+        self.mapping = {}
 
     def get_groups(self, group):
         if group['foreach']:
             return [group]
 
         for p in group['params']:
-            if self.key(p) not in self.fake_groups:
-                self.fake_groups[self.key(p)] = {**group, 'params': [p]}
+            if p not in self.fake_groups:
+                self.fake_groups[p] = {**group, 'params': [p]}
 
-        return [self.fake_groups[self.key(p)] for p in group['params']]
+        return [self.fake_groups[p] for p in group['params']]
 
     def state_(self, arg: Tensor):
-        return self.state[self.key(arg)]
+        return self.state[self.mapping.get(arg, arg)]
 
     def mars_correct_list(self, group, p_list, g_list, mars_gamma, beta):
         for p, g in zip(p_list, g_list):
@@ -538,6 +536,8 @@ class StatefulOptimizer(torch.optim.Optimizer):
             p_views = merge_group(group, p)
             if grad is not None:
                 grad = merge_group(group, grad)
+            for i, pv in enumerate(p_views):
+                self.mapping[pv] = (p, i)
             if isinstance(p_views, Tensor):
                 yield p_views, grad
                 continue
@@ -622,11 +622,14 @@ class StatefulOptimizer(torch.optim.Optimizer):
             for top_group in self.param_groups:
                 for group in self.get_groups(top_group):
                     self._step(group)
+                    self.mapping.clear()
                     if self.use_ema:
                         self.ema_update(group)
+
         return loss
 
 
+
 class ScheduleFree(StatefulOptimizer):
     def eval(self):
         for group in self.param_groups:

{heavyball-0.24.3 → heavyball-0.25.0}/heavyball.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: heavyball
-Version: 0.24.3
+Version: 0.25.0
 Summary: Efficient optimizers
 Home-page: https://github.com/clashluke/heavyball
 Author: Lucas Nestler

{heavyball-0.24.3 → heavyball-0.25.0}/setup.py

@@ -10,7 +10,7 @@ setuptools.setup(
     name='heavyball',
     license='BSD',
     description='Efficient optimizers',
-    version='0.24.3',
+    version='0.25.0',
     long_description=README,
     url='https://github.com/clashluke/heavyball',
     packages=setuptools.find_packages(),