heavyball 0.24.4__tar.gz → 0.25.1__tar.gz

{heavyball-0.24.4 → heavyball-0.25.1}/PKG-INFO

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

{heavyball-0.24.4 → heavyball-0.25.1}/heavyball/__init__.py

@@ -15,6 +15,10 @@ from .precond_schedule_sfpsoap import PrecondScheduleSFPaLMSOAP
 from .psgd_kron import ForeachPSGDKron
 from .pure_psgd import ForeachPurePSGD
 from .schedule_free_palm_foreach_soap import SFPaLMForeachSOAP
+from .foreach_solp import ForeachSOLP
+from .palm_foreach_solp import PaLMForeachSOLP
+from .precond_schedule_palm_foreach_solp import PrecondSchedulePaLMForeachSOLP
+from .precond_schedule_foreach_solp import PrecondScheduleForeachSOLP
 
 PalmForEachSoap = PaLMForeachSOAP
 
@@ -35,12 +39,18 @@ PaLMPAdam = ForeachPaLMPAdam
 DelayedPSGD = ForeachDelayedPSGD
 CachedPSGDKron = ForeachCachedPSGDKron
 CachedDelayedPSGDKron = ForeachCachedDelayedPSGDKron
+SOLP = ForeachSOLP
+PaLMSOLP = PaLMForeachSOLP
+PrecondSchedulePaLMSOLP = PrecondSchedulePaLMForeachSOLP
+PrecondScheduleSOLP = PrecondScheduleForeachSOLP
 
 __all__ = ['PalmForEachSoap', 'PaLMForeachSFAdamW', 'PaLMForeachSOAP', 'SFPaLMForeachSOAP', 'PrecondScheduleSFPaLMSOAP',
            'ForeachSOAP', 'ForeachSFAdamW', 'ForeachLaProp', 'ForeachADOPT', 'PrecondScheduleForeachSOAP',
            'PrecondSchedulePaLMForeachSOAP', 'ForeachPSGDKron', 'ForeachAdamW', 'ForeachPurePSGD', 'ForeachPaLMPAdam',
-           'ForeachDelayedPSGD', 'ForeachCachedPSGDKron', 'ForeachCachedDelayedPSGDKron',
+           'ForeachDelayedPSGD', 'ForeachCachedPSGDKron', 'ForeachCachedDelayedPSGDKron', 'ForeachSOLP',
+           'PaLMForeachSOLP', 'PrecondSchedulePaLMForeachSOLP', 'PrecondScheduleForeachSOLP',
              #
            'PaLMSOAP', 'PaLMSFAdamW', 'PaLMSFSoap', 'PaLMSFAdamW', 'PrecondScheduleSFPaLMSOAP', 'SOAP', 'SFAdamW',
            'LaProp', 'ADOPT', 'PSGDKron', 'AdamW', 'PurePSGD', 'PaLMPAdam', 'DelayedPSGD', 'CachedPSGDKron',
-           'CachedDelayedPSGDKron', 'PrecondScheduleSOAP', 'PrecondSchedulePaLMSOAP']
+           'CachedDelayedPSGDKron', 'PrecondScheduleSOAP', 'PrecondSchedulePaLMSOAP', 'SOLP', 'PrecondScheduleSOLP',
+           'PrecondSchedulePaLMSOLP', 'PrecondScheduleSOLP']

{heavyball-0.24.4 → heavyball-0.25.1}/heavyball/foreach_soap.py

@@ -6,7 +6,7 @@ from .utils import init_preconditioner, update_preconditioner, project, beta_deb
 
 class ForeachSOAP(StatefulOptimizer):
     """
-    SFPaLMForeachSOAP
+    ForeachSOAP
 
     Sources:
         Baseline SOAP:
@@ -75,6 +75,8 @@ class ForeachSOAP(StatefulOptimizer):
         step_size = -group["lr"] * min(step / group['warmup_steps'], 1)
 
         for p, g, gp, ea, eas in zip(p_list, grad, grad_projected, exp_avg, exp_avg_sq):
+            state = self.state_(p)
+
             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']

heavyball-0.25.1/heavyball/foreach_solp.py

@@ -0,0 +1,89 @@
+import torch
+
+from .utils import init_preconditioner, update_preconditioner, project, beta_debias, exp_avg_sq_, update_param_, set_, \
+    StatefulOptimizer, laprop_exp_avg_
+
+
+class ForeachSOLP(StatefulOptimizer):
+    """
+    ForeachSOLP
+
+    Sources:
+        LaProp:
+            LaProp: Separating Momentum and Adaptivity in Adam
+            Liu Ziyin, Zhikang T.Wang, Masahito Ueda
+            https://arxiv.org/abs/2002.04839
+            https://github.com/ClashLuke/HeavyBall/blob/main/heavyball/foreach_laprop.py
+
+        Baseline SOAP:
+            SOAP: Improving and Stabilizing Shampoo using Adam
+            Nikhil Vyas, Depen Morwani, Rosie Zhao, Itai Shapira, David Brandfonbrener, Lucas Janson, Sham Kakade
+            https://arxiv.org/abs/2409.11321
+            https://github.com/nikhilvyas/SOAP
+
+        ScheduleFree:
+            The Road Less Scheduled
+            Aaron Defazio, Xingyu Alice Yang, Harsh Mehta, Konstantin Mishchenko, Ahmed Khaled, Ashok Cutkosky
+            https://arxiv.org/abs/2405.15682
+            https://github.com/facebookresearch/schedule_free
+    """
+
+    def __init__(self, params, lr: float = 3e-3, betas=(0.9, 0.95), shampoo_beta: float = 0.95, eps: float = 1e-8,
+                 weight_decay: float = 0.01, precondition_frequency: int = 2, max_precond_dim: int = 2048,  #
+                 merge_dims: bool = True, precondition_1d: bool = False, normalize_grads: bool = False,
+                 data_format: str = "channels_first", correct_bias: bool = True, warmup_steps: int = 1,
+                 split: bool = False, foreach: bool = True, mars: bool = False, caution: bool = False,
+                 mars_gamma: float = 0.0025):
+        defaults = {"lr": lr, "betas": betas, "shampoo_beta": shampoo_beta, "eps": eps, "weight_decay": weight_decay,
+                    "precondition_frequency": precondition_frequency, "max_precond_dim": max_precond_dim,
+                    "merge_dims": merge_dims, "precondition_1d": precondition_1d, "normalize_grads": normalize_grads,
+                    "correct_bias": correct_bias, 'warmup_steps': warmup_steps, 'split': split, 'mars': mars,
+                    'caution': caution, 'mars_gamma': mars_gamma}
+        super().__init__(params, defaults, foreach)
+        self._data_format = data_format
+
+    def _step(self, group):
+        vals = []
+        step = 0
+
+        max_precond_dim = group['max_precond_dim']
+        precondition_1d = group['precondition_1d']
+
+        for p, g in self.split_p_and_g_in_group(group, beta1=group['betas'][0]):
+            state = self.state_(p)
+            step = state['step'] = state.get("step", -1) + 1
+
+            if "exp_avg" not in state:
+                state["exp_avg"] = torch.zeros_like(g, dtype=torch.float32, memory_format=torch.preserve_format)
+                state["exp_avg_sq"] = torch.zeros_like(g, dtype=torch.float32, memory_format=torch.preserve_format)
+                init_preconditioner(g, state, max_precond_dim, precondition_1d)
+                update_preconditioner(g, state, max_precond_dim, precondition_1d, 0, True)
+                continue  # first step is skipped so that we never use the current gradients in the projection.
+
+            # Projecting gradients to the eigenbases of Shampoo's preconditioner
+            # i.e. projecting to the eigenbases of matrices in state['GG']
+            grad_projected = project(g, state['Q'], False)
+            exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+            vals.append((p, g, grad_projected, exp_avg, exp_avg_sq))
+
+        if not vals:
+            return
+
+        p_list, grad, grad_projected, exp_avg, exp_avg_sq = zip(*vals)
+        beta1, beta2 = group["betas"]
+
+        old_debiased2 = beta_debias(beta2, step)
+
+        # 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)
+
+        step_size = -group["lr"] * min(step / group['warmup_steps'], 1)
+
+        for p, g, gp, ea, eas in zip(p_list, grad, grad_projected, exp_avg, exp_avg_sq):
+            laprop_exp_avg_(ea, eas, gp, beta1, beta2, step_tensor)
+            state = self.state_(p)
+            precond = project(ea, state['Q'], True)
+
+            update_preconditioner(g, state, max_precond_dim, precondition_1d, old_debiased2, step > 0 and step % group['precondition_frequency'] == 0)
+            update_param_([p], [precond], step_size, group["weight_decay"], caution=group['caution'], grad=[g])

{heavyball-0.24.4 → heavyball-0.25.1}/heavyball/palm_foreach_soap.py

@@ -6,7 +6,7 @@ from .utils import init_preconditioner, update_preconditioner, project, beta_deb
 
 class PaLMForeachSOAP(StatefulOptimizer):
     """
-    SFPaLMForeachSOAP
+    PaLMForeachSOAP
 
     Sources:
         Baseline SOAP:
@@ -84,6 +84,7 @@ class PaLMForeachSOAP(StatefulOptimizer):
         step_size = -group["lr"] * min(step / group['warmup_steps'], 1)
 
         for p, g, gp, ea, eas in zip(p_list, grad, grad_projected, exp_avg, exp_avg_sq):
+            state = self.state_(p)
             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']

heavyball-0.25.1/heavyball/palm_foreach_solp.py

@@ -0,0 +1,98 @@
+import torch
+
+from .utils import init_preconditioner, update_preconditioner, project, beta_debias, exp_avg_sq_, update_param_, set_, \
+    StatefulOptimizer, laprop_exp_avg_
+
+
+class PaLMForeachSOLP(StatefulOptimizer):
+    """
+    PaLMForeachSOAP
+
+    Sources:
+        LaProp:
+            LaProp: Separating Momentum and Adaptivity in Adam
+            Liu Ziyin, Zhikang T.Wang, Masahito Ueda
+            https://arxiv.org/abs/2002.04839
+            https://github.com/ClashLuke/HeavyBall/blob/main/heavyball/foreach_laprop.py
+
+        Baseline SOAP:
+            SOAP: Improving and Stabilizing Shampoo using Adam
+            Nikhil Vyas, Depen Morwani, Rosie Zhao, Itai Shapira, David Brandfonbrener, Lucas Janson, Sham Kakade
+            https://arxiv.org/abs/2409.11321
+            https://github.com/nikhilvyas/SOAP
+
+        ScheduleFree:
+            The Road Less Scheduled
+            Aaron Defazio, Xingyu Alice Yang, Harsh Mehta, Konstantin Mishchenko, Ahmed Khaled, Ashok Cutkosky
+            https://arxiv.org/abs/2405.15682
+            https://github.com/facebookresearch/schedule_free
+
+        Beta2 Schedule:
+            PaLM: Scaling Language Modeling with Pathways
+            Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Levskaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowycz, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier-Hellstern, Douglas Eck, Jeff Dean, Slav Petrov, Noah Fiedel
+            https://arxiv.org/abs/2204.02311
+    """
+
+    def __init__(self, params, lr: float = 3e-3, beta=0.9, betas=(None, None), shampoo_beta: float = 0.95,
+                 eps: float = 1e-8, weight_decay: float = 0.01, precondition_frequency: int = 2,
+                 max_precond_dim: int = 2048,  #
+                 merge_dims: bool = True, precondition_1d: bool = False, normalize_grads: bool = False,
+                 data_format: str = "channels_first", correct_bias: bool = True, warmup_steps: int = 1,
+                 beta2_scale: float = 0.8, split: bool = False, foreach: bool = True, mars: bool = False,
+                 caution: bool = False, mars_gamma: float = 0.0025):
+        if betas[0] is not None:
+            beta = betas[0]
+        defaults = {"lr": lr, "beta": beta, "shampoo_beta": shampoo_beta, "eps": eps, "weight_decay": weight_decay,
+                    "precondition_frequency": precondition_frequency, "max_precond_dim": max_precond_dim,
+                    "merge_dims": merge_dims, "precondition_1d": precondition_1d, "normalize_grads": normalize_grads,
+                    "correct_bias": correct_bias, 'warmup_steps': warmup_steps, 'beta2_scale': beta2_scale,
+                    'split': split, 'mars': mars, 'caution': caution, 'mars_gamma': mars_gamma}
+        super().__init__(params, defaults, foreach)
+        self._data_format = data_format
+
+    def _step(self, group):
+        vals = []
+        step = 0
+
+        max_precond_dim = group['max_precond_dim']
+        precondition_1d = group['precondition_1d']
+
+        for p, g in self.split_p_and_g_in_group(group, beta1=group['beta']):
+            state = self.state_(p)
+            step = state['step'] = state.get("step", -1) + 1
+
+            if "exp_avg" not in state:
+                state["exp_avg"] = torch.zeros_like(g, dtype=torch.float32, memory_format=torch.preserve_format)
+                state["exp_avg_sq"] = torch.zeros_like(g, dtype=torch.float32, memory_format=torch.preserve_format)
+                init_preconditioner(g, state, max_precond_dim, precondition_1d)
+                update_preconditioner(g, state, max_precond_dim, precondition_1d, 0, True)
+                continue  # first step is skipped so that we never use the current gradients in the projection.
+
+            # Projecting gradients to the eigenbases of Shampoo's preconditioner
+            # i.e. projecting to the eigenbases of matrices in state['GG']
+            grad_projected = project(g, state['Q'], False)
+            exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+            vals.append((p, g, grad_projected, exp_avg, exp_avg_sq))
+
+        if not vals:
+            return
+
+        p_list, grad, grad_projected, exp_avg, exp_avg_sq = zip(*vals)
+        beta1 = group["beta"]
+
+        beta2 = 1 - step ** -group['beta2_scale']
+        old_debiased2 = beta_debias(beta2, step)
+
+        # Decay the first and second moment running average coefficient
+        # 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)
+        step_size = -group["lr"] * min(step / group['warmup_steps'], 1)
+
+        for p, g, gp, ea, eas in zip(p_list, grad, grad_projected, exp_avg, exp_avg_sq):
+            laprop_exp_avg_(ea, eas, gp, beta1, beta2, step_tensor)
+            state = self.state_(p)
+            precond = project(ea, state['Q'], True)
+
+            update_preconditioner(g, state, max_precond_dim, precondition_1d, old_debiased2, step > 0 and step % group['precondition_frequency'] == 0)
+            update_param_([p], [precond], step_size, group["weight_decay"], caution=group['caution'], grad=[g])

heavyball-0.25.1/heavyball/precond_schedule_foreach_solp.py

@@ -0,0 +1,95 @@
+
+
+import random
+
+import torch
+
+from .utils import init_preconditioner, update_preconditioner, project, beta_debias, update_param_, \
+    precond_schedule, set_, StatefulOptimizer, laprop_exp_avg_
+
+
+class PrecondScheduleForeachSOLP(StatefulOptimizer):
+    """
+    Sources:
+        LaProp:
+            LaProp: Separating Momentum and Adaptivity in Adam
+            Liu Ziyin, Zhikang T.Wang, Masahito Ueda
+            https://arxiv.org/abs/2002.04839
+            https://github.com/ClashLuke/HeavyBall/blob/main/heavyball/foreach_laprop.py
+
+        Preconditioner Schedules:
+            Preconditioned Stochastic Gradient Descent
+            Xi-Lin Li, Omead Pooladzandi, Evan Walters
+            https://arxiv.org/abs/1512.04202
+            https://github.com/evanatyourservice/kron_torch
+            https://github.com/lixilinx/psgd_torch
+
+        Baseline SOAP:
+            SOAP: Improving and Stabilizing Shampoo using Adam
+            Nikhil Vyas, Depen Morwani, Rosie Zhao, Itai Shapira, David Brandfonbrener, Lucas Janson, Sham Kakade
+            https://arxiv.org/abs/2409.11321
+            https://github.com/nikhilvyas/SOAP
+    """
+
+    def __init__(self, params, lr: float = 3e-3, betas=(0.9, 0.95), shampoo_beta: float = 0.95, eps: float = 1e-8,
+                 weight_decay: float = 0.01, precondition_frequency: int = 2, max_precond_dim: int = 2048,  #
+                 merge_dims: bool = True, precondition_1d: bool = False, normalize_grads: bool = False,
+                 data_format: str = "channels_first", correct_bias: bool = True, warmup_steps: int = 1,
+                 precond_scheduler=(1 / 3, 9), split: bool = False, foreach: bool = True, mars: bool = False,
+                 caution: bool = False, mars_gamma: float = 0.0025):
+        defaults = {"lr": lr, "betas": betas, "shampoo_beta": shampoo_beta, "eps": eps, "weight_decay": weight_decay,
+                    "precondition_frequency": precondition_frequency, "max_precond_dim": max_precond_dim,
+                    "merge_dims": merge_dims, "precondition_1d": precondition_1d, "normalize_grads": normalize_grads,
+                    "correct_bias": correct_bias, 'warmup_steps': warmup_steps, 'precond_scheduler': precond_scheduler,
+                    'split': split, 'mars': mars, 'caution': caution, 'mars_gamma': mars_gamma}
+        super().__init__(params, defaults, foreach)
+        self._data_format = data_format
+        self.rng = random.Random(0x120983109)
+
+    def _step(self, group):
+        vals = []
+        step = 0
+
+        max_precond_dim = group['max_precond_dim']
+        precondition_1d = group['precondition_1d']
+
+        for p, g in self.split_p_and_g_in_group(group, beta1=group['betas'][0]):
+            state = self.state_(p)
+            step = state['step'] = state.get("step", -1) + 1
+
+            if "exp_avg" not in state:
+                state["exp_avg"] = torch.zeros_like(g, dtype=torch.float32, memory_format=torch.preserve_format)
+                state["exp_avg_sq"] = torch.zeros_like(g, dtype=torch.float32, memory_format=torch.preserve_format)
+                init_preconditioner(g, state, max_precond_dim, precondition_1d)
+                update_preconditioner(g, state, max_precond_dim, precondition_1d, 0, True)
+                continue  # first step is skipped so that we never use the current gradients in the projection.
+
+            # Projecting gradients to the eigenbases of Shampoo's preconditioner
+            # i.e. projecting to the eigenbases of matrices in state['GG']
+            grad_projected = project(g, state['Q'], False)
+            exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+            vals.append((p, g, grad_projected, exp_avg, exp_avg_sq))
+
+        if not vals:
+            return
+
+        p_list, grad, grad_projected, exp_avg, exp_avg_sq = zip(*vals)
+        beta1, beta2 = group["betas"]
+
+        old_debiased2 = beta_debias(beta2, step)
+
+        # 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)
+
+        update_precond = precond_schedule(step, group['precond_scheduler'], self.rng)
+        step_size = -group["lr"] * min(step / group['warmup_steps'], 1)
+
+
+        for p, g, gp, ea, eas in zip(p_list, grad, grad_projected, exp_avg, exp_avg_sq):
+            laprop_exp_avg_(ea, eas, gp, beta1, beta2, step_tensor)
+            state = self.state_(p)
+            precond = project(ea, state['Q'], True)
+
+            update_preconditioner(g, state, max_precond_dim, precondition_1d, old_debiased2, update_precond)
+            update_param_([p], [precond], step_size, group["weight_decay"], caution=group['caution'], grad=[g])

heavyball-0.25.1/heavyball/precond_schedule_palm_foreach_solp.py

@@ -0,0 +1,103 @@
+import random
+
+import torch
+
+from .utils import init_preconditioner, update_preconditioner, project, beta_debias, laprop_exp_avg_, update_param_, \
+    precond_schedule, set_, StatefulOptimizer
+
+
+class PrecondSchedulePaLMForeachSOLP(StatefulOptimizer):
+    """
+    Sources:
+        LaProp:
+            LaProp: Separating Momentum and Adaptivity in Adam
+            Liu Ziyin, Zhikang T.Wang, Masahito Ueda
+            https://arxiv.org/abs/2002.04839
+            https://github.com/ClashLuke/HeavyBall/blob/main/heavyball/foreach_laprop.py
+
+        Preconditioner Schedules:
+            Preconditioned Stochastic Gradient Descent
+            Xi-Lin Li, Omead Pooladzandi, Evan Walters
+            https://arxiv.org/abs/1512.04202
+            https://github.com/evanatyourservice/kron_torch
+            https://github.com/lixilinx/psgd_torch
+
+        Baseline SOAP:
+            SOAP: Improving and Stabilizing Shampoo using Adam
+            Nikhil Vyas, Depen Morwani, Rosie Zhao, Itai Shapira, David Brandfonbrener, Lucas Janson, Sham Kakade
+            https://arxiv.org/abs/2409.11321
+            https://github.com/nikhilvyas/SOAP
+
+        Beta2 Schedule:
+            PaLM: Scaling Language Modeling with Pathways
+            Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Levskaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowycz, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier-Hellstern, Douglas Eck, Jeff Dean, Slav Petrov, Noah Fiedel
+            https://arxiv.org/abs/2204.02311
+    """
+
+    def __init__(self, params, lr: float = 3e-3, beta=0.9, shampoo_beta: float = 0.95, eps: float = 1e-8,
+                 weight_decay: float = 0.01, precondition_frequency: int = 2, max_precond_dim: int = 2048,  #
+                 merge_dims: bool = True, precondition_1d: bool = False, normalize_grads: bool = False,
+                 data_format: str = "channels_first", correct_bias: bool = True, warmup_steps: int = 1,
+                 precond_scheduler=(1 / 3, 9), betas=(None, None), beta2_scale: float = 0.8, split: bool = False,
+                 foreach: bool = True, mars: bool = False, caution: bool = False, mars_gamma: float = 0.0025):
+        if betas[0] is not None:
+            beta = betas[0]
+        defaults = {"lr": lr, "beta": beta, "shampoo_beta": shampoo_beta, "eps": eps, "weight_decay": weight_decay,
+                    "precondition_frequency": precondition_frequency, "max_precond_dim": max_precond_dim,
+                    "merge_dims": merge_dims, "precondition_1d": precondition_1d, "normalize_grads": normalize_grads,
+                    "correct_bias": correct_bias, 'warmup_steps': warmup_steps, 'precond_scheduler': precond_scheduler,
+                    'beta2_scale': beta2_scale, 'split': split, 'mars': mars, 'caution': caution,
+                    'mars_gamma': mars_gamma}
+        super().__init__(params, defaults, foreach)
+        self._data_format = data_format
+        self.rng = random.Random(0x120983109)
+
+    def _step(self, group):
+        vals = []
+        step = 0
+
+        max_precond_dim = group['max_precond_dim']
+        precondition_1d = group['precondition_1d']
+
+        for p, g in self.split_p_and_g_in_group(group, beta1=group['beta']):
+            state = self.state_(p)
+            step = state['step'] = state.get("step", -1) + 1
+
+            if "exp_avg" not in state:
+                state["exp_avg"] = torch.zeros_like(g, dtype=torch.float32, memory_format=torch.preserve_format)
+                state["exp_avg_sq"] = torch.zeros_like(g, dtype=torch.float32, memory_format=torch.preserve_format)
+                init_preconditioner(g, state, max_precond_dim, precondition_1d)
+                update_preconditioner(g, state, max_precond_dim, precondition_1d, 0, True)
+                continue  # first step is skipped so that we never use the current gradients in the projection.
+
+            # Projecting gradients to the eigenbases of Shampoo's preconditioner
+            # i.e. projecting to the eigenbases of matrices in state['GG']
+            grad_projected = project(g, state['Q'], False)
+            exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+            vals.append((p, g, grad_projected, exp_avg, exp_avg_sq))
+
+        if not vals:
+            return
+
+        p_list, grad, grad_projected, exp_avg, exp_avg_sq = zip(*vals)
+        beta1 = group["beta"]
+
+        beta2 = 1 - max(step, 1) ** -group['beta2_scale']
+        old_debiased1 = beta_debias(beta1, step)
+        old_debiased2 = beta_debias(beta2, step)
+
+        # Decay the first and second moment running average coefficient
+        # 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)
+
+        update_precond = precond_schedule(step, group['precond_scheduler'], self.rng)
+        step_size = -group["lr"] * min(step / group['warmup_steps'], 1)
+
+        for p, g, gp, ea, eas in zip(p_list, grad, grad_projected, exp_avg, exp_avg_sq):
+            laprop_exp_avg_(ea, eas, gp, beta1, beta2, step_tensor)
+            state = self.state_(p)
+            precond = project(ea, state['Q'], True)
+
+            update_preconditioner(g, state, max_precond_dim, precondition_1d, old_debiased2, update_precond)
+            update_param_([p], [precond], step_size, group["weight_decay"], caution=group['caution'], grad=[g])

{heavyball-0.24.4 → heavyball-0.25.1}/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:
@@ -690,6 +693,29 @@ def exp_avg_(exp_avg: List[Tensor], exp_avg_sq: List[Tensor], grad: List[Tensor]
     return denom
 
 
+
+@decorator_knowngood
+def _compilable_laprop_exp_avg_(exp_avg: List[Tensor], exp_avg_sq: List[Tensor],
+                         grad_projected: List[Tensor], beta1: Tensor, beta2: Tensor, step: Tensor):
+    beta1 = beta_debias(beta1, step)
+    beta2 = beta_debias(beta2, step)
+
+    gp32, exp_avg_sq32 = [list(map(promote, x)) for x in [grad_projected, exp_avg_sq]]
+
+    denom = exp_avg_sq_(exp_avg_sq32, gp32, beta2, 1e-8)
+    gp32 = torch._foreach_div(gp32, denom)
+    stochastic_lerp_(exp_avg, gp32, 1 - beta1)
+
+    copy_stochastic_list_(exp_avg_sq, exp_avg_sq32)
+
+
+def laprop_exp_avg_(exp_avg: List[Tensor], exp_avg_sq: List[Tensor], grad_projected: List[Tensor],
+             beta1: float, beta2: float, step: int):
+    exp_avg, exp_avg_sq, grad_projected = list_guard(exp_avg), list_guard(exp_avg_sq), list_guard(grad_projected)
+    beta1, beta, step = scalar_guard(beta1, exp_avg[0]), scalar_guard(beta2, exp_avg[0]), scalar_guard(step, exp_avg[0])
+    _compilable_laprop_exp_avg_(exp_avg, exp_avg_sq, grad_projected, beta1, beta2, step)
+
+
 @decorator_knowngood
 def _compilable_copy_stochastic_(target: Tensor, source: Tensor):
     """Taken as-is from https://github.com/pytorch/pytorch/issues/120376#issuecomment-1974828905"""

{heavyball-0.24.4 → heavyball-0.25.1}/heavyball.egg-info/PKG-INFO

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

{heavyball-0.24.4 → heavyball-0.25.1}/heavyball.egg-info/SOURCES.txt

@@ -10,11 +10,15 @@ heavyball/foreach_adopt.py
 heavyball/foreach_laprop.py
 heavyball/foreach_sfadamw.py
 heavyball/foreach_soap.py
+heavyball/foreach_solp.py
 heavyball/p_adam.py
 heavyball/palm_foreach_sfadamw.py
 heavyball/palm_foreach_soap.py
+heavyball/palm_foreach_solp.py
 heavyball/precond_schedule_foreach_soap.py
+heavyball/precond_schedule_foreach_solp.py
 heavyball/precond_schedule_palm_foreach_soap.py
+heavyball/precond_schedule_palm_foreach_solp.py
 heavyball/precond_schedule_sfpsoap.py
 heavyball/psgd_kron.py
 heavyball/pure_psgd.py
@@ -40,4 +44,5 @@ test/test_merge.py
 test/test_no_grad.py
 test/test_psgd.py
 test/test_soap.py
+test/test_solp.py
 test/test_stochastic_updates.py

{heavyball-0.24.4 → heavyball-0.25.1}/setup.py

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

heavyball-0.25.1/test/test_solp.py

@@ -0,0 +1,50 @@
+import heavyball
+import heavyball.utils
+import pytest
+import torch
+from benchmark.utils import get_optim
+from heavyball.utils import clean, set_torch, ScheduleFree
+from torch import nn
+from torch._dynamo import config
+
+config.cache_size_limit = 128
+
+@pytest.mark.parametrize("opt", heavyball.__all__)
+@pytest.mark.parametrize("size,depth", [(128, 2)])
+def test_solp(opt, size, depth: int, iterations: int = 65536, outer_iterations: int = 2):
+    set_torch()
+    if 'SOAP' not in opt:
+        raise pytest.skip('This test is only for SOAP')
+
+    opt_name = opt
+    peaks = []
+    losses = []
+
+    for use_solp in [True, False]:
+        try:
+            opt = getattr(heavyball, opt_name.replace("SOAP", "SOLP") if use_solp else opt_name)
+        except AttributeError:
+            raise pytest.skip(f'{opt_name} does not have a SOLP variant')
+        print(opt, opt_name.replace("SOAP", "SOLP"))
+
+        torch.manual_seed(0x2131290)
+        peaks.append([])
+        losses.append([])
+
+        for i in range(outer_iterations):
+            model = nn.Sequential(*[nn.Linear(size, size) for _ in range(depth)]).cuda()
+            o = get_optim(opt, model.parameters(), lr=1e-5)
+
+            for _ in range(iterations):
+                loss = model(torch.randn((1024, size), device='cuda')).square().mean()
+                loss.backward()
+                o.step()
+                o.zero_grad()
+                losses[-1].append(loss.detach())
+
+            del model, o
+            clean()
+
+    for i, (l0, l1) in enumerate(zip(*losses)):
+        print(i, l0.item(), l1.item())
+        assert l0.item() <= l1.item() * 1.1