Foundry/docs/quantization.md

# quantization

Tensor quantization for model compression. Python counterpart: Bitsandbytes.

**Header:** `quantization/quantization.h`
**Source:** `quantization/quantization.c`
**Status:** API Complete

---

## Overview

The quantization module provides:
- INT8 quantization (symmetric, per-tensor or per-channel)
- INT4 per-group quantization
- Ternary quantization (2-bit, BitNet b1.58 absmean)
- Quantized matrix multiplication (with native ternary kernel)
- Dequantization
- Offline model export (see [quantize-model](quantize-model.md))

---

## Quick Start

```c
#include "quantization/quantization.h"

// Quantize weights to INT8
int64_t shape[] = {4096, 4096};
float* weights = load_weights();

fq_tensor* qw = fq_quantize_int8(weights, shape, 2);
printf("Compressed to %.1f MB (was %.1f MB)\n",
       fq_memory_usage(qw) / 1e6,
       4096 * 4096 * 4 / 1e6);

// Quantized matmul
float* activations = get_activations();
float* output = fq_matmul(activations, qw, batch, 4096, 4096);

free(output);
fq_free(qw);
```

### Ternary Quantization

```c
// Quantize to {-1, 0, +1} using absmean threshold
fq_tensor* qt = fq_quantize_ternary(weights, shape, 2);

// Or with manual threshold
fq_tensor* qt2 = fq_quantize_ternary_threshold(weights, shape, 2, 0.05f);

// Per-channel ternary (different threshold per output channel)
fq_tensor* qt3 = fq_quantize_ternary_per_channel(weights, shape, 2, 0);

// Dequantize: ternary_value * scale
float* restored = fq_dequantize(qt);
```

---

## Ternary Format

### Encoding

Each value is encoded in 2 bits: `-1 = 0b00`, `0 = 0b01`, `+1 = 0b10`. Four values are packed per byte, LSB-first:

```
byte = (val3 << 6) | (val2 << 4) | (val1 << 2) | val0
```

Packed size: `(numel + 3) / 4` bytes.

### Absmean Thresholding (BitNet b1.58)

```
threshold = mean(|w|)
w > threshold   ->  +1
w < -threshold  ->  -1
otherwise       ->   0
scale = threshold  (for reconstruction: ternary_value * scale)
```

### Compression

| Dtype | Bits/value | Compression vs F32 |
|-------|------------|-------------------|
| INT8 | 8 | 4x |
| INT4 | 4 | 8x |
| Ternary | 2 | 16x |

Real-world result on RWKV-X 0.2B: 837 MB -> 53 MB (15.8x).

---

## Ternary Quantization (BitNet b1.58)

2-bit ternary quantization constrains weights to `{-1, 0, +1}`, enabling integer-only matmul (no floating-point multiply in the inner loop). This is the foundation for 1-bit LLM architectures like BitNet b1.58.

### How It Works

1. **Absmean scale**: `gamma = mean(|weights|)` over all elements
2. **Quantize**: `q = round(w / gamma)`, clamped to `{-1, 0, +1}`
3. **Pack**: 4 values per byte (2 bits each), encoding `00=-1, 01=0, 10=+1`
4. **Matmul**: inner loop uses only add/subtract/skip (no multiply), then scales result by `gamma`

### Compression

| Dtype | Bits/weight | Compression vs FP32 |
|-------|------------|---------------------|
| INT8 | 8 | 4x |
| INT4 | 4 | 8x |
| Ternary | 2 | 16x |

### Usage

```c
#include "quantization/quantization.h"

// Quantize a weight matrix to ternary
int64_t shape[] = {4096, 4096};
float* weights = load_weights();

fq_tensor* qw = fq_quantize_ternary(weights, shape, 2);
// 4096x4096 = 16M weights -> 4 MB (was 64 MB in FP32)

// Native ternary matmul (no dequantization, integer-only inner loop)
float* activations = get_activations();  // [batch, 4096]
float* output = fq_matmul_ternary(activations, qw, batch, 4096, 4096);

// Also works through the generic fq_matmul / fq_matmul_acc interface
// (auto-dispatches to ternary kernel when B is FQ_TERNARY)
float* output2 = fq_matmul(activations, qw, batch, 4096, 4096);

free(output);
free(output2);
fq_free(qw);
```

### Native Matmul Kernel

The ternary matmul kernel unpacks weights on-the-fly during the inner loop. Each 2-bit weight maps to `{-1, 0, +1}` via a lookup table, so the inner loop body is effectively:

```
sum += A[k] * w   // where w is -1, 0, or +1
```

The compiler optimizes this to conditional add/subtract/nop. No temporary dequantized buffer is allocated.

---

## Ternary Training (Backward Pass + Stochastic Flip)

The quantization module supports training ternary models end-to-end without floating-point weight storage. Weight updates use stochastic flipping (below).

### Backward Pass

Two backward functions propagate gradients through a ternary matmul layer where the forward pass was `Y = scale * X @ W_ternary`:

- **`fq_matmul_ternary_backward_input`** — computes dL/dX by multiplying d_output against the packed ternary weights (same unpacking as forward, sums over N). Uses `+=` accumulation (caller pre-zeros d_input).
- **`fq_matmul_ternary_backward_weight`** — computes dL/dW as a standard float outer product `X^T @ dL/dY`. No ternary involvement, no scale factor — the float gradient is what the flip function consumes.

### Stochastic Weight Flip

`fq_ternary_stochastic_flip` replaces the optimizer entirely. For each weight:

1. Compute flip probability: `p = clamp(|grad| * temperature, 0, 1)`
2. If `rand() < p`: flip one step toward the negative gradient direction
3. States: `-1 ↔ 0 ↔ +1`, clamped at boundaries (can't go below -1 or above +1)

Linear clamp is used instead of sigmoid because `sigmoid(0) = 0.5` would cause random flipping with zero gradient. Linear maps `0 → 0` correctly.

The PRNG is caller-owned (`uint64_t*` state, xorshift64) for reproducibility and thread safety in parallel training.

### Training Loop

```c
#include "quantization/quantization.h"

// Quantize weight matrix to ternary
int64_t shape[] = {K, N};
fq_tensor* W = fq_quantize_ternary(float_weights, shape, 2);

uint64_t rng;
fq_seed(&rng, 42);

for (int step = 0; step < n_steps; step++) {
    // Forward: Y[M,N] = X[M,K] @ W[K,N] (ternary kernel)
    float* Y = fq_matmul_ternary(X, W, M, K, N);

    // ... compute loss and d_output[M,N] ...

    // Backward input: dX[M,K] (for upstream layers)
    float dX[M * K];
    memset(dX, 0, sizeof(dX));
    fq_matmul_ternary_backward_input(d_output, W, dX, M, K, N);

    // Backward weight: dW[K,N] = X^T @ d_output
    float dW[K * N];
    memset(dW, 0, sizeof(dW));
    fq_matmul_ternary_backward_weight(X, d_output, dW, M, K, N);

    // Update: stochastic flip (replaces optimizer)
    fq_ternary_stochastic_flip(W, dW, temperature, &rng);

    free(Y);
}

fq_free(W);
```

### Memory

Training with stochastic flipping uses **2 bits per parameter** — no optimizer state, no FP32 weight shadow. For Reed-150M (111M params), that's ~28 MB total training weight memory vs ~1.8 GB for FP32 STE + Adam.

---

## API Reference

### Quantization

```c
// INT8
fq_tensor* fq_quantize_int8(const float* data, int64_t* shape, int ndim);
fq_tensor* fq_quantize_int8_per_channel(const float* data, int64_t* shape,
                                         int ndim, int channel_dim);

// INT4
fq_tensor* fq_quantize_int4(const float* data, int64_t* shape,
                             int ndim, int group_size);

// Ternary
fq_tensor* fq_quantize_ternary(const float* data, int64_t* shape, int ndim);
fq_tensor* fq_quantize_ternary_threshold(const float* data, int64_t* shape,
                                          int ndim, float threshold);
fq_tensor* fq_quantize_ternary_per_channel(const float* data, int64_t* shape,
                                            int ndim, int channel_dim);

// Generic dispatcher
fq_tensor* fq_quantize(const float* data, int64_t* shape, int ndim,
                       fq_dtype dtype, fq_mode mode, fq_granularity granularity,
                       int channel_dim, int group_size);

void fq_free(fq_tensor* qt);
```

### Dequantization

```c
float* fq_dequantize(fq_tensor* qt);
float* fq_dequantize_range(fq_tensor* qt, int64_t start, int64_t end);
```

### Quantized Operations

```c
float* fq_matmul(const float* A, fq_tensor* B, int M, int K, int N);
void fq_matmul_acc(const float* A, fq_tensor* B, float* C, int M, int K, int N);

// Native ternary matmul (integer-only, no dequant)
float* fq_matmul_ternary(const float* A, fq_tensor* B, int M, int K, int N);
void fq_matmul_ternary_acc(const float* A, fq_tensor* B, float* C, int M, int K, int N);

// Ternary backward pass
void fq_matmul_ternary_backward_input(const float* d_output, fq_tensor* W,
                                       float* d_input, int M, int K, int N);
void fq_matmul_ternary_backward_weight(const float* input, const float* d_output,
                                        float* d_weight, int M, int K, int N);

// Stochastic ternary weight flip (replaces optimizer)
void fq_ternary_stochastic_flip(fq_tensor* W, const float* d_weight,
                                 float temperature, uint64_t* rng_state);
void fq_seed(uint64_t* rng_state, uint64_t seed);
```

### Types

```c
typedef enum {
    FQ_INT8,    // 8-bit signed
    FQ_INT4,    // 4-bit signed (packed)
    FQ_UINT8,   // 8-bit unsigned
    FQ_UINT4,   // 4-bit unsigned (packed)
    FQ_TERNARY, // 2-bit ternary {-1, 0, +1} (packed, 4 per byte)
} fq_dtype;

typedef enum {
    FQ_SYMMETRIC,   // Zero-point is always 0
    FQ_ASYMMETRIC,  // Zero-point can be non-zero
} fq_mode;

typedef enum {
    FQ_PER_TENSOR,  // Single scale for entire tensor
    FQ_PER_CHANNEL, // Different scale for each channel
    FQ_PER_GROUP,   // Group quantization (like GPTQ)
} fq_granularity;
```

### Utility

```c
const char* fq_dtype_name(fq_dtype dtype);
float fq_dtype_bytes(fq_dtype dtype);  // 1 for int8, 0.5 for int4, 0.25 for ternary
void fq_compute_scale(const float* data, int64_t n, fq_mode mode,
                      float* scale, int32_t* zero_point);
size_t fq_memory_usage(fq_tensor* qt);
void fq_print(fq_tensor* qt);
float fq_quantization_error(const float* original, fq_tensor* quantized);
```

---

## Export API

The export module provides offline model quantization — loads a safetensors model, quantizes all weight tensors, and saves a new safetensors file.

**Header:** `quantization/export.h`
**Source:** `quantization/export.c`

### One-shot Export

```c
#include "quantization/export.h"

fq_export_config cfg = fq_export_config_defaults();
cfg.target_dtype = FQ_TERNARY;

fq_export_quantized("model.safetensors", "model-ternary.safetensors", &cfg);
```

### Step-by-step Export (for GUI)

```c
fq_export_config cfg = fq_export_config_defaults();
cfg.target_dtype = FQ_INT8;

fq_export_state* s = fq_export_begin("model.safetensors", &cfg);

while (fq_export_step(s)) {
    printf("%.0f%% - %s\n",
           fq_export_progress(s) * 100,
           fq_export_current_tensor(s));
}

fq_export_finish(s, "model-q8.safetensors");
```

### Export Config

```c
typedef struct {
    fq_dtype target_dtype;        // FQ_INT8, FQ_INT4, FQ_TERNARY
    fq_granularity granularity;   // PER_TENSOR, PER_CHANNEL, PER_GROUP
    int group_size;               // for per-group (default 128)
    float ternary_threshold;      // 0.0 = auto absmean, >0 = manual

    bool skip_embeddings;         // don't quantize embed layers (default true)
    bool skip_norms;              // don't quantize norm layers (default true)
    bool skip_head;               // don't quantize lm_head (default false)
    bool skip_biases;             // don't quantize biases (default true)
    bool skip_1d;                 // don't quantize 1D tensors (default true)
    int min_elements;             // skip tensors < this size (default 1024)

    void (*progress_cb)(int current, int total, const char* name, void* ud);
    void* userdata;
} fq_export_config;
```

### Export Functions

```c
fq_export_config fq_export_config_defaults(void);

// One-shot
bool fq_export_quantized(const char* input, const char* output, fq_export_config* cfg);

// Step-by-step
fq_export_state* fq_export_begin(const char* input, fq_export_config* cfg);
int fq_export_tensor_count(fq_export_state* s);
bool fq_export_step(fq_export_state* s);
float fq_export_progress(fq_export_state* s);
const char* fq_export_current_tensor(fq_export_state* s);
bool fq_export_finish(fq_export_state* s, const char* output);
void fq_export_abort(fq_export_state* s);
```

### Output Format

Quantized tensors are stored in safetensors format with auxiliary metadata tensors:

| Tensor | Dtype | Description |
|--------|-------|-------------|
| `layer.weight` | `I8` / `U8` | Quantized data (INT8 raw, INT4/ternary packed) |
| `layer.weight.__scale` | `F32` | Scale factors for reconstruction |
| `layer.weight.__shape` | `I64` | Original dimensions |
| `layer.weight.__zp` | `I32` | Zero points (asymmetric only) |

Skipped tensors (norms, biases, embeddings) are passed through unchanged.

File-level metadata:

```json
{
    "quantization.dtype": "TERNARY",
    "quantization.granularity": "per_tensor",
    "quantization.group_size": "128",
    "quantization.format_version": "1"
}
```

---

## Testing

```bash
cd foundry

# Quantization core (INT8, INT4, ternary)
make test-quantization    # 28 tests

# Export pipeline (round-trip, skip logic, metadata, step API)
make test-export          # 14 tests
```

---

## See Also

- [quantize-model](quantize-model.md) - CLI tool for offline model quantization
- [tensor](tensor.md) - Tensor operations
- [safetensors](safetensors.md) - Model file I/O
- [lora](lora.md) - LoRA for efficient fine-tuning