adv-optm 0.1.8__tar.gz → 1.0.0__tar.gz

adv_optm-1.0.0/PKG-INFO

@@ -0,0 +1,174 @@
+Metadata-Version: 2.4
+Name: adv_optm
+Version: 1.0.0
+Summary: A family of highly efficient, lightweight yet powerful optimizers.
+Home-page: https://github.com/Koratahiu/Advanced_Optimizers
+Author: Koratahiu
+Author-email: hiuhonor@gmail.com
+License: Apache 2.0
+Keywords: llm,fine-tuning,memory-efficient,low-rank,compression,pytorch,optimizer,adam
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.0
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: keywords
+Dynamic: license
+Dynamic: license-file
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# Advanced Optimizers (AIO)
+
+A comprehensive, all-in-one collection of optimization algorithms for deep learning, designed for maximum efficiency, minimal memory footprint, and superior performance across diverse model architectures and training scenarios.
+
+[![PyPI](https://img.shields.io/pypi/v/adv_optm)](https://pypi.org/project/adv_optm/)
+
+---
+
+## 📦 Installation
+
+```bash
+pip install adv_optm
+```
+
+---
+
+## 🧠 Core Innovations
+
+This library integrates multiple state-of-the-art optimization techniques validated through extensive research and practical training, with 1-bit compression for optimizer states:
+
+### **Memory-Efficient Optimization (SMMF-inspired)**
+- **Paper**: [SMMF: Square-Matricized Momentum Factorization](https://arxiv.org/abs/2412.08894)
+- **Approach**: Uses rank-1 non-negative matrix factorization with reconstruction cycle (factor → reconstruct → update → factor)
+- **Innovation**: 
+  - First moment split into **1-bit sign + absolute value**
+  - Final storage: **four factored vectors + one 1-bit sign state**
+  - Preserves Adam-like update quality with drastically reduced memory
+
+---
+
+## ⚡ Performance Characteristics
+
+### Memory Efficiency (SDXL Model - 6.5GB)
+| Optimizer | Memory Usage | Description |
+|-----------|--------------|-------------|
+| `Adopt_Factored` | 328 MB | 4 small vectors + 1-bit state |
+| `Adopt_Factored + AdEMAMix` | 625 MB | 6 small vectors + two 1-bit states |
+| `Simplified_AdEMAMix` | 328 MB | Same as standard factored (no extra state) |
+
+### Speed Comparison (SDXL, Batch Size 4)
+| Optimizer | Speed | Notes |
+|-----------|-------|-------|
+| `Adafactor` | ~8.5s/it | Baseline |
+| `Adopt_Factored` | ~10s/it | +18% overhead from compression |
+| `Adopt_Factored + AdEMAMix` | ~12s/it | +41% overhead (3 factored states) |
+
+---
+
+## 🧪 Available Optimizers
+
+### Standard Optimizers (All support `factored=True/False`)
+| Optimizer | Description | Best For |
+|-----------|-------------|----------|
+| `Adam_Adv` | Advanced Adam implementation | General purpose |
+| `Adopt_Adv` | Adam-variant with independent beta2 | Stable training for small batch size regimes |
+| `Prodigy_Adv` | Prodigy with D-Adaptation | Adam with automatic LR tuning |
+| `Simplified_AdEMAMix` | Adam variant with accumulator momentum | Small/large batch training when tuned correctly |
+| `Lion_Adv` | Advanced Lion implementation | Memory-constrained environments |
+| `Prodigy_Lion_Adv` | Prodigy + Lion combination | Lion with automatic LR tuning |
+
+### Feature Matrix
+| Feature | Adam_Adv | Adopt_Adv | Prodigy_Adv | Simplified_AdEMAMix | Lion_Adv |
+|---------|----------|-----------|-------------|---------------------|----------|
+| Factored | ✓ | ✓ | ✓ | ✓ | ✓ |
+| AdEMAMix | ✓ | ✓ | ✓ | ✗ | ✗ |
+| Simplified_AdEMAMix | ✗ | ✗ | ✓ | ✓ | ✗ |
+| OrthoGrad | ✓ | ✓ | ✓ | ✓ | ✓ |
+| Grams | ✓ | ✓ | ✓ | ✗ | ✗ |
+| Cautious | ✓ | ✓ | ✓ | ✗ | ✓ |
+| atan2 | ✓ | ✓ | ✓ | ✗ | ✗ |
+| Stochastic Rounding | ✓ | ✓ | ✓ | ✓ | ✓ |
+| Fused Backward Pass | ✓ | ✓ | ✓ | ✓ | ✓ |
+
+---
+
+## ⚙️ Key Features & Parameters
+
+### Comprehensive Feature Guide
+
+| Feature | Description | Recommended Usage | Performance Impact | Theoretical Basis | Compatibility |
+|---------|-------------|-------------------|--------------------|-------------------|--------------|
+| **Factored** | Memory-efficient optimization using rank-1 factorization | Enable for large models (>1B params) or limited VRAM | +12-41% time overhead, 1-bit memory usage | [SMMF](https://arxiv.org/abs/2412.08894) | All optimizers |
+| **AdEMAMix** | Dual EMA system for momentum | Use for long training runs (10k+ steps) | +1 state memory. | [AdEMAMix](https://arxiv.org/abs/2409.03137) | Adam/Adopt/Prodigy |
+| **Simplified_AdEMAMix** | Accumulator-based momentum | Small batch training (≤32) | Same memory as standard, no extra overhead | [Schedule-Free Connections](https://arxiv.org/abs/2502.02431) | Adam/Prodigy |
+| **OrthoGrad** | Removes gradient component parallel to weights | Full finetuning without weight decay | +33% time overhead, no memory impact | [Grokking at Edge](https://github.com/LucasPrietoAl/grokking-at-the-edge-of-numerical-stability) | All optimizers |
+| **Stochastic Rounding** | Improves precision for BF16 training | BF16 training | Minimal overhead (<5%) | [Revisiting BFloat16 Training](https://arxiv.org/abs/2010.06192) | All optimizers |
+| **atan2** | Robust eps replacement + built-in clipping | Use with Adopt or unstable training | No overhead | [Adam-atan2](https://github.com/lucidrains/adam-atan2-pytorch) | Adam/Adopt/prodigy |
+| **Cautious** | Update only when the direction align with the gradients | should faster the convergence | No overhead | [C-Optim](https://github.com/kyleliang919/C-Optim) | Adam/Adopt/prodigy |
+| **Grams** | Update direction from the gradients | should have a stronger effect than cautious | No overhead | [Grams](https://github.com/Gunale0926/Grams) | Adam/Adopt/prodigy |
+
+---
+
+## Simplified_AdEMAMix Parameters
+Simplified_AdEMAMix replaces standard momentum with an accumulator for better small-large batch performance.
+
+| Parameter | Recommended Values | Description |
+|-----------|---------------------|-------------|
+| `beta1` | 0.9 (large BS), 0.99-0.9999 (small BS) | Determines memory length of accumulator |
+| `alpha` | 100-10 (small BS), 1-0 (large BS) | Gradient smoothing factor |
+
+**Alpha Tuning Guide**:
+| Batch Size | Recommended α | Rationale |
+|------------|---------------|-----------|
+| Small (≤32) | 100, 50, 20, 10 | Emphasizes recent gradients for quick adaptation |
+| Medium (32-512) | 10, 5, 2, 1 | Balanced approach |
+| Large (≥512) | 1, 0.5, 0 | Emphasizes historical gradients for stability |
+
+⚠️ **Important**: Use **~100x smaller learning rate** with Simplified_AdEMAMix compared to AdamW (e.g., 1e-6 instead of 1e-4)
+
+### 📊 Performance Validation
+Small Batch Training (SDXL, BS=2, 1.8K steps)
+![Training Comparison](https://github.com/user-attachments/assets/7eff0671-cc59-47fc-8b63-d5205456d649)
+
+- **🟢 Prodigy_adv** (beta1=0.9, d0=1e-5): Final LR=2.9e-4
+- **🔵 Prodigy_adv + Simplified_AdEMAMix** (beta1=0.99, α=100, d0=1e-7): Final LR=5.8e-6
+
+**Results**:
+- Simplified_AdEMAMix shows faster convergence and better final performance
+- D-Adaptation automatically handles aggressive updates (50x smaller LR)
+- Generated samples show significantly better quality with Simplified_AdEMAMix
+
+---
+
+## ⚠️ Known Limitations
+
+### 1. Prodigy_Adv Sensitivity
+- Highly sensitive to gradient modifications (Adopt normalization, low-rank factorization)
+- May fail to increase learning rate in some LoRA scenarios
+- **Fix**: Disable factorization or set beta1=0
+
+### 2. Aggressive Learning Rates
+- Can destabilize factored first moment
+- **Recommendation**: Check Prodigy learning rate as reference for safe LR threshold
+
+---
+
+## 📚 References
+
+1. [SMMF: Square-Matricized Momentum Factorization](https://arxiv.org/abs/2412.08894)
+2. [The AdEMAMix Optimizer: Better, Faster, Older](https://arxiv.org/abs/2409.03137)
+3. [Connections between Schedule-Free Optimizers, AdEMAMix, and Accelerated SGD Variants](https://arxiv.org/abs/2502.02431)
+
+---

adv_optm-1.0.0/README.md

@@ -0,0 +1,143 @@
+# Advanced Optimizers (AIO)
+
+A comprehensive, all-in-one collection of optimization algorithms for deep learning, designed for maximum efficiency, minimal memory footprint, and superior performance across diverse model architectures and training scenarios.
+
+[![PyPI](https://img.shields.io/pypi/v/adv_optm)](https://pypi.org/project/adv_optm/)
+
+---
+
+## 📦 Installation
+
+```bash
+pip install adv_optm
+```
+
+---
+
+## 🧠 Core Innovations
+
+This library integrates multiple state-of-the-art optimization techniques validated through extensive research and practical training, with 1-bit compression for optimizer states:
+
+### **Memory-Efficient Optimization (SMMF-inspired)**
+- **Paper**: [SMMF: Square-Matricized Momentum Factorization](https://arxiv.org/abs/2412.08894)
+- **Approach**: Uses rank-1 non-negative matrix factorization with reconstruction cycle (factor → reconstruct → update → factor)
+- **Innovation**: 
+  - First moment split into **1-bit sign + absolute value**
+  - Final storage: **four factored vectors + one 1-bit sign state**
+  - Preserves Adam-like update quality with drastically reduced memory
+
+---
+
+## ⚡ Performance Characteristics
+
+### Memory Efficiency (SDXL Model - 6.5GB)
+| Optimizer | Memory Usage | Description |
+|-----------|--------------|-------------|
+| `Adopt_Factored` | 328 MB | 4 small vectors + 1-bit state |
+| `Adopt_Factored + AdEMAMix` | 625 MB | 6 small vectors + two 1-bit states |
+| `Simplified_AdEMAMix` | 328 MB | Same as standard factored (no extra state) |
+
+### Speed Comparison (SDXL, Batch Size 4)
+| Optimizer | Speed | Notes |
+|-----------|-------|-------|
+| `Adafactor` | ~8.5s/it | Baseline |
+| `Adopt_Factored` | ~10s/it | +18% overhead from compression |
+| `Adopt_Factored + AdEMAMix` | ~12s/it | +41% overhead (3 factored states) |
+
+---
+
+## 🧪 Available Optimizers
+
+### Standard Optimizers (All support `factored=True/False`)
+| Optimizer | Description | Best For |
+|-----------|-------------|----------|
+| `Adam_Adv` | Advanced Adam implementation | General purpose |
+| `Adopt_Adv` | Adam-variant with independent beta2 | Stable training for small batch size regimes |
+| `Prodigy_Adv` | Prodigy with D-Adaptation | Adam with automatic LR tuning |
+| `Simplified_AdEMAMix` | Adam variant with accumulator momentum | Small/large batch training when tuned correctly |
+| `Lion_Adv` | Advanced Lion implementation | Memory-constrained environments |
+| `Prodigy_Lion_Adv` | Prodigy + Lion combination | Lion with automatic LR tuning |
+
+### Feature Matrix
+| Feature | Adam_Adv | Adopt_Adv | Prodigy_Adv | Simplified_AdEMAMix | Lion_Adv |
+|---------|----------|-----------|-------------|---------------------|----------|
+| Factored | ✓ | ✓ | ✓ | ✓ | ✓ |
+| AdEMAMix | ✓ | ✓ | ✓ | ✗ | ✗ |
+| Simplified_AdEMAMix | ✗ | ✗ | ✓ | ✓ | ✗ |
+| OrthoGrad | ✓ | ✓ | ✓ | ✓ | ✓ |
+| Grams | ✓ | ✓ | ✓ | ✗ | ✗ |
+| Cautious | ✓ | ✓ | ✓ | ✗ | ✓ |
+| atan2 | ✓ | ✓ | ✓ | ✗ | ✗ |
+| Stochastic Rounding | ✓ | ✓ | ✓ | ✓ | ✓ |
+| Fused Backward Pass | ✓ | ✓ | ✓ | ✓ | ✓ |
+
+---
+
+## ⚙️ Key Features & Parameters
+
+### Comprehensive Feature Guide
+
+| Feature | Description | Recommended Usage | Performance Impact | Theoretical Basis | Compatibility |
+|---------|-------------|-------------------|--------------------|-------------------|--------------|
+| **Factored** | Memory-efficient optimization using rank-1 factorization | Enable for large models (>1B params) or limited VRAM | +12-41% time overhead, 1-bit memory usage | [SMMF](https://arxiv.org/abs/2412.08894) | All optimizers |
+| **AdEMAMix** | Dual EMA system for momentum | Use for long training runs (10k+ steps) | +1 state memory. | [AdEMAMix](https://arxiv.org/abs/2409.03137) | Adam/Adopt/Prodigy |
+| **Simplified_AdEMAMix** | Accumulator-based momentum | Small batch training (≤32) | Same memory as standard, no extra overhead | [Schedule-Free Connections](https://arxiv.org/abs/2502.02431) | Adam/Prodigy |
+| **OrthoGrad** | Removes gradient component parallel to weights | Full finetuning without weight decay | +33% time overhead, no memory impact | [Grokking at Edge](https://github.com/LucasPrietoAl/grokking-at-the-edge-of-numerical-stability) | All optimizers |
+| **Stochastic Rounding** | Improves precision for BF16 training | BF16 training | Minimal overhead (<5%) | [Revisiting BFloat16 Training](https://arxiv.org/abs/2010.06192) | All optimizers |
+| **atan2** | Robust eps replacement + built-in clipping | Use with Adopt or unstable training | No overhead | [Adam-atan2](https://github.com/lucidrains/adam-atan2-pytorch) | Adam/Adopt/prodigy |
+| **Cautious** | Update only when the direction align with the gradients | should faster the convergence | No overhead | [C-Optim](https://github.com/kyleliang919/C-Optim) | Adam/Adopt/prodigy |
+| **Grams** | Update direction from the gradients | should have a stronger effect than cautious | No overhead | [Grams](https://github.com/Gunale0926/Grams) | Adam/Adopt/prodigy |
+
+---
+
+## Simplified_AdEMAMix Parameters
+Simplified_AdEMAMix replaces standard momentum with an accumulator for better small-large batch performance.
+
+| Parameter | Recommended Values | Description |
+|-----------|---------------------|-------------|
+| `beta1` | 0.9 (large BS), 0.99-0.9999 (small BS) | Determines memory length of accumulator |
+| `alpha` | 100-10 (small BS), 1-0 (large BS) | Gradient smoothing factor |
+
+**Alpha Tuning Guide**:
+| Batch Size | Recommended α | Rationale |
+|------------|---------------|-----------|
+| Small (≤32) | 100, 50, 20, 10 | Emphasizes recent gradients for quick adaptation |
+| Medium (32-512) | 10, 5, 2, 1 | Balanced approach |
+| Large (≥512) | 1, 0.5, 0 | Emphasizes historical gradients for stability |
+
+⚠️ **Important**: Use **~100x smaller learning rate** with Simplified_AdEMAMix compared to AdamW (e.g., 1e-6 instead of 1e-4)
+
+### 📊 Performance Validation
+Small Batch Training (SDXL, BS=2, 1.8K steps)
+![Training Comparison](https://github.com/user-attachments/assets/7eff0671-cc59-47fc-8b63-d5205456d649)
+
+- **🟢 Prodigy_adv** (beta1=0.9, d0=1e-5): Final LR=2.9e-4
+- **🔵 Prodigy_adv + Simplified_AdEMAMix** (beta1=0.99, α=100, d0=1e-7): Final LR=5.8e-6
+
+**Results**:
+- Simplified_AdEMAMix shows faster convergence and better final performance
+- D-Adaptation automatically handles aggressive updates (50x smaller LR)
+- Generated samples show significantly better quality with Simplified_AdEMAMix
+
+---
+
+## ⚠️ Known Limitations
+
+### 1. Prodigy_Adv Sensitivity
+- Highly sensitive to gradient modifications (Adopt normalization, low-rank factorization)
+- May fail to increase learning rate in some LoRA scenarios
+- **Fix**: Disable factorization or set beta1=0
+
+### 2. Aggressive Learning Rates
+- Can destabilize factored first moment
+- **Recommendation**: Check Prodigy learning rate as reference for safe LR threshold
+
+---
+
+## 📚 References
+
+1. [SMMF: Square-Matricized Momentum Factorization](https://arxiv.org/abs/2412.08894)
+2. [The AdEMAMix Optimizer: Better, Faster, Older](https://arxiv.org/abs/2409.03137)
+3. [Connections between Schedule-Free Optimizers, AdEMAMix, and Accelerated SGD Variants](https://arxiv.org/abs/2502.02431)
+
+---

{adv_optm-0.1.8 → adv_optm-1.0.0}/adv_optm/__init__.py

@@ -16,4 +16,4 @@ __all__ = [
     "Lion_Prodigy_adv",
 ]
 
-__version__ = "0.1.8"
+__version__ = "1.0.0"

{adv_optm-0.1.8 → adv_optm-1.0.0}/adv_optm/optim/AdamW_adv.py

@@ -86,6 +86,9 @@ class AdamW_adv(torch.optim.Optimizer):
             raise ValueError(f"Epsilon should be >= 0.0. Got {eps}")
         if not (weight_decay >= 0.0):
             raise ValueError(f"Weight-decay should be >= 0.0. Got {weight_decay}")
+        if use_cautious and use_grams:
+            print("Warning: use_cautious is incompatible with use_grams, Disabling use_cautious.")
+            use_cautious = False
 
         defaults = {
             "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay,

{adv_optm-0.1.8 → adv_optm-1.0.0}/adv_optm/optim/Adopt_adv.py

@@ -62,6 +62,16 @@ class Adopt_adv(torch.optim.Optimizer):
             the warmup, `alpha` ramps from 0 to its target value. If `None`,
             the scheduler is disabled and the full `alpha` value is used from
             the start. (default: None)
+        Simplified_AdEMAMix (bool): whether to use the Simplified AdEMAMix update rule.
+            This changes the EMA to accumulator and the update numerator to `alpha_grad * grad + mt`, which can be
+            more responsive, especially for small batch sizes. Enabling this will
+            automatically disable `use_AdEMAMix`, `use_cautious`, `use_grams`,
+            and `use_atan2`. (default: False)
+        alpha_grad (float): Mixing coefficient for the Simplified AdEMAMix update rule
+            (only used when `Simplified_AdEMAMix` is `True`). Controls the weight of the
+            current gradient. For small batch sizes, use high values (e.g., 10-100) to be
+            more responsive. For large batch sizes, use low values (e.g., 0-1) for
+            stability. (default: 100.0)
         factored (bool): whether to use the factorization or disable it to use
             the uncompressed optimizer. (default: False)
     """
@@ -77,13 +87,15 @@ class Adopt_adv(torch.optim.Optimizer):
         vector_reshape: bool = True,
         stochastic_rounding: bool = True,
         use_atan2: bool = False,
-        use_cautious: bool = True,
+        use_cautious: bool = False,
         use_grams: bool = False,
         use_orthograd: bool = False,
         use_AdEMAMix: bool = False,
         beta3_ema: float = 0.9999,
         alpha: float = 5.0,
         t_alpha: int | None = None,
+        Simplified_AdEMAMix: bool = False,
+        alpha_grad: float = 100.0,
         factored: bool = False,
     ):
         if not (lr >= 0.0):
@@ -94,19 +106,34 @@ class Adopt_adv(torch.optim.Optimizer):
             raise ValueError(f"Epsilon should be >= 0.0. Got {eps}")
         if not (weight_decay >= 0.0):
             raise ValueError(f"Weight-decay should be >= 0.0. Got {weight_decay}")
+        if use_cautious and use_grams:
+            print("Warning: use_cautious is incompatible with use_grams, Disabling use_cautious.")
+            use_cautious = False
+        if betas[0] == 0.0 and Simplified_AdEMAMix:
+            raise ValueError(f"Beta1 cannot be 0.0 when using Simplified_AdEMAMix. Got {betas[0]}")
+        if use_AdEMAMix and Simplified_AdEMAMix:
+            print("Warning: use_AdEMAMix is incompatible with Simplified_AdEMAMix, Disabling use_AdEMAMix.")
+        if use_grams and Simplified_AdEMAMix:
+            print("Warning: use_grams is incompatible with Simplified_AdEMAMix, Disabling use_grams.")
+        if use_cautious and Simplified_AdEMAMix:
+            print("Warning: use_cautious is incompatible with Simplified_AdEMAMix, Disabling use_cautious.")
+        if use_atan2 and Simplified_AdEMAMix:
+            print("Warning: use_atan2 is incompatible with Simplified_AdEMAMix. Disabling use_atan2.")
+            use_atan2 = False
 
         defaults = {
             "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay,
             "vector_reshape": vector_reshape, "beta3_ema": beta3_ema, "alpha": alpha,
-            "t_alpha": t_alpha,
+            "t_alpha": t_alpha, "alpha_grad": alpha_grad, 
         }
         self.clip_lambda = clip_lambda
         self.stochastic_rounding = stochastic_rounding
-        self.use_atan2 = use_atan2
-        self.use_cautious = use_cautious
-        self.use_grams = use_grams
+        self.use_atan2 = use_atan2 and not Simplified_AdEMAMix
+        self.use_cautious = use_cautious and not Simplified_AdEMAMix
+        self.use_grams = use_grams and not Simplified_AdEMAMix
         self.use_orthograd = use_orthograd
-        self.use_AdEMAMix = use_AdEMAMix
+        self.use_AdEMAMix = use_AdEMAMix and not Simplified_AdEMAMix
+        self.Simplified_AdEMAMix = Simplified_AdEMAMix
         self.factored = factored
         super().__init__(params, defaults)
 
@@ -185,6 +212,8 @@ class Adopt_adv(torch.optim.Optimizer):
             alpha_t = alpha
             if t_alpha is not None and t_alpha > 0 and current_step < t_alpha:
                 alpha_t = min(current_step * alpha / t_alpha, alpha)
+        if self.Simplified_AdEMAMix:
+            alpha_grad = group["alpha_grad"]
 
         if state['factored']:
             d1, d2 = state['effective_shape']
@@ -224,7 +253,10 @@ class Adopt_adv(torch.optim.Optimizer):
             del denom
 
             # ADOPT Step B: Update momentum m_t using normalized gradient
-            mt.mul_(beta1).add_(normalized_grad, alpha=1.0 - beta1)
+            if self.Simplified_AdEMAMix:
+                mt.mul_(beta1).add_(normalized_grad, alpha=1.0)
+            else:
+                mt.mul_(beta1).add_(normalized_grad, alpha=1.0 - beta1)
             if self.use_grams:
                 mt = grad_reshaped.sign() * mt.abs()
             elif self.use_cautious:
@@ -237,6 +269,8 @@ class Adopt_adv(torch.optim.Optimizer):
                 mt_slow.mul_(beta3_ema).add_(normalized_grad, alpha=1.0 - beta3_ema)
                 update = torch.add(mt, m_slow, alpha=alpha_t)
                 update = update.view(p.shape)
+            elif self.Simplified_AdEMAMix:
+                update = torch.add(mt, grad_reshaped, alpha=alpha_grad)
             else:
                 update = mt.view(p.shape)
 
@@ -283,7 +317,10 @@ class Adopt_adv(torch.optim.Optimizer):
             del denom
 
             # ADOPT Step B: Update momentum m_t
-            m.mul_(beta1).add_(normalized_grad, alpha=1.0 - beta1)
+            if self.Simplified_AdEMAMix:
+                m.mul_(beta1).add_(normalized_grad, alpha=1.0)
+            else:
+                m.mul_(beta1).add_(normalized_grad, alpha=1.0 - beta1)
 
             if self.use_grams:
                 m = grad.sign() * m.abs()
@@ -296,6 +333,8 @@ class Adopt_adv(torch.optim.Optimizer):
             if self.use_AdEMAMix:
                 m_slow.mul_(beta3_ema).add_(normalized_grad, alpha=1.0 - beta3_ema)
                 update = torch.add(m, m_slow, alpha=alpha_t)
+            elif self.Simplified_AdEMAMix:
+                update = torch.add(m, grad, alpha=alpha_grad)
             else:
                 update = m.clone()
 

{adv_optm-0.1.8 → adv_optm-1.0.0}/adv_optm/optim/Prodigy_adv.py

@@ -127,8 +127,11 @@ class Prodigy_adv(torch.optim.Optimizer):
             raise ValueError(f"Weight-decay should be >= 0.0. Got {weight_decay}")
         if not (prodigy_steps >= 0):
             raise ValueError(f"prodigy_steps should be >= 0. Got {prodigy_steps}")
+        if use_cautious and use_grams:
+            print("Warning: use_cautious is incompatible with use_grams, Disabling use_cautious.")
+            use_cautious = False
         if betas[0] == 0.0 and Simplified_AdEMAMix:
-            raise ValueError(f"Beta 1 cannot be 0.0 when using Simplified_AdEMAMix. Got {betas[0]}")
+            raise ValueError(f"Beta1 cannot be 0.0 when using Simplified_AdEMAMix. Got {betas[0]}")
         if use_AdEMAMix and Simplified_AdEMAMix:
             print("Warning: use_AdEMAMix is incompatible with Simplified_AdEMAMix, Disabling use_AdEMAMix.")
         if use_grams and Simplified_AdEMAMix:

adv_optm-1.0.0/adv_optm.egg-info/PKG-INFO

@@ -0,0 +1,174 @@
+Metadata-Version: 2.4
+Name: adv_optm
+Version: 1.0.0
+Summary: A family of highly efficient, lightweight yet powerful optimizers.
+Home-page: https://github.com/Koratahiu/Advanced_Optimizers
+Author: Koratahiu
+Author-email: hiuhonor@gmail.com
+License: Apache 2.0
+Keywords: llm,fine-tuning,memory-efficient,low-rank,compression,pytorch,optimizer,adam
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.0
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: keywords
+Dynamic: license
+Dynamic: license-file
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# Advanced Optimizers (AIO)
+
+A comprehensive, all-in-one collection of optimization algorithms for deep learning, designed for maximum efficiency, minimal memory footprint, and superior performance across diverse model architectures and training scenarios.
+
+[![PyPI](https://img.shields.io/pypi/v/adv_optm)](https://pypi.org/project/adv_optm/)
+
+---
+
+## 📦 Installation
+
+```bash
+pip install adv_optm
+```
+
+---
+
+## 🧠 Core Innovations
+
+This library integrates multiple state-of-the-art optimization techniques validated through extensive research and practical training, with 1-bit compression for optimizer states:
+
+### **Memory-Efficient Optimization (SMMF-inspired)**
+- **Paper**: [SMMF: Square-Matricized Momentum Factorization](https://arxiv.org/abs/2412.08894)
+- **Approach**: Uses rank-1 non-negative matrix factorization with reconstruction cycle (factor → reconstruct → update → factor)
+- **Innovation**: 
+  - First moment split into **1-bit sign + absolute value**
+  - Final storage: **four factored vectors + one 1-bit sign state**
+  - Preserves Adam-like update quality with drastically reduced memory
+
+---
+
+## ⚡ Performance Characteristics
+
+### Memory Efficiency (SDXL Model - 6.5GB)
+| Optimizer | Memory Usage | Description |
+|-----------|--------------|-------------|
+| `Adopt_Factored` | 328 MB | 4 small vectors + 1-bit state |
+| `Adopt_Factored + AdEMAMix` | 625 MB | 6 small vectors + two 1-bit states |
+| `Simplified_AdEMAMix` | 328 MB | Same as standard factored (no extra state) |
+
+### Speed Comparison (SDXL, Batch Size 4)
+| Optimizer | Speed | Notes |
+|-----------|-------|-------|
+| `Adafactor` | ~8.5s/it | Baseline |
+| `Adopt_Factored` | ~10s/it | +18% overhead from compression |
+| `Adopt_Factored + AdEMAMix` | ~12s/it | +41% overhead (3 factored states) |
+
+---
+
+## 🧪 Available Optimizers
+
+### Standard Optimizers (All support `factored=True/False`)
+| Optimizer | Description | Best For |
+|-----------|-------------|----------|
+| `Adam_Adv` | Advanced Adam implementation | General purpose |
+| `Adopt_Adv` | Adam-variant with independent beta2 | Stable training for small batch size regimes |
+| `Prodigy_Adv` | Prodigy with D-Adaptation | Adam with automatic LR tuning |
+| `Simplified_AdEMAMix` | Adam variant with accumulator momentum | Small/large batch training when tuned correctly |
+| `Lion_Adv` | Advanced Lion implementation | Memory-constrained environments |
+| `Prodigy_Lion_Adv` | Prodigy + Lion combination | Lion with automatic LR tuning |
+
+### Feature Matrix
+| Feature | Adam_Adv | Adopt_Adv | Prodigy_Adv | Simplified_AdEMAMix | Lion_Adv |
+|---------|----------|-----------|-------------|---------------------|----------|
+| Factored | ✓ | ✓ | ✓ | ✓ | ✓ |
+| AdEMAMix | ✓ | ✓ | ✓ | ✗ | ✗ |
+| Simplified_AdEMAMix | ✗ | ✗ | ✓ | ✓ | ✗ |
+| OrthoGrad | ✓ | ✓ | ✓ | ✓ | ✓ |
+| Grams | ✓ | ✓ | ✓ | ✗ | ✗ |
+| Cautious | ✓ | ✓ | ✓ | ✗ | ✓ |
+| atan2 | ✓ | ✓ | ✓ | ✗ | ✗ |
+| Stochastic Rounding | ✓ | ✓ | ✓ | ✓ | ✓ |
+| Fused Backward Pass | ✓ | ✓ | ✓ | ✓ | ✓ |
+
+---
+
+## ⚙️ Key Features & Parameters
+
+### Comprehensive Feature Guide
+
+| Feature | Description | Recommended Usage | Performance Impact | Theoretical Basis | Compatibility |
+|---------|-------------|-------------------|--------------------|-------------------|--------------|
+| **Factored** | Memory-efficient optimization using rank-1 factorization | Enable for large models (>1B params) or limited VRAM | +12-41% time overhead, 1-bit memory usage | [SMMF](https://arxiv.org/abs/2412.08894) | All optimizers |
+| **AdEMAMix** | Dual EMA system for momentum | Use for long training runs (10k+ steps) | +1 state memory. | [AdEMAMix](https://arxiv.org/abs/2409.03137) | Adam/Adopt/Prodigy |
+| **Simplified_AdEMAMix** | Accumulator-based momentum | Small batch training (≤32) | Same memory as standard, no extra overhead | [Schedule-Free Connections](https://arxiv.org/abs/2502.02431) | Adam/Prodigy |
+| **OrthoGrad** | Removes gradient component parallel to weights | Full finetuning without weight decay | +33% time overhead, no memory impact | [Grokking at Edge](https://github.com/LucasPrietoAl/grokking-at-the-edge-of-numerical-stability) | All optimizers |
+| **Stochastic Rounding** | Improves precision for BF16 training | BF16 training | Minimal overhead (<5%) | [Revisiting BFloat16 Training](https://arxiv.org/abs/2010.06192) | All optimizers |
+| **atan2** | Robust eps replacement + built-in clipping | Use with Adopt or unstable training | No overhead | [Adam-atan2](https://github.com/lucidrains/adam-atan2-pytorch) | Adam/Adopt/prodigy |
+| **Cautious** | Update only when the direction align with the gradients | should faster the convergence | No overhead | [C-Optim](https://github.com/kyleliang919/C-Optim) | Adam/Adopt/prodigy |
+| **Grams** | Update direction from the gradients | should have a stronger effect than cautious | No overhead | [Grams](https://github.com/Gunale0926/Grams) | Adam/Adopt/prodigy |
+
+---
+
+## Simplified_AdEMAMix Parameters
+Simplified_AdEMAMix replaces standard momentum with an accumulator for better small-large batch performance.
+
+| Parameter | Recommended Values | Description |
+|-----------|---------------------|-------------|
+| `beta1` | 0.9 (large BS), 0.99-0.9999 (small BS) | Determines memory length of accumulator |
+| `alpha` | 100-10 (small BS), 1-0 (large BS) | Gradient smoothing factor |
+
+**Alpha Tuning Guide**:
+| Batch Size | Recommended α | Rationale |
+|------------|---------------|-----------|
+| Small (≤32) | 100, 50, 20, 10 | Emphasizes recent gradients for quick adaptation |
+| Medium (32-512) | 10, 5, 2, 1 | Balanced approach |
+| Large (≥512) | 1, 0.5, 0 | Emphasizes historical gradients for stability |
+
+⚠️ **Important**: Use **~100x smaller learning rate** with Simplified_AdEMAMix compared to AdamW (e.g., 1e-6 instead of 1e-4)
+
+### 📊 Performance Validation
+Small Batch Training (SDXL, BS=2, 1.8K steps)
+![Training Comparison](https://github.com/user-attachments/assets/7eff0671-cc59-47fc-8b63-d5205456d649)
+
+- **🟢 Prodigy_adv** (beta1=0.9, d0=1e-5): Final LR=2.9e-4
+- **🔵 Prodigy_adv + Simplified_AdEMAMix** (beta1=0.99, α=100, d0=1e-7): Final LR=5.8e-6
+
+**Results**:
+- Simplified_AdEMAMix shows faster convergence and better final performance
+- D-Adaptation automatically handles aggressive updates (50x smaller LR)
+- Generated samples show significantly better quality with Simplified_AdEMAMix
+
+---
+
+## ⚠️ Known Limitations
+
+### 1. Prodigy_Adv Sensitivity
+- Highly sensitive to gradient modifications (Adopt normalization, low-rank factorization)
+- May fail to increase learning rate in some LoRA scenarios
+- **Fix**: Disable factorization or set beta1=0
+
+### 2. Aggressive Learning Rates
+- Can destabilize factored first moment
+- **Recommendation**: Check Prodigy learning rate as reference for safe LR threshold
+
+---
+
+## 📚 References
+
+1. [SMMF: Square-Matricized Momentum Factorization](https://arxiv.org/abs/2412.08894)
+2. [The AdEMAMix Optimizer: Better, Faster, Older](https://arxiv.org/abs/2409.03137)
+3. [Connections between Schedule-Free Optimizers, AdEMAMix, and Accelerated SGD Variants](https://arxiv.org/abs/2502.02431)
+
+---

{adv_optm-0.1.8 → adv_optm-1.0.0}/setup.py

@@ -5,7 +5,7 @@ with open("README.md", "r", encoding="utf-8") as fh:
 
 setup(
     name="adv_optm",
-    version="0.1.8",
+    version="1.0.0",
     author="Koratahiu",
     author_email="hiuhonor@gmail.com",
     license='Apache 2.0',

adv_optm-0.1.8/PKG-INFO

@@ -1,130 +0,0 @@
-Metadata-Version: 2.4
-Name: adv_optm
-Version: 0.1.8
-Summary: A family of highly efficient, lightweight yet powerful optimizers.
-Home-page: https://github.com/Koratahiu/Advanced_Optimizers
-Author: Koratahiu
-Author-email: hiuhonor@gmail.com
-License: Apache 2.0
-Keywords: llm,fine-tuning,memory-efficient,low-rank,compression,pytorch,optimizer,adam
-Classifier: Programming Language :: Python :: 3
-Classifier: License :: OSI Approved :: Apache Software License
-Classifier: Operating System :: OS Independent
-Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
-Classifier: Topic :: Software Development :: Libraries :: Python Modules
-Requires-Python: >=3.8
-Description-Content-Type: text/markdown
-License-File: LICENSE
-Requires-Dist: torch>=2.0
-Dynamic: author
-Dynamic: author-email
-Dynamic: classifier
-Dynamic: description
-Dynamic: description-content-type
-Dynamic: home-page
-Dynamic: keywords
-Dynamic: license
-Dynamic: license-file
-Dynamic: requires-dist
-Dynamic: requires-python
-Dynamic: summary
-
-# Advanced Optimizers
-
-This repo introduces a new family of highly efficient, lightweight yet powerful optimizers, born from extensive research into recent academic literature and validated through practical training runs across diverse models.
-
----
-
-### Install
-
-`pip install adv_optm`
-
----
-
-### Theory (Inspired by SMMF)
-
-Based primarily on:  
-**[SMMF: Square-Matricized Momentum Factorization for Memory-Efficient Optimization](https://arxiv.org/abs/2412.08894)**
-
-The core innovation:
-- Uses fast, non-negative matrix factorization (NNMF - rank 1), but **reconstructs the full state before each update** to preserve momentum accuracy, then re-factors afterward (factor → reconstruct → update → factor cycle).
-- For the *signed first moment*, we split into **sign + absolute value**:
-  - Sign is stored as **1-bit state** via bitwise ops (SMMF originally used 8-bit with 7 bits wasted).
-  - Absolute value goes through the factor/reconstruct cycle using two factored vectors + the signed state.
-- Final storage: **four factored vectors + one 1-bit sign**.
-- Updates behave like full-state Adam but with drastically reduced memory.
-
-> ✅ **TL;DR**: Lightweight, strong, memory-efficient optimizer.
-
----
-
-### Memory Cost
-
-- **Adopt_Factored** for full SDXL finetune: **328 MB** (4 small vectors + 1-bit state)
-- **Adopt_Factored with AdEMAMix** for full SDXL finetune: **625 MB** (6 small vectors + two 1-bit states)
-> SDXL is 6.5GB model.
-
----
-
-### ⏱️ Speed (my tests in SDXL - BS 4)
-
-- **Adopt_Factored**: ~10s/it
-- **Adopt_Factored with AdEMAMix**: ~12s/it
-- **Adafactor**: ~8.5s/it  
-→ Overhead from compression/reconstruction cycles.
-→ It's faster than [MLorc](https://arxiv.org/abs/2506.01897) (~12s/it), which uses RSVD compression, and should be the fastest momentum compression (AFAIK).
-
----
-
-### 📈 Performance
-
-- **Better than Adafactor, and CAME factorzation methods**
-- **Comparable or identical to Adam** (see SMMF paper results)
-
----
-
-### Available Optimizers (all support `Factored` toggle)
-
-Set `Factored=False` to disable factorization and run as a full uncompressed optimizer (like vanilla Adam).
-
-1. **Adam**
-2. **Prodigy**
-3. **Adopt**
-
----
-
-### Bonus Features (Built-in)
-
-- **Fused Backward Pass**
-
-- **Stochastic Rounding (SR)**: Improves quality and convergence for **BF16 training**.
-
-- **[AdEMAMix](https://arxiv.org/abs/2409.03137)**  
-  → This adds a second, slow-moving EMA, which is combined with the primary momentum to stabilize updates, especially during long runs of full finetuning.
-  → A higher value of beta3 (e.g., 0.9999) gives the EMA a longer memory, making it more stable but slower to adapt. A lower value (e.g., 0.999) is often better for shorter training runs (2k-4k steps).
-  → When `factored` is true, it compresses the new momentum in the same way as the first moment (1-bit state + 2 vectors). However, this introduces noticeable overhead as we are compressing/reconstructing a third state each step.
-
-  ⚠️ **Note**: AdEMAMix updates are more aggressive than normal Adam/Adopt, so use a x2-x5 smaller LR than usual (or use Prodigy).
-
-- **[`atan2` smoothing & scaling](https://github.com/lucidrains/adam-atan2-pytorch)**  
-  → Robust `eps` replacement (no tuning!) + built-in gradient clipping  
-  → *Ideal for ADOPT* (which normally needs higher `eps` and clipping), so `use_atan2` is all-in-one for it.
-
-- **[OrthoGrad](https://github.com/LucasPrietoAl/grokking-at-the-edge-of-numerical-stability)**  
-  → Removes gradient component parallel to weights → prevents "naïve loss minimization" (NLM) → reduces natural overfitting  
-  → Perfect for fine-tuning the direction of existing features (e.g., full finetune or training a trained LoRA) without weight decay erasing prior knowledge.
-
-  ⚠️ **Note**: OrthoGrad introduces **~33% time overhead**, so take this into account.
-
-- **[Grams: Gradient Descent with Adaptive Momentum Scaling](https://github.com/Gunale0926/Grams)**  
-  → Eliminates the need for 1-bit momentum sign storage by using the **sign of gradients** for the first moment.
-
-  ⚠️ **Not recommended for small batch sizes**: gradients are too noisy, which can destabilize momentum (tested for Prodigy and it made the optimizer slower to find the LR or converge in BS 4).
-
-### Other Notes
-
-- **Adopt** skips the first step (only initializes the states) and has built-in clipping (sticking to the original optimizer), but we skip both of these when you enable `use_atan2`; as the optimizer becomes scale-invariant and the values of the states won't cause any issues or instability.
-
-- When `use_atan2` is True, `eps` will be ignored and you should also disable any gradient clipping.
-
----

adv_optm-0.1.8/README.md

@@ -1,99 +0,0 @@
-# Advanced Optimizers
-
-This repo introduces a new family of highly efficient, lightweight yet powerful optimizers, born from extensive research into recent academic literature and validated through practical training runs across diverse models.
-
----
-
-### Install
-
-`pip install adv_optm`
-
----
-
-### Theory (Inspired by SMMF)
-
-Based primarily on:  
-**[SMMF: Square-Matricized Momentum Factorization for Memory-Efficient Optimization](https://arxiv.org/abs/2412.08894)**
-
-The core innovation:
-- Uses fast, non-negative matrix factorization (NNMF - rank 1), but **reconstructs the full state before each update** to preserve momentum accuracy, then re-factors afterward (factor → reconstruct → update → factor cycle).
-- For the *signed first moment*, we split into **sign + absolute value**:
-  - Sign is stored as **1-bit state** via bitwise ops (SMMF originally used 8-bit with 7 bits wasted).
-  - Absolute value goes through the factor/reconstruct cycle using two factored vectors + the signed state.
-- Final storage: **four factored vectors + one 1-bit sign**.
-- Updates behave like full-state Adam but with drastically reduced memory.
-
-> ✅ **TL;DR**: Lightweight, strong, memory-efficient optimizer.
-
----
-
-### Memory Cost
-
-- **Adopt_Factored** for full SDXL finetune: **328 MB** (4 small vectors + 1-bit state)
-- **Adopt_Factored with AdEMAMix** for full SDXL finetune: **625 MB** (6 small vectors + two 1-bit states)
-> SDXL is 6.5GB model.
-
----
-
-### ⏱️ Speed (my tests in SDXL - BS 4)
-
-- **Adopt_Factored**: ~10s/it
-- **Adopt_Factored with AdEMAMix**: ~12s/it
-- **Adafactor**: ~8.5s/it  
-→ Overhead from compression/reconstruction cycles.
-→ It's faster than [MLorc](https://arxiv.org/abs/2506.01897) (~12s/it), which uses RSVD compression, and should be the fastest momentum compression (AFAIK).
-
----
-
-### 📈 Performance
-
-- **Better than Adafactor, and CAME factorzation methods**
-- **Comparable or identical to Adam** (see SMMF paper results)
-
----
-
-### Available Optimizers (all support `Factored` toggle)
-
-Set `Factored=False` to disable factorization and run as a full uncompressed optimizer (like vanilla Adam).
-
-1. **Adam**
-2. **Prodigy**
-3. **Adopt**
-
----
-
-### Bonus Features (Built-in)
-
-- **Fused Backward Pass**
-
-- **Stochastic Rounding (SR)**: Improves quality and convergence for **BF16 training**.
-
-- **[AdEMAMix](https://arxiv.org/abs/2409.03137)**  
-  → This adds a second, slow-moving EMA, which is combined with the primary momentum to stabilize updates, especially during long runs of full finetuning.
-  → A higher value of beta3 (e.g., 0.9999) gives the EMA a longer memory, making it more stable but slower to adapt. A lower value (e.g., 0.999) is often better for shorter training runs (2k-4k steps).
-  → When `factored` is true, it compresses the new momentum in the same way as the first moment (1-bit state + 2 vectors). However, this introduces noticeable overhead as we are compressing/reconstructing a third state each step.
-
-  ⚠️ **Note**: AdEMAMix updates are more aggressive than normal Adam/Adopt, so use a x2-x5 smaller LR than usual (or use Prodigy).
-
-- **[`atan2` smoothing & scaling](https://github.com/lucidrains/adam-atan2-pytorch)**  
-  → Robust `eps` replacement (no tuning!) + built-in gradient clipping  
-  → *Ideal for ADOPT* (which normally needs higher `eps` and clipping), so `use_atan2` is all-in-one for it.
-
-- **[OrthoGrad](https://github.com/LucasPrietoAl/grokking-at-the-edge-of-numerical-stability)**  
-  → Removes gradient component parallel to weights → prevents "naïve loss minimization" (NLM) → reduces natural overfitting  
-  → Perfect for fine-tuning the direction of existing features (e.g., full finetune or training a trained LoRA) without weight decay erasing prior knowledge.
-
-  ⚠️ **Note**: OrthoGrad introduces **~33% time overhead**, so take this into account.
-
-- **[Grams: Gradient Descent with Adaptive Momentum Scaling](https://github.com/Gunale0926/Grams)**  
-  → Eliminates the need for 1-bit momentum sign storage by using the **sign of gradients** for the first moment.
-
-  ⚠️ **Not recommended for small batch sizes**: gradients are too noisy, which can destabilize momentum (tested for Prodigy and it made the optimizer slower to find the LR or converge in BS 4).
-
-### Other Notes
-
-- **Adopt** skips the first step (only initializes the states) and has built-in clipping (sticking to the original optimizer), but we skip both of these when you enable `use_atan2`; as the optimizer becomes scale-invariant and the values of the states won't cause any issues or instability.
-
-- When `use_atan2` is True, `eps` will be ignored and you should also disable any gradient clipping.
-
----

adv_optm-0.1.8/adv_optm.egg-info/PKG-INFO

@@ -1,130 +0,0 @@
-Metadata-Version: 2.4
-Name: adv_optm
-Version: 0.1.8
-Summary: A family of highly efficient, lightweight yet powerful optimizers.
-Home-page: https://github.com/Koratahiu/Advanced_Optimizers
-Author: Koratahiu
-Author-email: hiuhonor@gmail.com
-License: Apache 2.0
-Keywords: llm,fine-tuning,memory-efficient,low-rank,compression,pytorch,optimizer,adam
-Classifier: Programming Language :: Python :: 3
-Classifier: License :: OSI Approved :: Apache Software License
-Classifier: Operating System :: OS Independent
-Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
-Classifier: Topic :: Software Development :: Libraries :: Python Modules
-Requires-Python: >=3.8
-Description-Content-Type: text/markdown
-License-File: LICENSE
-Requires-Dist: torch>=2.0
-Dynamic: author
-Dynamic: author-email
-Dynamic: classifier
-Dynamic: description
-Dynamic: description-content-type
-Dynamic: home-page
-Dynamic: keywords
-Dynamic: license
-Dynamic: license-file
-Dynamic: requires-dist
-Dynamic: requires-python
-Dynamic: summary
-
-# Advanced Optimizers
-
-This repo introduces a new family of highly efficient, lightweight yet powerful optimizers, born from extensive research into recent academic literature and validated through practical training runs across diverse models.
-
----
-
-### Install
-
-`pip install adv_optm`
-
----
-
-### Theory (Inspired by SMMF)
-
-Based primarily on:  
-**[SMMF: Square-Matricized Momentum Factorization for Memory-Efficient Optimization](https://arxiv.org/abs/2412.08894)**
-
-The core innovation:
-- Uses fast, non-negative matrix factorization (NNMF - rank 1), but **reconstructs the full state before each update** to preserve momentum accuracy, then re-factors afterward (factor → reconstruct → update → factor cycle).
-- For the *signed first moment*, we split into **sign + absolute value**:
-  - Sign is stored as **1-bit state** via bitwise ops (SMMF originally used 8-bit with 7 bits wasted).
-  - Absolute value goes through the factor/reconstruct cycle using two factored vectors + the signed state.
-- Final storage: **four factored vectors + one 1-bit sign**.
-- Updates behave like full-state Adam but with drastically reduced memory.
-
-> ✅ **TL;DR**: Lightweight, strong, memory-efficient optimizer.
-
----
-
-### Memory Cost
-
-- **Adopt_Factored** for full SDXL finetune: **328 MB** (4 small vectors + 1-bit state)
-- **Adopt_Factored with AdEMAMix** for full SDXL finetune: **625 MB** (6 small vectors + two 1-bit states)
-> SDXL is 6.5GB model.
-
----
-
-### ⏱️ Speed (my tests in SDXL - BS 4)
-
-- **Adopt_Factored**: ~10s/it
-- **Adopt_Factored with AdEMAMix**: ~12s/it
-- **Adafactor**: ~8.5s/it  
-→ Overhead from compression/reconstruction cycles.
-→ It's faster than [MLorc](https://arxiv.org/abs/2506.01897) (~12s/it), which uses RSVD compression, and should be the fastest momentum compression (AFAIK).
-
----
-
-### 📈 Performance
-
-- **Better than Adafactor, and CAME factorzation methods**
-- **Comparable or identical to Adam** (see SMMF paper results)
-
----
-
-### Available Optimizers (all support `Factored` toggle)
-
-Set `Factored=False` to disable factorization and run as a full uncompressed optimizer (like vanilla Adam).
-
-1. **Adam**
-2. **Prodigy**
-3. **Adopt**
-
----
-
-### Bonus Features (Built-in)
-
-- **Fused Backward Pass**
-
-- **Stochastic Rounding (SR)**: Improves quality and convergence for **BF16 training**.
-
-- **[AdEMAMix](https://arxiv.org/abs/2409.03137)**  
-  → This adds a second, slow-moving EMA, which is combined with the primary momentum to stabilize updates, especially during long runs of full finetuning.
-  → A higher value of beta3 (e.g., 0.9999) gives the EMA a longer memory, making it more stable but slower to adapt. A lower value (e.g., 0.999) is often better for shorter training runs (2k-4k steps).
-  → When `factored` is true, it compresses the new momentum in the same way as the first moment (1-bit state + 2 vectors). However, this introduces noticeable overhead as we are compressing/reconstructing a third state each step.
-
-  ⚠️ **Note**: AdEMAMix updates are more aggressive than normal Adam/Adopt, so use a x2-x5 smaller LR than usual (or use Prodigy).
-
-- **[`atan2` smoothing & scaling](https://github.com/lucidrains/adam-atan2-pytorch)**  
-  → Robust `eps` replacement (no tuning!) + built-in gradient clipping  
-  → *Ideal for ADOPT* (which normally needs higher `eps` and clipping), so `use_atan2` is all-in-one for it.
-
-- **[OrthoGrad](https://github.com/LucasPrietoAl/grokking-at-the-edge-of-numerical-stability)**  
-  → Removes gradient component parallel to weights → prevents "naïve loss minimization" (NLM) → reduces natural overfitting  
-  → Perfect for fine-tuning the direction of existing features (e.g., full finetune or training a trained LoRA) without weight decay erasing prior knowledge.
-
-  ⚠️ **Note**: OrthoGrad introduces **~33% time overhead**, so take this into account.
-
-- **[Grams: Gradient Descent with Adaptive Momentum Scaling](https://github.com/Gunale0926/Grams)**  
-  → Eliminates the need for 1-bit momentum sign storage by using the **sign of gradients** for the first moment.
-
-  ⚠️ **Not recommended for small batch sizes**: gradients are too noisy, which can destabilize momentum (tested for Prodigy and it made the optimizer slower to find the LR or converge in BS 4).
-
-### Other Notes
-
-- **Adopt** skips the first step (only initializes the states) and has built-in clipping (sticking to the original optimizer), but we skip both of these when you enable `use_atan2`; as the optimizer becomes scale-invariant and the values of the states won't cause any issues or instability.
-
-- When `use_atan2` is True, `eps` will be ignored and you should also disable any gradient clipping.
-
----