NanoJev-mlx-4bit / README.md
zerodegress's picture
Upload README.md with huggingface_hub
b4e52f3 verified
|
Raw History Blame Contribute Delete
7.26 kB
---
base_model: C-Tianyu/NanoJev
library_name: mlx
tags:
- nanojev
- jev
- decision-model
- probability-distributions
- mlx
- 4-bit
- quantized
---
# NanoJev-mlx-4bit
**An MLX affine 4-bit quantization of the root checkpoint of [C-Tianyu/NanoJev](https://huggingface.co/C-Tianyu/NanoJev)** β€” a nano replica of TypeSafe's Jev System One decision model: a Qwen3-0.6B backbone (28 layers, hidden 1024) with structured decision heads, 596M parameters in total. It returns complete probability distributions over dynamic candidates (`choice` 2–255 options, `boolean`, `score` 2–10 levels) in one batched backbone forward, with zero output-token decoding.
No training was performed and no weight changed its value: this repository re-encodes the upstream fp32 checkpoint in a lower-precision weight format.
## What this is, and what it is not
- **It is** a post-training **quantization**: every quantized tensor is mapped onto 4-bit integer values inside groups of 64 input features, with one fp16 scale and one fp16 bias per group.
- **It is not a fine-tuned model.** Zero gradients, zero data, zero optimizer steps were involved.
- **It is not a new architecture or an independent Jev implementation.** The upstream checkpoint is the trained artifact; this is the same decision function at lower weight precision.
- **License**: the upstream repository declares no license for the weights (its code is MIT). This derivative therefore claims **no license** either, and points back to the upstream checkpoint for any licensing question.
## Quantization configuration
| | |
|---|---|
| scheme | MLX affine, `bits=4`, `group_size=64` |
| quantized (595,984,384 params) | the 196 backbone linear projections (`self_attn.{q,k,v,o}_proj`, `mlp.{gate,up,down}_proj`, 440,401,920 params) **and** `backbone.embed_tokens` (155,582,464 params) |
| kept at fp16 (no quantization) | all RMSNorm / LayerNorm weights, all biases, and the **entire decision head** (`norm`, `scalar`, `set_project`, `set_attention`, `set_output` β€” 200,578 params, 0.03 % of the model) |
| scales / biases dtype | `float16`; a bf16 runtime should cast them to the activation dtype at load time |
The decision head is deliberately left unquantized: it is 0.03 % of the parameters and it is where the decision boundary lives, so there is nothing to gain and something to lose.
Measured weight-space error of the quantization (per-tensor max / mean relative error vs the upstream fp32 weights): `mlp` 12.45 % / 6.79 %, `self_attn` 11.75 % / 7.74 %, `embed_tokens` 6.70 % / 6.70 %.
## Size
| file | bytes | relative |
|---|---|---|
| upstream fp32 `best.safetensors` | 2,385,039,280 | 100 % |
| [NanoJev-bf16](https://huggingface.co/ZeroDegress/NanoJev-bf16) | 1,192,538,284 | 50.0 % |
| **this artifact** | **335,857,708** | **14.1 %** |
## Measured quality cost
Reference: the upstream fp32 checkpoint. Split: the frozen `stage2/dev.jsonl` (120 states / 360 questions / 898 candidate paths), Apple M4, one backbone forward per 8-state batch, zero decode steps.
Method, stated plainly: this artifact was dequantized back to fp32 and evaluated through the **unmodified upstream inference path**, so the numbers below isolate the quantization error itself. A real MLX runtime also computes in a half-precision activation dtype, which adds a smaller second term that is **not** included here.
| metric | fp32 reference | this artifact |
|---|---|---|
| teacher-target soft CE | 0.499212 | **0.498122** (Ξ” βˆ’0.00109) |
| gold-target soft CE | 0.606506 | 0.601031 |
| argmax hit rate vs teacher | 92.16 % | 92.44 % |
| argmax agreement with the fp32 model | β€” | **99.17 %** (3 of 360 questions flip) |
| mean total variation | β€” | 0.00660 (median 0.00365, max 0.0489) |
| mean KL divergence | β€” | 2.91e-4 |
| `boolean` polarity flips | β€” | **0 / 120** |
Per question type, argmax agreement with fp32: `boolean` 100 %, `score` 100 %, `choice` 97.5 %. All three flips are `choice` questions decided near a tie; two of the three move *closer* to the teacher target after quantization. The cross-entropy does not rise β€” the quantization error is smaller than the model's remaining distance to its own targets.
A variant that leaves `embed_tokens` unquantized was measured too: identical flip set, mean TV 0.00635 instead of 0.00660. Quantizing the embedding (26 % of all parameters) therefore costs essentially nothing, which is why it is included here.
## This does not make inference faster
Measured on the same machine with MLX 0.32 at realistic shapes, a 4-bit matmul is **at parity** with fp32/bf16 weight matmul (2.96 vs 2.99 / 3.60 TFLOPS). This workload is one batched prefill per call, so it is compute-bound rather than weight-bandwidth-bound, and dequantization is pure overhead in it.
**The gain here is footprint β€” 3.5Γ— smaller than the bf16 conversion β€” not throughput.** If the goal is speed, the levers that actually pay are prefix sharing across the candidate paths of a question, length-bucketed batching, and a fused/compiled runtime.
## How to load it
The tensor layout follows the MLX convention:
- `<name>` β€” `uint32`, packed 4-bit weights
- `<name>.scales`, `<name>.biases` β€” `float16`, one entry per group of 64 input features
- every other tensor is a plain fp16 tensor
```python
import mlx.core as mx
w = mx.load("best.safetensors")
print(w["backbone.layers.0.self_attn.q_proj.weight"].shape, # (2048, 16) packed
w["backbone.layers.0.self_attn.q_proj.weight"].dtype, # uint32
w["backbone.layers.0.self_attn.q_proj.weight.scales"].shape) # (2048, 16)
```
This repository ships weights only. The architecture keys (`backbone.*` plus the decision head) are not part of any public MLX model registry, so a runtime needs an MLX implementation of that forward pass; the upstream reference implementation is at [TianyuCodings/NanoJev](https://github.com/TianyuCodings/NanoJev) (MIT, PyTorch/CUDA), and the same checkpoint in unquantized form is available upstream and as a bf16 conversion from this account.
## Verification and reproducibility
- source weights sha256 `fff62d1412685c1714eaa386acb603f9690371fb3cc8ad03dc41319302597c28` (the upstream root checkpoint, `v3_teacher_coords_multi_seed17`)
- this artifact's sha256 `332d47eb23c81b0370f580e87e0e55ef737ba1baa206a8d300eadf9617307b62`
- `MANIFEST.json` records the source hash, the quantization configuration, the parameter counts per tensor group, the per-group weight-space error, and the hash of every produced file
- produced by `scripts/quantize_mlx.py` (MLX `mx.quantize`, then `mx.dequantize` for the evaluation copy);
the evaluation ran through the frozen-split harness kept with the bf16 conversion workspace
(`compare_precision.py`), against the upstream fp32 checkpoint as reference
## Attribution
- Upstream model: [C-Tianyu/NanoJev](https://huggingface.co/C-Tianyu/NanoJev) β€” code MIT, weights with no declared license.
- Upstream data: [C-Tianyu/NanoJev-Data](https://huggingface.co/datasets/C-Tianyu/NanoJev-Data) (per-record `metadata.license = CC0-1.0`).
- Quantization and evaluation: performed locally on the upstream checkpoint; no upstream weights were retrained or altered in value beyond the quantization itself.