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

heavyball/pure_psgd.py

@@ -1,105 +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
-"""
-
-import torch
-from heavyball.utils import identity
-
-from .utils import update_param_, warmup, psgd_precond_grad, init_Q_exprs, PSGDBase, \
-    line_to_triu, triu_to_line, promote
-
-
-class ForeachPurePSGD(PSGDBase):
-    """
-    Kronecker Factorized PSGD WITHOUT Momentum
-
-    Args:
-        params (iterable): Iterable of parameters to optimize or dicts defining
-            parameter groups.
-        lr (float): Learning rate.
-        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, 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: callable = None, store_triu_as_line: bool = True, foreach: bool = True,
-                 q_dtype='float32', stochastic_schedule: bool = True, 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 <= weight_decay:
-            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
-
-        if clip_fn is None:
-            clip_fn = identity
-
-        assert not mars, "MARS is not supported in this optimizer"
-
-        defaults = dict(lr=lr, 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, 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)
-        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']
-        store_triu_as_line = group['store_triu_as_line']
-        q_dtype = getattr(torch, group['q_dtype'])
-
-        vals = []
-
-        for p, g in self.split_p_and_g_in_group(group, should_promote=False, beta1=0.0):
-            state = self.state_(p)
-
-            if 'Q' not in state:
-                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["Q"]))
-
-        if not vals:
-            return
-
-        p_list, grad_list, Q_list = zip(*vals)
-        del vals
-
-        group["step"] += 1
-
-        Q_list = list(Q_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)
-            q = line_to_triu(q_orig) if store_triu_as_line else q_orig
-
-            if group:
-                q32 = [promote(q_) for q_ in q]
-                self.do_update(group, [p], [g], [q32], precond_lr, [q_orig], store_triu_as_line)
-            psgd_precond_grad(True, self.state_(p)["exprs"][-1], g, *q)
-            update_param_([p], self.clip_fn([g]), lr, weight_decay, caution=group['caution'], grad=[g])

heavyball/schedule_free_palm_foreach_soap.py

@@ -1,136 +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, copy_stochastic_list_, promote, decorator_knowngood, \
-    mars_correction
-
-
-@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 SFPaLMForeachSOAP(ScheduleFree):
-    """
-    SFPaLMForeachSOAP
-
-    Sources:
-        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, 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,
-                    'gradient_clip_val': gradient_clip_val, 'weight_sum': 0, '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']
-        mars = group['mars']
-
-        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)
-
-        adaptive_gradient_clipping_(p_list, grad, group["gradient_clip_val"], eps=group["eps"])
-
-        vals = []
-
-        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).float()
-                state["exp_avg_sq"] = torch.zeros_like(g, dtype=torch.float32, memory_format=torch.preserve_format)
-                if mars:
-                    state['mars_prev_grad'] = g.clone()
-                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)
-
-        beta2 = 1 - max(step, 1) ** -group['beta2_scale']
-        new_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_(new_debiased2)
-        _compilable_exp_avg_sq_(exp_avg_sq, grad_projected, old_debiased_tensor, group["eps"])
-
-        update_precond = group['step'] > 0 and group['step'] % group['precondition_frequency'] == 0
-
-        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
-            # CANT DO /= HERE AS EXP_AVG MAY POINT TO THE BUFFER
-            set_(gp, project(gp, state['Q'], back=True))
-
-            update_preconditioner(g, state, max_precond_dim, precondition_1d, 1 - new_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-0.25.1.dist-info/RECORD

@@ -1,28 +0,0 @@
-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
-heavyball/foreach_adamw.py,sha256=K4xTes4drylAqaqWky8O_Bg_mmbAmcHZ5DEBs5vMD-s,2860
-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=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=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=_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,,