Metadata-Version: 2.4
Name: galore-adamw
Version: 0.1.0
Summary: GaLore AdamW — Memory-efficient optimizer using Gradient Low-Rank Projection for training LLMs on consumer GPUs
Author: Agustin
License-Expression: MIT
Project-URL: Homepage, https://github.com/agustin/galore-adamw
Project-URL: Issues, https://github.com/agustin/galore-adamw/issues
Keywords: optimizer,galore,adamw,low-rank,deep-learning,pytorch,memory-efficient,llm
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.0.0
Dynamic: license-file

# 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

```bash
pip install galore-adamw
```

## Quick Start

```python
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

| Parameter | Default | Description |
|---|---|---|
| `lr` | `1e-3` | Learning rate |
| `rank` | `128` | Max rank of gradient projection |
| `update_proj_gap` | `200` | Steps between SVD re-projection |
| `scale` | `1.0` | Scale factor for projected update |
| `proj_type` | `"std"` | `"std"` (rows) or `"reverse"` (cols) |
| `adaptive_rank` | `True` | Auto-scale rank per layer |
| `use_8bit_states` | `False` | Quantize m,v to int8 |
| `use_sophia_clip` | `False` | Diagonal Hessian clipping |
| `sophia_rho` | `0.03` | Clipping threshold |
| `weight_decay` | `0.01` | Decoupled weight decay |
| `max_grad_norm` | `1.0` | Gradient clipping (0=disabled) |

## Memory Stats

```python
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
