Metadata-Version: 2.4
Name: turboadam
Version: 0.1.0
Summary: Drop-in Adam/AdamW replacement with 6.5× optimizer-state memory reduction. One line change, no model modifications. Compresses both moments in-place during training with bounded per-element error guarantees.
Author-email: David Kogan <davidkny22@gmail.com>
License: MIT
Project-URL: Homepage, https://github.com/davidkny22/turboadam
Project-URL: Repository, https://github.com/davidkny22/turboadam
Project-URL: Issues, https://github.com/davidkny22/turboadam/issues
Keywords: optimizer,adam,adamw,pytorch,memory-efficient,quantization,deep-learning,training
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.2.0
Requires-Dist: numpy>=1.26.0
Provides-Extra: triton
Requires-Dist: triton>=2.2.0; extra == "triton"
Provides-Extra: dev
Requires-Dist: pytest>=8.0.0; extra == "dev"
Requires-Dist: matplotlib>=3.8.0; extra == "dev"
Requires-Dist: datasets>=2.18.0; extra == "dev"
Requires-Dist: transformers>=4.40.0; extra == "dev"
Dynamic: license-file

# TurboAdam

[![Tests](https://img.shields.io/badge/tests-151%2F151-brightgreen)]() [![Python](https://img.shields.io/badge/python-3.10%2B-blue)]() [![PyTorch](https://img.shields.io/badge/PyTorch-2.2%2B-orange)]() [![License](https://img.shields.io/badge/license-MIT-green)]()

**Drop-in Adam/AdamW replacement with 6.5× optimizer-state memory reduction.**

One line change. No model modifications. No training-loop changes.

```python
from turboadam import TurboAdam

optimizer = TurboAdam(model.parameters(), lr=1e-3)
```

---

## Why TurboAdam?

Adam stores two full-precision copies of every parameter (first and second moments). For a 7B model that is **28 GB** of optimizer state alone — often the memory bottleneck that forces smaller batch sizes or shorter context lengths.

TurboAdam compresses both moments in-place during training, cutting optimizer-state memory from **64 bits/param → 9.9 bits/param** (6.5× reduction). On GPT-2 124M it converges within **+0.25 loss points** of full-precision AdamW (**1.2% relative** — within run-to-run noise).

| Model size | AdamW optimizer state | TurboAdam | Savings |
|-----------|----------------------|-----------|---------|
| 125M (GPT-2) | 0.50 GB | 0.08 GB | **0.42 GB** |
| 7B | 28.0 GB | **4.3 GB** | **23.7 GB** |
| 70B | 280.0 GB | **43.0 GB** | **237.0 GB** |

---

## Quick start

### Install

```bash
pip install turboadam
```

For the latest source version:

```bash
pip install git+https://github.com/davidkny22/turboadam.git
```

Requirements: Python >=3.10, PyTorch >=2.2, Triton (optional, for CUDA speed-ups).

### Use

```python
from turboadam import TurboAdam

# Drop-in replacement for torch.optim.AdamW
optimizer = TurboAdam(
    model.parameters(),
    lr=6e-4,
    betas=(0.9, 0.999),
    weight_decay=0.01,
    v_bits=4,          # 2, 3, 4, 6, or 8
    compress_m=True,   # CoState first-moment compression
    compress_v=True,   # Log-scale second-moment compression
)
```

---

## How it works

TurboAdam combines two independent, separable compression techniques. You can enable either or both.

### 1Q — Second-moment (v) compression

v is stored as n-bit **log-scale quantized** values per 128-element block:

1. **Decompress** block min/max → reconstruct v via exp interpolation
2. **EMA update**: `v_new = β₂·v_old + (1-β₂)·g²`
3. **Bias-correct** denominator: `denom = √(v / (1-β₂ᵗ)) + ε`
4. **Re-compress** with **stochastic rounding** (unbiased — prevents systematic EMA drift)

Storage per block: `n_bits` uint8 indices + 2× fp16 scales.  
Default **4-bit** = **4.25 bits/param**.

**Key insight:** Theoretical analysis predicted 4-bit would fail due to accumulated quantization noise (22× amplification from β₂=0.999 EMA). In practice it works because quantization errors are correlated — same elements map to the same buckets step-to-step.

### CoState — First-moment (m) compression

Gradient-residual decomposition: `m = α·g + δ`

- `α = (m·g) / (g·g)` — scalar projection onto current gradient
- `δ = m - α·g` — residual orthogonal to gradient

δ is partitioned into 128-element blocks and classified into three **costates**:

| Costate | Condition | Storage | Typical share |
|---------|-----------|---------|---------------|
| **Null** | `r < P₁₀` | 1-bit flag | ~10% |
| **Phase** | `P₁₀ ≤ r < P₉₀` | 1-bit sign per element | ~80% |
| **Amplitude** | `r ≥ P₉₀` | 1-bit sign + fp16 block scale | ~10% |

**Key insight:** For Adam, direction matters more than magnitude because `m/√v` normalizes per-element. Sign-only encoding preserves direction for 80% of components. This is why CoState works at ~2 bits/param while low-rank approaches fail — they preserve magnitude for few directions but lose direction for many.

---

## Results

### Memory

Measured on one GPT-2 layer (9 parameter tensors, CUDA).

| Configuration | Persistent optimizer memory | vs AdamW |
|--------------|----------------------------|----------|
| AdamW (baseline) | 56.6 MB | 1.00× |
| TurboAdam (v only, 4-bit) | 35.6 MB | **0.63×** |
| TurboAdam (m only, CoState) | 29.6 MB | **0.52×** |
| **TurboAdam (m + v, default)** | **8.6 MB** | **0.15×** |

### Speed

Measured on one GPT-2 layer, RTX 4070, 200-step average.

| Configuration | Time/step | vs AdamW |
|--------------|-----------|----------|
| AdamW (baseline) | 12.0 ms | 1.00× |
| TurboAdam (v only) | 8.4 ms | **0.70×** |
| TurboAdam (m + v, default) | 17.0 ms | **1.41×** |

The v-only path is actually **faster** than AdamW because 4-bit log-scale decompression is cheaper than full fp32 EMA updates on small tensors. The m+v path adds ~40% overhead from CoState encode/decode.

### Convergence — GPT-2 124M on WikiText-103

| Configuration | Loss @ step 500 | Gap vs AdamW |
|--------------|-----------------|--------------|
| AdamW (full fp32) | 19.28 | — |
| TurboAdam (8-bit v + CoState) | 19.79 | +0.51 |
| **TurboAdam (4-bit v + CoState, default)** | **19.58** | **+0.25** |
| TurboAdam (CoState only, fp32 v) | 19.80 | +0.52 |
| TurboAdam (v only, fp32 m) | 19.28 | ~0.00 |

The +0.25 gap is structural to CoState's sign-only encoding and shrinks as training progresses (+2.94 at step 50, +0.25 at step 500). Threshold tuning and error feedback do not reduce it. For workloads where every tenth of a point matters, run with `compress_m=False` for v-only compression at zero convergence cost.

---

## API

```python
TurboAdam(
    params,                    # iterable of parameters or param groups
    lr=1e-3,                   # learning rate
    betas=(0.9, 0.999),        # (β₁, β₂) EMA decay coefficients
    eps=1e-8,                  # numerical stability
    weight_decay=0.0,          # AdamW-style decoupled weight decay
    block_size=128,            # quantization block size (elements)
    v_bits=4,                  # bits per element for v: 2, 3, 4, 6, or 8
    compress_m=True,           # enable CoState m compression
    compress_v=True,           # enable v compression
    null_pct=0.10,             # CoState null threshold percentile
    amp_pct=0.90,              # CoState amplitude threshold percentile
    error_feedback=False,      # CoState error feedback (tested, no improvement)
    capturable=False,          # CUDA graph capture (not yet supported)
    min_m_compress_elements=4096,  # minimum param size for CoState m compression
)
```

All arguments are standard PyTorch Optimizer kwargs plus TurboAdam-specific compression controls. State dicts are fully compatible with `torch.save` / `torch.load`.

**Notes:**
- `torch.compile` will graph-break at `opt.step()` (expected for Python-loop optimizers; does not affect correctness).
- FSDP / DeepSpeed ZeRO compatibility is on the [roadmap](ROADMAP.md) for v0.2.0.

---

## Validation

```bash
# Full test suite (151 tests)
python -m pytest tests/ -q

# Quick convergence smoke test
python -c "
import torch, torch.nn as nn
from turboadam import TurboAdam

torch.manual_seed(0)
x = nn.Parameter(torch.randn(50, device='cuda'))
opt = TurboAdam([x], lr=1e-2)
for _ in range(200):
    opt.zero_grad()
    loss = (x**2).sum()
    loss.backward()
    opt.step()
print(f'Final loss: {loss.item():.6f}')  # < 5% of initial
"

# GPT-2 124M training run (~36 min on RTX 4070)
python experiments/train_turboadam.py --steps 500 --log_every 50

# Speed benchmark
python scripts/benchmark_speed.py

# Memory profiler
python scripts/profile_memory.py
```

---

## Design decisions

1. **Compress-every-step (not freeze-refresh).** The original design froze v for 1000 steps and refreshed periodically. This caused a +3.75 loss gap from v staleness. Compress-every-step with stochastic rounding eliminates staleness — the EMA runs continuously on the compressed state.

2. **4-bit default.** 4-bit gives 6.5× compression with +0.25 gap. 8-bit gives 4.1× with +0.51. The sweet spot is 4-bit — going higher barely improves precision, going lower risks noise accumulation.

3. **Stochastic rounding.** Unbiased rounding prevents systematic drift in the EMA. Without it, deterministic rounding accumulates a bias of ~1000× the per-step error (for β₂=0.999).

4. **Sign-only for CoState (not low-rank).** We tested LoRA-Pre style low-rank projection (rank 8–512). It fails for Adam because momentum is NOT low-rank — rank-8 captures only 4% of energy. Sign-only encoding captures direction for ALL elements, which is what Adam's per-coordinate denominator normalization needs.

5. **P10/P90 thresholds.** Extensive testing showed threshold changes (P5/P85, P5/P80, P10/P95, etc.) produce identical convergence. The gap is structural to sign encoding, not the null/phase/amplitude split.

---

## Project status

- **Phase 1** (current): RTX 4070 8GB, models ≤ 125M — **complete**. Correctness validated, speed optimized, Triton kernels production-ready.
- **Phase 2** (next): DGX Spark 128GB, models up to 7B — pending hardware.

---

## Citation

```bibtex
@misc{kogan2026turboadam,
  title={TurboAdam: Memory-Efficient Adam via In-Place Optimizer State Compression},
  author={Kogan, David},
  year={2026},
  howpublished={\url{https://github.com/davidkogan/turboadam}}
}
```

---

## License

MIT
