heavyball 1.4.0__py3-none-any.whl → 1.4.3__py3-none-any.whl

heavyball/chainable.py

@@ -477,7 +477,7 @@ class ChainOpt(utils.StatefulOptimizer):
         else:
             chain(self.state_, group, g, p, *self.fns)
 
-        group['lr'] = None
+        group['lr'] = group['prev_lr']
         group['step'] = None
 
 

heavyball/utils.py

@@ -195,10 +195,8 @@ def scale_by_exp_avg_sq_(exp_avg_sq, grad, beta2, eps):
 
 @decorator_knowngood
 def _compilable_exp_avg_(state, grad, beta):
-    for s, g in zip(state, grad):
-        lerped = s.lerp(g, 1 - beta)
-        copy_stochastic_(s, lerped)
-        copy_stochastic_(g, lerped)
+    lerped = _lerp32(state, grad, beta)
+    copy_stochastic_list_(grad, lerped)
 
 
 def scale_by_exp_avg_(state, grad, beta):
@@ -746,6 +744,7 @@ def copy_stochastic_list_(target: List[Tensor], source: List[Tensor]):
 def _lerp32(state: List[Tensor], grad: List[Tensor], beta):
     ea32 = list(map(promote, state))
     grad = list(map(promote, grad))
+    beta = promote(beta)
 
     ea32 = [e.lerp(g, 1 - beta) for e, g in zip(ea32, grad)]
     copy_stochastic_list_(state, ea32)
@@ -1032,26 +1031,19 @@ def psgd_balance_Q(Q_in):
 
 
 def psgd_calc_A_and_conjB(exprA, G, Q):
-    V = torch.randn(G.shape, dtype=G.dtype, device=G.device)
     eps = scalar_guard(math.sqrt(torch.finfo(torch.float32).eps), G)
     eps *= G.norm() / G.numel()
-    G = G + V * eps
+    G = G + torch.randn_like(G) * eps
     md = min_dtype(Q + [G])
     A = torch.einsum(exprA, *[q.to(md) for q in Q], G.to(md)).to(G.dtype)
     order = G.dim()
-    p = list(range(order))
-    conjB = torch.permute(V, p[1:] + p[:1]).to(promote(G.dtype))
+    conjB = torch.randn(G.shape[1:] + G.shape[:1], dtype=promote(G.dtype), device=G.device)
     Q = [promote(q) for q in Q]
     for i, q in enumerate(Q):
         if q.dim() <= 1:
             conjB /= q
         else:
-            unsqueeze = conjB.dim() <= 1
-            if unsqueeze:
-                conjB = conjB.unsqueeze(0)
-            conjB = torch.linalg.solve_triangular(q, conjB, upper=True, left=False)
-            if unsqueeze:
-                conjB = conjB.squeeze(0)
+            conjB = torch.linalg.solve_triangular(q, conjB.reshape(-1, q.size(0)), upper=True, left=False).reshape_as(conjB)
         if i < order - 1:
             conjB = torch.transpose(conjB, i, order - 1)
     return A, conjB
@@ -1092,7 +1084,7 @@ def psgd_update_precond(Q, exprs, G, precond_lr, oq, store_triu_as_line):
         else:
             torch.triu(term1, out=term1)
             term1 /= torch.where(norm > 0, psgd_lb(term2, norm), norm).clamp_(tiny_bf16)
-            term1 = torch.mm(term1, q)
+            term1 = torch.mm(term1, q.to(term1.dtype))
         if store_triu_as_line:
             term1 = triu_to_line([term1])[0][1]
             o = o[1]

heavyball-1.4.3.dist-info/METADATA

@@ -0,0 +1,934 @@
+Metadata-Version: 2.1
+Name: heavyball
+Version: 1.4.3
+Summary: Efficient optimizers
+Home-page: https://github.com/clashluke/heavyball
+Author: Lucas Nestler
+Author-email: github.heavyball@nestler.sh
+License: BSD
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: License :: OSI Approved :: BSD License
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3.7
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Topic :: Software Development :: Libraries
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Classifier: Intended Audience :: Developers
+Requires-Python: >=3.7
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: opt-einsum
+Requires-Dist: torch
+Requires-Dist: numpy
+
+# `heavyball`: Efficient Optimizers
+
+* [Public API](#Public-API)
+    - [Foreach Optimizers](#Foreach-Optimizers)
+    - [`heavyball.utils`](#heavyball.utils)
+    - [Example Usage](#Example-Usage)
+
+* [`heavyball.chainable`](##heavyball.chainable)
+    - [Core Concept](#Core-Concept)
+    - [`FunctionTransform` and Guards](#FunctionTransform-and-Guards)
+    - [Chaining Transformations](#Chaining-Transformations)
+    - [Building Optimizers](#Building-Optimizers)
+    - [Creating New Transformations](#Creating-New-Transformations)
+
+* [Optimizer Recommendations](#Optimizer-Recommendations)
+    - [Choosing the Right Optimizer](#Choosing-the-Right-Optimizer)
+
+---
+
+The `heavyball` library provides a collection of efficient optimizers designed for deep learning. It leverages
+techniques like preconditioning, momentum, and adaptive learning rates to accelerate training and improve convergence.
+The library's core strength lies in its `chainable` API, which allows for flexible composition of optimizers, enabling
+users to build custom optimization strategies.
+
+## Public API
+
+The `heavyball` library exposes the following optimizers through its main namespace:
+
+### Foreach Optimizers
+
+These optimizers are designed to be efficient by operating on batches of parameters simultaneously using `foreach`
+operations whenever possible.
+
+#### `ForeachAdamW`
+
+```python
+class ForeachAdamW(C.BaseOpt):
+    def __init__(self, params, lr=0.0025, betas=(0.9, 0.99), eps=1e-8, weight_decay=0, warmup_steps=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):
+# ...
+```
+
+A foreach implementation of the AdamW optimizer. It incorporates weight decay into the update rule and uses adaptive
+learning rates based on the first and second moments of the gradients.
+
+**Key Parameters:**
+
+* **`lr`**: Learning rate.
+* **`betas`**: Coefficients used for computing running averages of the gradient and its square.
+* **`eps`**: A small constant for numerical stability.
+* **`weight_decay`**: Weight decay coefficient.
+* **`warmup_steps`**: Number of steps for linear learning rate warmup.
+* **`foreach`**: Enables/disables the use of `foreach` operations.
+* **`storage_dtype`**: The floating-point type to be used for internal state. `"float32"` or `"bfloat16"`.
+* **`mars`**: Enables/disables Mars correction.
+* **`caution`**: Enables/disables the use of a cautious update rule, avoiding updates that point in the opposite
+  direction to the gradients.
+* **`mars_gamma`**: Mars correction coefficient.
+* **`gradient_clipping`**: Gradient clipping function or method. See `heavyball.utils` for available options.
+* **`update_clipping`**: Update clipping function or method. See `heavyball.utils` for available options.
+* **`palm`**: Enables/disables PaLM's beta2 schedule.
+* **`beta2_scale`**: if we're using the PaLM schedule, `beta2 = step ** -beta2_scale`
+
+#### `ForeachRMSprop`
+
+```python
+class ForeachRMSprop(C.BaseOpt):
+    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):
+# ...
+```
+
+A foreach implementation of a debiased RMSprop optimizer (Note: this is different from `torch.optim.RMSprop`). It uses
+adaptive learning rates based on the second moment of the gradients.
+
+**Key Parameters:**
+
+* **`lr`**: Learning rate.
+* **`betas`**: Coefficients used for computing running averages of the squared gradient.
+* **`eps`**: A small constant for numerical stability.
+* **`weight_decay`**: Weight decay coefficient.
+* **`warmup_steps`**: Number of steps for linear learning rate warmup.
+* **`r`**: Schedule-Free coefficient that controls dependence of the learning rate on step count.
+* **`weight_lr_power`**: Schedule-Free coefficient that controls the sensitivity of `r` to the learning rate.
+* **`foreach`**: Enables/disables the use of `foreach` operations.
+* **`storage_dtype`**: The floating-point type to be used for internal state. `"float32"` or `"bfloat16"`.
+* **`mars`**: Enables/disables Mars correction.
+* **`caution`**: Enables/disables the use of a cautious update rule, avoiding updates that point in the opposite
+  direction to the gradients.
+* **`mars_gamma`**: Mars correction coefficient.
+* **`gradient_clipping`**: Gradient clipping function or method. See `heavyball.utils` for available options.
+* **`update_clipping`**: Update clipping function or method. See `heavyball.utils` for available options.
+* **`palm`**: Enables/disables PaLM's beta2 schedule.
+* **`beta2_scale`**: if we're using the PaLM schedule, `beta2 = step ** -beta2_scale`
+
+#### `ForeachSFAdamW`
+
+```python
+class ForeachSFAdamW(C.ScheduleFree):
+    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):
+# ...
+```
+
+A foreach implementation of the Schedule-Free AdamW optimizer. It combines the benefits of AdamW with the Schedule-Free
+approach, which dynamically adjusts the learning rate based on the current state of optimization.
+
+**Key Parameters:**
+
+* **`lr`**: Base learning rate. The effective learning rate at each step depends on `lr`, `r`, and `weight_lr_power`.
+* **`betas`**: Coefficients used for computing running averages of the gradient and its square.
+* **`eps`**: A small constant for numerical stability.
+* **`weight_decay`**: Weight decay coefficient.
+* **`warmup_steps`**: Number of steps for linear learning rate warmup.
+* **`r`**: Schedule-Free coefficient that controls dependence of the learning rate on step count.
+* **`weight_lr_power`**: Schedule-Free coefficient that controls the sensitivity of `r` to the learning rate.
+* **`foreach`**: Enables/disables the use of `foreach` operations.
+* **`storage_dtype`**: The floating-point type to be used for internal state. `"float32"` or `"bfloat16"`.
+* **`mars`**: Enables/disables Mars correction.
+* **`caution`**: Enables/disables the use of a cautious update rule, avoiding updates that point in the opposite
+  direction to the gradients.
+* **`mars_gamma`**: Mars correction coefficient.
+* **`gradient_clipping`**: Gradient clipping function or method. See `heavyball.utils` for available options.
+* **`update_clipping`**: Update clipping function or method. See `heavyball.utils` for available options.
+* **`palm`**: Enables/disables PaLM's beta2 schedule.
+* **`beta2_scale`**: if we're using the PaLM schedule, `beta2 = step ** -beta2_scale`
+
+#### `PaLMForeachSFAdamW`
+
+```python
+class PaLMForeachSFAdamW(ForeachSFAdamW):
+    palm: bool = True
+```
+
+A specialized version of `ForeachSFAdamW` with PaLM's beta2 schedule enabled by default.
+
+#### `ForeachADOPT`
+
+```python
+class ForeachADOPT(C.BaseOpt):
+    def __init__(self, params, lr=0.0025, betas=(0.9, 0.99), eps=1e-8, weight_decay=0, warmup_steps=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):
+# ...
+```
+
+A foreach implementation of the ADOPT optimizer, which uses a debiased estimate of the second moment of the gradients.
+
+**Key Parameters:**
+
+* **`lr`**: Learning rate.
+* **`betas`**: Coefficients used for computing running averages of the gradient and its square.
+* **`eps`**: A small constant for numerical stability.
+* **`weight_decay`**: Weight decay coefficient.
+* **`warmup_steps`**: Number of steps for linear learning rate warmup.
+* **`foreach`**: Enables/disables the use of `foreach` operations.
+* **`storage_dtype`**: The floating-point type to be used for internal state. `"float32"` or `"bfloat16"`.
+* **`mars`**: Enables/disables Mars correction.
+* **`caution`**: Enables/disables the use of a cautious update rule, avoiding updates that point in the opposite
+  direction to the gradients.
+* **`mars_gamma`**: Mars correction coefficient.
+* **`gradient_clipping`**: Gradient clipping function or method. See `heavyball.utils` for available options.
+* **`update_clipping`**: Update clipping function or method. See `heavyball.utils` for available options.
+* **`palm`**: Enables/disables PaLM's beta2 schedule.
+* **`beta2_scale`**: if we're using the PaLM schedule, `beta2 = step ** -beta2_scale`
+
+#### `ForeachMuon`
+
+```python
+class ForeachMuon(C.BaseOpt):
+    def __init__(self, params, lr=0.0025, betas=(0.9, 0.99), eps=1e-8, weight_decay=0, warmup_steps=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,
+                 nesterov: bool = True):
+# ...
+```
+
+A foreach implementation of the Muon optimizer, incorporating orthogonal updates via the `orthogonalize_update`
+transformation.
+
+**Key Parameters:**
+
+* **`lr`**: Learning rate.
+* **`betas`**: Coefficients used for computing running averages of the gradient and its square.
+* **`eps`**: A small constant for numerical stability.
+* **`weight_decay`**: Weight decay coefficient.
+* **`warmup_steps`**: Number of steps for linear learning rate warmup.
+* **`foreach`**: Enables/disables the use of `foreach` operations.
+* **`storage_dtype`**: The floating-point type to be used for internal state. `"float32"` or `"bfloat16"`.
+* **`mars`**: Enables/disables Mars correction.
+* **`caution`**: Enables/disables the use of a cautious update rule, avoiding updates that point in the opposite
+  direction to the gradients.
+* **`mars_gamma`**: Mars correction coefficient.
+* **`gradient_clipping`**: Gradient clipping function or method. See `heavyball.utils` for available options.
+* **`update_clipping`**: Update clipping function or method. See `heavyball.utils` for available options.
+* **`palm`**: Enables/disables PaLM's beta2 schedule.
+* **`beta2_scale`**: if we're using the PaLM schedule, `beta2 = step ** -beta2_scale`
+* **`nesterov`**: Enables/disables Nesterov momentum.
+
+#### `ForeachLaProp`
+
+```python
+class ForeachLaProp(C.BaseOpt):
+    def __init__(self, params, lr=0.0025, betas=(0.9, 0.99), eps=1e-8, weight_decay=0, warmup_steps=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):
+# ...
+```
+
+A foreach implementation of the LaProp optimizer.
+
+**Key Parameters:**
+
+* **`lr`**: Learning rate.
+* **`betas`**: Coefficients used for computing running averages of the gradient and its square.
+* **`eps`**: A small constant for numerical stability.
+* **`weight_decay`**: Weight decay coefficient.
+* **`warmup_steps`**: Number of steps for linear learning rate warmup.
+* **`foreach`**: Enables/disables the use of `foreach` operations.
+* **`storage_dtype`**: The floating-point type to be used for internal state. `"float32"` or `"bfloat16"`.
+* **`mars`**: Enables/disables Mars correction.
+* **`caution`**: Enables/disables the use of a cautious update rule, avoiding updates that point in the opposite
+  direction to the gradients.
+* **`mars_gamma`**: Mars correction coefficient.
+* **`gradient_clipping`**: Gradient clipping function or method. See `heavyball.utils` for available options.
+* **`update_clipping`**: Update clipping function or method. See `heavyball.utils` for available options.
+* **`palm`**: Enables/disables PaLM's beta2 schedule.
+* **`beta2_scale`**: if we're using the PaLM schedule, `beta2 = step ** -beta2_scale`
+
+#### `MuonLaProp`
+
+```python
+class MuonLaProp(C.BaseOpt):
+    def __init__(self, params, lr=0.0025, betas=(0.9, 0.99), eps=1e-8, weight_decay=0, warmup_steps=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):
+# ...
+```
+
+A variant of LaProp that incorporates orthogonal updates via the `orthogonalize_update` transformation.
+
+**Key Parameters:**
+
+* **`lr`**: Learning rate.
+* **`betas`**: Coefficients used for computing running averages of the gradient and its square.
+* **`eps`**: A small constant for numerical stability.
+* **`weight_decay`**: Weight decay coefficient.
+* **`warmup_steps`**: Number of steps for linear learning rate warmup.
+* **`foreach`**: Enables/disables the use of `foreach` operations.
+* **`storage_dtype`**: The floating-point type to be used for internal state. `"float32"` or `"bfloat16"`.
+* **`mars`**: Enables/disables Mars correction.
+* **`caution`**: Enables/disables the use of a cautious update rule, avoiding updates that point in the opposite
+  direction to the gradients.
+* **`mars_gamma`**: Mars correction coefficient.
+* **`gradient_clipping`**: Gradient clipping function or method. See `heavyball.utils` for available options.
+* **`update_clipping`**: Update clipping function or method. See `heavyball.utils` for available options.
+* **`palm`**: Enables/disables PaLM's beta2 schedule.
+* **`beta2_scale`**: if we're using the PaLM schedule, `beta2 = step ** -beta2_scale`
+
+#### `ForeachSOAP`
+
+```python
+class ForeachSOAP(C.BaseOpt):
+    use_precond_schedule: bool = False
+
+    def __init__(self, params, lr: float = 3e-3, betas=(0.9, 0.95), shampoo_beta: float = 0.95, eps: float = 1e-8,
+                 weight_decay: float = 0.01, precondition_frequency: int = 2, max_precond_dim: int = 2048,  #
+                 merge_dims: bool = True, precondition_1d: bool = False, normalize_grads: bool = False,
+                 data_format: str = "channels_first", correct_bias: bool = True, warmup_steps: int = 1,
+                 split: bool = False, foreach: bool = True, mars: bool = False, caution: bool = False,
+                 mars_gamma: float = 0.0025, palm: bool = C.use_default, precond_scheduler=(1 / 3, 9),
+                 beta2_scale: float = 0.8, use_precond_schedule: bool = C.use_default,
+                 gradient_clipping: C.str_or_fn = C.use_default, update_clipping: C.str_or_fn = C.use_default):
+# ...
+```
+
+A foreach implementation of the SOAP (Second-Order Adaptive Preconditioner) optimizer. It uses a preconditioner based on
+the second-order statistics of the gradients to accelerate convergence.
+
+**Key Parameters:**
+
+* **`lr`**: Learning rate.
+* **`betas`**: Coefficients used for computing running averages of the gradient.
+* **`shampoo_beta`**: Coefficient used for computing running average of the preconditioner.
+* **`eps`**: A small constant for numerical stability.
+* **`weight_decay`**: Weight decay coefficient.
+* **`precondition_frequency`**: Frequency of preconditioner updates. If using `use_precond_schedule`, this parameter is
+  ignored.
+* **`max_precond_dim`**: Maximum dimension of the preconditioner.
+* **`merge_dims`**: Whether to merge dimensions when forming the preconditioner.
+* **`precondition_1d`**: Whether to use a 1D preconditioner for 1D parameters.
+* **`normalize_grads`**: Whether to normalize gradients before applying SOAP.
+* **`data_format`**: `"channels_first"` or `"channels_last"`. Specifies the data format of the input tensors.
+* **`correct_bias`**: Enables/disables bias correction for the running averages.
+* **`warmup_steps`**: Number of steps for linear learning rate warmup.
+* **`split`**: Whether to split large dimensions when forming the preconditioner.
+* **`foreach`**: Enables/disables the use of `foreach` operations.
+* **`mars`**: Enables/disables Mars correction.
+* **`caution`**: Enables/disables the use of a cautious update rule, avoiding updates that point in the opposite
+  direction to the gradients.
+* **`mars_gamma`**: Mars correction coefficient.
+* **`palm`**: Enables/disables PaLM's beta2 schedule.
+* **`precond_scheduler`**: A tuple `(power, log_base)` specifying the preconditioner update schedule, where the update
+  probability is `1 / (step ** power * log_base)`. This parameter is only used if `use_precond_schedule` is `True`.
+* **`beta2_scale`**: if we're using the PaLM schedule, `beta2 = step ** -beta2_scale`
+* **`use_precond_schedule`**: Whether to use a dynamic preconditioner update schedule instead of a fixed frequency.
+* **`gradient_clipping`**: Gradient clipping function or method. See `heavyball.utils` for available options.
+* **`update_clipping`**: Update clipping function or method. See `heavyball.utils` for available options.
+
+#### `PaLMForeachSOAP`
+
+```python
+class PaLMForeachSOAP(ForeachSOAP):
+    use_precond_schedule: bool = False
+    palm: bool = True
+```
+
+A specialized version of `ForeachSOAP` with PaLM's beta2 schedule enabled by default.
+
+#### `PrecondScheduleForeachSOAP`
+
+```python
+class PrecondScheduleForeachSOAP(ForeachSOAP):
+    use_precond_schedule: bool = True
+```
+
+A specialized version of `ForeachSOAP` that uses a dynamic preconditioner update schedule.
+
+#### `PrecondSchedulePaLMForeachSOAP`
+
+```python
+class PrecondSchedulePaLMForeachSOAP(ForeachSOAP):
+    use_precond_schedule: bool = True
+    palm: bool = True
+```
+
+A specialized version of `ForeachSOAP` with both PaLM-specific modifications and a dynamic preconditioner update
+schedule enabled by default.
+
+#### `ForeachPSGDKron`
+
+```python
+class ForeachPSGDKron(C.BaseOpt):
+    delayed: bool = False
+    cached: bool = False
+    exp_avg_input: bool = True
+
+    def __init__(self, params, lr=0.001, beta=0.9, weight_decay=0.0, preconditioner_update_probability=None,
+                 max_size_triangular=2048, min_ndim_triangular=2, memory_save_mode=None,
+                 momentum_into_precond_update=True, warmup_steps: int = 1, merge_dims: bool = False,
+                 split: bool = False, store_triu_as_line: bool = True, foreach: bool = True, q_dtype='float32',
+                 stochastic_schedule: bool = True, storage_dtype: str = 'float32', mars: bool = False,
+                 caution: bool = False, mars_gamma: float = 0.0025, delayed: Optional[bool] = C.use_default,
+                 cached: Optional[bool] = C.use_default, 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,  #
+                 # expert parameters
+                 precond_init_scale=1.0, precond_lr=0.1):
+# ...
+```
+
+A foreach implementation of the PSGD (Preconditioned Stochastic Gradient Descent) optimizer with Kronecker-factored
+preconditioners.
+
+**Key Parameters:**
+
+* **`lr`**: Learning rate.
+* **`beta`**: Coefficient used for computing running average of the gradient.
+* **`weight_decay`**: Weight decay coefficient.
+* **`preconditioner_update_probability`**: Probability of updating the preconditioner at each step. If `None`, a default
+  schedule is used.
+* **`max_size_triangular`**: Maximum size of triangular matrices used in the preconditioner.
+* **`min_ndim_triangular`**: Minimum number of dimensions for a tensor to be considered for triangular preconditioner.
+* **`memory_save_mode`**: Memory saving mode for the preconditioner. Can be `None`, `"one_diag"`, or `"all_diag"`.
+* **`momentum_into_precond_update`**: Whether to use momentum in the preconditioner update.
+* **`warmup_steps`**: Number of steps for linear learning rate warmup.
+* **`merge_dims`**: Whether to merge dimensions when forming the preconditioner.
+* **`split`**: Whether to split large dimensions when forming the preconditioner.
+* **`store_triu_as_line`**: Whether to store the upper triangular part of the preconditioner as a 1D vector.
+* **`foreach`**: Enables/disables the use of `foreach` operations.
+* **`q_dtype`**: The floating-point type to be used for the preconditioner. `"float32"` or `"bfloat16"`.
+* **`stochastic_schedule`**: Whether to use a stochastic schedule for updating the preconditioner.
+* **`storage_dtype`**: The floating-point type to be used for internal state. `"float32"` or `"bfloat16"`.
+* **`mars`**: Enables/disables Mars correction.
+* **`caution`**: Enables/disables the use of a cautious update rule, avoiding updates that point in the opposite
+  direction to the gradients.
+* **`mars_gamma`**: Mars correction coefficient.
+* **`delayed`**: Enables/disables delayed preconditioner updates.
+* **`cached`**: Enables/disables caching of preconditioner-related computations.
+* **`exp_avg_input`**: Whether to apply `exp_avg` to the input before calculating the preconditioner.
+* **`gradient_clipping`**: Gradient clipping function or method. See `heavyball.utils` for available options.
+* **`update_clipping`**: Update clipping function or method. See `heavyball.utils` for available options.
+* **`precond_init_scale`**: Initial scale of the preconditioner.
+* **`precond_lr`**: Learning rate for preconditioner updates.
+
+#### `ForeachPurePSGD`
+
+```python
+class ForeachPurePSGD(ForeachPSGDKron):
+    exp_avg_input: bool = False
+```
+
+A specialized version of `ForeachPSGDKron` that does not apply `exp_avg` to the input before calculating the
+preconditioner.
+
+#### `ForeachCachedDelayedPSGDKron`
+
+```python
+class ForeachCachedDelayedPSGDKron(ForeachPSGDKron):
+    delayed: bool = True
+    cached: bool = True
+```
+
+A specialized version of `ForeachPSGDKron` with both delayed preconditioner updates and caching enabled by default.
+
+#### `ForeachCachedPSGDKron`
+
+```python
+class ForeachCachedPSGDKron(ForeachPSGDKron):
+    cached: bool = True
+```
+
+A specialized version of `ForeachPSGDKron` with caching enabled by default.
+
+#### `ForeachDelayedPSGD`
+
+```python
+class ForeachDelayedPSGD(ForeachPSGDKron):
+    delayed: bool = True
+```
+
+A specialized version of `ForeachPSGDKron` with delayed preconditioner updates enabled by default.
+
+## `heavyball.utils`
+
+The `heavyball.utils` module provides several important functions and settings that users may find useful:
+
+### Settings
+
+* **`compile_mode`**:  (defaults to `"max-autotune-no-cudagraphs"`) Controls the compilation mode used by
+  `torch.compile`. Setting this to `"default"` or `"max-autotune-no-cudagraphs"` improves performance at the cost of
+  increasd compile time. Setting it to `None` disables compilation.
+* **`dynamic`**: (defaults to `False`) Enables/disables dynamic shapes during compilation. Enabling this reduces
+  compilation time but may lead to slower execution.
+* **`zeroth_power_mode`**: (defaults to `"qr"`) Controls the method used for computing the zeroth power of a matrix (
+  orthogonalization) in certain preconditioners. Options include:
+    * `"qr"`: Uses QR decomposition.
+    * `"svd"`: Uses singular value decomposition.
+    * `"newtonschulz"`: Uses Newton-Schulz iteration.
+
+### Gradient/Update Clipping
+
+The following functions are used for gradient and update clipping. They can be passed to the `gradient_clipping` or
+`update_clipping` arguments of the optimizers:
+
+* **`l2_clip_`**: Clips the gradient/update by its L2 norm.
+* **`rmsnorm_clip_`**: Clips the gradient/update by its RMS norm.
+* **`trust_region_clip_`**: Clips the gradient/update using a trust region method.
+* **`mu_law_compress`**: Compresses the gradient/update using the µ-law algorithm.
+* **`a_law_compress`**: Compresses the gradient/update using the A-law algorithm.
+* **`identity`**: Does not modify the gradient/update (no clipping).
+
+### Other Utilities
+
+* **`set_torch`**: Sets recommended PyTorch settings for performance, including enabling cuDNN benchmark mode, disabling
+  deterministic algorithms, setting the precision of float32 matrix multiplications, and enabling opt-einsum with the "
+  auto-hq" strategy.
+* **`clean`**: Clears the CUDA cache.
+* **`hook_optimizer_into_model`**: Hooks an optimizer into a model's `post_accumulate_grad_hook`.
+* **`fused_hook`**: Hooks an optimizer into a model's `post_accumulate_grad_hook`, fusing multiple parameter updates
+  into a single step.
+* **`disable_caution_scaling`**: Disables the scaling factor applied when `caution` is enabled in optimizers.
+
+## Example Usage
+
+```python
+import torch
+from torch import nn
+import heavyball
+
+# Define a simple model
+model = nn.Linear(10, 2)
+
+# Create an optimizer
+optimizer = heavyball.ForeachAdamW(model.parameters(), lr=1e-3, weight_decay=1e-2)
+# alternative:
+optimizer = heavyball.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-2)
+
+# Generate some dummy data
+input = torch.randn(1, 10)
+target = torch.randn(1, 2)
+
+# Training loop
+for _ in range(100):
+    # Forward pass
+    output = model(input)
+    loss = (output - target).sum()
+
+    # Backward pass
+    loss.backward()
+
+    # Optimizer step
+    optimizer.step()
+
+    # optional: zero gradients; optimizer.step() already does this, which is different from torch.optim
+    optimizer.zero_grad()
+```
+
+This example demonstrates how to create an `AdamW` optimizer and use it to train a simple linear model. You can easily
+replace `AdamW` with any other optimizer from the `heavyball` library and customize its behavior using the various
+available parameters and settings.
+
+By using `heavyball`'s optimizers and understanding the options in `heavyball.utils`, users can achieve better
+performance, control over training, and easier experimentation with advanced optimization techniques.
+
+
+---
+
+# `heavyball.chainable`: A Composable Optimizer API
+
+The `heavyball.chainable` module provides a powerful and flexible way to build optimizers through function composition,
+similar to Optax. It allows you to chain together a sequence of transformations to create custom optimization algorithms
+tailored to your specific needs. This modular approach makes it easy to experiment with different optimization
+strategies and build complex optimizers from simple, reusable components.
+
+## Core Concept
+
+At the heart of `heavyball.chainable` lies the concept of gradient transformations. A gradient transformation is simply
+a function that takes a state dictionary, a group dictionary, an update tensor, a gradient tensor, and a parameter
+tensor as input, and returns a new (or modified) update tensor. These transformations can be chained together to form an
+optimization algorithm.
+
+The state dictionary stores any persistent state needed by the transformation, such as momentum buffers or
+preconditioners. The group dictionary contains hyperparameters specific to a group of parameters. The update tensor is
+the current update being processed, the gradient tensor is the gradient of the loss with respect to the parameter, and
+the parameter tensor is the parameter itself.
+
+### Function Signature
+
+A typical gradient transformation function has the following signature:
+
+```python
+
+def my_transformation(state: dict, group: dict, update: List[torch.Tensor], grad: List[torch.Tensor],
+                      param: List[torch.Tensor]) -> torch.Tensor:
+    # ... transformation logic ...
+    return update
+```
+
+or
+
+```python
+@C.no_state_no_foreach
+def my_transformation(group: dict, update: torch.Tensor, grad: torch.Tensor, param: torch.Tensor, *args,
+                      **kwargs) -> torch.Tensor:
+    # ... transformation logic ...
+    return update
+```
+
+Note that the second version has no state and processes updates one by one, while the first version processes updates
+in parallel.
+
+These functions modify the `update` in place or return a new tensor.
+
+### Example: Scaling by the learning rate
+
+```python
+from heavyball import chainable as C
+
+
+@C.no_state_no_foreach
+def scale_by_learning_rate(group: dict, update: torch.Tensor, grad: torch.Tensor, param: torch.Tensor) -> torch.Tensor:
+    return update * group["lr"]
+```
+
+## `FunctionTransform` and Guards
+
+To make it easier to create gradient transformations, `heavyball.chainable` provides the `FunctionTransform` class and a
+set of "guard" decorators.
+
+### `FunctionTransform`
+
+`FunctionTransform` is a base class for gradient transformations that provides a common interface and helper methods. It
+takes a function `fn` as input and stores it along with its name.
+
+```python
+class FunctionTransform:
+    def __init__(self, fn):
+        self.fn = fn
+        self.fn_name = self.get_fn().__name__
+
+    def __call__(self, state, group, update, grad, param, *args, **kwargs):
+        raise NotImplementedError
+
+    def get_fn(self):
+        if hasattr(self.fn, 'get_fn'):
+            return self.fn.get_fn()
+        return self.fn
+
+    def val_name(self, name):
+        return f"{self.fn_name}_{name}"
+```
+
+### Guards
+
+Guards are decorators that help manage the state dictionary and ensure that transformations are applied correctly. They
+handle common tasks like initializing state variables and preventing redundant computations.
+
+#### `zero_guard`
+
+The `zero_guard` decorator ensures that a specific variable in the state dictionary is initialized to zero if it doesn't
+exist.
+
+```python
+@C.zero_guard("momentum")
+def my_transformation(state, group, update, grad, param, momentum):
+    # ... momentum will be initialized to zero if it doesn't exist in state ...
+    return update
+```
+
+#### `copy_guard`
+
+The `copy_guard` decorator creates a copy of a specified input (update, grad, or param) and stores it in the state
+dictionary.
+
+```python
+@C.copy_guard(0, "update_copy")  # 0 refers to the 'update' argument
+def my_transformation(state, group, update, grad, param, update_copy):
+    # ... update_copy will be a copy of the update tensor ...
+    return update
+```
+
+#### `general_guard`
+
+The `general_guard` decorator provides a more flexible way to manage state. It allows you to specify a custom
+initialization function that is called if a specific variable is not found in the state.
+
+```python
+def init_preconditioner(state, group, update, grad, param, **kwargs):
+
+
+# ... initialize preconditioner ...
+
+@C.general_guard("precond", init_fn=init_preconditioner)
+def my_transformation(state, group, update, grad, param, precond):
+    # ... precond will be initialized using init_preconditioner if it doesn't exist ...
+    return update
+```
+
+#### `no_state`
+
+The `no_state` decorator indicates that a transformation does not use or modify any state.
+
+#### `no_state_no_foreach`
+
+The `no_state_no_foreach` decorator indicates that a transformation does not use or modify any state and also does not
+support `foreach` implementations.
+
+## Chaining Transformations
+
+The power of `heavyball.chainable` comes from its ability to chain transformations together. This is achieved through
+the `chain` function.
+
+```python
+def chain(state: Union[callable, dict], group, grad, param, *fns):
+    update = [torch.clone(g, memory_format=torch.preserve_format) for g in grad]
+    skip_update = False
+    for fn in fns:
+        try:
+            update = fn(state, group, update, grad, param)
+        except SkipUpdate:
+            skip_update = True
+            continue
+        if update is None:
+            break
+    if not skip_update and update is not None:
+        utils.update_param_(param, update, group['lr'], group['weight_decay'], caution=group['caution'], grad=grad)
+```
+
+The `chain` function takes a state dictionary, a group dictionary, a gradient tensor, a parameter tensor, and a sequence
+of gradient transformations as input. It applies each transformation in order, passing the output of one transformation
+as the input to the next.
+
+## Building Optimizers
+
+The `ChainOpt` class provides a convenient way to build optimizers from chained transformations.
+
+```python
+class ChainOpt(utils.StatefulOptimizer):
+    # ...
+    def __init__(self, params, defaults, foreach: bool, *fns):
+        # ...
+        self.fns = tuple(fns)
+
+    def _step(self, group):
+        # ...
+        if not group['foreach'] or len(p) == 1:
+            for param, grad in zip(p, g):
+                chain(self.state_, group, [grad], [param], *self.fns)
+        else:
+            chain(self.state_, group, g, p, *self.fns)
+        # ...
+```
+
+### BaseOpt
+
+The `BaseOpt` class extends `ChainOpt` and provides additional features like gradient clipping, update clipping, and
+optional PaLM beta2 schedule.
+
+```python
+class BaseOpt(ChainOpt):
+    # ...
+    def __init__(self, params, defaults, foreach: bool, gradient_clipping: str_or_fn, update_clipping: str_or_fn,
+                 palm: bool = use_default, *fns, compile_step: bool = use_default, promote: bool = use_default):
+# ...
+```
+
+### `ScheduleFree`
+
+The `ScheduleFree` class provides a convenient interface for using the `update_by_schedule_free` transformation.
+
+### Predefined Transformations
+
+`heavyball.chainable` provides a number of predefined gradient transformations, including:
+
+* `exp_avg`: Calculates the exponential moving average of the gradients.
+* `scale_by_exp_avg_sq`: Scales the updates by the inverse square root of the exponential moving average of squared
+  gradients.
+* `scale_by_adam`: Scales the updates using the Adam algorithm.
+* `update_by_adam`: Updates the parameters using the Adam algorithm.
+* `scale_by_laprop`: Scales the updates using the LaProp algorithm.
+* `update_by_laprop`: Updates the parameters using the LaProp algorithm.
+* `update_by_schedule_free`: Updates the parameters using the Schedule-Free algorithm.
+* `update_by_adopt`: Updates the parameters using the ADOPT algorithm.
+* `scale_by_adopt`: Scales the updates using the ADOPT algorithm.
+* `orthogonalize_update`: Orthogonalizes the update tensor.
+* `nesterov_momentum`: Applies Nesterov momentum to the updates.
+* `heavyball_momentum`: Applies heavy-ball momentum to the updates.
+* `scale_by_soap`: Scales the updates using the SOAP preconditioner.
+* `scale_by_psgd`: Scales the updates using the PSGD preconditioner.
+* `scale_by_delayed_psgd`: Scales the updates using the delayed PSGD preconditioner.
+* `update_by_psgd`: Updates the parameters using the PSGD preconditioner.
+* `update_by_delayed_psgd`: Updates the parameters using the delayed PSGD preconditioner.
+* `palm_beta2`: Modifies the beta2 parameter for PaLM optimizers.
+
+## Creating New Transformations
+
+You can easily create new gradient transformations by following the function signature and using the provided guards and
+`FunctionTransform` class.
+
+### Example: Clipping gradients by norm
+
+```python
+from heavyball import chainable as C
+from heavyball import utils
+
+
+@C.no_state
+def clip_by_global_norm(group: dict, update: torch.Tensor, grad: torch.Tensor, param: torch.Tensor,
+                        max_norm: float) -> torch.Tensor:
+    """Clips the gradient by its global norm."""
+    total_norm = torch.norm(torch.stack([torch.norm(g) for g in grad]))
+    clip_coef = max_norm / (total_norm + 1e-6)
+    if clip_coef < 1:
+        return [u * clip_coef for u in update]
+    return update
+```
+
+### Example: L2-Normalization of updates
+
+```python
+from heavyball import chainable as C
+from heavyball import utils
+
+
+@C.no_state_no_foreach
+def l2_normalize_updates(group: dict, update: torch.Tensor, grad: torch.Tensor, param: torch.Tensor) -> torch.Tensor:
+    """L2-normalizes the updates."""
+    norm = update.norm()
+    if norm > 0:
+        return update / norm
+    return update
+```
+
+---
+
+## Optimizer Recommendations
+
+This hierarchy ranks optimizers from most recommended (top) to least recommended (bottom) for general deep learning
+tasks. However, the best choice always depends on your specific model, dataset, and computational resources.
+
+**1. Preconditioned Optimizers (SOAP and PSGD):**
+
+- **Recommendation:** **Start here.** These are generally the most powerful and efficient optimizers in `heavyball`.
+- **`ForeachSOAP`** (and its variants: `PaLMForeachSOAP`, `PrecondScheduleForeachSOAP`,
+  `PrecondSchedulePaLMForeachSOAP`):
+    - **Strengths:**
+        - **Adaptive Preconditioning:** SOAP dynamically adapts to the curvature of the loss landscape using
+          second-order information, leading to faster convergence, especially in ill-conditioned problems.
+        - **Robustness:** Less sensitive to hyperparameter choices compared to Adam.
+        - **Strong Empirical Performance:** Often outperforms other optimizers across various tasks and architectures.
+    - **Weaknesses:**
+        - **Computational Cost:** Higher per-step cost due to preconditioner computation and updates.
+        - **Memory Usage:** Can use more memory than simpler optimizers, particularly for large models.
+        - **`precondition_frequency` or `precond_scheduler`:** Needs to be tuned, though the default schedule usually
+          works well.
+    - **When to use:**
+        - **Complex models and datasets:** Where optimization is challenging.
+        - **When training stability is crucial.**
+        - **When you can't retune hyperparameters.**
+    - **Variants:**
+        - `PaLMForeachSOAP`: Enables PaLM's beta2 schedule by default.
+        - `PrecondScheduleForeachSOAP`: Uses a dynamic schedule for preconditioner updates.
+        - `PrecondSchedulePaLMForeachSOAP`: Combines the PaLM schedule with a dynamic preconditioner schedule.
+
+- **`ForeachPSGDKron`** (and its variants: `ForeachPurePSGD`, `ForeachCachedDelayedPSGDKron`, `ForeachCachedPSGDKron`,
+  `ForeachDelayedPSGD`):
+    - **Strengths:**
+        - **Preconditioning:** Uses Kronecker-factored approximations to capture curvature information, providing many
+          of the benefits of second-order methods at a lower cost than full curvature methods.
+        - **Efficiency:** Relatively efficient in terms of computation.
+        - **Tunability:** Offers many options for customization.
+        - **Convergence:** Tends to converge faster than SOAP.
+    - **Weaknesses:**
+        - **No baseline:** SOAP can copy Adam's hyperparameters - PSGD requires more tuning.
+        - **Complexity:** Has many hyperparameters to tune.
+    - **When to use:**
+        - **Large models:** Where memory is a constraint.
+        - **When `ForeachSOAP` is too computationally expensive.**
+        - **When you want potentially the best performance regardless of computational cost.**
+    - **Variants:**
+        - `ForeachPurePSGD`: Disables exponential averaging of the input when calculating the preconditioner.
+        - `ForeachCachedDelayedPSGDKron`: Caches preconditioner-related computations and uses delayed preconditioner
+          updates.
+        - `ForeachCachedPSGDKron`: Caches preconditioner-related computations.
+        - `ForeachDelayedPSGD`: Uses delayed preconditioner updates.
+
+**2. Muon:**
+
+- **`ForeachMuon`** (and `MuonLaProp`):
+    - **Strengths:**
+        - **Momentum with Orthogonal Updates:** Combines momentum with orthogonalized updates, which can
+          improve stability and exploration.
+        - **Good Generalization:** Often leads to better generalization performance compared to Adam.
+    - **Weaknesses:**
+        - **Performance:** Generally outperformed by SOAP and PSGD.
+        - **Computational Cost:** Higher overheads than SOAP and PSGD.
+    - **When to use:**
+        - **When generalization is a primary concern.**
+        - **When you want an optimizer less prone to finding sharp minima.**
+
+**3. Adam-Based Optimizers:**
+
+- **`ForeachLaProp`**:
+    - **Strengths:**
+        - **Backward Compatibility:** Can use Adam's hyperparameters, but allows a larger range of betas.
+        - **Stability:** More stable than Adam.
+    - **Weaknesses:**
+        - **Performance:** Generally outperformed by SOAP, PSGD, and Muon.
+    - **When to use:**
+        - **When you want less risk or better losses than Adam, but can't run advanced methods.**
+
+- **`ForeachAdamW`** (and `ForeachSFAdamW`, `PaLMForeachSFAdamW`):
+    - **Strengths:**
+        - **Widely Used:** A popular and well-established optimizer.
+    - **Weaknesses:**
+        - **Performance:** Often outperformed by preconditioned optimizers (SOAP, PSGD) and Muon.
+        - **Sensitivity to Hyperparameters:** Can be sensitive to the choice of learning rate and beta parameters.
+    - **When to use:**
+        - **As a strong baseline.**
+        - **When you are familiar with Adam and want a robust starting point.**
+        - **When computational cost is a major concern (compared to second-order methods).**
+    - **Variants:**
+        - `ForeachSFAdamW`: A Schedule-Free version of AdamW that dynamically adjusts the learning rate.
+        - `PaLMForeachSFAdamW`: A PaLM version of Schedule-Free AdamW.
+
+## Choosing the Right Optimizer
+
+1. **Start with Preconditioning:** Begin with either `ForeachSOAP` or `ForeachPSGDKron`. If computational resources are
+   a major constraint, lean towards `ForeachPSGDKron`. If performance is paramount, try `ForeachSOAP` first.
+
+2. **Consider Muon:** If preconditioned optimizers are not feasible or if you want to explore alternatives that
+   incorporate momentum and orthogonal updates, try `ForeachMuon`.
+
+3. **Use LaProp or Adam as Baselines:** `ForeachLaProp` can serve as a simple adaptive baseline. `ForeachAdamW` is a
+   strong and widely used baseline that you should always compare against.
+
+4. **Experiment and Tune:** The best optimizer ultimately depends on your specific problem. It's crucial to experiment
+   with different optimizers and carefully tune their hyperparameters (especially the learning rate).
+
+## Important Notes
+
+* **Learning Rate:** The learning rate is the most important hyperparameter. You'll likely need to adjust it when
+  switching between optimizers.
+* **Warmup:** Consider using a learning rate warmup, especially for more complex optimizers like SOAP and PSGD.
+* **Weight Decay:** Weight decay can improve generalization for many optimizers, especially AdamW.
+* **`foreach`:** Use `foreach` versions of the optimizers when possible for better performance.
+* **`heavyball.utils`:** Remember to utilize the settings and functions in `heavyball.utils` (e.g., `set_torch`,
+  `compile_mode`, `zeroth_power_mode`, clipping functions) to optimize performance and experiment with different
+  configurations.
+

heavyball-1.4.3.dist-info/RECORD

@@ -0,0 +1,8 @@
+heavyball/__init__.py,sha256=miRgcXlzLWTNzojeRF5hEcg-x_GqfMHjRzOaiR_zO3U,10981
+heavyball/chainable.py,sha256=-5ovRa7yD7V41_cgaBJtO5fBrnBemAILl4YKjQmeuns,24183
+heavyball/utils.py,sha256=x0rSU8lko7ACdI9GuTLC0wP6HwIZxwB8f8tukBOR0xA,48129
+heavyball-1.4.3.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
+heavyball-1.4.3.dist-info/METADATA,sha256=RM_pOme3dsQL-drKcKD6FJ0qE3SSh4JdPM-kC9vpbeU,43584
+heavyball-1.4.3.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+heavyball-1.4.3.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
+heavyball-1.4.3.dist-info/RECORD,,

heavyball-1.4.0.dist-info/METADATA

@@ -1,136 +0,0 @@
-Metadata-Version: 2.1
-Name: heavyball
-Version: 1.4.0
-Summary: Efficient optimizers
-Home-page: https://github.com/clashluke/heavyball
-Author: Lucas Nestler
-Author-email: github.heavyball@nestler.sh
-License: BSD
-Classifier: Development Status :: 5 - Production/Stable
-Classifier: License :: OSI Approved :: BSD License
-Classifier: Programming Language :: Python
-Classifier: Programming Language :: Python :: 3.7
-Classifier: Programming Language :: Python :: 3.8
-Classifier: Programming Language :: Python :: 3.9
-Classifier: Topic :: Software Development :: Libraries
-Classifier: Topic :: Software Development :: Libraries :: Python Modules
-Classifier: Intended Audience :: Developers
-Requires-Python: >=3.7
-Description-Content-Type: text/markdown
-License-File: LICENSE
-Requires-Dist: opt-einsum
-Requires-Dist: torch
-Requires-Dist: numpy
-
-# HeavyBall
-
-> [!IMPORTANT]  
-> It's recommended to use `heavyball.utils.set_torch()` for faster training and less memory usage.
-
-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-12-07, 1.0.0), the recommended stable optimizer is `PrecondSchedulePaLMSOAP` (see below). The
-recommended experimental optimizer is `DelayedPSGDKron` ([tuning guide](docs/psgd_efficiency.md)).
-
-## Features
-
-* **Optax-like API**: `C = heavyball.chainable; grokfast = C.ChainOpt(p, lr, C.exp_avg, C.scale_by_adam)`
-* **Stochastic Rounding**: [FP32 convergence with BF16 parameters](https://github.com/pytorch/pytorch/issues/120376)
-* **Inplace EMA**: Same math, but less memory, less compute and higher stability
-* **Foreach**: Fast multi-tensor application (turn it off to save memory via `foreach=False`)
-* **PaLM Beta2**: Fast initial
-  convergence, [stable late convergence](https://x.com/_clashluke/status/1820810798693818761)
-* **ScheduleFree**: No learning rate schedule, but better convergence
-* [**Preconditioner Schedule**](https://github.com/lixilinx/psgd_torch/): Improved loss-per-step in early convergence,
-  better step-per-second in late convergence (explained below)
-* **Memory-efficient storage** PSGD supports `store_triu_as_line` (default: `True`) and `q_dtype` to trade off memory
-  usage for memory
-  bandwidth; Other optimizers have `storage_dtype`, supporting lower-precision EMAs at no(?) performance drop via
-  stochastic rounding
-
-## Getting started
-
-```bash
-pip install heavyball
-```
-
-```python
-import torch
-import heavyball
-
-# Create a model
-model = torch.nn.Linear(16, 1)
-
-# Create an optimizer
-optimizer = heavyball.PrecondSchedulePaLMSOAP(model.parameters(), lr=1e-3)
-
-x = torch.randn(128, 16)
-y = torch.randn(128, 1)
-
-for _ in range(1000):
-    optimizer.zero_grad()
-    loss = torch.nn.functional.mse_loss(model(x), y)
-    loss.backward()
-    optimizer.step()
-```
-
-## Optimizers
-
-| Name                          | Description                                                                                                                                                       | Advantages / Disadvantages                                                                                                                                                                                                                                                                                                            |
-|-------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
-| **AdamW**                     | More efficient (speed, memory) [AdamW](https://arxiv.org/abs/1711.05101)                                                                                          | + Faster than AdamW<br>+ Possibly more (numerically) stable                                                                                                                                                                                                                                                                           
-| **LaProp**                    | More efficient (speed, memory) [LaProp](https://arxiv.org/abs/2002.04839)                                                                                         | + Same cost as AdamW<br>+ Marginally better converence (better proofs)<br>+ Higher hyperparameter stability<br>- Not a guaranteed win (can be neutral)<br>- No "Slingshot"                                                                                                                                                            |
-| **ADOPT**                     | More efficient (speed, memory) [ADOPT](https://arxiv.org/abs/2411.02853)                                                                                          | + Same cost as AdamW<br>+ Rigorous mathematical convergence proofs, even for challenging models (GANs)<br>- Empirically underperforms LaProp<br>- no bf16                                                                                                                                                                             |
-| **SFAdamW**                   | More efficient (speed, memory) [ScheduleFree AdamW](https://arxiv.org/abs/2405.15682)                                                                             | + Same cost as AdamW, but better eval perf<br>+ Full control over hyperparameters                                                                                                                                                                                                                                                     |
-| **PaLMSFAdamW**               | ForeachSFAdamW with [PaLM's beta2 schedule](https://arxiv.org/abs/2204.02311)                                                                                     | + Same cost as AdamW, but better eval perf<br>+ Less control, but faster early and more stable late convergence<br>+ ScheduleFree<br>- slow early convergence                                                                                                                                                                         |
-| **SOAP**                      | More efficient (speed, memory) [SOAP](https://arxiv.org/abs/2409.11321)                                                                                           | + Faster convergence (loss-at-step)<br>+ Full control over hyperparameters<br>- more memory usage<br>- more hyperparameters<br>- higher overhead than AdamW (can be ammortized; better loss-at-second)                                                                                                                                |
-| **PaLMSOAP**                  | ForeachSOAP with [PaLM's beta2 schedule](https://arxiv.org/abs/2204.02311)                                                                                        | + Faster convergence (loss-at-step)<br>+ Less control, but faster early and more stable late convergence<br>- more memory usage<br>- more hyperparameters<br>- higher overhead than AdamW (can be ammortized; better loss-at-second)                                                                                                  |
-| **SFPaLMSOAP**                | ScheduleFree PaLMForeachSOAP                                                                                                                                      | + Fast convergence (loss-at-step)<br>+ less memory usage than PaLMForeachSOAP (more tham AdamW)<br>- slower initial convergence than PaLMForeachSOAP (but allows higher LRs)<br>- higher overhead than AdamW (can be ammortized)                                                                                                      |
-| **PrecondScheduleSFPaLMSOAP** | SFPaLMForeachSOAP with [preconditioner schedule](https://github.com/lixilinx/psgd_torch/), matching the error of PrecondEvery=2 with the cost of PrecondEvery=512 | + Better initial convergence than SFPaLMForeachSOAP<br>+ Significantly faster (sec/it) later<br>+ less memory usage than PaLMForeachSOAP (more tham AdamW)<br>- slower initial convergence than PaLMForeachSOAP (but allows higher LRs)<br>- higher overhead than AdamW (can be ammortized), goes to 0 with increasing number of step |
-| **PrecondSchedulePaLMSOAP**   | PrecondScheduleSFPaLMForeachSOAP without schedule-free                                                                                                            | + Best initial convergence<br>+ Significantly faster (sec/it) later<br>+ high stability<br>- more memory usage than PrecondScheduleSFPaLMForeachSOAP<br>- higher overhead than AdamW (can be ammortized), goes to 0 with increasing number of steps                                                                                   |
-| **PrecondScheduleSOAP**       | PrecondScheduleSFPaLMForeachSOAP without PaLM's beta2 schedule                                                                                                    | + Better initial convergence<br>+ Significantly faster (sec/it) later<br>- more memory usage than PrecondScheduleSFPaLMForeachSOAP<br>- higher overhead than AdamW (can be ammortized), goes to 0 with increasing number of steps                                                                                                     |
-
-## Precond Schedule
-
-The default preconditioner schedule (`f`) would yield the following update intervals:
-
-| Steps     | Interval, `f` | Total (schedule) | Total (constant, every 2) | Total (constant, every 16) |
-|-----------|---------------|------------------|---------------------------|----------------------------|
-| 10        | 1.00005       | 10               | 5 (0.5x)                  | 0 (0.0x)                   |
-| 100       | 1.026         | 99               | 50 (0.5x)                 | 6 (0.1x)                   |
-| 1,000     | 2.0           | 738              | 500 (0.7x)                | 62 (0.1x)                  |
-| 10,000    | 14.3          | 2,168            | 5,000 (2.3x)              | 625 (0.3x)                 |
-| 100,000   | 100.2         | 4,049            | 50,000 (12.3x)            | 6,250 (1.5x)               |
-| 1,000,000 | 513           | 7,245            | 500,000 (69.0x)           | 62,500 (8.6x)              |
-
-## Memory
-
-Second order optimizers make it difficult to estimate memory usage, as it depends on shapes and hyperparameters. To
-estimate your memory usage, you may use `test/test_memory.py` which attempts to ensure there are no regressions.\
-Furthermore, you can find real-world memory usage of a 300M parameters video diffusion model below:
-![img.png](assets/memory.png)
-
-## PSGD
-
-HeavyBall offers various configurations of PSGD:
-
-* "PSGDKron" is the baseline, equivalent to [kron_torch](https://github.com/evanatyourservice/kron_torch/), but with
-  lower compute and memory
-  overhead.
-* "PurePSGD" has no momentum, further reducing memory and compute
-* "DelayedPSGD" implements SOAP/ADOPT-style off-by-one momentum, which has worse initial convergence but higher
-  stability
-  ![img.png](assets/delayed_psgd.png)
-
-## Utils
-
-To access `heavyball.utils`, you need to explicitly `import heavyball.utils`.\
-It has several handy functions:
-
-* `set_torch()` sets pytorch optimization settings (TF32, opt_einsum, benchmark, ...)
-* `compile_mode`, a string passed as-is to `torch.compile(mode=compile_mode)` in all compiled heavyball calls; `compile_mode=None` disables torch_compile
-* `zeroth_power_mode`, a string determining whether to use QR, newtonschulz, or svd or eigh to approximate
-  the eigenvectors.

heavyball-1.4.0.dist-info/RECORD

@@ -1,8 +0,0 @@
-heavyball/__init__.py,sha256=miRgcXlzLWTNzojeRF5hEcg-x_GqfMHjRzOaiR_zO3U,10981
-heavyball/chainable.py,sha256=OYlCVe06SjpxUM8tBBJUIOrmU3uMYwVYwPzkaQMwN98,24171
-heavyball/utils.py,sha256=djwaSLZOB8B-xD2jJxZfXTJpJrcWp-mWTmKxC2F5Sh0,48330
-heavyball-1.4.0.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
-heavyball-1.4.0.dist-info/METADATA,sha256=Pbb_2JAZevdXVpjkJajuBNrTvrLnpDYBL5NLx1SZxHg,12022
-heavyball-1.4.0.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-heavyball-1.4.0.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
-heavyball-1.4.0.dist-info/RECORD,,