GaLore AdamW

Memory-efficient PyTorch optimizer using Gradient Low-Rank Projection (GaLore) for training LLMs on consumer GPUs.

Projects gradients to a low-rank subspace before computing optimizer states, reducing memory by 8-32x for large weight matrices.

Features

• GaLore projection: Randomized SVD projects gradients to low-rank subspace
• 8-bit quantized states: Optional int8 quantization of Adam m/v (4x extra savings)
• Sophia-style clipping: Optional diagonal Hessian clipping for stability
• Adaptive rank: Automatically scales projection rank per layer size

Installation

pip install galore-adamw

Quick Start

from galore_adamw import GaLoreAdamW, GaLoreConfig

# Basic usage
cfg = GaLoreConfig(lr=1e-3, rank=128)
optimizer = GaLoreAdamW(model.parameters(), cfg)

# With 8-bit states + Sophia clipping (maximum memory savings)
cfg = GaLoreConfig(
    lr=1e-3,
    rank=128,
    use_8bit_states=True,
    use_sophia_clip=True,
)
optimizer = GaLoreAdamW(model.parameters(), cfg)

# Training loop (standard PyTorch)
for batch in dataloader:
    loss = model(batch)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

Configuration

Memory Stats

stats = optimizer.get_memory_stats()
print(f"GaLore: {stats['galore_state_mb']:.1f} MB")
print(f"AdamW:  {stats['adamw_state_mb']:.1f} MB")
print(f"Savings: {stats['savings_ratio']:.1f}x")

License

MIT