heavyball 1.7.1__tar.gz → 1.7.2__tar.gz

{heavyball-1.7.1 → heavyball-1.7.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: heavyball
-Version: 1.7.1
+Version: 1.7.2
 Summary: Efficient Optimizers
 Author-email: HeavyBall Authors <github.heavyball@nestler.sh>
 Project-URL: source, https://github.com/HomebrewML/HeavyBall

{heavyball-1.7.1 → heavyball-1.7.2}/heavyball/utils.py

@@ -50,7 +50,7 @@ def decorator(func):
     return _fn
 
 
-def decorator_knowngood(func: Callable):
+def decorator_knowngood(func: Callable, fullgraph: bool = True):
     compiled = None
 
     @functools.wraps(func)
@@ -59,7 +59,7 @@ def decorator_knowngood(func: Callable):
             return func(*args, **kwargs)
         nonlocal compiled
         if compiled is None:
-            compiled = torch.compile(fullgraph=True, dynamic=dynamic, mode=compile_mode)(func)
+            compiled = torch.compile(fullgraph=fullgraph, dynamic=dynamic, mode=compile_mode)(func)
         return compiled(*args, **kwargs)
 
     return _fn
@@ -1332,6 +1332,7 @@ def divided_root(x: torch.Tensor, y: torch.Tensor, pow0: float, pow1: float, eps
 def precond_init_scale(scale, scale_scale, grad, hessian_vector, vector, scale_max: float = 1e6):
     automatic_scale = True
     manual_hint = " Set it manually using `precond_init_scale=0.1`"
+
     if scale is not None:
         automatic_scale = False
         warn_once(
@@ -1345,12 +1346,16 @@ def precond_init_scale(scale, scale_scale, grad, hessian_vector, vector, scale_m
         scale = mean_root(grad, 4) * scale_scale
     else:
         scale = divided_root(vector, hessian_vector, 2, 4) * scale_scale
+
     if isinstance(scale, torch.Tensor):
         scale = scale.item()  # slow, but necessary
+
     if np.isfinite(scale):
-        if scale > scale_max or scale < 1 / scale_max:
+        if scale > scale_max or scale < 1 / scale_max:  # fallthrough to later checks
             warn_once(f"The computed precond_init_scale {scale} is outside of the expected range.{manual_hint}")
-        return scale
+        else:
+            return scale
+
     if not automatic_scale:
         raise ValueError("The manually set precond_init_scale is not finite")
 
@@ -1361,6 +1366,10 @@ def precond_init_scale(scale, scale_scale, grad, hessian_vector, vector, scale_m
             raise ValueError(f"Grad or HVP is all 0s, causing NaNs in precond_init_scale computation.{manual_hint}")
         if not torch.isfinite(x).all().item():
             raise ValueError("Grad or HVP is not finite")
+
+    if np.isfinite(scale):
+        return scale
+
     raise ValueError(f"Computed precond_init_scale is not finite.{manual_hint}")
 
 
@@ -1634,35 +1643,58 @@ def dampen_multiple(g: List[Tensor], damp: float = 2**-13):
     return flatten(vs), flatten(gs)
 
 
-@decorator_knowngood
 def casted_einsum(expr: str, *args: Tensor) -> Tensor:
     md = min_dtype(args)
     return torch.einsum(expr, *[a.to(md) for a in args]).to(args[-1].dtype)
 
 
+@decorator_knowngood
+def _psgd_calc_scalars_(Qs: List[Tensor], conjB: Tensor):
+    triangular_qs = []
+    for i, q in enumerate(Qs):
+        q = promote(q)
+        if q.dim() <= 1:
+            shape = [1] * conjB.ndim
+            shape[i] = -1
+            conjB /= q.view(shape)
+        else:
+            triangular_qs.append((i, q))
+    return triangular_qs
+
+
+@decorator_knowngood
+def _reshape_conjB(solved: Tensor, original_shape: List[int], last_dim: int, new_shape: int):
+    solved = solved.reshape(original_shape)
+    solved.transpose(last_dim, -1)
+    return solved.reshape(new_shape).contiguous()
+
+
 def psgd_calc_A_and_conjB(exprA, G, Q, conjB):  # conjB ("V", "vector") == randn during hvp/whitening
     order = G.dim()
     if order > 1:
         conjB = conjB.view_as(G).permute(*range(1, order), 0)
     conjB = conjB.to(promote(G.dtype))
     A = casted_einsum(exprA, *Q, G)
-    for i, q in enumerate(Q):
-        q = promote(q)
-        if q.dim() <= 1:
-            conjB /= q
-        else:
-            solved = torch.linalg.solve_triangular(q, conjB.reshape(-1, q.size(0)).contiguous(), upper=True, left=False)
-            conjB = solved.reshape_as(conjB)
-        if i < order - 1:
-            conjB = conjB.transpose(i, -1)
+    solve = torch.compiler.disable(torch.linalg.solve_triangular)
+    original_shape = conjB.shape
+    prev_i = -1
+    for i, tri_q in _psgd_calc_scalars_(Q, conjB):
+        conjB = _reshape_conjB(conjB, original_shape, prev_i, [-1, tri_q.size(0)])
+        prev_i = i
+        conjB = solve(tri_q, conjB, upper=True, left=False)
+    conjB = _reshape_conjB(conjB, original_shape, prev_i, original_shape)
     return A, conjB
 
 
-def psgd_lb(A, max_abs):
+@decorator_knowngood
+def _max_select(to_index: Tensor, to_argmax: Tensor):
+    idx = to_argmax.argmax()
+    return to_index.index_select(1, idx).flatten().contiguous()
+
+
+def psgd_lb(A: Tensor, max_abs: Tensor):
     A /= max_abs
-    a0 = torch.einsum("ij,ij->j", A, A)
-    i = torch.argmax(a0)
-    x = torch.index_select(A, 1, i).flatten().contiguous()
+    x = _max_select(A, torch.einsum("ij,ij->j", A, A))
     x = torch.einsum("i,ij->j", x, A)
     x /= x.norm()
     x = torch.einsum("j,kj->k", x, A)
@@ -1671,30 +1703,51 @@ def psgd_lb(A, max_abs):
     return x
 
 
+@decorator_knowngood
+def _subtract_from_line_(state: Tensor, term: Tensor):
+    stochastic_add_([state], [triu_to_line([term])[0][1]], -1)
+
+
+@decorator_knowngood
+def _prescale_term_(term1: Tensor, fac: Tensor, norm: Tensor, lower_bound: Tensor):
+    out = term1.float().triu() * fac
+    out = out / torch.where(norm > 0, lower_bound, norm).clamp(tiny_bf16)
+    copy_stochastic_(term1, out)
+
+
+@decorator_knowngood
+def _compilable_stochastic_multiply_div_(x: Tensor, fac: Tensor, y: Tensor, z: Tensor):
+    copy_stochastic_(x, promote(x) * promote(fac) * promote(y) / promote(z).clamp(min=tiny_bf16))
+
+
+@decorator_knowngood
+def _compilable_add_sub_(x: Tensor, y: Tensor):
+    x = promote(x)
+    y = promote(y)
+    return x - y, x + y
+
+
 @decorator
 def psgd_update_precond(Q, exprs, G, precond_lr, oq, store_triu_as_line, V):
     """Update Kronecker product preconditioner Q with pair (V, G)."""
     exprA, exprGs, _ = exprs
     A, conjB = psgd_calc_A_and_conjB(exprA, G, Q, V)
+    precond_lr = scalar_guard(precond_lr, G)
 
     for q, exprG, o in zip(Q, exprGs, oq):
-        term1 = promote(torch.einsum(exprG, A, A))
-        term2 = promote(torch.einsum(exprG, conjB, conjB))
-        term1, term2 = term1 - term2, term1 + term2
-        term1 *= precond_lr
+        term1 = torch.einsum(exprG, A, A)
+        term2 = torch.einsum(exprG, conjB, conjB)
+        term1, term2 = _compilable_add_sub_(term1, term2)
         norm = term2.norm(float("inf"))
         if q.dim() < 2:
-            term1 *= q.to(term1.dtype) / norm.clamp_(min=tiny_bf16)
+            _compilable_stochastic_multiply_div_(term1, precond_lr, q, norm)
         else:
-            torch.triu(term1, out=term1)
-            term1 /= torch.where(norm > 0, psgd_lb(term2, norm), norm).clamp_(tiny_bf16)
-            term1 = torch.mm(term1, q.to(term1.dtype))
+            lower_bound = psgd_lb(term2, norm)
+            _prescale_term_(term1, precond_lr, lower_bound, norm)
+            torch.mm(term1, q.to(term1.dtype), out=term1)
         if store_triu_as_line:
-            term1 = triu_to_line([term1])[0][1]  # Convert update to line format
-            # Apply update directly to the tensor part of the state tuple o[1]
-            stochastic_add_(o[1], term1, -1)
+            _subtract_from_line_(q, term1)
         else:
-            # Apply update to the state tensor o
             stochastic_add_(o, term1, -1)
 
 

{heavyball-1.7.1 → heavyball-1.7.2}/heavyball.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: heavyball
-Version: 1.7.1
+Version: 1.7.2
 Summary: Efficient Optimizers
 Author-email: HeavyBall Authors <github.heavyball@nestler.sh>
 Project-URL: source, https://github.com/HomebrewML/HeavyBall

{heavyball-1.7.1 → heavyball-1.7.2}/heavyball.egg-info/SOURCES.txt

@@ -9,9 +9,6 @@ heavyball.egg-info/SOURCES.txt
 heavyball.egg-info/dependency_links.txt
 heavyball.egg-info/requires.txt
 heavyball.egg-info/top_level.txt
-heavyball/optimizations/__init__.py
-heavyball/optimizations/integrator.py
-heavyball/optimizations/optimizations.py
 test/test_bf16_params.py
 test/test_bf16_q.py
 test/test_bf16_storage.py
@@ -27,10 +24,4 @@ test/test_memory_leak.py
 test/test_merge.py
 test/test_no_grad.py
 test/test_soap.py
-test/test_stochastic_updates.py
-tests/test_psgd_kron_line_optim.py
-tests/test_psgd_kron_regression.py
-tests/test_psgd_lra_regression.py
-tests/test_psgd_optimization.py
-tests/test_psgd_optimizations.py
-tests/test_psgd_training_performance.py
+test/test_stochastic_updates.py

{heavyball-1.7.1 → heavyball-1.7.2}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 [project]
 name = "heavyball"
 description = "Efficient Optimizers"
-version = "1.7.1"
+version = "1.7.2"
 authors = [{ name = "HeavyBall Authors", email = "github.heavyball@nestler.sh" }]
 classifiers = ["Intended Audience :: Developers",
     "Intended Audience :: Science/Research",

heavyball-1.7.1/heavyball/optimizations/__init__.py

@@ -1,38 +0,0 @@
-"""
-PSGD optimization module - optimized implementations of PSGD functions
-to improve execution speed while maintaining numerical equivalence.
-"""
-
-# Import optimized functions
-# Import integrator API
-from .integrator import (
-    enable_optimizations,
-    get_optimization_status,
-    restore_original_functions,
-)
-from .optimizations import (
-    # LRA optimizations
-    low_rank_mm_optimized,
-    lra_precond_optimized,
-    precond_grad_cached_optimized,
-    # KRON optimizations
-    psgd_calc_A_and_conjB_optimized,
-    psgd_precond_grad_optimized,
-    psgd_update_precond_optimized,
-    update_lra_precond_optimized,
-)
-
-__all__ = [
-    # Optimized functions
-    "low_rank_mm_optimized",
-    "update_lra_precond_optimized",
-    "lra_precond_optimized",
-    "psgd_calc_A_and_conjB_optimized",
-    "psgd_update_precond_optimized",
-    "psgd_precond_grad_optimized",
-    "precond_grad_cached_optimized",
-    # Integrator API
-    "enable_optimizations",
-    "restore_original_functions",
-    "get_optimization_status",
-]

heavyball-1.7.1/heavyball/optimizations/integrator.py

@@ -1,169 +0,0 @@
-"""
-Integration module to selectively enable optimized implementations
-of PSGD functions while maintaining API compatibility.
-"""
-
-import os
-import sys
-from typing import Any, Dict
-
-import torch
-
-from . import optimizations
-from .. import utils
-
-# Store original function references
-_original_functions = {}
-_optimized_functions = {}
-
-# Mapping of original functions to their optimized versions
-OPTIMIZATION_MAP = {
-    # LRA functions
-    utils.update_lra_precond_: optimizations.update_lra_precond_optimized,
-    utils.lra_precond: optimizations.lra_precond_optimized,
-    # KRON functions
-    utils.psgd_update_precond: optimizations.psgd_update_precond_optimized,
-    utils.psgd_precond_grad: optimizations.psgd_precond_grad_optimized,
-    utils.precond_grad_cached_: optimizations.precond_grad_cached_optimized,
-}
-
-# Config for enabling/disabling optimizations
-_config = {
-    "enabled": os.environ.get("HEAVYBALL_OPTIMIZE", "1") == "1",
-    "torch_compile_allowed": os.environ.get("HEAVYBALL_USE_COMPILE", "1") == "1",
-    "enable_lra": True,
-    "enable_kron": True,
-    "verbose": os.environ.get("HEAVYBALL_VERBOSE", "0") == "1",
-}
-
-
-def _apply_monkey_patch(original_func, optimized_func):
-    """Monkey patch a function with its optimized version."""
-    if original_func not in _original_functions:
-        _original_functions[original_func] = original_func
-
-    # Store reference to the optimized function
-    _optimized_functions[original_func] = optimized_func
-
-    # Get the module where the original function is defined
-    module = original_func.__module__
-    func_name = original_func.__name__
-
-    # Replace the function in its module
-    if hasattr(sys.modules[module], func_name):
-        setattr(sys.modules[module], func_name, optimized_func)
-
-        if _config["verbose"]:
-            print(f"Replaced {module}.{func_name} with optimized version")
-    else:
-        if _config["verbose"]:
-            print(f"Warning: Could not find {func_name} in module {module}")
-
-
-def enable_optimizations(
-    enable: bool = True, lra: bool = True, kron: bool = True, torch_compile: bool = True, verbose: bool = False
-):
-    """
-    Enable or disable PSGD optimizations.
-
-    Args:
-        enable: Whether to enable optimizations at all
-        lra: Whether to enable LRA-specific optimizations
-        kron: Whether to enable Kron-specific optimizations
-        torch_compile: Whether to allow torch.compile optimizations
-        verbose: Whether to print optimization status messages
-    """
-    _config["enabled"] = enable
-    _config["enable_lra"] = lra
-    _config["enable_kron"] = kron
-    _config["torch_compile_allowed"] = torch_compile
-    _config["verbose"] = verbose
-
-    if verbose:
-        print(f"PSGD Optimizations: {'enabled' if enable else 'disabled'}")
-        print(f"  - LRA optimizations: {'enabled' if lra else 'disabled'}")
-        print(f"  - KRON optimizations: {'enabled' if kron else 'disabled'}")
-        print(f"  - torch.compile: {'allowed' if torch_compile else 'disabled'}")
-
-    if not enable:
-        # Restore original functions
-        restore_original_functions()
-        return
-
-    # Apply optimizations based on config
-    for orig_func, opt_func in OPTIMIZATION_MAP.items():
-        # Skip LRA functions if disabled
-        if not _config["enable_lra"] and orig_func in [utils.update_lra_precond_, utils.lra_precond]:
-            continue
-
-        # Skip KRON functions if disabled
-        if not _config["enable_kron"] and orig_func in [
-            utils.psgd_update_precond,
-            utils.psgd_precond_grad,
-            utils.precond_grad_cached_,
-        ]:
-            continue
-
-        _apply_monkey_patch(orig_func, opt_func)
-
-    # Disable torch.compile if not allowed
-    if not _config["torch_compile_allowed"]:
-        # Monkey patch torch.compile to be a no-op
-        def _noop_compile(fn, **kwargs):
-            return fn
-
-        if not hasattr(torch, "_original_compile"):
-            torch._original_compile = torch.compile
-            torch.compile = _noop_compile
-            if verbose:
-                print("Disabled torch.compile (replaced with no-op)")
-    else:
-        # Restore original torch.compile
-        if hasattr(torch, "_original_compile"):
-            torch.compile = torch._original_compile
-            del torch._original_compile
-            if verbose:
-                print("Restored original torch.compile")
-
-
-def restore_original_functions():
-    """Restore all original function implementations."""
-    for orig_func, func_ref in _original_functions.items():
-        module = orig_func.__module__
-        func_name = orig_func.__name__
-
-        if hasattr(sys.modules[module], func_name):
-            setattr(sys.modules[module], func_name, func_ref)
-
-            if _config["verbose"]:
-                print(f"Restored original implementation of {module}.{func_name}")
-
-    # Also restore torch.compile if it was modified
-    if hasattr(torch, "_original_compile"):
-        torch.compile = torch._original_compile
-        del torch._original_compile
-        if _config["verbose"]:
-            print("Restored original torch.compile")
-
-
-def get_optimization_status() -> Dict[str, Any]:
-    """Get current optimization status."""
-    return {
-        "enabled": _config["enabled"],
-        "lra_enabled": _config["enable_lra"],
-        "kron_enabled": _config["enable_kron"],
-        "torch_compile_allowed": _config["torch_compile_allowed"],
-        "optimized_functions": list(_optimized_functions.keys()),
-        "original_functions": list(_original_functions.keys()),
-    }
-
-
-# Auto-initialize optimizations based on environment
-if os.environ.get("HEAVYBALL_AUTO_OPTIMIZE", "1") == "1":
-    enable_optimizations(
-        enable=_config["enabled"],
-        lra=_config["enable_lra"],
-        kron=_config["enable_kron"],
-        torch_compile=_config["torch_compile_allowed"],
-        verbose=_config["verbose"],
-    )