heavyball 2.1.3__tar.gz → 2.1.4__tar.gz

{heavyball-2.1.3 → heavyball-2.1.4}/PKG-INFO

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

{heavyball-2.1.3 → heavyball-2.1.4}/heavyball/__init__.py

@@ -1,6 +1,6 @@
 import functools
 import math
-from typing import Optional
+from typing import Optional, Type, Union
 
 import torch.optim
 
@@ -8,39 +8,6 @@ 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 SGD(C.BaseOpt):
     def __init__(
         self,
@@ -778,7 +745,7 @@ class ForeachPSGDKron(C.BaseOpt):
         beta=None,
         betas=(0.9, 0.999),
         weight_decay=0.0,
-        preconditioner_update_probability=None,
+        preconditioner_update_probability=C.use_default,
         max_size_triangular=2048,
         min_ndim_triangular=2,
         memory_save_mode=None,
@@ -830,8 +797,8 @@ class ForeachPSGDKron(C.BaseOpt):
         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()
+        self.precond_schedule = C.default(
+            defaults.pop("preconditioner_update_probability"), utils.precond_update_prob_schedule()
         )
         params = defaults.pop("params")
 
@@ -890,7 +857,7 @@ class ForeachPSGDLRA(C.BaseOpt):
         lr=0.001,
         beta=0.9,
         weight_decay=0.0,
-        preconditioner_update_probability=None,
+        preconditioner_update_probability=C.use_default,
         momentum_into_precond_update=True,
         rank: Optional[int] = None,
         warmup_steps: int = 0,
@@ -924,8 +891,8 @@ class ForeachPSGDLRA(C.BaseOpt):
         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()
+        self.precond_schedule = C.default(
+            defaults.pop("preconditioner_update_probability"), utils.precond_update_prob_schedule()
         )
         params = defaults.pop("params")
 
@@ -960,6 +927,54 @@ class NewtonHybrid2PSGDLRA(ForeachNewtonPSGDLRA):
     hvp_interval = 2
 
 
+class SAMWrapper(torch.optim.Optimizer):
+    def __init__(
+        self,
+        params,
+        wrapped_optimizer: Union[utils.StatefulOptimizer, Type[utils.StatefulOptimizer]] = ForeachAdamW,
+        ball: float = 0.1,
+    ):
+        params = list(params)
+        super().__init__(params, {"ball": ball})
+
+        if isinstance(wrapped_optimizer, type):
+            if not issubclass(wrapped_optimizer, utils.StatefulOptimizer):
+                raise ValueError(f"{wrapped_optimizer.__name__} is not a HeavyBall optimizer")
+            wrapped_optimizer = wrapped_optimizer(params)
+        elif not isinstance(wrapped_optimizer, utils.StatefulOptimizer):
+            raise ValueError(f"{wrapped_optimizer.__class__.__name__} is not a HeavyBall optimizer")
+
+        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):
+            try:
+                _loss = original_handle_closure(closure)
+            finally:
+                for group, old in zip(self.param_groups, old_params):
+                    utils.copy_stochastic_list_(group["params"], old)
+            return _loss
+
+        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(set_to_none=set_to_none)
+
+
 PalmForEachSoap = PaLMForeachSOAP
 PaLMSOAP = PaLMForeachSOAP
 PaLMSFAdamW = PaLMForeachSFAdamW
@@ -983,52 +998,4 @@ PSGDLRA = ForeachPSGDLRA
 NewtonPSGDLRA = ForeachNewtonPSGDLRA
 NewtonPSGDKron = ForeachCachedNewtonPSGD
 
-__all__ = [
-    "Muon",
-    "RMSprop",
-    "PrecondSchedulePaLMSOAP",
-    "PSGDKron",
-    "PurePSGD",
-    "DelayedPSGD",
-    "CachedPSGDKron",
-    "CachedDelayedPSGDKron",
-    "PalmForEachSoap",
-    "PaLMSOAP",
-    "PaLMSFAdamW",
-    "LaProp",
-    "ADOPT",
-    "PrecondScheduleSOAP",
-    "PrecondSchedulePaLMSOAP",
-    "RMSprop",
-    "MuonLaProp",
-    "ForeachSignLaProp",
-    "ForeachDelayedPSGDLRA",
-    "ForeachPSGDLRA",
-    "ForeachPSGDLRA",
-    "ForeachNewtonPSGDLRA",  #
-    "ForeachAdamW",
-    "ForeachSFAdamW",
-    "ForeachLaProp",
-    "ForeachADOPT",
-    "ForeachSOAP",
-    "ForeachPSGDKron",
-    "ForeachPurePSGD",
-    "ForeachDelayedPSGD",
-    "ForeachCachedPSGDKron",
-    "ForeachCachedDelayedPSGDKron",
-    "ForeachRMSprop",
-    "ForeachMuon",
-    "ForeachCachedNewtonPSGD",
-    "OrthoLaProp",
-    "LaPropOrtho",
-    "SignLaProp",
-    "DelayedPSGD",
-    "PSGDLRA",
-    "NewtonPSGDLRA",
-    "NewtonHybrid2PSGDLRA",
-    "NewtonHybrid2PSGDKron",
-    "MSAMLaProp",
-    "NewtonPSGDKron",
-    "ForeachAdamC",
-    "SGD",
-]
+__all__ = [k for k, v in globals().items() if isinstance(v, type) and issubclass(v, torch.optim.Optimizer)]

{heavyball-2.1.3 → heavyball-2.1.4}/heavyball/chainable.py

@@ -62,8 +62,6 @@ class FunctionTransform:
                     self._init(st, group, *a, **kwargs)
                 except SkipUpdate:
                     skip_update = True
-                except:
-                    raise
                 finally:
                     if "is_initialized" not in st:
                         st["is_initialized"] = set()
@@ -499,6 +497,7 @@ def scale_by_suds(group, update, grad, param, exp_avg, exp_avg_sq, fisher_approx
     precond_update, _ = utils.eigvecs_product_rank1(precond_update.flatten(), fisher_approx.flatten(), w)
 
     new_approx = utils.oja_update(fisher_approx.flatten().to(update.dtype), update.flatten(), group["precond_lr"])
+    new_approx = new_approx.view_as(fisher_approx)
     utils.copy_stochastic_(fisher_approx, new_approx)
     return precond_update
 
@@ -565,7 +564,7 @@ def _init_psgd_kron(state, group, update, grad, param, cached: bool = False, pro
     )
     state["Q"] = utils.triu_to_line(Q) if group["store_triu_as_line"] else Q
     state["running_lower_bound"] = [torch.zeros((1,), device=q.device, dtype=torch.float64) for q in Q]
-    state["step"] = torch.zeros((), device=param.device, dtype=torch.int64)
+    state["step"] = torch.zeros((), device=param.device, dtype=torch.float64)  # torch casts int to float in ckpt load
     if group["adaptive"]:
         state["velocity"] = [torch.zeros((), device=q.device, dtype=q.dtype) for q in Q]
     if not cached:
@@ -750,7 +749,9 @@ def _update_psgd_precond(
     if isinstance(prob, float):
         float_prob = prob
     else:
-        float_prob = prob(group.get(f"cumulative_prob_{id(Q)}_prob_step", 1))
+        prob_step = group.get(f"cumulative_prob_{id(Q)}_prob_step", 1)
+        float_prob = prob(prob_step)
+        group[f"cumulative_prob_{id(Q)}_prob_step"] = prob_step + 1
     group["is_cached"] = should_use_cache = cached and float_prob < 0.5
 
     if precond is not None:
@@ -1086,6 +1087,7 @@ class ChainOpt(utils.StatefulOptimizer):
         if not group["foreach"] or len(p) == 1:
             for param, grad in zip(p, g):
                 chain(self.state_, group, [grad], [param], *self.fns)
+                group["caution"] = caution
         else:
             chain(self.state_, group, g, p, *self.fns)
 

{heavyball-2.1.3 → heavyball-2.1.4}/heavyball/helpers.py

@@ -1,5 +1,4 @@
-from __future__ import annotations
-
+import copy
 import functools
 import math
 import threading
@@ -14,7 +13,7 @@ import pandas as pd
 import torch
 from hebo.design_space.design_space import DesignSpace
 from hebo.optimizers.hebo import HEBO
-from optuna._transform import _SearchSpaceTransform
+from optuna._transform import _SearchSpaceTransform as SearchSpaceTransform
 from optuna.distributions import BaseDistribution, CategoricalDistribution, FloatDistribution, IntDistribution
 from optuna.samplers import BaseSampler, CmaEsSampler, RandomSampler
 from optuna.samplers._lazy_random_state import LazyRandomState
@@ -60,9 +59,9 @@ def manual_seed(seed: int | None = None) -> Generator[None, None, None]:
 class SimpleAPIBaseSampler(BaseSampler):
     def __init__(
         self,
-        search_space: dict[str, BaseDistribution] = None,
+        search_space: Optional[dict[str, BaseDistribution]] = None,
     ):
-        self.search_space = search_space
+        self.search_space = {} if search_space is None else dict(search_space)
 
     def suggest_all(self, trial: FrozenTrial):
         return {k: trial._suggest(k, dist) for k, dist in self.search_space.items()}
@@ -154,7 +153,7 @@ def _untransform_numerical_param_torch(
 
 
 @torch.no_grad()
-def untransform(self: _SearchSpaceTransform, trans_params: Tensor) -> dict[str, Any]:
+def untransform(self: SearchSpaceTransform, trans_params: Tensor) -> dict[str, Any]:
     assert trans_params.shape == (self._raw_bounds.shape[0],)
 
     if self._transform_0_1:
@@ -182,29 +181,31 @@ class BoTorchSampler(SimpleAPIBaseSampler):
 
     def __init__(
         self,
-        search_space: dict[str, BaseDistribution] = None,
+        search_space: Optional[dict[str, BaseDistribution]] = None,
         *,
-        candidates_func: None = None,
-        constraints_func: None = None,
+        candidates_func: Optional[Callable[..., Tensor]] = None,
+        constraints_func: Optional[Callable[..., Tensor]] = None,
         n_startup_trials: int = 10,
         consider_running_trials: bool = False,
-        independent_sampler: None = None,
+        independent_sampler: Optional[BaseSampler] = None,
         seed: int | None = None,
         device: torch.device | str | None = None,
         trial_chunks: int = 128,
     ):
-        assert constraints_func is None
-        assert candidates_func is None
-        assert consider_running_trials is False
-        assert independent_sampler is None
-        self._candidates_func = None
-        self._independent_sampler = RandomSampler(seed=seed)
+        if constraints_func is not None:
+            raise NotImplementedError("constraints_func is currently not supported by BoTorchSampler.")
+        if consider_running_trials:
+            raise NotImplementedError("consider_running_trials is currently not supported by BoTorchSampler.")
+        if candidates_func is not None and not callable(candidates_func):
+            raise TypeError("candidates_func must be callable.")
+        self._candidates_func = candidates_func
+        self._independent_sampler = independent_sampler or RandomSampler(seed=seed)
         self._n_startup_trials = n_startup_trials
         self._seed = seed
         self.trial_chunks = trial_chunks
 
         self._study_id: int | None = None
-        self.search_space = search_space
+        self.search_space = {} if search_space is None else dict(search_space)
         if isinstance(device, str):
             device = torch.device(device)
         self._device = device or torch.device("cpu")
@@ -212,14 +213,24 @@ class BoTorchSampler(SimpleAPIBaseSampler):
         self._values = None
         self._params = None
         self._index = 0
+        self._bounds_dim: int | None = None
 
     def infer_relative_search_space(self, study: Study, trial: FrozenTrial) -> dict[str, BaseDistribution]:
         return self.search_space
 
     @torch.no_grad()
     def _preprocess_trials(
-        self, trans: _SearchSpaceTransform, study: Study, trials: list[FrozenTrial]
+        self, trans: SearchSpaceTransform, study: Study, trials: list[FrozenTrial]
     ) -> Tuple[int, Tensor, Tensor]:
+        bounds_dim = trans.bounds.shape[0]
+        if self._bounds_dim is not None and self._bounds_dim != bounds_dim:
+            self._values = None
+            self._params = None
+            self._index = 0
+            self.seen_trials = set()
+        if self._bounds_dim is None:
+            self._bounds_dim = bounds_dim
+
         new_trials = []
         for trial in trials:
             tid: int = trial._trial_id
@@ -230,6 +241,10 @@ class BoTorchSampler(SimpleAPIBaseSampler):
 
         n_objectives = len(study.directions)
         if not new_trials:
+            if self._values is None or self._params is None:
+                empty_values = torch.zeros((0, n_objectives), dtype=torch.float64, device=self._device)
+                empty_params = torch.zeros((0, bounds_dim), dtype=torch.float64, device=self._device)
+                return n_objectives, empty_values, empty_params
             return n_objectives, self._values[: self._index], self._params[: self._index]
 
         n_completed_trials = len(trials)
@@ -246,18 +261,28 @@ class BoTorchSampler(SimpleAPIBaseSampler):
             if direction == StudyDirection.MINIMIZE:  # BoTorch always assumes maximization.
                 values[:, obj_idx] *= -1
 
-        if self._values is None:
+        bounds_dim = trans.bounds.shape[0]
+        cache_stale = (
+            self._values is None
+            or self._params is None
+            or self._values.size(1) != n_objectives
+            or self._params.size(1) != bounds_dim
+        )
+        if cache_stale:
             self._values = torch.zeros((self.trial_chunks, n_objectives), dtype=torch.float64, device=self._device)
-            self._params = torch.zeros(
-                (self.trial_chunks, trans.bounds.shape[0]), dtype=torch.float64, device=self._device
-            )
+            self._params = torch.zeros((self.trial_chunks, bounds_dim), dtype=torch.float64, device=self._device)
+            self._index = 0
+            self.seen_trials = set()
+            self._bounds_dim = bounds_dim
         spillage = (self._index + n_completed_trials) - self._values.size(0)
         if spillage > 0:
             pad = int(math.ceil(spillage / self.trial_chunks) * self.trial_chunks)
             self._values = F.pad(self._values, (0, 0, 0, pad))
             self._params = F.pad(self._params, (0, 0, 0, pad))
-        self._values[self._index : self._index + n_completed_trials] = torch.from_numpy(values)
-        self._params[self._index : self._index + n_completed_trials] = torch.from_numpy(params)
+        values_tensor = torch.from_numpy(values).to(self._device)
+        params_tensor = torch.from_numpy(params).to(self._device)
+        self._values[self._index : self._index + n_completed_trials] = values_tensor
+        self._params[self._index : self._index + n_completed_trials] = params_tensor
         self._index += n_completed_trials
 
         return n_objectives, self._values[: self._index], self._params[: self._index]
@@ -276,7 +301,7 @@ class BoTorchSampler(SimpleAPIBaseSampler):
         if n_completed_trials < self._n_startup_trials:
             return {}
 
-        trans = _SearchSpaceTransform(search_space)
+        trans = SearchSpaceTransform(search_space)
         n_objectives, values, params = self._preprocess_trials(trans, study, completed_trials)
 
         if self._candidates_func is None:
@@ -379,10 +404,10 @@ class HEBOSampler(optunahub.samplers.SimpleBaseSampler, SimpleAPIBaseSampler):
         independent_sampler: BaseSampler | None = None,
     ) -> None:
         super().__init__(search_space, seed)
-        assert constant_liar is False
-        assert independent_sampler is None
+        if constant_liar:
+            raise NotImplementedError("constant_liar is not supported by HEBOSampler.")
         self._hebo = HEBO(_convert_to_hebo_design_space(search_space), scramble_seed=self._seed)
-        self._independent_sampler = optuna.samplers.RandomSampler(seed=seed)
+        self._independent_sampler = independent_sampler or optuna.samplers.RandomSampler(seed=seed)
         self._rng = np.random.default_rng(seed)
 
     def sample_relative(
@@ -451,10 +476,12 @@ class FastINGO:
         learning_rate: Optional[float] = None,
         last_n: int = 4096,
         loco_step_size: float = 0.1,
-        device="cuda",
+        device: str | None = None,
         batchnorm_decay: float = 0.99,
         score_decay: float = 0.99,
     ) -> None:
+        if device is None:
+            device = _use_cuda()
         n_dimension = len(mean)
         if population_size is None:
             population_size = 4 + int(np.floor(3 * np.log(n_dimension)))
@@ -521,8 +548,14 @@ class FastINGO:
         if y.numel() <= 2:
             return
 
-        y = y + torch.where(y.min() <= 0, 1e-8 - y.min(), 0)
-        y = y.log()
+        min_y = y.min()
+        max_y = y.max()
+        if torch.isclose(max_y, min_y, rtol=0.0, atol=1e-12):
+            return
+
+        if min_y <= 0:
+            y = y + (1e-8 - min_y)
+        y = y.clamp_min_(1e-8).log()
 
         ema = -torch.arange(y.size(0), device=y.device, dtype=y.dtype)
         weight = self.batchnorm_decay**ema
@@ -583,7 +616,7 @@ class ImplicitNaturalGradientSampler(BaseSampler):
     def reseed_rng(self) -> None:
         self._independent_sampler.reseed_rng()
         if self._optimizer:
-            self._optimizer._rng.seed()
+            self._optimizer.generator.seed()
 
     def infer_relative_search_space(
         self, study: "optuna.Study", trial: "optuna.trial.FrozenTrial"
@@ -633,14 +666,11 @@ class ImplicitNaturalGradientSampler(BaseSampler):
             self._warn_independent_sampling = False
             return {}
 
-        trans = _SearchSpaceTransform(search_space)
+        trans = SearchSpaceTransform(search_space)
 
-        if self._optimizer is None:
+        if self._optimizer is None or self._optimizer.dim != len(trans.bounds):
             self._optimizer = self._init_optimizer(trans, population_size=self._population_size)
-
-        if self._optimizer.dim != len(trans.bounds):
-            self._warn_independent_sampling = False
-            return {}
+            self._param_queue.clear()
 
         solution_trials = [t for t in completed_trials if self._check_trial_is_generation(t)]
         for t in solution_trials:
@@ -651,7 +681,7 @@ class ImplicitNaturalGradientSampler(BaseSampler):
 
     def _init_optimizer(
         self,
-        trans: _SearchSpaceTransform,
+        trans: SearchSpaceTransform,
         population_size: Optional[int] = None,
     ) -> FastINGO:
         lower_bounds = trans.bounds[:, 0]
@@ -705,6 +735,7 @@ class ThreadLocalSampler(threading.local):
 
 
 def init_cmaes(study, seed, trials, search_space):
+    trials = copy.deepcopy(trials)
     trials.sort(key=lambda trial: trial.datetime_complete)
     return CmaEsSampler(seed=seed, source_trials=trials, lr_adapt=True)
 
@@ -716,8 +747,14 @@ def init_hebo(study, seed, trials, search_space):
     return sampler
 
 
+def _use_cuda():
+    return "cuda" if torch.cuda.is_available() else "cpu"
+
+
 def init_botorch(study, seed, trials, search_space):
-    return BoTorchSampler(search_space=search_space, seed=seed, device="cuda")  # will automatically pull in latest data
+    return BoTorchSampler(
+        search_space=search_space, seed=seed, device=_use_cuda()
+    )  # will automatically pull in latest data
 
 
 def init_nsgaii(study, seed, trials, search_space):
@@ -739,17 +776,20 @@ class AutoSampler(BaseSampler):
     def __init__(
         self,
         samplers: Iterable[Tuple[int, Callable]] | None = None,
-        search_space: dict[str, BaseDistribution] = None,
+        search_space: Optional[dict[str, BaseDistribution]] = None,
         *,
         seed: int | None = None,
-        constraints_func: None = None,
+        constraints_func: Optional[Callable[..., Any]] = None,
     ) -> None:
-        assert constraints_func is None
+        if constraints_func is not None:
+            raise NotImplementedError("constraints_func is not supported by AutoSampler.")
         if samplers is None:
+            if search_space is None:
+                raise ValueError("AutoSampler requires a search_space when using the default sampler schedule.")
             samplers = ((0, init_hebo), (100, init_nsgaii))
         self.sampler_indices = np.sort(np.array([x[0] for x in samplers], dtype=np.int32))
         self.samplers = [x[1] for x in sorted(samplers, key=lambda x: x[0])]
-        self.search_space = search_space
+        self.search_space = {} if search_space is None else dict(search_space)
         self._rng = LazyRandomState(seed)
         self._random_sampler = RandomSampler(seed=seed)
         self._thread_local_sampler = ThreadLocalSampler()
@@ -792,7 +832,7 @@ class AutoSampler(BaseSampler):
         complete_trials = study._get_trials(deepcopy=False, states=(TrialState.COMPLETE,), use_cache=True)
         self._completed_trials = max(self._completed_trials, len(complete_trials))
         new_index = (self._completed_trials >= self.sampler_indices).sum() - 1
-        if new_index == self._current_index:
+        if new_index == self._current_index or new_index < 0:
             return
         self._current_index = new_index
         self._sampler = self.samplers[new_index](
@@ -805,7 +845,7 @@ class AutoSampler(BaseSampler):
     def sample_relative(
         self, study: Study, trial: FrozenTrial, search_space: dict[str, BaseDistribution]
     ) -> dict[str, Any]:
-        return self._sampler.sample_relative(study, trial, self.search_space)
+        return self._sampler.sample_relative(study, trial, search_space or self.search_space)
 
     def sample_independent(
         self,
@@ -834,5 +874,6 @@ class AutoSampler(BaseSampler):
         state: TrialState,
         values: Sequence[float] | None,
     ) -> None:
-        assert state in [TrialState.COMPLETE, TrialState.FAIL, TrialState.PRUNED]
+        if state not in (TrialState.COMPLETE, TrialState.FAIL, TrialState.PRUNED):
+            raise ValueError(f"Unsupported trial state: {state}.")
         self._sampler.after_trial(study, trial, state, values)

{heavyball-2.1.3 → heavyball-2.1.4}/heavyball/utils.py

@@ -343,7 +343,8 @@ def set_(dst: Tensor, src: Tensor):
 
 
 def clean():
-    torch.cuda.empty_cache()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
     gc.collect()
 
 
@@ -470,7 +471,7 @@ def msign(G: torch.Tensor, steps: int = 10, eps: float = 1e-7) -> torch.Tensor:
 
     if should_transpose:
         x = x.mT
-    return x.float()
+    return x.to(G.dtype)
 
 
 ###### END
@@ -670,9 +671,9 @@ def get_orthogonal_matrix(mat, max_eps: float = 1e-3, min_eps: float = 1e-30):
             final.append(None)
             continue
 
+        device, dtype = m.device, m.dtype
         m = promote(m.data)
 
-        device, dtype = m.device, m.dtype
         eps = min_eps
         while True:
             try:
@@ -700,7 +701,6 @@ def get_orthogonal_matrix(mat, max_eps: float = 1e-3, min_eps: float = 1e-30):
                     raise
             clean()
 
-        eigvec = eigvec.to(device=m.device, dtype=m.dtype)
         eigvec = torch.flip(eigvec, [1])
         final.append(eigvec)
 
@@ -1053,13 +1053,15 @@ class StatefulOptimizer(torch.optim.Optimizer):
     def get_groups(self, group):
         return [group]
 
-    @functools.lru_cache(maxsize=None)
     def state_(self, arg: Tensor, fail: bool = True):
-        if not fail and arg not in self.mapping:
-            return {}
-        if _tensor_key(arg) not in self.mapping_inverse:
+        key = _tensor_key(arg)
+        if key not in self.mapping_inverse:
             self._init_mapping()
-        state_param, index = self.mapping_inverse[_tensor_key(arg)]
+        if key not in self.mapping_inverse:
+            if not fail:
+                return {}
+            raise KeyError("Tensor has no tracked state.")
+        state_param, index = self.mapping_inverse[key]
         if state_param not in self.state:
             self.state[state_param] = collections.defaultdict(dict)
         return self.state[state_param][index]
@@ -1147,7 +1149,7 @@ class StatefulOptimizer(torch.optim.Optimizer):
                 active_p = [p for p in group["params"]]
 
                 if not active_p:
-                    return
+                    continue
 
                 k = group["ema_step"] = group.get("ema_step", -1) + 1
 
@@ -1164,7 +1166,7 @@ class StatefulOptimizer(torch.optim.Optimizer):
                 active_p = [p for p in group["params"]]
 
                 if not active_p:
-                    return
+                    continue
 
                 for p in active_p:
                     if "param_ema" in self.state_(p):
@@ -1178,7 +1180,7 @@ class StatefulOptimizer(torch.optim.Optimizer):
                 active_p = [p for p in group["params"]]
 
                 if not active_p:
-                    return
+                    continue
 
                 for p in active_p:
                     if "param_ema" in self.state_(p):
@@ -1207,7 +1209,7 @@ class StatefulOptimizer(torch.optim.Optimizer):
         for group in self.param_groups:
             for p, g in self.split_p_and_g_in_group(group, skip_none=True, raw=True):
                 p.grad = grads.pop(0)
-                stochastic_add_(g, p.grad, -1)  # technically, we have to divide by the scale here
+                stochastic_add_divide_(g, p.grad, -1, torch.finfo(p.dtype).eps ** 0.5)
                 p.hessian_vector = g
                 p.data.copy_(p.orig)
                 del p.orig
@@ -1297,6 +1299,8 @@ class StatefulOptimizer(torch.optim.Optimizer):
             self._is_preconditioning = psgd_should_update(self.inner_group, self.precond_schedule, self.precond_rng)
         loss = self._handle_closure(closure)
 
+        if self.use_ema:
+            self.ema_update()
         # we assume that parameters are constant and that there are no excessive recompiles
         with torch.no_grad(), torch._dynamo.utils.disable_cache_limit():
             for group in self.param_groups:
@@ -1304,8 +1308,6 @@ class StatefulOptimizer(torch.optim.Optimizer):
                     group["param_count"] = sum(p.numel() for p in group["params"])
                 group["is_preconditioning"] = self._is_preconditioning
                 self._step(group)
-                if self.use_ema:
-                    self.ema_update()
                 for real, views in self.mapping.items():
                     for tensor in (real, *views):
                         for key in ("grad", "vector", "hessian_vector", "orig"):
@@ -1569,18 +1571,20 @@ def stochastic_round_list_(ref: List[Tensor], source: List[Tensor]):
 
 @decorator_knowngood
 def stochastic_round_(ref: Tensor, source: Tensor | None = None):
-    if source is not None:
-        if source.dtype == torch.bfloat16 or ref.dtype == source.dtype:
-            return source
-        if ref.dtype != torch.bfloat16:
-            return source.to(ref.dtype)
-    else:
+    if source is None:
         source = ref
-    source = source.float()
-    result = torch.randint_like(source, dtype=torch.int32, low=0, high=(1 << 16))
-    result.add_(source.view(dtype=torch.int32))
-    result.bitwise_and_(-65536)  # -65536 = FFFF0000 as a signed int32
-    return result.view(dtype=torch.float32).bfloat16()
+    if ref.dtype != torch.bfloat16:
+        return source.to(ref.dtype)
+    if source.dtype == torch.bfloat16:
+        return source
+    if source.dtype in (torch.float16, torch.float32, torch.float64):
+        source = source.to(torch.float32)
+        noise = torch.randint_like(source, dtype=torch.int32, low=0, high=(1 << 16))
+        bits = source.view(dtype=torch.int32)
+        bits.add_(noise)
+        bits.bitwise_and_(-65536)  # FFFF0000 mask, preserves sign+exp+7 mantissa bits
+        return bits.view(dtype=torch.float32).bfloat16()
+    return source.to(ref.dtype)
 
 
 @decorator_knowngood
@@ -1913,7 +1917,8 @@ def update_lra_precond_(
 
     # LU factorization to reuse computation
     try:
-        LU, pivots = torch.linalg.lu_factor(IpVtU)
+        lu_matrix = promote(IpVtU)  # operate in fp32 when inputs are bf16/half
+        LU, pivots = torch.linalg.lu_factor(lu_matrix)
     except RuntimeError:
         # Error:
         # U[2,2] is zero and using it on lu_solve would result in a division by zero.
@@ -1923,8 +1928,13 @@ def update_lra_precond_(
         # So, we skip this step and reattempt on the next one
         return U.to(U_orig[0].dtype), V.to(V_orig[0].dtype), d.to(d_orig[0].dtype)
 
-    invQtv = invQtv - V @ torch.linalg.lu_solve(LU, pivots, (U.T @ invQtv).view(-1, 1), adjoint=True).flatten()
-    invPv = U @ torch.linalg.lu_solve(LU, pivots, (V.T @ invQtv).view(-1, 1)).flatten()
+    solve_dtype = LU.dtype
+    rhs = (U.T @ invQtv).view(-1, 1).to(solve_dtype)
+    correction = torch.linalg.lu_solve(LU, pivots, rhs, adjoint=True).to(V.dtype)
+    invQtv = invQtv - (V @ correction).flatten()
+    rhs = (V.T @ invQtv).view(-1, 1).to(solve_dtype)
+    solution = torch.linalg.lu_solve(LU, pivots, rhs).to(U.dtype)
+    invPv = (U @ solution).flatten()
 
     eps, step = scalar_guard(eps, step, vector)
     _compilable_d_step(d, d_orig, invQtv, vector, invPv, hessian_vector, Ph, eps, step, delayed)
@@ -2044,7 +2054,10 @@ def extract_from_flat_update(params: List[Tensor], update: Tensor):
 @decorator_knowngood
 def flatten(x: List[Tensor], remaining: int = 0) -> Tensor:
     last_dim = x[0].shape[-remaining:] if remaining else []
-    return torch.cat([i.reshape(-1, *last_dim) for i in x if i.numel()], 0)
+    tensors = [i.reshape(-1, *last_dim) for i in x if i.numel()]
+    if not tensors:
+        return torch.zeros((), dtype=x[0].device, device=x[0].device)
+    return torch.cat(tensors, 0)
 
 
 @decorator_knowngood
@@ -2116,16 +2129,6 @@ def psgd_calc_A_and_conjB(G: Tensor, Q, conjB: Tensor | None):  # conjB ("V", "v
     return A, conjB
 
 
-@decorator_knowngood
-def _random_projection(x: Tensor, scale: Optional[Tensor]):
-    if scale is None:
-        scale = x.norm(float("inf")).clamp(min=1e-8)
-    k = 2 ** math.ceil(math.log2(math.log2(min(x.shape))))  # next-largest-power-of-2 of log2-of-size
-    norm = x.square().sum(0)
-    indices = torch.topk(norm, k, largest=True).indices
-    return x.index_select(1, indices).contiguous() / scale, scale
-
-
 def max_singular_value_exact(A, use_lobpcg: bool = False):
     try:
         if use_lobpcg:
@@ -2169,7 +2172,15 @@ def max_singular_value_cholesky(A: Tensor, max_abs: Optional[Tensor] = None):
     """
     Adapted from @evanatyourservice
     """
-    Y, max_abs = _random_projection(A, max_abs)
+    if max_abs is None:
+        max_abs = A.norm(float("inf")).clamp(min=1e-8)
+
+    # cholesky uses random projection, but this uses topk -- topk is a warm start, which may converge to a biased result
+    k = 2 ** math.ceil(math.log2(math.log2(min(A.shape))))  # next-largest-power-of-2 of log2-of-size
+    norm = A.square().sum(0)
+    indices = torch.topk(norm, k, largest=True).indices
+    Y = A.index_select(1, indices).contiguous() / max_abs
+
     Q = inplace_orthogonal_(Y, precise_zeroth_power_mode)
     Q = Q / max_abs
     Z = A.T @ Q
@@ -2557,7 +2568,7 @@ def _psgd_quad_preconditioner_grad(GG: List[Tensor], Q: List[Tensor], numel: int
         else:
             scale = gg.size(0) / numel
             gg = 2 * torch.eye(gg.size(0), device=gg.device, dtype=gg.dtype) - gg * scale
-            update = q - gg @ q @ gg
+            update = q - casted_einsum("ab,cd,bc", gg, gg, q)
             out.append(update + update.T)  # make matrix symmetric
     return out
 
@@ -3111,7 +3122,7 @@ def pointwise_lr_adaptation(
 ):
     grads, update, state, delta = list_guard(grads, update, state, delta)
     lr_lr = scalar_guard(lr_lr, grads[0])
-    _compilable_lr_adapt_(grads, update, state, delta, lr_lr)
+    _compilable_pointwise_lr_adapt_(grads, update, state, delta, lr_lr)
 
 
 def hook_optimizer_into_model(model, optimizer, *args, **kwargs):
@@ -3131,8 +3142,6 @@ def hook_optimizer_into_model(model, optimizer, *args, **kwargs):
 
 def fused_hook(parameters, optimizer, *args, **kwargs):
     parameters = list(parameters)
-    param_count = len(parameters)
-    seen_params = set()
 
     o = optimizer(parameters, *args, **kwargs)
     step_fn = o.step
@@ -3141,12 +3150,8 @@ def fused_hook(parameters, optimizer, *args, **kwargs):
     )
 
     def _step(p: Tensor):
-        seen_params.add(p)
-
-        if len(seen_params) < param_count:
-            step_fn()
-            o.zero_grad()
-            seen_params.clear()
+        step_fn()
+        o.zero_grad()
 
     for p in parameters:
         p.register_post_accumulate_grad_hook(_step)
@@ -3171,6 +3176,8 @@ def sam_step(parameters, ball_size, adaptive: bool = True):
     old_params = []
     for p in parameters:
         old_params.append(p.detach().clone())
+        if not hasattr_none(p, "grad"):
+            continue
         grad = promote(p.grad)
         if adaptive:
             grad = grad * promote(p).square()

{heavyball-2.1.3 → heavyball-2.1.4}/heavyball.egg-info/PKG-INFO

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

{heavyball-2.1.3 → heavyball-2.1.4}/heavyball.egg-info/SOURCES.txt

@@ -36,6 +36,7 @@ test/test_save_restore.py
 test/test_singular_values.py
 test/test_soap.py
 test/test_stochastic_updates.py
+test/test_stochastic_utils_cpu.py
 test/test_toy_training.py
 test/test_utils_cpu.py
 test/test_utils_property.py

{heavyball-2.1.3 → heavyball-2.1.4}/pyproject.toml

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

{heavyball-2.1.3 → heavyball-2.1.4}/test/test_bf16_params.py

@@ -14,6 +14,11 @@ os.environ["TORCH_LOGS"] = "+recompiles"
 
 config.cache_size_limit = 128
 
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="CUDA is required to run bf16 foreach parameter tests.",
+)
+
 
 @pytest.mark.parametrize("opt", heavyball.__all__)
 @pytest.mark.parametrize("size,depth", [(256, 1)])

{heavyball-2.1.3 → heavyball-2.1.4}/test/test_foreach.py

@@ -1,3 +1,5 @@
+import os
+
 import pytest
 import torch
 from lightbench.utils import get_optim
@@ -15,13 +17,40 @@ def get_memory():
     return torch.cuda.memory_allocated()
 
 
+def _read_int(name: str, default: int, *, minimum: int) -> int:
+    raw = os.environ.get(name)
+    if raw is None:
+        return default
+    try:
+        return max(minimum, int(raw))
+    except ValueError:
+        return default
+
+
+DEFAULT_SIZE = _read_int("HB_FOREACH_TEST_SIZE", 128, minimum=1)
+DEFAULT_DEPTH = _read_int("HB_FOREACH_TEST_DEPTH", 16, minimum=1)
+DEFAULT_ITERATIONS = _read_int("HB_FOREACH_TEST_ITERATIONS", 64, minimum=1)
+DEFAULT_OUTER = _read_int("HB_FOREACH_TEST_OUTER", 1, minimum=1)
+DEFAULT_WARMUP = _read_int("HB_FOREACH_TEST_WARMUP", 1, minimum=0)
+
+
 @pytest.mark.parametrize("opt", heavyball.__all__)
-@pytest.mark.parametrize("size,depth", [(256, 128)])
-def test_foreach(opt, size, depth: int, iterations: int = 4096, outer_iterations: int = 2):
+@pytest.mark.parametrize("size,depth", [(DEFAULT_SIZE, DEFAULT_DEPTH)])
+def test_foreach(
+    opt,
+    size,
+    depth: int,
+    iterations: int = DEFAULT_ITERATIONS,
+    outer_iterations: int = DEFAULT_OUTER,
+    warmup_runs: int = DEFAULT_WARMUP,
+):
     set_torch()
 
     opt = getattr(heavyball, opt)
 
+    total_runs = warmup_runs + outer_iterations
+    assert total_runs >= 1
+
     peaks = []
     losses = []
 
@@ -30,7 +59,7 @@ def test_foreach(opt, size, depth: int, iterations: int = 4096, outer_iterations
         peaks.append([])
         losses.append([])
 
-        for i in range(outer_iterations):
+        for i in range(total_runs):
             clean()
             model = nn.Sequential(*[nn.Linear(size, size) for _ in range(depth)]).cuda()
             clean()
@@ -56,8 +85,14 @@ def test_foreach(opt, size, depth: int, iterations: int = 4096, outer_iterations
 
             peak = torch.cuda.memory_stats()["allocated_bytes.all.peak"]
 
-            if i > 0:
-                peaks[-1].append(peak)
+            if i < warmup_runs:
+                continue
+
+            peaks[-1].append(peak)
+
+    if warmup_runs:
+        cutoff = warmup_runs * iterations
+        losses = [loss_list[cutoff:] for loss_list in losses]
 
     for p0, p1 in zip(*peaks):
         assert p0 > p1

{heavyball-2.1.3 → heavyball-2.1.4}/test/test_memory.py

@@ -19,7 +19,6 @@ expected_memory = {
     "adamw": {"after": 4, "peak": 5.1},
     "soap": {"after": 7, "peak": 14},
     "psgd": {"after": 4, "peak": 11.5},
-    "padam": {"after": 5, "peak": 11.4},
 }
 
 

heavyball-2.1.4/test/test_stochastic_utils_cpu.py

@@ -0,0 +1,49 @@
+import os
+
+os.environ.setdefault("TORCH_COMPILE_DISABLE", "1")
+
+import torch
+
+from heavyball.utils import copy_stochastic_, stochastic_add_, stochastic_divide_with_eps_
+
+
+def _average_stochastic_round(source: torch.Tensor, trials: int = 512) -> torch.Tensor:
+    dest = torch.empty_like(source, dtype=torch.bfloat16)
+    totals = torch.zeros_like(source, dtype=torch.float64)
+    for _ in range(trials):
+        copy_stochastic_(dest, source)
+        totals += dest.to(dtype=torch.float64)
+    return totals / trials
+
+
+def test_copy_stochastic_round_is_close_to_source_mean():
+    torch.manual_seed(0x5566AA)
+    values = torch.randn(2048, dtype=torch.float32) * 3.0
+    averaged = _average_stochastic_round(values, trials=256)
+    delta = averaged - values.double()
+
+    # Stochastic round should stay close to the original float32 values.
+    assert delta.abs().mean().item() < 5e-3
+    assert delta.abs().max().item() < 2.5e-2
+
+
+def test_stochastic_add_broadcasts_partner_lists():
+    torch.manual_seed(0x172893)
+    targets = [torch.zeros(4, dtype=torch.bfloat16) for _ in range(2)]
+    partner = [torch.linspace(-1.0, 1.0, 4, dtype=torch.float32)]
+
+    stochastic_add_(targets, partner, alpha=0.25)
+    expected = partner[0] * 0.25
+    for tensor in targets:
+        assert torch.allclose(tensor.float(), expected, atol=5e-3, rtol=0)
+
+
+def test_stochastic_divide_with_eps_matches_float_result():
+    torch.manual_seed(0xABCDEF)
+    numerator = torch.randn(32, dtype=torch.bfloat16)
+    denominator = torch.rand(32, dtype=torch.bfloat16) + 0.05
+    result = numerator.clone()
+
+    stochastic_divide_with_eps_(result, denominator, eps=1e-3)
+    expected = numerator.float() / (denominator.float() + 1e-3)
+    assert torch.allclose(result.float(), expected, atol=2e-2, rtol=2e-2)

{heavyball-2.1.3 → heavyball-2.1.4}/test/test_utils_cpu.py

@@ -3,7 +3,6 @@ import random
 import warnings
 from copy import deepcopy
 
-import numpy as np
 import pytest
 import torch
 from torch import Tensor, nn
@@ -222,29 +221,31 @@ def test_stochastic_math_helpers_match_expected_results():
     assert torch.allclose(a, torch.full_like(a, expected), atol=1e-6)
 
 
-def test_stochastic_math_accuracy(steps: int = 100, items: int = 32, target_shift: float = 1.0):
-    """
-    TODO: Rework this test or stochastic rounding.
-    With target_shift=1, it passes. With target_shift != 1, it does not pass -- this is unexpected
-    """
+def test_stochastic_math_accuracy():
     torch.manual_seed(0x172893)
-    rng = np.random.default_rng(0x213112)
+    items = 8
+    steps = 2048
+    increments = torch.full((items,), 1e-3, dtype=torch.float32)
 
-    accum_baseline = torch.zeros(items, dtype=torch.bfloat16) + target_shift
-    accum_stochastic = torch.zeros(items, dtype=torch.bfloat16) + target_shift
-    accum_groundtruth = torch.zeros(items, dtype=torch.float64) + target_shift
+    baseline = torch.zeros(items, dtype=torch.bfloat16)
+    stochastic = torch.zeros(items, dtype=torch.bfloat16)
+    ground_truth = torch.zeros(items, dtype=torch.float64)
 
-    add = 1 / (1 + 2 * torch.arange(items, dtype=torch.float32))
-    alphas = np.exp(-2 - 2 * rng.random((steps // 2,)))
-    for _ in range(2):
-        for alpha in alphas:
-            accum_baseline.add_(add, alpha=alpha)
-            accum_groundtruth.add_(add, alpha=alpha)
-            stochastic_add_(accum_stochastic, add, alpha=alpha)
-        assert (accum_baseline.double() - accum_groundtruth).norm().item() > (
-            accum_stochastic.double() - accum_groundtruth
-        ).norm().item()
-        alphas = -alphas
+    for _ in range(steps):
+        baseline.add_(increments)
+        ground_truth.add_(increments)
+        stochastic_add_(stochastic, increments)
+
+    baseline_error = torch.abs(baseline.float() - ground_truth.float()).mean().item()
+    stochastic_error = torch.abs(stochastic.float() - ground_truth.float()).mean().item()
+
+    assert baseline_error > 1.0
+    assert stochastic_error < 0.2
+    assert stochastic_error < baseline_error * 0.2
+
+    baseline_bias = abs(baseline.float().mean().item() - ground_truth.float().mean().item())
+    stochastic_bias = abs(stochastic.float().mean().item() - ground_truth.float().mean().item())
+    assert stochastic_bias < baseline_bias
 
 
 def test_disable_caution_scaling_toggles_behavior():