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

heavyball/utils.py

@@ -44,10 +44,10 @@ def _compilable_schedule_free_(p, z, ckp1, grad, lr, beta1):
     z32 = z.float()
     p32.lerp_(end=z32, weight=1 - ckp1)
     p32.add_(grad, alpha=lr * (beta1 * (1 - ckp1) - 1))
-    _compilable_copy_stochastic_(p, p32)
+    _guarded_copy_stochastic(p, p32)
 
     z32.add_(grad, alpha=-lr)
-    _compilable_copy_stochastic_(z, z32)
+    _guarded_copy_stochastic(z, z32)
 
 
 def schedule_free_(lr: float, weight_lr_power: float, weight_sum: float, beta1: float, parameters: List[torch.Tensor],
@@ -157,11 +157,11 @@ def adaptive_gradient_clipping_(parameters: List[torch.Tensor], gradients: List[
 
 
 def set_(dst: torch.Tensor, src: torch.Tensor):
-    if src.data_ptr() == dst.data_ptr():
+    if not torch.compiler.is_compiling() and src.data_ptr() == dst.data_ptr():
         return
     if src.shape != dst.shape:
         src = src.reshape_as(dst)
-    if src.is_contiguous() and dst.is_contiguous() and src.dtype == dst.dtype:
+    if not torch.compiler.is_compiling() and src.is_contiguous() and dst.is_contiguous() and src.dtype == dst.dtype:
         dst.set_(src)
     else:
         dst.copy_(src)
@@ -486,6 +486,12 @@ 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)
 def _compilable_copy_stochastic_(target: torch.Tensor, source: torch.Tensor):
     """Taken as-is from https://github.com/pytorch/pytorch/issues/120376#issuecomment-1974828905"""
@@ -505,10 +511,7 @@ 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():
         return
-    if target.dtype != torch.bfloat16 or source.dtype not in (torch.float16, torch.float32, torch.float64):
-        set_(target, source)
-        return
-    _compilable_copy_stochastic_(target, source)
+    _guarded_copy_stochastic(target, source)
 
 
 @torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=True)
@@ -521,7 +524,7 @@ def _compilable_update_one_(p, u, decay, add_fn, lr):
         p32.add_(u32, alpha=lr)
     else:
         add_fn(p32, u32, lr)
-    _compilable_copy_stochastic_(p, p32)
+    _guarded_copy_stochastic(p, p32)
 
 
 def update_param_(param: List[torch.Tensor], update: List[torch.Tensor], lr: float, decay: float,
@@ -840,12 +843,13 @@ class PSGDBase(StatefulOptimizer):
             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, original_q if original_q else q_list):
-            if g.dim() > 1:
-                if store_triu_as_line:
-                    psgd_balance_Q([q_ for _, q_ in q])
-                else:
-                    psgd_balance_Q(q)
+        if self.should_update(group, self.balance_probability, "balance_prob"):
+            for g, q in zip(grad_list, original_q if original_q else q_list):
+                if g.dim() > 1:
+                    if store_triu_as_line:
+                        psgd_balance_Q([q_ for _, q_ in q])
+                    else:
+                        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.5.dist-info → heavyball-0.18.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: heavyball
-Version: 0.18.5
+Version: 0.18.7
 Summary: Efficient optimizers
 Home-page: https://github.com/clashluke/heavyball
 Author: Lucas Nestler
@@ -32,7 +32,7 @@ A simple package of efficient optimizers
 The goal is not to thrive for completeness, full maintenance or abstraction, but instead to provide a simple
 largely static alternative to `torch.optim` with more and better optimizers.
 
-Currently (2024-11-20, 0.17.0), the recommended stable optimizer is `PrecondSchedulePaLMSOAP` (see below). The
+Currently (2024-11-21, 0.18.6), the recommended stable optimizer is `PrecondSchedulePaLMSOAP` (see below). The
 recommended experimental optimizer is `DelayedPSGDKron` ([tuning guide](docs/psgd_efficiency.md)).
 
 ## Features

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

@@ -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=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,,
+heavyball/utils.py,sha256=lKIV11qvlHITK7lwaScGbP1ryCmInse9Fe64t0OBmQQ,31408
+heavyball-0.18.7.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
+heavyball-0.18.7.dist-info/METADATA,sha256=KUYpVlwytyMmQBuby0Jf1WaklYdc2GPddiMAqyGKzsM,11810
+heavyball-0.18.7.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+heavyball-0.18.7.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
+heavyball-0.18.7.dist-info/RECORD,,