Metadata-Version: 2.4
Name: nanofold
Version: 0.2.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 x 512-token chunks of held-out prose, with calibration on 8 chunks from a **disjoint** part of the corpus. fp32 baseline perplexity **29.82**; fp16 model footprint **250.5 MB**. Reproduce with `benchmarks/opt125m.py`.

Size is the **whole model**, not just the compressed part -- a weight left unfolded still has to be paid for at fp16, and reporting only what was compressed flatters the result.

| method | bits/weight | total size | perplexity | vs fp32 |
|---|---:|---:|---:|---:|
| RTN 4-bit, group 64 | 4.39 | 90.0 MB | 32.95 | +10.5% |
| RTN 4-bit, group 32 | 4.77 | 97.6 MB | 31.79 | +6.6% |
| RTN 4-bit, group 16 | 5.52 | 113.0 MB | 31.48 | +5.6% |
| **FOLD 4-bit, group 64** | **4.63** | **94.8 MB** | **31.19** | **+4.6%** |
| RTN 3-bit, group 64 | 3.39 | 69.5 MB | 45.62 | +53.0% |
| RTN 3-bit, group 32 | 3.77 | 77.2 MB | 39.85 | +33.6% |
| RTN 3-bit, group 16 | 4.52 | 92.5 MB | 36.55 | +22.5% |
| **FOLD 3-bit, group 64** | **3.63** | **74.4 MB** | **36.75** | **+23.2%** |
| RTN 2-bit, group 64 | 2.39 | 49.0 MB | 4183.95 | +13929% |
| **FOLD 2-bit** | **2.87** | **58.8 MB** | **149.02** | **+400%** |

Reading the table:

- **At 4 bits FOLD wins outright.** 94.8 MB at +4.6% beats RTN's best 4-bit setting, which needs 113.0 MB -- **19% more memory** -- to land at +5.6%. Against RTN at a comparable size, quality loss drops from +6.6% to +4.6%.
- **At 3 bits FOLD trades quality for size.** 74.4 MB at +23.2% against RTN group-16's 92.5 MB at +22.5%: statistically the same quality for **20% less memory**, 18 MB on a 125M model. The win here is the bytes, not the perplexity.
- **At 2 bits plain quantisation collapses** to unusable perplexity while FOLD stays finite. Neither is usable; FOLD just degrades gracefully instead of catastrophically.

These numbers fold everything -- embeddings included. Folding the embedding costs about 0.8 points of relative perplexity and saves 54 MB, which is the right trade when memory is the binding constraint. Pass `targets=["linear"]` to leave it dense: that scores +3.6% at 148.9 MB.

### 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.89** | **+3.6%** |

(Linear layers only, so the stages are compared like for like.)

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.
- **There is no fused dequantisation kernel.** Weights are unpacked and expanded in plain PyTorch, then handed to a normal matmul. Storage is genuinely 3-4x smaller, which is what decides whether a model fits at all — but on hardware where the uncompressed model already fits, this will be *slower* than just running it uncompressed. The measurements here are compression quality, never throughput.
- For a fused expert stack the transient expansion is the whole stack for that layer, not just the experts a token routes to. Resident memory still drops by the full ratio; per-token compute does not.

## What it can compress

FOLD is not limited to `nn.Linear`. Most of the weight in a modern checkpoint is not in one:

| weight | example | how |
|---|---|---|
| `nn.Linear` | any transformer projection | swapped for `FoldLinear` |
| **fused MoE expert stacks** | GLM-4.5, DeepSeek — all experts in one 3-D parameter | folded per expert, each with its own scales, outliers and delta |
| transposed matrices | GPT-2's `Conv1D`, stored `(in, out)` | `layout="transposed"` |
| convolutions | `nn.Conv1d/2d/3d` | spatial dims flattened into the input axis |
| embeddings | `nn.Embedding` | folded like any matrix |
| **anything else** | architectures this library has never seen | catch-all: any float parameter with 2+ dims |

The catch-all is the important one. It folds the parameter away and puts a dense tensor back on the attribute via forward hooks, so **the owning module's `forward` is never touched** — which is why a fused expert block written for GLM works without nanofold knowing anything about GLM.

```python
nanofold.fold_model(model, nanofold.FoldConfig(bits=4))          # everything
nanofold.fold_model(model, targets=["linear", "embedding"])      # or narrow it
```

Per-expert folding matters: experts specialise, so their weight distributions diverge. Sharing one set of scales across a stack would waste range on every expert but the widest.

## Memory during compression

Folding holds several full-precision copies of whatever it is given at once. A 277M-parameter expert block folded in one pass needs several GB and will simply die on a small machine. `chunk_elements` (default 16M) bounds it by processing a few slices at a time:

```python
FoldConfig(bits=4, chunk_elements=4_194_304)   # tighter memory, same output
```

Measured on a 25M-parameter stack — peak RSS 0.77 GB / 1.67 GB / 2.34 GB at 4M / 16M / 128M `chunk_elements`, producing byte-identical results. The floor is one slice, so the largest single matrix sets the minimum requirement.

## 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. Folded parameters do the same through their hooks. The expansion dtype follows the activations by default; `nanofold.set_compute_dtype(model, torch.float16)` pins it, which halves the transient scratch.

### 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                 # 197 tests
ruff check src tests
```

## License

MIT — see [LICENSE](LICENSE).
