heavyball 0.18.3__py3-none-any.whl → 0.18.5__py3-none-any.whl

heavyball/foreach_adamw.py

@@ -34,9 +34,9 @@ class ForeachAdamW(StatefulOptimizer):
 
         # Decay the first and second moment running average coefficient
         torch._foreach_lerp_(exp_avg, grad, 1 - beta_debias(group['betas'][0], k + 1))
-        denom = exp_avg_sq_(exp_avg_sq, grad, beta_debias(group['betas'][1], k + 1), eps)
+        denom = list(exp_avg_sq_(exp_avg_sq, grad, beta_debias(group['betas'][1], k + 1), eps))
 
         # 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: torch._foreach_addcdiv_(p, e, denom, l))
+        update_param_(y, exp_avg, lr, decay, lambda p, e, l: p.addcdiv_(e, denom.pop(0), value=l))
         group['k'] = k + 1

heavyball/utils.py

@@ -38,6 +38,18 @@ def warmup(lr: float, step: int, warmup_steps: int):
     return lr * step / warmup_steps
 
 
+@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.add_(grad, alpha=lr * (beta1 * (1 - ckp1) - 1))
+    _compilable_copy_stochastic_(p, p32)
+
+    z32.add_(grad, alpha=-lr)
+    _compilable_copy_stochastic_(z, z32)
+
+
 def schedule_free_(lr: float, weight_lr_power: float, weight_sum: float, beta1: float, parameters: List[torch.Tensor],
                    z: List[torch.Tensor], grad: list[torch.Tensor], r: float = 0.0, step: int = 0):
     weight = lr ** weight_lr_power * max(step, 1) ** r
@@ -50,15 +62,10 @@ def schedule_free_(lr: float, weight_lr_power: float, weight_sum: float, beta1:
 
     # These operations update y in-place,
     # without computing x explicitly.
-    p32 = [promote(p) for p in parameters]
-    z32 = [promote(z_) for z_ in z]
-    torch._foreach_lerp_(p32, z32, weight=ckp1)
-    torch._foreach_add_(p32, grad, alpha=lr * (beta1 * (1 - ckp1) - 1))
-    copy_stochastic_list_(parameters, p32)
-
-    # z step
-    torch._foreach_sub_(z, grad, alpha=lr)
-    copy_stochastic_list_(z, z32)
+    lr_tensor = torch.empty((), dtype=torch.float32, device=parameters[0].device).fill_(lr)
+    ckp1_tensor = torch.empty((), dtype=torch.float32, device=parameters[0].device).fill_(ckp1)
+    for p, z_, g in zip(parameters, z, grad):
+        _compilable_schedule_free_(p, z_, ckp1_tensor, g, lr_tensor, beta1)
     return weight_sum
 
 
@@ -479,7 +486,7 @@ def copy_stochastic_list_(target: List[torch.Tensor], source: List[torch.Tensor]
         copy_stochastic_(t, s)
 
 
-@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=False)
+@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=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
@@ -504,17 +511,24 @@ def copy_stochastic_(target: torch.Tensor, source: torch.Tensor):
     _compilable_copy_stochastic_(target, source)
 
 
-def update_param_(param: List[torch.Tensor], update: List[torch.Tensor], lr: float, decay: float,
-                  add_fn: callable = None):
-    param32 = [promote(p) for p in param]
-    update32 = [promote(u.view(p.shape)) for u, p in zip(update, param)]
+@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()
     if decay > 0:
-        torch._foreach_mul_(param32, 1 - decay * lr)
+        p32.mul_(1 - decay * lr)
     if add_fn is None:
-        torch._foreach_add_(param32, update32, alpha=lr)
+        p32.add_(u32, alpha=lr)
     else:
-        add_fn(param32, update32, lr)
-    copy_stochastic_list_(param, param32)
+        add_fn(p32, u32, lr)
+    _compilable_copy_stochastic_(p, p32)
+
+
+def update_param_(param: List[torch.Tensor], update: List[torch.Tensor], lr: float, decay: float,
+                  add_fn: callable = None):
+    lr_tensor = torch.empty((), dtype=torch.float32, device=param[0].device).fill_(lr)
+    for p, u in zip(param, update):
+        _compilable_update_one_(p, u, decay, add_fn, lr_tensor)
 
 
 def precond_schedule(step, precond_scheduler, rng):
@@ -815,19 +829,23 @@ class PSGDBase(StatefulOptimizer):
 
     def do_update(self, group, p_list, grad_list, q_list, precond_lr, original_q: Optional[List] = None,
                   store_triu_as_line=False):
-        for i, (p, grad, Q) in enumerate(zip(p_list, grad_list, q_list)):
+        if original_q:
+            if store_triu_as_line:
+                update_fn = update_triu_
+            else:
+                update_fn = copy_stochastic_list_
+        else:
+            update_fn = lambda x, y: None
+        for i, (p, grad, Q, oq) in enumerate(zip(p_list, grad_list, q_list, original_q)):
             psgd_update_precond(Q, self.state_(p)["exprs"], torch.randn_like(grad), grad, precond_lr, self._tiny)
+            update_fn(oq, Q)
 
-        for g, q in zip(grad_list, q_list):
+        for g, q in zip(grad_list, original_q if original_q else q_list):
             if g.dim() > 1:
-                psgd_balance_Q(q)
-
-        if original_q:
-            for q in q_list:
                 if store_triu_as_line:
-                    update_triu_(original_q[i], Q)
+                    psgd_balance_Q([q_ for _, q_ in q])
                 else:
-                    copy_stochastic_list_(original_q[i], Q)
+                    psgd_balance_Q(q)
 
 
 def precond_update_prob_schedule(max_prob=1.0, min_prob=0.03, decay=0.001, flat_start=250):

{heavyball-0.18.3.dist-info → heavyball-0.18.5.dist-info}/METADATA

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

{heavyball-0.18.3.dist-info → heavyball-0.18.5.dist-info}/RECORD

@@ -2,7 +2,7 @@ heavyball/__init__.py,sha256=iqP428JWwwx-XDOZ0nUdbCkOLEyfoqVyWZLQLAcwxaw,2214
 heavyball/cached_delayed_psgd_kron.py,sha256=PQAER6UgVh5l87DGRZrJ8CVP9UhyCG5wJD9rPLnj_G8,6460
 heavyball/cached_psgd_kron.py,sha256=GaeneBp0irksCSBIrJY4D_0hCpZ-uSRPMhqVX_a-og8,6417
 heavyball/delayed_psgd.py,sha256=fhBWFLTSl1S2gHWCeYak-STaXRwpC56sWZGLFMKFEJM,5589
-heavyball/foreach_adamw.py,sha256=h_ar0ZZRM0q_wxuEkxEOEYe-2p-mB4OMgAHivrUnPl8,1777
+heavyball/foreach_adamw.py,sha256=CTg7rfUmlTSjihD5KY9xP0sT2dUKZyZ4-2V42Vlr28U,1780
 heavyball/foreach_adopt.py,sha256=ogOw2JjwEQNj7AKlweAphQFdMJ_GcMDm-RyDvEzugoc,1911
 heavyball/foreach_laprop.py,sha256=yGVmGqWiSw8Y2Xj70ndkR8ZMygakTB4_iRwV02Svkqg,1816
 heavyball/foreach_sfadamw.py,sha256=15-n6-lx4PAHYsKYmXbugxsR5MnqaPYy2vUudPRiitg,2087
@@ -16,9 +16,9 @@ heavyball/precond_schedule_sfpsoap.py,sha256=vq7jd302refKPa_9X2lkOTOtCCcTBVByPdo
 heavyball/psgd_kron.py,sha256=u46dorOUXx-do1IYeno2wj-6l1zYKMQQC-N2Zr2PzLI,5476
 heavyball/pure_psgd.py,sha256=iUy7mMKWxwNiVUMYrQ7SBnreu3t_XSbnhTW3a1yw4m0,4835
 heavyball/schedule_free_palm_foreach_soap.py,sha256=zkcikH5wWbzq4kOrmBjilvY3iWzuUddcv2HNEPKr3MI,6366
-heavyball/utils.py,sha256=qs_WfzJdS-3XyEuw-m6mWMEeR95r7bGFVC8wWCHtD48,30365
-heavyball-0.18.3.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
-heavyball-0.18.3.dist-info/METADATA,sha256=Cx8LM2g3BFOk8WJH3B8ve8kQ7HghMCIRLggdJp37x4g,11810
-heavyball-0.18.3.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-heavyball-0.18.3.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
-heavyball-0.18.3.dist-info/RECORD,,
+heavyball/utils.py,sha256=2VBEQhtQ4mwsD99JMu7iWbiYPkutspjG3hGwCbIHZ9U,31134
+heavyball-0.18.5.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
+heavyball-0.18.5.dist-info/METADATA,sha256=Zcc87BhCxDTX7bjJ3pGG7VdIRmpZuYLwmWBKDiLc3AU,11810
+heavyball-0.18.5.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+heavyball-0.18.5.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
+heavyball-0.18.5.dist-info/RECORD,,