heavyball 0.18.7__tar.gz → 0.18.8__tar.gz

{heavyball-0.18.7 → heavyball-0.18.8}/PKG-INFO

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

{heavyball-0.18.7 → heavyball-0.18.8}/heavyball/foreach_adamw.py

@@ -1,7 +1,21 @@
 import torch
 import torch.optim
+from heavyball.utils import copy_stochastic_list_
 
-from .utils import warmup, exp_avg_sq_, beta_debias, update_param_, StatefulOptimizer
+from .utils import warmup, exp_avg_sq_, beta_debias, update_param_, StatefulOptimizer, promote
+
+
+@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=True)
+def _compilable_step_(y, grad, exp_avg_sq, exp_avg, beta1, beta2, step, lr, eps, decay):
+    g32, exp_avg32, exp_avg_sq32 = [list(map(promote, x)) for x in [grad, exp_avg, exp_avg_sq]]
+
+    torch._foreach_lerp_(exp_avg32, g32, 1 - beta_debias(beta1, step + 1))
+    denom = list(exp_avg_sq_(exp_avg_sq32, g32, beta_debias(beta2, step + 1), eps))
+
+    update_param_(y, exp_avg32, lr, decay, lambda p, e, l: p.addcdiv_(e, denom.pop(0), value=l))
+
+    copy_stochastic_list_(exp_avg, exp_avg32)
+    copy_stochastic_list_(exp_avg_sq, exp_avg_sq32)
 
 
 class ForeachAdamW(StatefulOptimizer):
@@ -30,13 +44,11 @@ class ForeachAdamW(StatefulOptimizer):
                 self.state_(p)['exp_avg_sq'] = torch.zeros_like(p.data, dtype=torch.float32)
 
         y, grad, exp_avg_sq, exp_avg = zip(
-            *[(p.data, p.grad.float(), self.state_(p)['exp_avg_sq'], self.state_(p)['exp_avg']) for p in active_p])
-
-        # Decay the first and second moment running average coefficient
-        torch._foreach_lerp_(exp_avg, grad, 1 - beta_debias(group['betas'][0], k + 1))
-        denom = list(exp_avg_sq_(exp_avg_sq, grad, beta_debias(group['betas'][1], k + 1), eps))
+            *[(p.data, p.grad, self.state_(p)['exp_avg_sq'], self.state_(p)['exp_avg']) for p in active_p])
 
-        # Normalize grad in-place for memory efficiency
         lr = -warmup(group['lr'], k + 1, group['warmup_steps'])
-        update_param_(y, exp_avg, lr, decay, lambda p, e, l: p.addcdiv_(e, denom.pop(0), value=l))
+        lr = torch.empty((), dtype=torch.float32, device=y[0].device).fill_(lr)
+        step = torch.empty((), dtype=torch.int32, device=y[0].device).fill_(k)
+        _compilable_step_(y, grad, exp_avg_sq, exp_avg, group['betas'][0], group['betas'][1], step, lr, eps, decay)
+
         group['k'] = k + 1

{heavyball-0.18.7 → heavyball-0.18.8}/heavyball/foreach_adopt.py

@@ -1,7 +1,27 @@
 import torch
 import torch.optim
+from heavyball.utils import copy_stochastic_list_
 
-from .utils import warmup, beta_debias, update_param_, StatefulOptimizer
+from .utils import warmup, beta_debias, update_param_, StatefulOptimizer, promote
+
+
+@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=True)
+def _compilable_step_(y, grad, exp_avg_sq, exp_avg, beta1, beta2, step, lr, eps, decay):
+    g32, exp_avg32, exp_avg_sq32 = [list(map(promote, x)) for x in [grad, exp_avg, exp_avg_sq]]
+    update_param_(y, exp_avg, lr, decay)
+
+    beta1 = beta_debias(beta1, step)
+    denom = torch._foreach_sqrt(exp_avg_sq32)
+    torch._foreach_maximum_(denom, eps)
+    torch._foreach_mul_(exp_avg32, beta1)
+    [ea32.addcdiv_(g, d, value=1 - beta1) for ea32, g, d in zip(exp_avg32, g32, denom)]
+
+    beta2 = beta_debias(beta2, step + 1)
+    torch._foreach_mul_(exp_avg_sq32, beta2)
+    [eas32.addcmul_(g, g, value=1 - beta2) for eas32, g in zip(exp_avg_sq32, g32)]
+
+    copy_stochastic_list_(exp_avg, exp_avg32)
+    copy_stochastic_list_(exp_avg_sq, exp_avg_sq32)
 
 
 class ForeachADOPT(StatefulOptimizer):
@@ -31,12 +51,18 @@ class ForeachADOPT(StatefulOptimizer):
                 self.state_(p)['exp_avg_sq'] = torch.zeros_like(p.data, dtype=torch.float32)
 
         y, grad, exp_avg_sq, exp_avg = zip(
-            *[(p.data, p.grad.float(), self.state_(p)['exp_avg_sq'], self.state_(p)['exp_avg']) for p in active_p])
+            *[(p.data, p.grad, self.state_(p)['exp_avg_sq'], self.state_(p)['exp_avg']) for p in active_p])
+
+        group['k'] = k + 1
 
         if k > 1:
             lr = -warmup(group['lr'], k - 1, group['warmup_steps'])
+            lr = torch.empty((), dtype=torch.float32, device=y[0].device).fill_(lr)
+            k = torch.empty((), dtype=torch.int32, device=y[0].device).fill_(k)
+            _compilable_step_(y, grad, exp_avg_sq, exp_avg, group['betas'][0], group['betas'][1], k, lr, eps, decay)
+            return
 
-            update_param_(y, exp_avg, lr, decay)
+        grad = [promote(g) for g in grad]
         if k > 0:
             beta1 = beta_debias(group['betas'][0], k)
             denom = torch._foreach_sqrt(exp_avg_sq)
@@ -48,5 +74,3 @@ class ForeachADOPT(StatefulOptimizer):
         torch._foreach_mul_(exp_avg_sq, beta2)
         torch._foreach_addcmul_(exp_avg_sq, grad, grad, value=1 - beta2)
         del grad
-
-        group['k'] = k + 1

{heavyball-0.18.7 → heavyball-0.18.8}/heavyball/foreach_laprop.py

@@ -1,7 +1,20 @@
 import torch
 import torch.optim
 
-from .utils import warmup, exp_avg_sq_, beta_debias, update_param_, StatefulOptimizer
+from .utils import warmup, exp_avg_sq_, beta_debias, update_param_, StatefulOptimizer, promote
+
+
+@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=True)
+def _compilable_step_(y, grad, exp_avg_sq, exp_avg, beta1, beta2, step, lr, eps, decay):
+    g32, exp_avg32, exp_avg_sq32 = [list(map(promote, x)) for x in [grad, exp_avg, exp_avg_sq]]
+
+    denom = exp_avg_sq_(exp_avg_sq32, g32, beta_debias(beta2, step), eps)
+
+    beta1 = beta_debias(beta1, step)
+    torch._foreach_mul_(exp_avg32, beta1)
+    [ea32.addcdiv_(g, d, value=1 - beta1) for ea32, g, d in zip(exp_avg32, g32, denom)]
+
+    update_param_(y, exp_avg32, lr, decay)
 
 
 class ForeachLaProp(StatefulOptimizer):
@@ -31,17 +44,13 @@ class ForeachLaProp(StatefulOptimizer):
                 self.state_(p)['exp_avg_sq'] = torch.zeros_like(p.data, dtype=torch.float32)
 
         y, grad, exp_avg_sq, exp_avg = zip(
-            *[(p.data, p.grad.float(), self.state_(p)['exp_avg_sq'], self.state_(p)['exp_avg']) for p in active_p])
-
-        # Decay the first and second moment running average coefficient
-        denom = exp_avg_sq_(exp_avg_sq, grad, beta_debias(group['betas'][1], k + 1), eps)
-        beta1 = beta_debias(group['betas'][0], k + 1)
-        torch._foreach_mul_(exp_avg, beta1)
-        torch._foreach_addcdiv_(exp_avg, grad, denom, 1 - beta1)
-        del grad
+            *[(p.data, p.grad, self.state_(p)['exp_avg_sq'], self.state_(p)['exp_avg'])  #
+              for p in active_p])
 
-        # Normalize grad in-place for memory efficiency
         lr = -warmup(group['lr'], k + 1, group['warmup_steps'])
-        update_param_(y, exp_avg, lr, decay)
+        lr = torch.empty((), dtype=torch.float32, device=y[0].device).fill_(lr)
+        step = torch.empty((), dtype=torch.int32, device=y[0].device).fill_(k + 1)
+
+        _compilable_step_(y, grad, exp_avg_sq, exp_avg, group['betas'][0], group['betas'][1], step, lr, eps, decay)
 
         group['k'] = k + 1

heavyball-0.18.8/heavyball/foreach_sfadamw.py

@@ -0,0 +1,60 @@
+import torch
+import torch.optim
+from heavyball.utils import get_ckp1
+
+from .utils import warmup, ScheduleFree, exp_avg_sq_, beta_debias, promote, _compilable_schedule_free_
+
+
+@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=True)
+def _compilable_step_(y, grad, exp_avg_sq, z, beta1, beta2, step, ckp1, eps, decay, lr):
+    old_debiased2 = beta_debias(beta2, step)
+
+    g32 = [promote(g_) for g_ in grad]
+    exp_avg_sq32 = [promote(e_) for e_ in exp_avg_sq]
+
+    denom = exp_avg_sq_(exp_avg_sq32, g32, old_debiased2, eps)
+    torch._foreach_div_(g32, denom)
+    if decay != 0:
+        torch._foreach_add_(g32, y, alpha=decay)
+    for p, z_, g in zip(y, z, g32):
+        _compilable_schedule_free_(p, z_, ckp1, g, lr, beta1)
+
+
+class ForeachSFAdamW(ScheduleFree):
+    def __init__(self, params, lr=0.0025, betas=(0.9, 0.99), eps=1e-8, weight_decay=0, warmup_steps=0, r=0.0,
+                 weight_lr_power=2.0, foreach: bool = True):
+
+        defaults = dict(lr=lr, betas=betas, eps=eps, r=r, k=0, warmup_steps=warmup_steps, train_mode=True,
+                        weight_sum=0.0, lr_max=-1.0, weight_lr_power=weight_lr_power, weight_decay=weight_decay,
+                        foreach=foreach)
+        super().__init__(params, defaults, foreach)
+
+    def _step(self, group):
+        eps = group['eps']
+        decay = group['weight_decay']
+        k = group['k']
+
+        if not group['train_mode']:
+            raise Exception("Not in train mode!")
+
+        active_p = [p for p in group['params'] if p.grad is not None]
+
+        if not active_p:
+            return
+
+        for p in active_p:
+            if 'z' not in self.state_(p):
+                self.state_(p)['z'] = torch.clone(p.data)
+                self.state_(p)['exp_avg_sq'] = torch.zeros_like(p.data, dtype=torch.float32)
+
+        y, grad, exp_avg_sq, z = zip(*[(p.data, p.grad, self.state_(p)['exp_avg_sq'], self.state_(p)['z'])  #
+                                       for p in active_p])
+
+        lr = warmup(group['lr'], k + 1, group['warmup_steps'])
+        ckp1, group['weight_sum'] = get_ckp1(lr, group['weight_lr_power'], group['weight_sum'], group['r'], k + 1)
+
+        step = torch.empty((), dtype=torch.int32, device=y[0].device).fill_(k + 1)
+        ckp1 = torch.empty((), dtype=torch.float32, device=y[0].device).fill_(ckp1)
+        lr = torch.empty((), dtype=torch.float32, device=y[0].device).fill_(lr)
+        _compilable_step_(y, grad, exp_avg_sq, z, group['betas'][0], group['betas'][1], step, ckp1, eps, decay, lr)
+        group['k'] = k + 1

heavyball-0.18.8/heavyball/palm_foreach_sfadamw.py

@@ -0,0 +1,64 @@
+import torch
+import torch.optim
+
+from .utils import schedule_free_, warmup, ScheduleFree, exp_avg_sq_, beta_debias, get_ckp1, promote, _compilable_schedule_free_
+
+
+@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=True)
+def _compilable_step_(y, grad, exp_avg_sq, z, beta1, beta2, step, ckp1, eps, decay, lr):
+    old_debiased2 = beta_debias(beta2, step)
+
+    g32 = [promote(g_) for g_ in grad]
+    exp_avg_sq32 = [promote(e_) for e_ in exp_avg_sq]
+
+    denom = exp_avg_sq_(exp_avg_sq32, g32, old_debiased2, eps)
+    torch._foreach_div_(g32, denom)
+    if decay != 0:
+        torch._foreach_add_(g32, y, alpha=decay)
+    for p, z_, g in zip(y, z, g32):
+        _compilable_schedule_free_(p, z_, ckp1, g, lr, beta1)
+
+
+class PaLMForeachSFAdamW(ScheduleFree):
+    def __init__(self, params, lr=0.0025, beta=0.9, betas=(None, None), eps=1e-8, weight_decay=0, warmup_steps=0, r=0.0,
+                 weight_lr_power=2.0, beta2_scale: float = 0.8, foreach: bool = True):
+        if betas[0] is not None:
+            beta = betas[0]
+        defaults = dict(lr=lr, beta=beta, eps=eps, r=r, k=0, warmup_steps=warmup_steps, train_mode=True, weight_sum=0.0,
+                        lr_max=-1.0, weight_lr_power=weight_lr_power, weight_decay=weight_decay,
+                        beta2_scale=beta2_scale)
+        super().__init__(params, defaults, foreach)
+
+    def _step(self, group):
+        eps = group['eps']
+        decay = group['weight_decay']
+        k = group['k']
+
+        if not group['train_mode']:
+            raise Exception("Not in train mode!")
+
+        active_p = [p for p in group['params'] if p.grad is not None]
+
+        if not active_p:
+            return
+
+        for p in active_p:
+            if 'z' not in self.state_(p):
+                self.state_(p)['z'] = torch.clone(p.data)
+                self.state_(p)['exp_avg_sq'] = torch.zeros_like(p.data, dtype=torch.float32)
+
+        # Decay the first moment running average coefficient
+        beta2 = 1 - (k + 1) ** -group['beta2_scale']
+
+        y, grad, exp_avg_sq, z = zip(*[(p.data, p.grad, self.state_(p)['exp_avg_sq'], self.state_(p)['z'])  #
+                                       for p in active_p])
+
+        lr = warmup(group['lr'], k + 1, group['warmup_steps'])
+        ckp1, group['weight_sum'] = get_ckp1(lr, group['weight_lr_power'], group['weight_sum'], group['r'], k + 1)
+
+        step = torch.empty((), dtype=torch.int32, device=y[0].device).fill_(k + 1)
+        ckp1 = torch.empty((), dtype=torch.float32, device=y[0].device).fill_(ckp1)
+        beta2 = torch.empty((), dtype=torch.float32, device=y[0].device).fill_(beta2)
+        lr = torch.empty((), dtype=torch.float32, device=y[0].device).fill_(lr)
+        _compilable_step_(y, grad, exp_avg_sq, z, group['beta'], beta2, step, ckp1, eps, decay, lr)
+        group['k'] = k + 1

{heavyball-0.18.7 → heavyball-0.18.8}/heavyball/psgd_kron.py

@@ -104,7 +104,8 @@ class ForeachPSGDKron(PSGDBase):
 
             if should_update:
                 q32 = [promote(q_) for q_ in q]
-                self.do_update(group, [p], [g], [q32], precond_lr, [q_orig], store_triu_as_line)
+                self.do_update(group, [p], [ea if momentum_into_precond_update else g], [q32], precond_lr, [q_orig],
+                               store_triu_as_line)
             set_(g, psgd_precond_grad(q, self.state_(p)["exprs"], ea))
 
         grad_list = self.clip_fn(grad_list)

{heavyball-0.18.7 → heavyball-0.18.8}/heavyball/utils.py

@@ -40,14 +40,25 @@ def warmup(lr: float, step: int, warmup_steps: int):
 
 @torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=True)
 def _compilable_schedule_free_(p, z, ckp1, grad, lr, beta1):
-    p32 = p.float()
-    z32 = z.float()
-    p32.lerp_(end=z32, weight=1 - ckp1)
+    p32 = promote(p)
+    z32 = promote(z)
+    p32.lerp_(end=z32, weight=ckp1)
     p32.add_(grad, alpha=lr * (beta1 * (1 - ckp1) - 1))
-    _guarded_copy_stochastic(p, p32)
+    copy_stochastic_(p, p32)
 
     z32.add_(grad, alpha=-lr)
-    _guarded_copy_stochastic(z, z32)
+    copy_stochastic_(z, z32)
+
+
+def get_ckp1(lr, weight_lr_power, weight_sum, r, step):
+    weight = lr ** weight_lr_power * max(step, 1) ** r
+    weight_sum = weight_sum + weight
+
+    try:
+        ckp1 = weight / weight_sum
+    except ZeroDivisionError:
+        ckp1 = 0
+    return ckp1, weight_sum
 
 
 def schedule_free_(lr: float, weight_lr_power: float, weight_sum: float, beta1: float, parameters: List[torch.Tensor],
@@ -136,7 +147,7 @@ def exp_avg_sq_(state, grad, beta2, eps, out=None):
         return torch.sqrt(state, out=out).clamp_(min=eps)
 
     torch._foreach_mul_(state, beta2)
-    torch._foreach_addcmul_(state, grad, grad, value=1 - beta2)
+    [s.addcmul_(g, g, value=1 - beta2) for s, g in zip(state, grad)]
     denom = torch._foreach_sqrt(state)
     torch._foreach_maximum_(denom, eps)
     return denom
@@ -332,9 +343,9 @@ def compute_ggt(grad, GG, max_precond_dim, precondition_1d, beta):
 
 
 def promote(x):
-    if x in (torch.bfloat16, torch.float16):
+    if isinstance(x, torch.dtype) and x in (torch.bfloat16, torch.float16):
         return torch.float32
-    if hasattr(x, 'dtype') and x.dtype in (torch.bfloat16, torch.float16):
+    if isinstance(x, torch.Tensor) and x.dtype in (torch.bfloat16, torch.float16):
         return x.float()
     return x
 
@@ -486,13 +497,8 @@ def copy_stochastic_list_(target: List[torch.Tensor], source: List[torch.Tensor]
         copy_stochastic_(t, s)
 
 
-def _guarded_copy_stochastic(target: torch.Tensor, source: torch.Tensor):
-    if target.dtype != torch.bfloat16 or source.dtype not in (torch.float16, torch.float32, torch.float64):
-        set_(target, source)
-    _compilable_copy_stochastic_(target, source)
-
-
-@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=True)
+# this can be dynamic for most optimizers - just not for PSGD. So, it's disabled for all
+@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True)
 def _compilable_copy_stochastic_(target: torch.Tensor, source: torch.Tensor):
     """Taken as-is from https://github.com/pytorch/pytorch/issues/120376#issuecomment-1974828905"""
     # create a random 16 bit integer
@@ -509,22 +515,24 @@ def _compilable_copy_stochastic_(target: torch.Tensor, source: torch.Tensor):
 
 
 def copy_stochastic_(target: torch.Tensor, source: torch.Tensor):
-    if target.data_ptr() == source.data_ptr():
+    if not torch.compiler.is_compiling() and target.data_ptr() == source.data_ptr():
         return
-    _guarded_copy_stochastic(target, source)
+    if target.dtype != torch.bfloat16 or source.dtype not in (torch.float16, torch.float32, torch.float64):
+        set_(target, source)
+    _compilable_copy_stochastic_(target, source)
 
 
 @torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=True)
 def _compilable_update_one_(p, u, decay, add_fn, lr):
-    p32 = p.float()
-    u32 = u.view(p.shape).float()
+    p32 = promote(p)
+    u32 = promote(u.view(p.shape))
     if decay > 0:
         p32.mul_(1 - decay * lr)
     if add_fn is None:
         p32.add_(u32, alpha=lr)
     else:
         add_fn(p32, u32, lr)
-    _guarded_copy_stochastic(p, p32)
+    copy_stochastic_(p, p32)
 
 
 def update_param_(param: List[torch.Tensor], update: List[torch.Tensor], lr: float, decay: float,

{heavyball-0.18.7 → heavyball-0.18.8}/heavyball.egg-info/PKG-INFO

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

{heavyball-0.18.7 → heavyball-0.18.8}/setup.py

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

{heavyball-0.18.7 → heavyball-0.18.8}/test/test_bf16_params.py

@@ -22,7 +22,6 @@ def get_memory():
 @pytest.mark.parametrize("size,depth", [(256, 2)])
 def test_foreach(opt, size, depth: int, iterations: int = 128, outer_iterations: int = 3):
     set_torch()
-
     opt = getattr(heavyball, opt)
 
     peaks = []

heavyball-0.18.7/heavyball/foreach_sfadamw.py

@@ -1,54 +0,0 @@
-import torch
-import torch.optim
-
-from .utils import schedule_free_, warmup, ScheduleFree, exp_avg_sq_, beta_debias
-
-
-class ForeachSFAdamW(ScheduleFree):
-    def __init__(self, params, lr=0.0025, betas=(0.9, 0.99), eps=1e-8, weight_decay=0, warmup_steps=0, r=0.0,
-                 weight_lr_power=2.0, foreach: bool = True):
-
-        defaults = dict(lr=lr, betas=betas, eps=eps, r=r, k=0, warmup_steps=warmup_steps, train_mode=True,
-                        weight_sum=0.0, lr_max=-1.0, weight_lr_power=weight_lr_power, weight_decay=weight_decay,
-                        foreach=foreach)
-        super().__init__(params, defaults, foreach)
-
-    def _step(self, group):
-        eps = group['eps']
-        decay = group['weight_decay']
-        k = group['k']
-
-        if not group['train_mode']:
-            raise Exception("Not in train mode!")
-
-        active_p = [p for p in group['params'] if p.grad is not None]
-
-        if not active_p:
-            return
-
-        for p in active_p:
-            if 'z' not in self.state_(p):
-                self.state_(p)['z'] = torch.clone(p.data)
-                self.state_(p)['exp_avg_sq'] = torch.zeros_like(p.data, dtype=torch.float32)
-
-        y, grad, exp_avg_sq, z = zip(
-            *[(p.data, p.grad.float(), self.state_(p)['exp_avg_sq'], self.state_(p)['z']) for p in active_p])
-
-        # Decay the first moment running average coefficient
-        old_debiased = beta_debias(group['betas'][1], k + 1)
-
-        # Decay the first and second moment running average coefficient
-        denom = exp_avg_sq_(exp_avg_sq, grad, old_debiased, eps)
-
-        # Normalize grad in-place for memory efficiency
-        torch._foreach_div_(grad, denom)
-
-        # Weight decay calculated at y
-        if decay != 0:
-            torch._foreach_add_(grad, y, alpha=decay)
-
-        lr = warmup(group['lr'], k + 1, group['warmup_steps'])
-        group['weight_sum'] = schedule_free_(lr, group['weight_lr_power'], group['weight_sum'], group['betas'][0], y, z,
-                                             grad, group['r'], k + 1)
-
-        group['k'] = k + 1

heavyball-0.18.7/heavyball/palm_foreach_sfadamw.py

@@ -1,57 +0,0 @@
-import torch
-import torch.optim
-
-from .utils import schedule_free_, warmup, ScheduleFree, exp_avg_sq_, beta_debias
-
-
-class PaLMForeachSFAdamW(ScheduleFree):
-    def __init__(self, params, lr=0.0025, beta=0.9, betas=(None, None), eps=1e-8, weight_decay=0, warmup_steps=0, r=0.0,
-                 weight_lr_power=2.0, beta2_scale: float = 0.8,
-                 foreach: bool = True):
-        if betas[0] is not None:
-            beta = betas[0]
-        defaults = dict(lr=lr, beta=beta, eps=eps, r=r, k=0, warmup_steps=warmup_steps, train_mode=True, weight_sum=0.0,
-                        lr_max=-1.0, weight_lr_power=weight_lr_power, weight_decay=weight_decay,
-                        beta2_scale=beta2_scale)
-        super().__init__(params, defaults, foreach)
-
-    def _step(self, group):
-        eps = group['eps']
-        decay = group['weight_decay']
-        k = group['k']
-
-        if not group['train_mode']:
-            raise Exception("Not in train mode!")
-
-        active_p = [p for p in group['params'] if p.grad is not None]
-
-        if not active_p:
-            return
-
-        for p in active_p:
-            if 'z' not in self.state_(p):
-                self.state_(p)['z'] = torch.clone(p.data)
-                self.state_(p)['exp_avg_sq'] = torch.zeros_like(p.data, dtype=torch.float32)
-
-        y, grad, exp_avg_sq, z = zip(
-            *[(p.data, p.grad.float(), self.state_(p)['exp_avg_sq'], self.state_(p)['z']) for p in active_p])
-
-        # Decay the first moment running average coefficient
-        beta2 = 1 - (k + 1) ** -group['beta2_scale']
-        old_debiased = beta_debias(beta2, k + 1)
-
-        # Decay the first and second moment running average coefficient
-        denom = exp_avg_sq_(exp_avg_sq, grad, old_debiased, eps)
-
-        # Normalize grad in-place for memory efficiency
-        torch._foreach_div_(grad, denom)
-
-        # Weight decay calculated at y
-        if decay != 0:
-            torch._foreach_add_(grad, y, alpha=decay)
-
-        lr = warmup(group['lr'], k + 1, group['warmup_steps'])
-        group['weight_sum'] = schedule_free_(lr, group['weight_lr_power'], group['weight_sum'], group['beta'], y, z,
-                                             grad, group['r'], k + 1)
-
-        group['k'] = k + 1