heavyball 0.24.4__py3-none-any.whl → 0.25.1__py3-none-any.whl

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/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/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/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/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/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/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/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.dist-info → heavyball-0.25.1.dist-info}/METADATA

@@ -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.dist-info → heavyball-0.25.1.dist-info}/RECORD

@@ -1,4 +1,4 @@
-heavyball/__init__.py,sha256=icHYN-MGsmHkLUlHCMcZkOlwY7GT63_ayR_a5iPKmzM,2226
+heavyball/__init__.py,sha256=RdUfGDTXw-rtoQJNediWnhDseYyyWNPVsr6tRq_ucp8,2813
 heavyball/cached_delayed_psgd_kron.py,sha256=HEyT6vW6Le6FmWpf-vAEzgbAkPH2mByqXcVZn07KCMk,6866
 heavyball/cached_psgd_kron.py,sha256=rOgWAeVMENI7kdoBuRo3ywrCeatAnIqBdeYPHuVk2aU,6998
 heavyball/delayed_psgd.py,sha256=L6qRLPxJmJ_1e0Mk2zLYUEVxkt8NGHq6v3HKawlgFcU,6334
@@ -6,19 +6,23 @@ heavyball/foreach_adamw.py,sha256=K4xTes4drylAqaqWky8O_Bg_mmbAmcHZ5DEBs5vMD-s,28
 heavyball/foreach_adopt.py,sha256=fHnbEqvKKc5IKPDWC9Qo9PiISSjj1MEViy0Jb3BRgZQ,3582
 heavyball/foreach_laprop.py,sha256=EXkwFQ-H7hHWLmiNUsxUcmXhzNNLMjieHjfOlY_6kmo,2868
 heavyball/foreach_sfadamw.py,sha256=TeWf0nKXQEFcz02rADYRJenDM9mX1dGHhvILLks6OW8,3087
-heavyball/foreach_soap.py,sha256=Tgwg4_Sir9nI_3R85f8NMQagquUBJmAEMQqh0uD3b0Y,4771
+heavyball/foreach_soap.py,sha256=ntFqg0fbkZ8EzERGlypXB8JWoGJ1sAY59f0CuWh_d48,4801
+heavyball/foreach_solp.py,sha256=1r7x_FUZRaUsoSLSvi-Z_-pZNtZrMresVJGq9m1EREA,4563
 heavyball/p_adam.py,sha256=qEcuU8VEc35vaWAXjT0O65vfCuNn_3ttwL4RlJKN3Xw,6389
 heavyball/palm_foreach_sfadamw.py,sha256=1qOr-uniSmI1sNCJc1SnvyKH5iFu80Z6H5h93lDTwcE,3410
-heavyball/palm_foreach_soap.py,sha256=zSjpYYm1hfgIudjo_q3ozu3Vkfhz8w8im1c-ou1U3sI,6198
+heavyball/palm_foreach_soap.py,sha256=fbRL1Tx9YeQ16sHWFPtY5Kj60BFV2AMngOnTiE4muK0,6231
+heavyball/palm_foreach_solp.py,sha256=N3M3tnahOfSHvLu3en76JTI1yo-ISEbSliSKlpt8ZWw,5994
 heavyball/precond_schedule_foreach_soap.py,sha256=p7oD2bESyCPsdGkJYhHluraDb_1K5Q28RNL6fIvD5C8,4969
+heavyball/precond_schedule_foreach_solp.py,sha256=xGEQ6HHUTCKeT9-ppEbLTdXVAfE74P0tph0qS16USyg,4768
 heavyball/precond_schedule_palm_foreach_soap.py,sha256=Sb3Fhv-EG28_oXnbVpE0iHe5R8i5_hltqoi_DgPuoEU,6505
+heavyball/precond_schedule_palm_foreach_solp.py,sha256=gaoJwJo_ZBnYuMamgFepnV9iWpUCmbYrxMWiL1QkPh0,6253
 heavyball/precond_schedule_sfpsoap.py,sha256=KUKdZzd336w24zPRcqwRatj7IVmd1Us0a_VuzASluIo,7565
 heavyball/psgd_kron.py,sha256=PtTe6eR547Y-4CvgjpchgkQsr_kWr4AN-uY9L_JO_C8,6088
 heavyball/pure_psgd.py,sha256=344NdVNHwUFX3fU2R1S_Xh9SXAML3E4ryHr7xfMh9Cc,5076
 heavyball/schedule_free_palm_foreach_soap.py,sha256=KTQY37MZH7YnOSTLKY8uVySUXxWXbFVUA1QXN3iv8Ds,7244
-heavyball/utils.py,sha256=AxhcHzbFAvhTgTFyIcdxs9TJkH4AgVEaNeBRjOLzoBM,40095
-heavyball-0.24.4.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
-heavyball-0.24.4.dist-info/METADATA,sha256=oksy8cvcHSdMEs9Mzv7WDAVkkkUPcpA1uYNlUgZM_bk,11926
-heavyball-0.24.4.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-heavyball-0.24.4.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
-heavyball-0.24.4.dist-info/RECORD,,
+heavyball/utils.py,sha256=_KvCCCnsu_l4I_OhiRr4noAiwUvzctN05JAuYPkrxXY,41191
+heavyball-0.25.1.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
+heavyball-0.25.1.dist-info/METADATA,sha256=WWR7dX_i7dcF-73-VJ42qcRFwZRL3unOSEwO4EM96e0,11926
+heavyball-0.25.1.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+heavyball-0.25.1.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
+heavyball-0.25.1.dist-info/RECORD,,