Hey everyone
At OpenBlade, we just open-sourced Avalanche, a standalone, ultra-low VRAM matrix-factorized optimizer designed to break the memory wall in foundation model pretraining and fine-tuning.
In deep model training, optimizer states frequently consume significantly more VRAM than the model weights themselves.
Standard FP32 AdamW requires storing two states per parameter—the first momentum vector m and the second uncentered variance vector $v$—costing 8 bytes of VRAM per parameter:
Instead of storing an M \times N dense grid of running squared gradients, Avalanche shifts second-moment estimation to a Rank-1 marginal factorization:
where row marginals r \in \mathbb{R}^M and column marginals c \in \mathbb{R}^N track mean squared gradients across rows and columns.
| Model Scale | Parameters | Standard FP32 AdamW (m+vm + vm+v) | 8-bit Adam (bitsandbytes) | Avalanche State | Net Reduction |
|---|---|---|---|---|---|
| Small / Edge | 125M | ~1.00 GB | ~250 MB | ~1.8 MB | 99.82% |
| Foundation | 1.12B | ~8.96 GB | ~2.24 GB | ~9.6 MB | 99.89% |
| 7B Base | 7.00B | ~56.00 GB | ~14.00 GB | ~24.5 MB | 99.96% |
| 13B Scale | 13.00B | ~104.00 GB | ~26.00 GB | ~42.0 MB | 99.96% |
| 70B Scale | 70.00B | ~560.00 GB | ~140.00 GB | ~78.0 MB | 99.99% |
Avalanche supports multiple first-moment momentum modes depending on your available VRAM headroom while keeping second moments factorized:
Avalanche (mode="none"): Avalanche4Bit (mode="4bit"): Avalanche8Bit (mode="8bit"): AvalancheBF16 / AvalancheFP16: AvalancheFP32:
In quantized modes (4bit, 8bit), quantization residuals are stored in compact half-precision buffers (e_t = m_t - \hat{m}_t) and injected back into the next momentum step:
Because residuals accumulate over time, no gradient signal or update magnitude is discarded.
Avalanche is a drop-in replacement inheriting directly from torch.optim.Optimizer:
Python
import torch
import torch.nn as nn
from avalanche import Avalanche, Avalanche4Bit
model = nn.Transformer(d_model=1024, nhead=16).cuda()
optimizer = Avalanche(model.parameters(), lr=1e-4, weight_decay=0.01)
optimizer = Avalanche4Bit(model.parameters(), lr=1e-4, betas=(0.9, 0.999), lossless=True)
for x, y in data:
optimizer.zero_grad(set_to_none=True)
loss = criterion(model(x), y)
loss.backward()
optimizer.step()
Trainer
from transformers import Trainer
from avalanche import Avalanche
class CustomTrainer(Trainer):
def create_optimizer(self):
self.optimizer = Avalanche(
self.model.parameters(),
lr=self.args.learning_rate,
weight_decay=self.args.weight_decay,
betas=(0.0, 0.999),
)
return self.optimizer
You can install it via pip or drop optimizer.py directly into your repository:
Bash
pip install git+https://github.com/OpenBladeAI/Avalanche.git
Give it a run on your setups or clusters and share your training loss curves with us here on the forum!