heavyball 0.18.3__tar.gz → 0.18.5__tar.gz

{heavyball-0.18.3 → heavyball-0.18.5}/PKG-INFO

@@ -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 → heavyball-0.18.5}/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-0.18.3 → heavyball-0.18.5}/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 → heavyball-0.18.5}/heavyball.egg-info/PKG-INFO

@@ -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 → heavyball-0.18.5}/heavyball.egg-info/SOURCES.txt

@@ -25,6 +25,7 @@ heavyball.egg-info/SOURCES.txt
 heavyball.egg-info/dependency_links.txt
 heavyball.egg-info/requires.txt
 heavyball.egg-info/top_level.txt
+test/test_bf16_params.py
 test/test_bf16_q.py
 test/test_closure.py
 test/test_foreach.py

{heavyball-0.18.3 → heavyball-0.18.5}/setup.py

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

heavyball-0.18.5/test/test_bf16_params.py

@@ -0,0 +1,52 @@
+import pytest
+import torch
+from torch import nn
+from torch._dynamo import config
+
+import heavyball
+import heavyball.utils
+from benchmark.utils import get_optim
+from heavyball.utils import clean, set_torch
+
+config.cache_size_limit = 128
+
+def get_memory():
+    clean()
+    torch.cuda.synchronize()
+    clean()
+    torch.cuda.synchronize()
+    return torch.cuda.memory_allocated()
+
+
+@pytest.mark.parametrize("opt", heavyball.__all__)
+@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 = []
+    losses = []
+
+    for dtype in [torch.float32, torch.bfloat16]:
+        torch.manual_seed(0x2131290)
+        peaks.append([])
+        losses.append([])
+
+        for i in range(outer_iterations):
+            model = nn.Sequential(*[nn.Linear(size, size) for _ in range(depth)]).cuda().to(dtype)
+            o = get_optim(opt, model.parameters(), lr=1e-3)
+
+            for _ in range(iterations):
+                loss = model(torch.randn((1024, size), device='cuda', dtype=dtype)).square().mean()
+                loss.backward()
+                o.step()
+                o.zero_grad()
+                losses[-1].append(loss.detach())
+
+            del model, o
+            clean()
+
+    for i, (l0, l1) in enumerate(zip(*losses)):
+        print(i, l0.item(), l1.item())
+        assert torch.allclose(l0.float(), l1.float(), rtol=0.1)

{heavyball-0.18.3 → heavyball-0.18.5}/test/test_bf16_q.py

@@ -1,10 +1,14 @@
-import heavyball
-import heavyball.utils
 import pytest
 import torch
+from torch import nn
+from torch._dynamo import config
+
+import heavyball
+import heavyball.utils
 from benchmark.utils import get_optim
 from heavyball.utils import clean, set_torch, PSGDBase
-from torch import nn
+
+config.cache_size_limit = 128
 
 
 def get_memory():
@@ -37,7 +41,7 @@ def test_foreach(opt, size, depth: int, iterations: int = 128, outer_iterations:
             o = get_optim(opt, model.parameters(), lr=1e-3, q_dtype=q_dtype)
 
             for _ in range(iterations):
-                loss = model(torch.randn((1024, size)).cuda()).square().mean()
+                loss = model(torch.randn((1024, size), device='cuda')).square().mean()
                 loss.backward()
                 o.step()
                 o.zero_grad()
@@ -46,7 +50,6 @@ def test_foreach(opt, size, depth: int, iterations: int = 128, outer_iterations:
             del model, o
             clean()
 
-
     for i, (l0, l1) in enumerate(zip(*losses)):
         print(i, l0.item(), l1.item())
         assert torch.allclose(l0, l1, rtol=0.1)

{heavyball-0.18.3 → heavyball-0.18.5}/test/test_foreach.py

@@ -45,7 +45,7 @@ def test_foreach(opt, size, depth: int, iterations: int = 5, outer_iterations: i
             clean()
 
             for _ in range(iterations):
-                loss = model(torch.randn((1, size)).cuda()).sum()
+                loss = model(torch.randn((1, size), device='cuda')).sum()
                 loss.backward()
                 o.step()
                 o.zero_grad()

{heavyball-0.18.3 → heavyball-0.18.5}/test/test_memory.py

@@ -48,7 +48,7 @@ def test_memory(opt, method, size, depth: int, iterations: int = 5, outer_iterat
         model_allocated = get_memory()
         o = get_optim(opt, model.parameters(), lr=1e-3)
         for _ in range(iterations):
-            model(torch.randn((1, size)).cuda()).sum().backward()
+            model(torch.randn((1, size), device='cuda')).sum().backward()
             o.step()
 
             opt_allocated = get_memory()

{heavyball-0.18.3 → heavyball-0.18.5}/test/test_stochastic_updates.py

@@ -38,7 +38,7 @@ def test_foreach(opt, size, depth: int, iterations: int = 8192, outer_iterations
             o = get_optim(opt, model.parameters(), lr=1e-3, stochastic_schedule=stochastic)
 
             for _ in range(iterations):
-                loss = model(torch.randn((128, size)).cuda()).square().mean()
+                loss = model(torch.randn((128, size), device-'cuda')).square().mean()
                 loss.backward()
                 o.step()
                 o.zero_grad()