turboadam 0.1.0

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.

TurboAdam

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

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

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

pip install turboadam

For the latest source version:

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

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

Use

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

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 for v0.2.0.

---

Validation

# 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

@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