diff --git a/.claude/settings.json b/.claude/settings.json new file mode 100644 index 0000000..bcf43f6 --- /dev/null +++ b/.claude/settings.json @@ -0,0 +1,3 @@ +{ + "includeCoAuthoredBy": false +} diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 5ce1c41..37acd5c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -3,8 +3,9 @@ name: CI on: push: branches: [ master, feature/** ] + # guirguispierre 2026-08-10 - no branch filter, restricting to master meant a + # pr stacked on another pr's branch got no ci run pull_request: - branches: [ master ] jobs: test: @@ -44,6 +45,19 @@ jobs: test -f atomic_1bit/core/libatomic.so || (echo "libatomic.so not found!" && exit 1) ls -lh atomic_1bit/core/libatomic.so + - name: Build embedded C++ runner + run: | + cd embedded + g++ -O3 -std=c++17 atomic_runner.cpp -o runner + cd .. + + - name: Verify runner was built + # guirguispierre 2026-08-10 - without this binary the parity tests skip + # silently and the claim goes unverified + run: | + test -f embedded/runner || (echo "embedded/runner not found!" && exit 1) + ./embedded/runner --help + - name: Run linting - Black run: | black --check atomic_1bit/ tests/ benchmarks/ tools/ scripts/ --exclude='/(\.git|\.venv|build|dist)/' diff --git a/README.md b/README.md index fcecc1a..b668d7b 100644 --- a/README.md +++ b/README.md @@ -22,10 +22,12 @@ weight == 0 -> skip (free sparsity) The result is a full transformer that runs on **integer arithmetic only**. No CUDA required. No FP16. No matrix multiply units. Just `add` and `sub` instructions that work on any processor manufactured in the last 30 years. **This matters because:** -- A 1.33M parameter model drops from **5.3 MB to 2.0 MB** (-62%) +- A 6.85M parameter model drops from **27.4 MB to 10.1 MB** (-63%) - The C++ runtime has **zero external dependencies** -- it's a single binary - It runs on a **Raspberry Pi**, an **ESP32**, or a **2015 laptop** -- The Python training stack and C++ inference engine produce **bit-exact identical output** +- The ternary kernel is **bit-exact** against the NumPy reference, and the C++ + engine tracks the PyTorch model to float32 precision (see + [Numerical parity](#numerical-parity)) > This is experimental research software. It works, it's verified, and it's honest about what it is: a proof that useful AI inference doesn't require expensive hardware. @@ -33,15 +35,20 @@ The result is a full transformer that runs on **integer arithmetic only**. No CU ## Performance -Benchmarked on Apple M-series, single thread, sequence length 128, 50 generated tokens: +Benchmarked on Apple M-series, single thread, 50 generated tokens, using the +Quick Start config below (256d, 6 layers, 4 heads, 4096 vocab, 6.85M params): -| Metric | FP16 Baseline | Atomic-1Bit | Improvement | +| Metric | FP32 Baseline | Atomic-1Bit | Improvement | |:---|:---|:---|:---| -| **Model Size** | 5.3 MB | **2.0 MB** | **-62%** | -| **Parameters** | 1.33M | 1.33M | Same | -| **Precision** | Float16 | Ternary {-1, 0, 1} | -- | -| **Throughput (C++)** | N/A | **~160-170 TPS** | Portable runtime | -| **Throughput (Python)** | ~826 TPS | ~130 TPS | Unoptimized | +| **Model Size** | 27.4 MB | **10.1 MB** | **-63%** | +| **Parameters** | 6.85M | 6.85M | Same | +| **Precision** | Float32 | Ternary {-1, 0, 1} | -- | +| **Throughput (C++)** | N/A | **~170-185 TPS** | Portable runtime | + +The C++ runner recomputes the full prefix each step rather than keeping a +KV-cache, so throughput falls off as the sequence grows -- generation is +O(n^2) in the number of tokens. The Python path (`AtomicTransformer.generate`) +*does* cache. Reproduce with `benchmarks/run_suite.py`.
Visual benchmarks @@ -70,6 +77,14 @@ cd Atomic-1Bit pip install -r requirements.txt ``` +### Build the ternary kernel + +The Python inference path loads a shared library, so build it first: + +```bash +make -C atomic_1bit/core # produces atomic_1bit/core/libatomic.so +``` + ### Verify the kernel works This confirms the C++ ternary kernel matches the Python/NumPy reference exactly: @@ -117,10 +132,34 @@ The project has three components: 1. **Research stack** (`atomic_1bit/`) -- PyTorch training, evaluation, and model architecture. Train on TinyStories or Alpaca-cleaned datasets with thermal safety monitoring, gradient accumulation, and cosine scheduling. -2. **Bare-metal runtime** (`embedded/`) -- Standalone C++ inference engine with zero dependencies. Supports CPU, Metal (Apple Silicon), and CUDA backends through conditional compilation. Produces bit-exact output matching the Python reference. +2. **Bare-metal runtime** (`embedded/`) -- Standalone C++ inference engine with zero dependencies. Supports CPU, Metal (Apple Silicon), and CUDA backends through conditional compilation. 3. **Gist tokens** -- Pre-computed "thought vectors" that compress a system prompt into a single embedding, injected into the attention stream at zero inference cost. +### Numerical parity + +Two different guarantees are worth separating, because only one of them is exact: + +| Claim | Status | Verified by | +|:---|:---|:---| +| C++ ternary kernel == NumPy int32 reference | **Bit-exact** | `tests/test_kernel_parity.py` | +| C++ engine == PyTorch model, whole forward pass | Agrees to float32 precision | `tests/test_cpp_runtime_parity.py` | + +Whole-model bit-equality is *not* achievable by design, and the reason is worth +knowing before you go hunting for a bug. Activations are quantized with a +per-tensor scale derived from a max over float values. `round()` is +discontinuous, so a last-ULP difference anywhere upstream -- summation order in +an RMS norm, a `tanh` implementation -- can push one value across a `.5` +boundary. That shifts a single int8 by one step, which is a *discrete* jump in +the accumulator, not an epsilon. Deeper models give this more chances to happen. + +In practice a shallow model matches PyTorch to ~1e-7 relative, and deeper ones +stay within a few percent with the same next-token prediction. Greedy decoding +can still fork after enough steps once two logits are near-tied. The test suite +pins both bounds, so a genuine mismatch -- a different activation function, +scale convention, or weight layout -- fails loudly instead of hiding in the +noise. + ![Ternary Matmul Diagram](assets/diagram_ternary_matmul_1765658104682.png) For more details, see [docs/USAGE.md](docs/USAGE.md). @@ -143,7 +182,7 @@ atomic_1bit/ embedded/ Standalone C++ runner + ESP32 port guide configs/ Model presets (4K pocket to 12.5M flagship) benchmarks/ Reproducible benchmark suite vs FP16 baselines -tests/ 67 pytest tests for correctness verification +tests/ 149 pytest tests for correctness verification scripts/ Plotting, evaluation, and reproduction scripts docs/ Installation, usage, commands, benchmarking guides examples/ Runnable example scripts @@ -153,12 +192,22 @@ examples/ Runnable example scripts ## Model Configurations -| Config | Parameters | Dimensions | Use Case | -|:---|:---|:---|:---| -| [`pocket_4k`](configs/pocket_4k.yaml) | ~100K | 256d, 4L, 4H | ESP32 / microcontrollers | -| [`stories_base`](configs/stories_base.yaml) | ~1.33M | 256d, 6L, 4H | Development / testing | -| [`flagship_12m`](configs/flagship_12m.yaml) | ~12.5M | 320d, 8L, 5H | Quality demos | -| [`mixed_precision`](configs/mixed_precision.yaml) | Configurable | Hybrid 1.58/4-bit | Experimental | +| Config | Parameters | Exported size | Dimensions | Use Case | +|:---|:---|:---|:---|:---| +| [`pocket_4k`](configs/pocket_4k.yaml) | 5.28M | 8.5 MB | 256d, 4L, 4H, 4K vocab | Smallest preset | +| [`stories_base`](configs/stories_base.yaml) | 30.5M | 69.2 MB | 256d, 6L, 4H, 50257 vocab | Development / testing | +| [`flagship_12m`](configs/flagship_12m.yaml) | 12.5M | 16.7 MB | 320d, 8L, 5H, 4K vocab | Quality demos | +| [`mixed_precision`](configs/mixed_precision.yaml) | Configurable | -- | Hybrid 1.58/4-bit | Experimental | + +Parameter counts are dominated by the token embedding table, which stays float32 +(`vocab_size x dim x 4` bytes) -- so vocabulary size, not depth, drives file +size. `stories_base.yaml` ships with the 50257-entry GPT-2 vocabulary; the +Quick Start above passes `--vocab_size 4096` instead, giving a 6.85M-parameter +model that exports to 10.1 MB. Set the vocabulary deliberately for your target. + +> None of these presets fit in ESP32 SRAM as-is. See +> [embedded/ESP32_PORT_GUIDE.md](embedded/ESP32_PORT_GUIDE.md) for what actually +> fits on-device. Load any config with: @@ -191,6 +240,8 @@ config = config_to_atomic(load_config("configs/stories_base.yaml")) ## Running Tests ```bash +pip install pytest + # Run the full test suite pytest tests/ -v @@ -199,6 +250,12 @@ pytest tests/test_bitlinear.py -v pytest tests/test_kernel_parity.py -v ``` +Two modules need compiled artifacts and skip without them: +`test_kernel_parity.py` needs `make -C atomic_1bit/core`, and +`test_cpp_runtime_parity.py` needs the runner +(`cd embedded && g++ -O3 -std=c++17 atomic_runner.cpp -o runner`). Build both +before trusting a green run. + --- ## Contributing diff --git a/atomic_1bit/core/backends/cpu_kernel.cpp b/atomic_1bit/core/backends/cpu_kernel.cpp index 2e4d37b..a54a894 100644 --- a/atomic_1bit/core/backends/cpu_kernel.cpp +++ b/atomic_1bit/core/backends/cpu_kernel.cpp @@ -36,33 +36,36 @@ void ternary_matmul(const int8_t *A, const int8_t *B_transposed, int32_t *C, int32x4_t v_acc0 = vdupq_n_s32(0); int32x4_t v_acc1 = vdupq_n_s32(0); + // guirguispierre 2026-08-10 - vmull_s8 widens to int16 because vmulq_s8 + // truncates: activation -128 against weight -1 wraps and flips the sum + // Process 32 items at a time (2x unrolled from 16) for (; k <= K - 32; k += 32) { // First 16 elements int8x16_t va0 = vld1q_s8(val_a + k); int8x16_t vb0 = vld1q_s8(val_b + k); - int8x16_t vprod0 = vmulq_s8(va0, vb0); - int16x8_t v_pair0 = vpaddlq_s8(vprod0); - int32x4_t v_quad0 = vpaddlq_s16(v_pair0); - v_acc0 = vaddq_s32(v_acc0, v_quad0); + v_acc0 = vaddq_s32( + v_acc0, vpaddlq_s16(vmull_s8(vget_low_s8(va0), vget_low_s8(vb0)))); + v_acc0 = vaddq_s32( + v_acc0, vpaddlq_s16(vmull_s8(vget_high_s8(va0), vget_high_s8(vb0)))); // Second 16 elements int8x16_t va1 = vld1q_s8(val_a + k + 16); int8x16_t vb1 = vld1q_s8(val_b + k + 16); - int8x16_t vprod1 = vmulq_s8(va1, vb1); - int16x8_t v_pair1 = vpaddlq_s8(vprod1); - int32x4_t v_quad1 = vpaddlq_s16(v_pair1); - v_acc1 = vaddq_s32(v_acc1, v_quad1); + v_acc1 = vaddq_s32( + v_acc1, vpaddlq_s16(vmull_s8(vget_low_s8(va1), vget_low_s8(vb1)))); + v_acc1 = vaddq_s32( + v_acc1, vpaddlq_s16(vmull_s8(vget_high_s8(va1), vget_high_s8(vb1)))); } // Handle remaining 16-element block for (; k <= K - 16; k += 16) { int8x16_t va = vld1q_s8(val_a + k); int8x16_t vb = vld1q_s8(val_b + k); - int8x16_t vprod = vmulq_s8(va, vb); - int16x8_t v_pair = vpaddlq_s8(vprod); - int32x4_t v_quad = vpaddlq_s16(v_pair); - v_acc0 = vaddq_s32(v_acc0, v_quad); + v_acc0 = vaddq_s32( + v_acc0, vpaddlq_s16(vmull_s8(vget_low_s8(va), vget_low_s8(vb)))); + v_acc0 = vaddq_s32( + v_acc0, vpaddlq_s16(vmull_s8(vget_high_s8(va), vget_high_s8(vb)))); } // Combine accumulators diff --git a/atomic_1bit/model/transformer.py b/atomic_1bit/model/transformer.py index b321b50..7c48972 100644 --- a/atomic_1bit/model/transformer.py +++ b/atomic_1bit/model/transformer.py @@ -90,7 +90,9 @@ def __init__(self, config: AtomicConfig): hidden_dim = 4 * config.dim self.fc1 = BitLinear(config.dim, hidden_dim, bias=False) self.fc2 = BitLinear(hidden_dim, config.dim, bias=False) - self.act = nn.GELU() + # guirguispierre 2026-08-10 - must match the tanh gelu in + # embedded/atomic_runner.cpp or python and c++ output diverge + self.act = nn.GELU(approximate="tanh") def forward(self, x): return self.fc2(self.act(self.fc1(x))) diff --git a/docs/API_REFERENCE.md b/docs/API_REFERENCE.md index b41b6ba..239dc0c 100644 --- a/docs/API_REFERENCE.md +++ b/docs/API_REFERENCE.md @@ -128,7 +128,7 @@ Core BitLinear compute kernel. Implements one linear layer of the form: quantize 2. Integer accumulation: `acc[o] += x_q[i] * w[i * out_dim + o]` (uses actual `int32_t` multiplication; no branch-on-sign optimization at this layer). 3. Dequantize: `out[o] = acc[o] / scale`. -**Known limitation:** The weight scale factor is not stored in the binary format used by `atomic_lib.h` and is therefore not applied during dequantization. The runner (`atomic_runner.cpp`) corrects this by reading and applying per-layer weight scales from the extended binary format written by `export_to_cpp.py`. +Per-tensor `float32` weight scales are read from the binary and applied during dequantization, matching the Python `BitLinear` forward. `atomic_lib.h` and `atomic_runner.cpp` now consume the identical format written by `export_to_cpp.py`. --- @@ -265,4 +265,4 @@ Each layer block contains: | `4 + dim * 4*dim` | `float32 + int8[]` | FC1 weight (scale + weights) | | `4 + 4*dim * dim` | `float32 + int8[]` | FC2 weight (scale + weights) | -Note: `atomic_lib.h` uses an older format that does not include the per-layer scale prefix for weights. The `atomic_runner.cpp` binary uses the current format with scale values. +Note: `atomic_lib.h` and `atomic_runner.cpp` read the same format, including the per-tensor `float32` scale prefix on every quantized weight tensor. diff --git a/docs/INSTALL.md b/docs/INSTALL.md index 04a385c..c01e3c9 100644 --- a/docs/INSTALL.md +++ b/docs/INSTALL.md @@ -37,11 +37,28 @@ This installs: python3 -c "import torch; import tiktoken; print('Dependencies OK')" ``` -## C++ Runtime (Optional) +## C++ Components -The C++ inference engine is optional -- you only need it if you want to run models on bare metal or embedded devices. +There are two separate C++ artifacts, and the first one is **not** optional. -### Build the engine +### 1. Ternary kernel (`libatomic.so`) -- required for Python inference + +`atomic_1bit/python/wrapper.py` loads this shared library, so +`python3 atomic_1bit/python/inference.py` and the kernel parity tests fail +without it: + +```bash +cd atomic_1bit/core +make # defaults to BACKEND=CPU +``` + +Verify: + +```bash +test -f atomic_1bit/core/libatomic.so && echo "kernel built" +``` + +### 2. Bare-metal runner -- needed to run exported models standalone ```bash cd embedded diff --git a/embedded/atomic_lib.h b/embedded/atomic_lib.h index e30c2a8..d3e9d31 100644 --- a/embedded/atomic_lib.h +++ b/embedded/atomic_lib.h @@ -48,14 +48,15 @@ inline void softmax(vector &x) { // Affine RMSNorm (used for the layer's ln1 / ln2 / ln_f). inline void rms_norm(const vector &x, const vector &w, vector &out, int dim) { - float ss = 0.0f; + // guirguispierre 2026-08-10 - double sum and a divide rather than a + // reciprocal multiply, to match atomic_runner.cpp and the python reference + double ss = 0.0; for (float f : x) - ss += f * f; - ss /= dim; - float rms = 1.0f / sqrt(ss + 1e-5f); + ss += (double)f * (double)f; + float rms = sqrt((float)(ss / dim) + 1e-5f); for (int i = 0; i < dim; ++i) { - out[i] = x[i] * rms * w[i]; + out[i] = (x[i] / rms) * w[i]; } } @@ -79,10 +80,11 @@ inline void bit_linear(const vector &x, const vector &w, int out_dim = (int)out.size(); // SubLN - float ss = 0.0f; + // guirguispierre 2026-08-10 - double for the same reason as rms_norm + double ss = 0.0; for (float v : x) - ss += v * v; - float rms = sqrt(ss / in_dim + 1e-5f); + ss += (double)v * (double)v; + float rms = sqrt((float)(ss / in_dim) + 1e-5f); // Activation absmax over x/rms = max(|x|)/rms float max_abs = 0.0f; @@ -93,11 +95,12 @@ inline void bit_linear(const vector &x, const vector &w, max_val = 1e-5f; float scale_x = 127.0f / max_val; - float effective_scale = scale_x / rms; + // guirguispierre 2026-08-10 - divide before scaling to match + // activation_quant(); folding it into one multiply rounds differently vector x_q(in_dim); for (int i = 0; i < in_dim; ++i) { - float q = roundf(x[i] * effective_scale); + float q = roundf((x[i] / rms) * scale_x); if (q > 127.0f) q = 127.0f; if (q < -127.0f) @@ -105,18 +108,22 @@ inline void bit_linear(const vector &x, const vector &w, x_q[i] = (int8_t)q; } - float dequant = scale_w / scale_x; - for (int o = 0; o < out_dim; ++o) { - int32_t acc = 0; - for (int i = 0; i < in_dim; ++i) { - int8_t wv = w[i * out_dim + o]; - if (wv == 1) - acc += x_q[i]; - else if (wv == -1) - acc -= x_q[i]; - } - out[o] = (float)acc * dequant; + // guirguispierre 2026-08-10 - i-outer streams each weight row contiguously + // instead of striding by out_dim, worth ~14x; integer sums are associative so + // the totals are unchanged + vector acc(out_dim, 0); + for (int i = 0; i < in_dim; ++i) { + int32_t xv = x_q[i]; + if (xv == 0) + continue; // free sparsity: skip the whole row + const int8_t *wrow = &w[(size_t)i * out_dim]; + for (int o = 0; o < out_dim; ++o) + acc[o] += wrow[o] * xv; } + + float dequant = scale_w / scale_x; + for (int o = 0; o < out_dim; ++o) + out[o] = (float)acc[o] * dequant; } struct AtomicLayer { diff --git a/embedded/atomic_runner.cpp b/embedded/atomic_runner.cpp index 15c0dc2..c4a68cf 100644 --- a/embedded/atomic_runner.cpp +++ b/embedded/atomic_runner.cpp @@ -4,6 +4,7 @@ #include #include #include +#include #include #include #include @@ -63,10 +64,12 @@ struct AtomicModel { // RMSNorm void rms_norm(const vector &x, const vector &weight, vector &out, int dim) { - float sum_sq = 0.0f; + // guirguispierre 2026-08-10 - float32 summation drifts from torch's pairwise + // sum enough to shift the int8 quant_scale downstream + double sum_sq = 0.0; for (float v : x) - sum_sq += v * v; - float rms = sqrt(sum_sq / dim + 1e-5f); + sum_sq += (double)v * (double)v; + float rms = sqrt((float)(sum_sq / dim) + 1e-5f); for (int i = 0; i < dim; ++i) { out[i] = (x[i] / rms) * weight[i]; } @@ -78,16 +81,17 @@ void rms_norm(const vector &x, const vector &weight, // Out: Float output [Out] // Scratch: [In] (Pre-allocated buffer to avoid malloc) void bit_linear(const vector &x, const TensorI8 &w, vector &out, - vector &scratch) { + vector &scratch, vector &acc32) { int in_dim = w.rows; int out_dim = w.cols; // 0. Internal RMSNorm (Non-affine) // Per BitNet paper/impl: BitLinear(x) = Layernorm(x) ... - float sum_sq = 0.0f; + // guirguispierre 2026-08-10 - double for the same reason as rms_norm + double sum_sq_d = 0.0; for (float v : x) - sum_sq += v * v; - float rms = sqrt(sum_sq / in_dim + 1e-5f); + sum_sq_d += (double)v * (double)v; + float rms = sqrt((float)(sum_sq_d / in_dim) + 1e-5f); // 1. Quantize Input // We effectively quantize (x / rms). @@ -100,14 +104,16 @@ void bit_linear(const vector &x, const TensorI8 &w, vector &out, max_val = 1e-5f; float quant_scale = 127.0f / max_val; - // effective_scale = quant_scale / rms - float effective_scale = quant_scale / rms; - vector x_q(in_dim); + // guirguispierre 2026-08-10 - divide before scaling to match + // activation_quant(); folding it into one multiply rounds differently + // guirguispierre 2026-08-10 - caller's buffer avoids a malloc per call in the + // hot loop on microcontrollers + if ((int)scratch.size() < in_dim) + scratch.resize(in_dim); + int8_t *x_q = scratch.data(); for (int i = 0; i < in_dim; ++i) { - // x_norm = x[i] / rms - // val = x_norm * quant_scale - float val = round(x[i] * effective_scale); + float val = round((x[i] / rms) * quant_scale); if (val > 127) val = 127; if (val < -127) @@ -116,42 +122,40 @@ void bit_linear(const vector &x, const TensorI8 &w, vector &out, } // 2. Ternary Kernel - fill(out.begin(), out.end(), 0.0f); - - for (int j = 0; j < out_dim; ++j) { - int32_t acc = 0; - for (int i = 0; i < in_dim; ++i) { - int8_t xv = x_q[i]; - int8_t wv = w.data[i * out_dim + j]; - - if (wv == 1) - acc += xv; - else if (wv == -1) - acc -= xv; - } - // 3. Dequantize - // out = (acc / x_scale) * w_scale - float val = ((float)acc / quant_scale) * w.scale; - out[j] = val; - - // Debug first few outputs of first layer - static int debug_count = 0; - if (debug_count < 5) { - // print diagnostics for first few calls - if (j == 0) { - cout << "[BitLinear Debug] Dim In: " << in_dim << " Out: " << out_dim - << endl; - cout << " Input Max: " << max_val << " Quant Scale: " << quant_scale - << " RMS: " << rms << endl; - cout << " Weight Scale: " << w.scale << endl; // Check if this is 0 - cout << " Acc: " << acc << " Dequant: " << val << endl; - cout << " Sample W: " << (int)w.data[0] << " " << (int)w.data[1] - << endl; - cout << " Sample X_Q: " << (int)x_q[0] << " " << (int)x_q[1] << endl; - debug_count++; - } - } + // guirguispierre 2026-08-10 - i-outer streams each weight row contiguously + // instead of striding by out_dim, worth ~14x; integer sums are associative so + // the totals are unchanged + if ((int)acc32.size() < out_dim) + acc32.resize(out_dim); + int32_t *acc = acc32.data(); + memset(acc, 0, out_dim * sizeof(int32_t)); + + for (int i = 0; i < in_dim; ++i) { + int32_t xv = x_q[i]; + if (xv == 0) + continue; // free sparsity: skip the whole row + const int8_t *wrow = &w.data[(size_t)i * out_dim]; + // guirguispierre 2026-08-10 - w is {-1,0,1} so the multiply is the add/sub + // and vectorizes; the branchy form stalls on mispredicts + for (int j = 0; j < out_dim; ++j) + acc[j] += wrow[j] * xv; } + + // 3. Dequantize + // out = (acc / x_scale) * w_scale + for (int j = 0; j < out_dim; ++j) + out[j] = ((float)acc[j] / quant_scale) * w.scale; +} + +// guirguispierre 2026-08-10 - full float32 precision so tools/parity_check.py +// can diff values exactly rather than by eye +void print_tensor(const char *name, const vector &v, int count = 10) { + int n = min((int)v.size(), count); + cout << "[" << name << "] n=" << v.size() << " "; + cout << std::setprecision(9); + for (int i = 0; i < n; ++i) + cout << v[i] << (i + 1 < n ? " " : ""); + cout << endl; } // Softmax @@ -274,6 +278,7 @@ struct Workspace { vector fc_act, mlp_out; vector xt, xn; vector scratch_q; + vector acc32; void ensure_size(int max_seq_len, int dim) { int max_el = max_seq_len * dim; @@ -339,13 +344,10 @@ void forward(AtomicModel &model, const vector &tokens, } // Print Embedding Sum (last token) - // if (parity_mode) { - // int last_idx = (seq_len - 1) * dim; - // // Define a vector to copy so we can use print_tensor - // vector tmp(dim); - // memcpy(tmp.data(), &ws.X[last_idx], dim * sizeof(float)); - // print_tensor("Embedding_Sum", tmp); - // } + if (parity_mode) { + vector tmp(&ws.X[(seq_len - 1) * dim], &ws.X[seq_len * dim]); + print_tensor("Embedding_Sum", tmp, dim); + } int l_idx = 0; for (auto &layer : model.layers) { @@ -356,22 +358,19 @@ void forward(AtomicModel &model, const vector &tokens, memcpy(&ws.X_norm[t * dim], ws.xn.data(), dim * sizeof(float)); } - // if (parity_mode) { - // int last_idx = (seq_len - 1) * dim; - // vector tmp(dim); - // memcpy(tmp.data(), &ws.X_norm[last_idx], dim * sizeof(float)); - // // Construct string layer0_LN1 - // string name = "Layer" + to_string(l_idx) + "_LN1"; - // print_tensor(name.c_str(), tmp); - // } + if (parity_mode) { + vector tmp(&ws.X_norm[(seq_len - 1) * dim], + &ws.X_norm[seq_len * dim]); + print_tensor(("Layer" + to_string(l_idx) + "_LN1").c_str(), tmp, dim); + } // Attention for (int t = 0; t < seq_len; ++t) { memcpy(ws.xt.data(), &ws.X_norm[t * dim], dim * sizeof(float)); - bit_linear(ws.xt, layer.q_w, ws.q, ws.scratch_q); - bit_linear(ws.xt, layer.k_w, ws.k, ws.scratch_q); - bit_linear(ws.xt, layer.v_w, ws.v, ws.scratch_q); + bit_linear(ws.xt, layer.q_w, ws.q, ws.scratch_q, ws.acc32); + bit_linear(ws.xt, layer.k_w, ws.k, ws.scratch_q, ws.acc32); + bit_linear(ws.xt, layer.v_w, ws.v, ws.scratch_q, ws.acc32); memcpy(&ws.Q[t * dim], ws.q.data(), dim * sizeof(float)); memcpy(&ws.K[t * dim], ws.k.data(), dim * sizeof(float)); @@ -410,24 +409,18 @@ void forward(AtomicModel &model, const vector &tokens, // O Proj for (int t = 0; t < seq_len; ++t) { memcpy(ws.xt.data(), &ws.AttnOut[t * dim], dim * sizeof(float)); - bit_linear(ws.xt, layer.o_w, ws.layer_out, ws.scratch_q); + bit_linear(ws.xt, layer.o_w, ws.layer_out, ws.scratch_q, ws.acc32); // Residual for (int i = 0; i < dim; ++i) ws.X[t * dim + i] += ws.layer_out[i]; } - // if (parity_mode) { - // // We want to print Attn Output (Residual Added? Python script prints - // // `attn_out` BEFORE residual add) Wait, python script printed - // `attn_out = - // // layer.attn(x_norm)`. Here `ws.layer_out` is `attn_out`. But we just - // // added it to X. We have it in `ws.layer_out`. Assuming seq_len=1 for - // // parity check usually. If seq_len > 1, this only captures the LAST - // // token's execution in that loop. - // string name = "Layer" + to_string(l_idx) + "_Attn_Out"; - // print_tensor(name.c_str(), ws.layer_out); - // } + // guirguispierre 2026-08-10 - layer_out is pre-residual here, matching what + // tools/parity_check.py prints on the python side + if (parity_mode) + print_tensor(("Layer" + to_string(l_idx) + "_Attn_Out").c_str(), + ws.layer_out, dim); // MLP for (int t = 0; t < seq_len; ++t) { @@ -442,21 +435,19 @@ void forward(AtomicModel &model, const vector &tokens, // } int hidden = 4 * dim; - bit_linear(ws.xn, layer.fc1_w, ws.fc_act, ws.scratch_q); + bit_linear(ws.xn, layer.fc1_w, ws.fc_act, ws.scratch_q, ws.acc32); for (int i = 0; i < hidden; ++i) ws.fc_act[i] = gelu(ws.fc_act[i]); - bit_linear(ws.fc_act, layer.fc2_w, ws.mlp_out, ws.scratch_q); + bit_linear(ws.fc_act, layer.fc2_w, ws.mlp_out, ws.scratch_q, ws.acc32); for (int i = 0; i < dim; ++i) ws.X[t * dim + i] += ws.mlp_out[i]; } - // if (parity_mode) { - // string name = "Layer" + to_string(l_idx) + "_MLP_Out"; - // print_tensor(name.c_str(), ws.mlp_out); - // cout << endl; // Separator - // } + if (parity_mode) + print_tensor(("Layer" + to_string(l_idx) + "_MLP_Out").c_str(), + ws.mlp_out, dim); l_idx++; } @@ -466,15 +457,13 @@ void forward(AtomicModel &model, const vector &tokens, memcpy(ws.xt.data(), &ws.X[last_t * dim], dim * sizeof(float)); rms_norm(ws.xt, model.ln_f.data, ws.xn, dim); - // if (parity_mode) { - // print_tensor("Final_LN", ws.xn); - // } + if (parity_mode) + print_tensor("Final_LN", ws.xn, dim); - bit_linear(ws.xn, model.head_w, logits, ws.scratch_q); + bit_linear(ws.xn, model.head_w, logits, ws.scratch_q, ws.acc32); - // if (parity_mode) { - // print_tensor("Final_Logits", logits, 20); - // } + if (parity_mode) + print_tensor("Final_Logits", logits, (int)logits.size()); } // Sampling: temperature + optional top-k + optional top-p (nucleus). diff --git a/tests/test_cpp_runtime_parity.py b/tests/test_cpp_runtime_parity.py new file mode 100644 index 0000000..38be745 --- /dev/null +++ b/tests/test_cpp_runtime_parity.py @@ -0,0 +1,143 @@ +import os +import subprocess +import sys + +import pytest +import torch +import torch.nn as nn + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + +from atomic_1bit.model.transformer import AtomicTransformer, AtomicConfig +from atomic_1bit.utils.export_to_cpp import export_model + +RUNNER = os.path.join(os.path.dirname(__file__), "..", "embedded", "runner") + +pytestmark = pytest.mark.skipif( + not os.path.exists(RUNNER), + reason="C++ runner not built (cd embedded && g++ -O3 -std=c++17 " + "atomic_runner.cpp -o runner)", +) + +CFG = AtomicConfig(vocab_size=256, dim=64, depth=4, heads=4, context_length=32) + +# guirguispierre 2026-08-10 - loose because the int8 quantizer rounds, so one +# ULP of float difference can flip a value across a .5 boundary and move an +# accumulator a whole step; this catches gross mismatches, not single flips +AGREEMENT_TOL = 5e-2 + + +def _run_runner(model_bin, *args): + return subprocess.run( + [RUNNER, "--model", str(model_bin), *args], + capture_output=True, text=True, timeout=120, + ) + + +def _generated_tokens(res, n_prompt=1): + body = res.stdout.split("Generating:", 1)[1].split("Done.")[0] + fields = body.split() + for tok in fields: + assert tok.lstrip("-").isdigit(), ( + f"non-numeric token {tok!r} in generation stream: {body!r}" + ) + # guirguispierre 2026-08-10 - runner echoes the prompt before generating + return [int(t) for t in fields[n_prompt:]] + + +def _cpp_logits(model_bin, start_token): + res = _run_runner(model_bin, "--parity", "--temp", "0", + "--start_token", str(start_token)) + for line in res.stdout.splitlines(): + if line.startswith("[Final_Logits]"): + vals = line.split("] n=", 1)[1].split()[1:] + return torch.tensor([float(v) for v in vals]) + raise AssertionError( + f"runner emitted no Final_Logits.\nstdout:\n{res.stdout[-2000:]}\n" + f"stderr:\n{res.stderr[-2000:]}" + ) + + +def _py_logits(model, start_token): + with torch.no_grad(): + return model(torch.tensor([[start_token]], dtype=torch.long))[0, -1, :] + + +@pytest.fixture(scope="module") +def exported(tmp_path_factory): + torch.manual_seed(1234) + model = AtomicTransformer(CFG) + model.eval() + path = tmp_path_factory.mktemp("parity") / "model.bin" + export_model(model, str(path)) + return model, path + + +class TestActivationContract: + # guirguispierre 2026-08-10 - the numeric tests below cannot catch a wrong + # activation: erf-vs-tanh gelu perturbs logits by the same ~1e-2 that one + # quantizer boundary flip does, and x86 and arm already differ by that much + def test_mlp_uses_tanh_gelu(self): + act = AtomicTransformer(CFG).layers[0].mlp.act + assert isinstance(act, nn.GELU) + assert act.approximate == "tanh", ( + "embedded/atomic_runner.cpp implements the tanh gelu; exact-erf " + "nn.GELU() here diverges by ~4e-4 per layer and compounds" + ) + + def test_tanh_and_erf_gelu_actually_differ(self): + x = torch.linspace(-3, 3, 64) + delta = (nn.GELU()(x) - nn.GELU(approximate="tanh")(x)).abs().max() + assert delta > 1e-5 + + +class TestCppRuntimeParity: + def test_output_stream_is_not_polluted(self, exported): + _, path = exported + res = _run_runner(path, "--steps", "8", "--temp", "0", + "--start_token", "7") + tokens = _generated_tokens(res) + assert len(tokens) == 8, f"expected 8 generated tokens, got {tokens!r}" + + @pytest.mark.parametrize("start_token", [0, 7, 42, 128, 255]) + def test_logits_agree(self, exported, start_token): + model, path = exported + cpp = _cpp_logits(path, start_token) + py = _py_logits(model, start_token) + + assert cpp.shape == py.shape + rel = ((cpp - py).abs().max() / py.abs().max()).item() + assert rel < AGREEMENT_TOL, ( + f"relative logit error {rel:.3e} exceeds {AGREEMENT_TOL:.0e} -- an " + f"op has diverged between atomic_1bit/ and embedded/" + ) + + @pytest.mark.parametrize("start_token", [0, 7, 42, 128, 255]) + def test_disagreement_is_only_ever_a_tie_break(self, exported, start_token): + # guirguispierre 2026-08-10 - asserting equal argmax would be flaky, so + # require instead that pytorch scores the token c++ picked as good as + # its own; a genuinely wrong engine fails this, a near-tie does not + model, path = exported + cpp = _cpp_logits(path, start_token) + py = _py_logits(model, start_token) + + delta = (cpp - py).abs().max().item() + shortfall = (py.max() - py[int(cpp.argmax())]).item() + assert shortfall <= 2 * delta + 1e-6, ( + f"C++ chose token {int(cpp.argmax())}, which PyTorch scores " + f"{shortfall:.4f} below its own pick {int(py.argmax())} -- more " + f"than the {delta:.3e} numerical difference explains" + ) + + +class TestRunnerCLI: + def test_help_exits_cleanly(self): + res = subprocess.run([RUNNER, "--help"], capture_output=True, + text=True, timeout=60) + assert res.returncode == 0 + assert "--model" in res.stdout and "--steps" in res.stdout + + def test_unknown_flag_is_rejected(self): + res = subprocess.run([RUNNER, "--not-a-flag"], capture_output=True, + text=True, timeout=60) + assert res.returncode != 0 diff --git a/tests/test_kernel_parity.py b/tests/test_kernel_parity.py index 98910f1..f2e36ef 100644 --- a/tests/test_kernel_parity.py +++ b/tests/test_kernel_parity.py @@ -86,3 +86,26 @@ def test_all_neg_ones_weights(self): C = self.ternary_matmul(A, B) expected = -A.astype(np.int32).sum(axis=1, keepdims=True) * np.ones((1, 4), dtype=np.int32) np.testing.assert_array_equal(C, expected) + + def test_int8_min_activation(self): + # guirguispierre 2026-08-10 - -128 against -1 is +128, which int8 cannot + # hold; K=64 keeps this on the vectorized path + A = np.full((1, 64), -128, dtype=np.int8) + B = np.full((64, 8), -1, dtype=np.int8) + C = self.ternary_matmul(A, B) + expected = A.astype(np.int32) @ B.astype(np.int32) + np.testing.assert_array_equal(C, expected) + assert C[0, 0] == 8192, "int8 overflow in the widening multiply" + + def test_full_int8_activation_range(self): + rng = np.random.default_rng(0) + for _ in range(40): + K = int(rng.integers(1, 300)) + N = int(rng.integers(1, 64)) + A = rng.integers(-128, 128, size=(2, K)).astype(np.int8) + B = rng.integers(-1, 2, size=(K, N)).astype(np.int8) + np.testing.assert_array_equal( + self.ternary_matmul(A, B), + A.astype(np.int32) @ B.astype(np.int32), + err_msg=f"mismatch at K={K}, N={N}", + ) diff --git a/tools/parity_check.py b/tools/parity_check.py index 9b816eb..6d9a628 100644 --- a/tools/parity_check.py +++ b/tools/parity_check.py @@ -98,7 +98,7 @@ def run_parity_check(): print(f" Manual Dequant: {dequant}") q_out = layer.attn.q_proj(x_norm) - attn_out = layer.attn(x_norm) + attn_out, _ = layer.attn(x_norm) print_tensor_stats(f"Layer{i}_Attn_Out", attn_out) x = x + attn_out