Metadata-Version: 2.4
Name: nanofold
Version: 0.1.0
Summary: FOLD compression: fit large models into small VRAM at near-lossless quality
Author: Nathan
License: MIT
Project-URL: Homepage, https://github.com/wisk-inc/aimet
Project-URL: Source, https://github.com/wisk-inc/aimet/tree/main/nanofold
Project-URL: Issues, https://github.com/wisk-inc/aimet/issues
Keywords: quantization,compression,vram,llm,low-rank,int4,model-compression,inference,offloading
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
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
Requires-Dist: numpy>=1.21
Provides-Extra: dev
Requires-Dist: pytest>=7; extra == "dev"
Requires-Dist: build; extra == "dev"
Requires-Dist: twine; extra == "dev"
Dynamic: license-file

# nanofold

**FOLD** — Fused Outlier and Low-rank Delta. Fit large models into small VRAM without paying for it in quality.

```bash
pip install nanofold
```

*Made by Nathan.*

---

## Why quantisation loses accuracy, and what FOLD does about it

Below 4 bits, round-to-nearest quantisation fails in a *structured* way. The error it leaves behind is not white noise:

- it is **heavy-tailed** — a handful of weights per tensor set each group's dynamic range and cost every other weight in that group its precision;
- it is **strongly correlated across rows** — the residual of a round-to-nearest pass is close to low-rank;
- and it is **spread evenly across input channels**, even though only some of those channels carry signal the model actually uses.

Plain quantisation throws all three away. FOLD stores a weight matrix as three parts that are cheap in very different ways, sized to recover exactly that structure:

```
W  ≈   S   +   dequant(Q)   +   A @ B
       │           │              │
       │           │              └── dense low-rank delta, rank r, fp16   (~6% of bytes)
       │           └───────────────── group-wise b-bit integer codes       (~93% of bytes)
       └───────────────────────────── sparse full-precision outliers       (~1% of bytes)
```

Two things make this more than the sum of its parts.

**The three parts are fitted jointly, not in sequence.** After fitting `A @ B` to the quantisation residual, FOLD re-quantises `W − A@B` rather than `W`, and refits. Each sweep hands the integer grid an easier target. Three sweeps converge to an error well below what any stage reaches alone.

**The residual fit is weighted by activation salience.** A plain SVD minimises `‖E − AB‖_F`, treating every input channel as equally important. What actually matters is error at the layer's *output*. FOLD weights each column by how strongly that channel fires on calibration data, so the objective approximates output error — and rank is never spent on channels the model does not excite.

## Measured results

`facebook/opt-125m`, perplexity on 24 × 512-token chunks of held-out prose. Calibration used 8 chunks from a **disjoint** part of the corpus. fp32 baseline perplexity **29.82**. Reproduce with `benchmarks/opt125m.py`.

| method | bits/weight | size | perplexity | vs fp32 |
|---|---:|---:|---:|---:|
| RTN 4-bit, group 64 | 4.40 | 67.9 MB | 32.40 | +8.7% |
| RTN 4-bit, group 32 | 4.77 | 73.7 MB | 31.61 | +6.0% |
| RTN 4-bit, group 16 | 5.52 | 85.3 MB | 31.39 | +5.3% |
| **FOLD 4-bit, group 64** | **4.65** | **71.7 MB** | **30.91** | **+3.6%** |
| RTN 3-bit, group 64 | 3.40 | 52.5 MB | 42.66 | +43.0% |
| RTN 3-bit, group 32 | 3.77 | 58.2 MB | 38.38 | +28.7% |
| RTN 3-bit, group 16 | 4.52 | 69.8 MB | 35.26 | +18.2% |
| **FOLD 3-bit, group 64** | **3.65** | **56.3 MB** | **35.17** | **+17.9%** |
| RTN 2-bit, group 64 | 2.40 | 37.0 MB | 3352.99 | +11143% |
| **FOLD 2-bit** | **2.89** | **44.7 MB** | **91.08** | **+205%** |

Reading the table:

- **At 4 bits, FOLD beats plain quantisation at every budget tested** — including RTN with group size 16, which costs **19% more bytes** (5.52 vs 4.65 bits/weight) and is still worse. Against RTN at the nearest matched size, FOLD cuts the quality loss from +6.0% to +3.6%.
- **At 3 bits, FOLD matches RTN group-16 quality using 19% fewer bits** — 35.17 at 3.65 bits/weight against 35.26 at 4.52, saving 13.5 MB on a 125M model. On quality this is a tie; the win is the 13.5 MB. Unlike the 4-bit result, it depends on having enough calibration data: with 2 calibration chunks instead of 8 it slips to +22.1% and RTN group-16 wins.
- **At 2 bits, plain quantisation collapses entirely** and FOLD survives. Neither is usable; FOLD just degrades gracefully instead of catastrophically.

### Where the gain comes from

Same model, 4 bits, each stage added on its own:

| configuration | perplexity | vs fp32 |
|---|---:|---:|
| RTN 4-bit, group 64 | 32.40 | +8.7% |
| + sparse outliers only | 31.87 | +6.9% |
| + low-rank delta only | 32.18 | +7.9% |
| + both, no calibration | 31.89 | +6.9% |
| **+ both, calibrated** | **30.91** | **+3.6%** |

Calibration is the single biggest contributor — the low-rank delta and the outliers each help, but weighting them by what the model actually uses is what roughly halves the remaining loss.

### Honest caveats

- Measured on a 125M-parameter model. Larger models generally quantise *better*, so these numbers are likely pessimistic — but they are what was measured, not extrapolated.
- Calibration quality matters, and more so the lower the bit width. The 4-bit advantage held under every split tried; the 3-bit one needs a reasonable calibration set to hold up.
- The advantage scales with how much structure a weight has. On a synthetic iid Gaussian with no channel spread and no heavy tail, the gain drops to about 6%, because there is nothing left to exploit. Trained weights have both in abundance; freshly initialised ones do not.
- FOLD costs more to *apply* than round-to-nearest: three alternating sweeps and a randomised SVD per layer. This is a one-time offline cost.
- 2-bit is not a usable operating point for a model this size under either method.

## Quick start

```python
import nanofold

model = ...  # any nn.Module

# optional but worth it: a few hundred real tokens
salience = nanofold.collect_salience(model, calibration_batches)

nanofold.fold_model(model, nanofold.FoldConfig(bits=4), salience=salience)

report = nanofold.compression_report(model)
print(f"{report['ratio']:.2f}x smaller, {report['bits_per_weight']:.2f} bits/weight")

nanofold.save_folded(model, "model.nfold")
```

Loading costs almost no resident memory — the container is memory-mapped, so weights are paged in by the OS on first touch:

```python
model = nanofold.load_folded(MyModel(), "model.nfold")
```

The `.nfold` container is a JSON header plus aligned raw blobs. There is no pickle anywhere in the load path, so opening a checkpoint cannot execute code from it.

### Fit a model to a VRAM budget

Layers differ enormously in how much they resent being quantised, so a uniform bit width either wastes budget on robust layers or starves fragile ones. The planner starts every layer at the cheapest setting and repeatedly spends the next byte where it buys the most error reduction:

```python
plan = nanofold.plan_budget(model, budget_bytes=6 * 1024**3, salience=salience)
nanofold.fold_model(model, per_layer=plan)
```

### Stream weights from host memory

Classic layer offloading is bottlenecked on PCIe because it ships fp16 weights across the bus on every forward pass. Folded weights stay *packed* all the way across and are expanded on the accelerator, so bus traffic falls by the full compression ratio. Transfers for layer *i+1* are issued on a side stream while layer *i* computes:

```python
paged = nanofold.page_model(model, device="cuda", prefetch=2)
out = paged.model(x)   # resident VRAM ≈ compressed size + one layer of scratch
```

`FoldLinear` never holds a dense parameter: each forward expands the codes into a temporary that is released as soon as the matmul is done.

### Command line

```bash
nanofold compress model.pt -o model.nfold --bits 4   # fold a checkpoint
nanofold inspect model.nfold -v                      # what's inside
nanofold expand model.nfold -o dense.pt              # back to dense
nanofold bench                                       # FOLD vs RTN, reproducible
```

## Configuration

| option | default | what it does |
|---|---|---|
| `bits` | `4` | integer width, 1–8. Sub-byte widths are packed with no padding — 3 bits really costs 3 bits |
| `group_size` | `64` | weights per affine scale, along the input dimension |
| `rank` | *auto* | absolute rank of the delta; leave unset to derive it from `rank_ratio` |
| `rank_ratio` | `1/128` | rank as a fraction of `min(out, in)` |
| `outlier_fraction` | `0.001` | share of weights kept at full precision |
| `iters` | `3` | alternating quantise/refit sweeps |
| `alpha` | `0.5` | salience equalisation strength; `0` disables it |

**Why `rank` defaults to a ratio.** The factors cost `(out + in) · rank`, while the codes cost `out · in · bits / 8`. A rank that is 6% overhead on a 4096-wide matrix is 25% on a 1024-wide one. Scaling rank with the matrix keeps that overhead flat, which is what lets one config work across a whole model.

## Development

```bash
pip install -e ".[dev]"
pytest                 # 161 tests
ruff check src tests
```

## License

MIT — see [LICENSE](LICENSE).
