# QuarterBit AXIOM

**High-Performance Quantized Optimizer for PyTorch**

Drop-in replacement for AdamW with significant memory savings.

[![PyPI](https://img.shields.io/pypi/v/quarterbit)](https://pypi.org/project/quarterbit/)

---

## Installation

```bash
pip install quarterbit
```

**Requirements:**
- Python 3.8+
- PyTorch 1.8+
- NVIDIA GPU with CUDA

**Supported GPUs:**
- Consumer: GTX 1650+, RTX 20/30/40 series
- Data Center: T4, V100, A10, A100, L4, L40, H100, H200

---

## Quick Start

```python
from quarterbit import Axiom

optimizer = Axiom(model.parameters())

for batch in dataloader:
    loss = model(batch)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
```

---

## Optimizer Parameters

AXIOM comes with tuned defaults. Advanced users can adjust:

```python
from quarterbit import Axiom

optimizer = Axiom(
    params,                    # Model parameters
    lr=5e-4,                   # Learning rate
    betas=(0.9, 0.999),        # Momentum coefficients
    eps=1e-8,                  # Numerical stability
    weight_decay=0.1,          # Decoupled weight decay
    total_steps=10000,         # Training steps (for scheduling)
    warmup_ratio=0.1,          # Warmup fraction
)
```

### Per-Layer Learning Rates

```python
optimizer = Axiom([
    {'params': model.backbone.parameters(), 'lr': 1e-4},
    {'params': model.head.parameters(), 'lr': 5e-4},
], weight_decay=0.1)
```

---

## Activation Checkpointing (Pro)

Custom compressed checkpointing with slot-based storage.

```python
from quarterbit.torch.utils import ActivationCheckpoint

# Create checkpoint storage
actcp = ActivationCheckpoint(
    max_slots=24,              # Number of layers/checkpoints
    max_elements=2**20,        # Max elements per activation
)

# During forward pass - store activations
actcp.store(activation, slot=layer_idx)

# During backward pass - restore activations
restored = actcp.restore(slot=layer_idx)

# Check memory savings
info = actcp.memory_info()
print(f"Memory saved: {info['savings_pct']:.1f}%")

# Clear when done
actcp.clear()
```

**Parameters:**
| Parameter | Description |
|-----------|-------------|
| `max_slots` | Number of checkpoint slots (typically num_layers) |
| `max_elements` | Maximum tensor size per slot |

**Methods:**
| Method | Description |
|--------|-------------|
| `store(tensor, slot)` | Compress and store activation |
| `restore(slot)` | Decompress and return activation |
| `clear(slot=None)` | Clear one slot or all |
| `memory_info()` | Get memory usage statistics |

---

## Gradient Compression (Pro)

Drift-free compression with error feedback for distributed training.

```python
from quarterbit.torch.utils import GradientCompressor

# Create compressor for each parameter
compressor = GradientCompressor(num_elements=param.numel())

# Compress before all-reduce
compressed = compressor.compress(param.grad)

# ... communicate compressed gradients ...

# Decompress after communication
param.grad = compressor.decompress(compressed)

# Reset error feedback (optional, between epochs)
compressor.reset()
```

**Parameters:**
| Parameter | Description |
|-----------|-------------|
| `num_elements` | Number of gradient elements |

**Methods:**
| Method | Description |
|--------|-------------|
| `compress(grads)` | Compress with error feedback |
| `decompress(compressed)` | Decompress to FP32 |
| `reset()` | Clear error feedback accumulator |

---

## Multi-GPU Training

### DataParallel

```python
model = torch.nn.DataParallel(model)
optimizer = Axiom(model.parameters())
```

### DistributedDataParallel

```python
from torch.nn.parallel import DistributedDataParallel as DDP

model = DDP(model, device_ids=[local_rank])
optimizer = Axiom(model.parameters())
```

### DeepSpeed

```python
import deepspeed

model, optimizer, _, _ = deepspeed.initialize(
    model=model,
    optimizer=Axiom(model.parameters()),
    config=ds_config
)
```

---

## Checkpointing

```python
# Save
torch.save({
    'model': model.state_dict(),
    'optimizer': optimizer.state_dict(),
    'step': step,
}, 'checkpoint.pt')

# Load
ckpt = torch.load('checkpoint.pt')
model.load_state_dict(ckpt['model'])
optimizer.load_state_dict(ckpt['optimizer'])
```

---

## Licensing

```bash
quarterbit activate <LICENSE_KEY>
```

| Tier | GPU Hours | Features |
|------|-----------|----------|
| **Free** | 10/month | Optimizer only |
| **Pro** | 1,000/month | + Checkpointing, Compression |
| **Enterprise** | Unlimited | + On-premise, Custom SLA |

Get your key: [quarterbit.dev/pricing](https://quarterbit.dev/pricing)

---

## Links

- [Website](https://quarterbit.dev)
- [Documentation](https://quarterbit.dev/docs)
- [Pricing](https://quarterbit.dev/pricing)
- [Support](mailto:info@quarterbit.dev)

---

Copyright (c) 2026 Clouthier Simulation Labs
