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

heavyball/precond_schedule_foreach_soap.py

@@ -1,95 +0,0 @@
-import random
-
-import torch
-
-from .utils import init_preconditioner, update_preconditioner, project, beta_debias, update_param_, \
-    precond_schedule, set_, StatefulOptimizer, exp_avg_
-
-
-class PrecondScheduleForeachSOAP(StatefulOptimizer):
-    """
-    Sources:
-        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):
-            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']
-            exp_avg_projected = project(ea, state['Q'], False)
-
-            # Projecting back the preconditioned (by Adam) exponential moving average of gradients
-            # to the original space
-            # CANT DO /= HERE AS EXP_AVG MAY POINT TO THE BUFFER
-            precond = project(exp_avg_projected / d, 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_foreach_solp.py

@@ -1,95 +0,0 @@
-
-
-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_soap.py

@@ -1,105 +0,0 @@
-import random
-
-import torch
-
-from .utils import init_preconditioner, update_preconditioner, project, beta_debias, exp_avg_, update_param_, \
-    precond_schedule, set_, StatefulOptimizer
-
-
-class PrecondSchedulePaLMForeachSOAP(StatefulOptimizer):
-    """
-    Sources:
-        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):
-            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']
-            exp_avg_projected = project(ea, state['Q'], False)
-
-            # Projecting back the preconditioned (by Adam) exponential moving average of gradients
-            # to the original space
-            # CANT DO /= HERE AS EXP_AVG MAY POINT TO THE BUFFER
-            exp_avg_projected = exp_avg_projected / d
-            precond = project(exp_avg_projected, 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

@@ -1,103 +0,0 @@
-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/precond_schedule_sfpsoap.py

@@ -1,141 +0,0 @@
-import random
-
-import torch
-
-from .utils import init_preconditioner, update_preconditioner, project, set_, adaptive_gradient_clipping_, exp_avg_sq_, \
-    beta_debias, schedule_free_, warmup, ScheduleFree, precond_schedule, copy_stochastic_list_, \
-    promote, decorator_knowngood
-
-
-@decorator_knowngood
-def _compilable_exp_avg_sq_(exp_avg_sq, grad_projected, old_debiased2, eps):
-    eas32, gp32 = [list(map(promote, x)) for x in (exp_avg_sq, grad_projected)]
-    denom = exp_avg_sq_(eas32, gp32, old_debiased2, eps)
-    torch._foreach_div_(gp32, denom)
-
-    copy_stochastic_list_(exp_avg_sq, eas32)
-    copy_stochastic_list_(grad_projected, gp32)
-
-
-class PrecondScheduleSFPaLMSOAP(ScheduleFree):
-    """
-    SFPaLMForeachSOAP with preconditioner schedules
-
-    Sources:
-        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
-
-        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, beta2_scale: float = 0.8, 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, r=0.0,
-                 weight_lr_power=2.0, gradient_clip_val: float = 0.1, precond_scheduler=(1 / 3, 9), betas=(None, None),
-                 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]
-
-        assert not caution, "Caution is not implemented in ScheduleFree optimizers"
-
-        defaults = {"lr": lr, "beta": beta, "beta2_scale": beta2_scale, "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, 'r': r,
-                    'weight_lr_power': weight_lr_power, 'train_mode': True, 'step': -1, 'weight_sum': 0,
-                    'gradient_clip_val': gradient_clip_val, '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 = []
-        max_precond_dim = group['max_precond_dim']
-        precondition_1d = group['precondition_1d']
-
-        step = group['step'] = group.get("step", 0) + 1
-
-        for p in group["params"]:
-            if p.grad is None:
-                continue
-            grad = p.grad.float()
-            vals.append((p, grad))
-
-        if not vals:
-            return
-
-        p_list, grad = zip(*vals)
-        vals = []
-
-        # adaptive gradient clipping
-        adaptive_gradient_clipping_(p_list, grad, group["gradient_clip_val"], eps=group["eps"])
-
-        for p, g in self.split_p_and_g_in_group(group, beta1=group['beta']):
-            state = self.state_(p)
-
-            if "z" not in state:
-                state["z"] = torch.clone(p.data, 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)
-            z, exp_avg_sq = state["z"], state["exp_avg_sq"]
-            vals.append((p, g, grad_projected, z, exp_avg_sq))
-
-        if not vals:
-            return
-
-        p_list, grad, grad_projected, z, exp_avg_sq = zip(*vals)
-        del vals
-
-        beta2 = 1 - max(step, 1) ** -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
-        old_debiased_tensor = torch.empty((), dtype=torch.float32, device=p_list[0].device).fill_(old_debiased2)
-        _compilable_exp_avg_sq_(exp_avg_sq, grad_projected, old_debiased_tensor, group["eps"])
-
-        update_precond = precond_schedule(step, group['precond_scheduler'], self.rng)
-
-        for p, g, gp in zip(p_list, grad, grad_projected):
-            state = self.state_(p)
-            # Projecting back the preconditioned (by Adam) exponential moving average of gradients
-            # to the original space
-            set_(gp, project(gp, state['Q'], back=True))
-
-            update_preconditioner(g, state, max_precond_dim, precondition_1d, old_debiased2, update_precond)
-
-        # Weight decay calculated at y
-        if group["weight_decay"] > 0:
-            torch._foreach_add_(grad, p_list, alpha=group["weight_decay"])
-
-        lr = warmup(group['lr'], step, group['warmup_steps'])
-        group['weight_sum'] = schedule_free_(lr, group['weight_lr_power'], group['weight_sum'], group['beta'], p_list,
-                                             z, grad_projected, group['r'], step)

heavyball/psgd_kron.py

@@ -1,120 +0,0 @@
-"""
-Originally from Evan Walters and Omead Pooladzandi, 2024
-Modified under Creative Commons Attribution 4.0 International
-Source available at https://github.com/evanatyourservice/kron_torch/blob/97a2b5ee8a1a4c29e4780bbf6c521e545189eff9/kron_torch/kron.py
-"""
-
-from typing import Optional
-
-import torch
-
-from .utils import update_param_, warmup, psgd_precond_grad, init_Q_exprs, trust_region_clip_, PSGDBase, \
-    line_to_triu, triu_to_line, promote, stochastic_lerp_, beta_debias
-
-
-class ForeachPSGDKron(PSGDBase):
-    """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.
-        b1 (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
-            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
-            update instead of raw gradients.
-    """
-
-    def __init__(self, params, lr=0.001, beta=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, warmup_steps: int = 1, merge_dims: bool = False,
-                 split: bool = False, clip_fn: Optional[callable] = None, store_triu_as_line: bool = True,
-                 foreach: bool = True, q_dtype='float32', stochastic_schedule: bool = True,
-                 storage_dtype: str = 'float32', mars: bool = False, caution: bool = False, mars_gamma: float = 0.0025,
-                 #
-                 # expert parameters
-                 precond_init_scale=1.0, precond_lr=0.1):
-        if not 0.0 <= lr:
-            raise ValueError(f"Invalid learning rate: {lr}")
-        if not 0.0 <= beta < 1.0:
-            raise ValueError(f"Invalid beta parameter: {beta}")
-        if not 0.0 <= weight_decay:
-            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
-
-        if clip_fn is None:
-            clip_fn = lambda x: trust_region_clip_(x, 0.9, 1.5)
-
-        defaults = dict(lr=lr, beta=beta, weight_decay=weight_decay, max_size_triangular=max_size_triangular,
-                        min_ndim_triangular=min_ndim_triangular, memory_save_mode=memory_save_mode,
-                        momentum_into_precond_update=momentum_into_precond_update, precond_lr=precond_lr,
-                        precond_init_scale=precond_init_scale, step=0, warmup_steps=warmup_steps, merge_dims=merge_dims,
-                        split=split, store_triu_as_line=store_triu_as_line, q_dtype=q_dtype,
-                        storage_dtype=storage_dtype,
-                        mars=mars, caution=caution, mars_gamma=mars_gamma)
-        super().__init__(params, defaults, foreach, stochastic_schedule, clip_fn, preconditioner_update_probability)
-
-    def _step(self, group):
-        should_update = self.should_update(group)
-        momentum_into_precond_update = group.get("momentum_into_precond_update", True)
-        precond_init_scale = group['precond_init_scale']
-        max_size_triangular = group['max_size_triangular']
-        min_ndim_triangular = group['min_ndim_triangular']
-        memory_save_mode = group['memory_save_mode']
-        precond_lr = group['precond_lr']
-        weight_decay = group['weight_decay']
-        lr = group['lr']
-        beta = group['beta']
-        store_triu_as_line = group['store_triu_as_line']
-        q_dtype = getattr(torch, group['q_dtype'])
-        storage_dtype = getattr(torch, group['storage_dtype'])
-
-        vals = []
-
-        for p, g in self.split_p_and_g_in_group(group, should_promote=False, beta1=beta):
-            state = self.state_(p)
-
-            if 'Q' not in state:
-                state["exp_avg"] = torch.zeros_like(g, dtype=storage_dtype, memory_format=torch.preserve_format)
-                Q, state["exprs"] = init_Q_exprs(p, precond_init_scale, max_size_triangular, min_ndim_triangular,
-                                                 memory_save_mode, dtype=q_dtype)
-                state['Q'] = triu_to_line(Q) if store_triu_as_line else Q
-
-            vals.append((p, g, state["exp_avg"], state["Q"]))
-
-        if not vals:
-            return
-
-        p_list, grad_list, exp_avg_list, Q_list = zip(*vals)
-        del vals
-
-        group["step"] += 1
-
-        beta = beta_debias(beta, group["step"])
-        beta = torch.empty((), dtype=torch.float32, device=grad_list[0].device).fill_(1 - beta)
-        stochastic_lerp_(exp_avg_list, grad_list, 1 - beta)
-
-        grad_list, Q_list, exp_avg_list = list(grad_list), list(Q_list), list(exp_avg_list)
-
-        lr = -warmup(lr, group['step'], group['warmup_steps'])
-
-        for i, (p, g) in enumerate(zip(p_list, grad_list)):
-            q_orig = Q_list.pop(0)
-            ea = exp_avg_list.pop(0)
-            q = line_to_triu(q_orig) if store_triu_as_line else q_orig
-
-            if should_update:
-                q32 = [promote(q_) for q_ in q]
-                self.do_update(group, [p], [ea if momentum_into_precond_update else g], [q32], precond_lr, [q_orig],
-                               store_triu_as_line)
-            g = psgd_precond_grad(False, self.state_(p)["exprs"][-1], ea, *q)
-            update_param_([p], self.clip_fn([g]), lr, weight_decay, caution=group['caution'], grad=[g])