heavyball 1.7.2__py3-none-any.whl → 2.0.0.dev0__py3-none-any.whl

heavyball/__init__.py

@@ -2,10 +2,45 @@ import functools
 import math
 from typing import Optional
 
+import torch.optim
+
 from . import chainable as C
 from . import utils
 
 
+class SAMWrapper(torch.optim.Optimizer):
+    def __init__(self, params, wrapped_optimizer: utils.StatefulOptimizer, ball: float = 0.1):
+        if not isinstance(wrapped_optimizer, utils.StatefulOptimizer):
+            raise ValueError(f"{wrapped_optimizer.__class__.__name__} is not a HeavyBall optimizer")
+        super().__init__(params, {"ball": ball})
+        self.wrapped_optimizer = wrapped_optimizer
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        if closure is None:
+            raise ValueError("SAM requires closure")
+        with torch.enable_grad():
+            closure()
+        old_params = [utils.sam_step(group["params"], group["ball"]) for group in self.param_groups]
+
+        originaL_handle_closure = self.wrapped_optimizer._handle_closure
+
+        def _handle_closure(closure):
+            originaL_handle_closure(closure)
+            for group, old in zip(self.param_groups, old_params):
+                utils.copy_stochastic_list_(group["params"], old)
+
+        try:
+            self.wrapped_optimizer._handle_closure = _handle_closure
+            loss = self.wrapped_optimizer.step(closure)
+        finally:
+            self.wrapped_optimizer._handle_closure = originaL_handle_closure
+        return loss
+
+    def zero_grad(self, set_to_none: bool = True):
+        self.wrapped_optimizer.zero_grad()
+
+
 class ForeachAdamW(C.BaseOpt):
     def __init__(
         self,
@@ -24,10 +59,16 @@ class ForeachAdamW(C.BaseOpt):
         update_clipping: C.str_or_fn = C.use_default,
         palm: bool = C.use_default,
         beta2_scale: float = 0.8,
+        **kwargs,
     ):
         defaults = locals()
         defaults.pop("self")
         params = defaults.pop("params")
+        defaults.update(defaults.pop("kwargs"))
+
+        if kwargs:
+            utils.warn_once(f"Working with uncaptured keyword arguments: {kwargs}")
+
         super().__init__(params, defaults, foreach, gradient_clipping, update_clipping, palm, C.update_by_adam)
 
 
@@ -55,10 +96,16 @@ class ForeachRMSprop(C.BaseOpt):
         update_clipping: C.str_or_fn = C.use_default,
         palm: bool = C.use_default,
         beta2_scale: float = 0.8,
+        **kwargs,
     ):
         defaults = locals()
         defaults.pop("self")
         params = defaults.pop("params")
+        defaults.update(defaults.pop("kwargs"))
+
+        if kwargs:
+            utils.warn_once(f"Working with uncaptured keyword arguments: {kwargs}")
+
         super().__init__(
             params,
             defaults,
@@ -90,10 +137,16 @@ class ForeachSFAdamW(C.ScheduleFree):
         update_clipping: C.str_or_fn = C.use_default,
         palm: bool = C.use_default,
         beta2_scale: float = 0.8,
+        **kwargs,
     ):
         defaults = locals()
         defaults.pop("self")
         params = defaults.pop("params")
+        defaults.update(defaults.pop("kwargs"))
+
+        if kwargs:
+            utils.warn_once(f"Working with uncaptured keyword arguments: {kwargs}")
+
         super().__init__(
             params,
             defaults,
@@ -106,6 +159,49 @@ class ForeachSFAdamW(C.ScheduleFree):
         )
 
 
+class MSAMLaProp(C.MSAM):
+    def __init__(
+        self,
+        params,
+        lr=0.0025,
+        betas=(0.9, 0.99),
+        eps=1e-6,
+        weight_decay=0,
+        warmup_steps=0,
+        r=0.0,
+        weight_lr_power=2.0,
+        foreach: bool = True,
+        storage_dtype: str = "float32",
+        mars: bool = False,
+        caution: bool = False,
+        mars_gamma: float = 0.0025,
+        gradient_clipping: C.str_or_fn = C.use_default,
+        update_clipping: C.str_or_fn = C.use_default,
+        palm: bool = C.use_default,
+        beta2_scale: float = 0.8,
+        sam_step_size: float = 0.1,
+        **kwargs,
+    ):
+        defaults = locals()
+        defaults.pop("self")
+        params = defaults.pop("params")
+        defaults.update(defaults.pop("kwargs"))
+
+        if kwargs:
+            utils.warn_once(f"Working with uncaptured keyword arguments: {kwargs}")
+
+        super().__init__(
+            params,
+            defaults,
+            foreach,
+            gradient_clipping,
+            update_clipping,
+            palm,
+            C.scale_by_exp_avg_sq,
+            C.update_by_msam,
+        )
+
+
 class PaLMForeachSFAdamW(ForeachSFAdamW):
     palm: bool = True
 
@@ -128,10 +224,16 @@ class ForeachADOPT(C.BaseOpt):
         update_clipping: C.str_or_fn = C.use_default,
         palm: bool = C.use_default,
         beta2_scale: float = 0.8,
+        **kwargs,
     ):
         defaults = locals()
         defaults.pop("self")
         params = defaults.pop("params")
+        defaults.update(defaults.pop("kwargs"))
+
+        if kwargs:
+            utils.warn_once(f"Working with uncaptured keyword arguments: {kwargs}")
+
         super().__init__(params, defaults, foreach, gradient_clipping, update_clipping, palm, C.update_by_adopt)
 
 
@@ -154,10 +256,16 @@ class ForeachMuon(C.BaseOpt):
         palm: bool = C.use_default,
         beta2_scale: float = 0.8,
         nesterov: bool = True,
+        **kwargs,
     ):
         defaults = locals()
         defaults.pop("self")
         params = defaults.pop("params")
+        defaults.update(defaults.pop("kwargs"))
+
+        if kwargs:
+            utils.warn_once(f"Working with uncaptured keyword arguments: {kwargs}")
+
         super().__init__(
             params,
             defaults,
@@ -165,7 +273,7 @@ class ForeachMuon(C.BaseOpt):
             gradient_clipping,
             update_clipping,
             palm,
-            C.nesterov_momentum if nesterov else C.heavyball_momentum,
+            C.nesterov_ema if nesterov else C.exp_avg,
             C.orthogonalize_update,
         )
 
@@ -188,10 +296,16 @@ class ForeachLaProp(C.BaseOpt):
         update_clipping: C.str_or_fn = C.use_default,
         palm: bool = C.use_default,
         beta2_scale: float = 0.8,
+        **kwargs,
     ):
         defaults = locals()
         defaults.pop("self")
         params = defaults.pop("params")
+        defaults.update(defaults.pop("kwargs"))
+
+        if kwargs:
+            utils.warn_once(f"Working with uncaptured keyword arguments: {kwargs}")
+
         super().__init__(params, defaults, foreach, gradient_clipping, update_clipping, palm, C.update_by_laprop)
 
 
@@ -213,10 +327,16 @@ class MuonLaProp(C.BaseOpt):
         update_clipping: C.str_or_fn = C.use_default,
         palm: bool = C.use_default,
         beta2_scale: float = 0.8,
+        **kwargs,
     ):
         defaults = locals()
         defaults.pop("self")
         params = defaults.pop("params")
+        defaults.update(defaults.pop("kwargs"))
+
+        if kwargs:
+            utils.warn_once(f"Working with uncaptured keyword arguments: {kwargs}")
+
         super().__init__(
             params,
             defaults,
@@ -271,12 +391,18 @@ class ForeachSOAP(C.BaseOpt):
         update_clipping: C.str_or_fn = C.use_default,
         storage_dtype: str = "float32",
         stochastic_schedule: bool = False,
+        precond_grad_accum: bool = False,
+        **kwargs,
     ):
         use_precond_schedule = C.default(use_precond_schedule, self.use_precond_schedule)
 
         defaults = locals()
         defaults.pop("self")
         params = defaults.pop("params")
+        defaults.update(defaults.pop("kwargs"))
+
+        if kwargs:
+            utils.warn_once(f"Working with uncaptured keyword arguments: {kwargs}")
 
         if use_precond_schedule:
             del defaults["precondition_frequency"]
@@ -313,10 +439,16 @@ class ForeachSignLaProp(C.BaseOpt):
         update_clipping: C.str_or_fn = C.use_default,
         palm: bool = C.use_default,
         beta2_scale: float = 0.8,
+        **kwargs,
     ):
         defaults = locals()
         defaults.pop("self")
         params = defaults.pop("params")
+        defaults.update(defaults.pop("kwargs"))
+
+        if kwargs:
+            utils.warn_once(f"Working with uncaptured keyword arguments: {kwargs}")
+
         super().__init__(
             params,
             defaults,
@@ -371,12 +503,17 @@ class ForeachSOLP(C.BaseOpt):
         update_clipping: C.str_or_fn = C.use_default,
         storage_dtype: str = "float32",
         stochastic_schedule: bool = False,
+        **kwargs,
     ):
         use_precond_schedule = C.default(use_precond_schedule, self.use_precond_schedule)
 
         defaults = locals()
         defaults.pop("self")
         params = defaults.pop("params")
+        defaults.update(defaults.pop("kwargs"))
+
+        if kwargs:
+            utils.warn_once(f"Working with uncaptured keyword arguments: {kwargs}")
 
         if use_precond_schedule:
             del defaults["precondition_frequency"]
@@ -427,10 +564,15 @@ class OrthoLaProp(C.BaseOpt):
         update_clipping: C.str_or_fn = C.use_default,
         palm: bool = C.use_default,
         beta2_scale: float = 0.8,
+        **kwargs,
     ):
         defaults = locals()
         defaults.pop("self")
         params = defaults.pop("params")
+        defaults.update(defaults.pop("kwargs"))
+
+        if kwargs:
+            utils.warn_once(f"Working with uncaptured keyword arguments: {kwargs}")
         super().__init__(
             params,
             defaults,
@@ -461,10 +603,15 @@ class LaPropOrtho(C.BaseOpt):
         update_clipping: C.str_or_fn = C.use_default,
         palm: bool = C.use_default,
         beta2_scale: float = 0.8,
+        **kwargs,
     ):
         defaults = locals()
         defaults.pop("self")
         params = defaults.pop("params")
+        defaults.update(defaults.pop("kwargs"))
+
+        if kwargs:
+            utils.warn_once(f"Working with uncaptured keyword arguments: {kwargs}")
         super().__init__(
             params,
             defaults,
@@ -487,12 +634,14 @@ class ForeachPSGDKron(C.BaseOpt):
     delayed: bool = False
     cached: bool = False
     exp_avg_input: bool = True
+    quad: bool = False
 
     def __init__(
         self,
         params,
         lr=0.001,
-        beta=0.9,
+        beta=None,
+        betas=(0.9, 0.999),
         weight_decay=0.0,
         preconditioner_update_probability=None,
         max_size_triangular=2048,
@@ -515,23 +664,38 @@ class ForeachPSGDKron(C.BaseOpt):
         exp_avg_input: Optional[bool] = C.use_default,
         gradient_clipping: C.str_or_fn = C.use_default,
         update_clipping: C.str_or_fn = C.use_default,  #
+        adaptive: bool = False,
+        ortho_method: Optional[str] = None,  # If None, no orthogonalization
+        precond_grad_accum: bool = False,
+        lower_bound_beta: float = 0.9,  # 0.0 recovers pre-2.0.0 PSGD
+        inverse_free: bool = C.use_default,
+        dampening: float = 2**-13,
+        precond_update_power_iterations: int = 2,
         # expert parameters
         precond_init_scale=None,
-        precond_init_scale_scale=1,
-        precond_lr=0.1,
+        precond_init_scale_scale: float = 1,
+        precond_init_scale_power: Optional[float] = None,
+        precond_lr: float = 0.1,
+        **kwargs,
     ):
+        delayed = C.default(delayed, self.delayed)
+        cached = C.default(cached, self.cached)
+        exp_avg_input = C.default(exp_avg_input, self.exp_avg_input)
+        update_clipping = C.default(update_clipping, utils.trust_region_clip_)
+        inverse_free = C.default(inverse_free, self.quad)
+
         defaults = locals()
         defaults.pop("self")
+        defaults.update(defaults.pop("kwargs"))
+
+        if kwargs:
+            utils.warn_once(f"Working with uncaptured keyword arguments: {kwargs}")
+
         self.precond_schedule = (
             defaults.pop("preconditioner_update_probability") or utils.precond_update_prob_schedule()
         )
         params = defaults.pop("params")
 
-        delayed = C.default(delayed, self.delayed)
-        cached = C.default(cached, self.cached)
-        exp_avg_input = C.default(exp_avg_input, self.exp_avg_input)
-        update_clipping = C.default(update_clipping, utils.trust_region_clip_)
-
         super().__init__(
             params,
             defaults,
@@ -569,6 +733,11 @@ class NewtonHybrid2PSGDKron(ForeachCachedNewtonPSGD):
     hvp_interval = 2
 
 
+class QUAD(ForeachPSGDKron):
+    quad = True
+    cached = True
+
+
 class ForeachPSGDLRA(C.BaseOpt):
     """
     Originally from Evan Walters and Omead Pooladzandi, 2024
@@ -601,13 +770,24 @@ class ForeachPSGDLRA(C.BaseOpt):
         gradient_clipping: C.str_or_fn = C.use_default,
         update_clipping: C.str_or_fn = C.use_default,
         eps: float = 1e-8,  #
-        # expert parameters
+        precond_grad_accum: bool = False,  # expert parameters
         precond_init_scale=None,
-        precond_init_scale_scale=1,
-        precond_lr=0.1,
+        precond_init_scale_scale: float = 1,
+        precond_init_scale_power: Optional[float] = None,
+        precond_lr: float = 0.1,
+        **kwargs,
     ):
+        delayed = C.default(delayed, self.delayed)
+        exp_avg_input = C.default(exp_avg_input, self.exp_avg_input)
+        update_clipping = C.default(update_clipping, utils.trust_region_clip_)
+
         defaults = locals()
         defaults.pop("self")
+        defaults.update(defaults.pop("kwargs"))
+
+        if kwargs:
+            utils.warn_once(f"Working with uncaptured keyword arguments: {kwargs}")
+
         self.precond_schedule = (
             defaults.pop("preconditioner_update_probability") or utils.precond_update_prob_schedule()
         )
@@ -621,10 +801,6 @@ class ForeachPSGDLRA(C.BaseOpt):
             defaults["rank"] = round(math.log2(sum(p.numel() for p in params)))
             utils.warn_once(f"rank was set to {defaults['rank']}")
 
-        delayed = C.default(delayed, self.delayed)
-        exp_avg_input = C.default(exp_avg_input, self.exp_avg_input)
-        update_clipping = C.default(update_clipping, utils.trust_region_clip_)
-
         super().__init__(
             params,
             defaults,
@@ -715,4 +891,5 @@ __all__ = [
     "NewtonPSGDLRA",
     "NewtonHybrid2PSGDLRA",
     "NewtonHybrid2PSGDKron",
+    "MSAMLaProp",
 ]