Repository navigation
pack ternary weights four per byte - #6
Merged
Merged
Conversation
Ternary has three states but the exporter spent a whole byte on each, so the
format wasted 4x on the thing the project exists to compress. Weights are now
written as w + 1, giving 2-bit codes {0, 1, 2} with 3 unused, filling each byte
from the low-order pair upward in the tensor's flat row-major order.
Bumps the format to v2. Existing .bin files must be re-exported.
Both unpack at load rather than in the inner loop, so bit_linear still sees a plain int8 array and throughput is unchanged at ~185 tps. Verified the unpacked weights are byte-identical to what quantize_and_transpose_weight produces.
Both ports carry their own loaders and would have read packed bytes as if unpacked. The esp32 path streams from flash, so it now reads one packed row per weight row instead of seeking once per weight, which is the 4x reduction in reads its own TODO asked for. Neither port checked the format version before; both do now.
Asserts the round trip is lossless across element counts either side of a byte boundary, that a weight costs a quarter byte, that code 3 never reaches the stream, and that padding bits are zero. Updates the layout parser and byte budget in test_export_roundtrip for v2.
The ESP32 guide called 2-bit packing a solution nobody had built and sized the model at 28MB against 16MB of flash. pocket_4k is now 5.4 MB and flagship_12m 8.4 MB, so both fit. The float32 embedding table is what dominates a v2 export.
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
Auto-approved: pull request authored by repository owner @guirguispierre.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Ternary weights have three states but the exporter spent a whole byte on each, so the binary format wasted 4x on precisely the thing the project exists to compress. This writes them as w + 1, giving 2-bit codes with one unused, four to a byte. pocket_4k goes 8.5 MB to 5.4 MB and flagship_12m 16.7 MB to 8.4 MB, so both now fit ESP32 flash, which is what the port guide has been asking for while describing the packing as something nobody had built yet.
Both C++ loaders unpack at load rather than in the inner loop, so bit_linear still sees a plain int8 array and throughput is unchanged at ~185 tps. I checked losslessness directly rather than inferring it from output similarity: the loader's unpacked weights are byte-identical to what quantize_and_transpose_weight produces, for a 4096-element projection and a 32768-element head. The esp32 and wasm ports carry their own loaders and would have silently read packed bytes as unpacked, so both are updated; the esp32 one now reads a packed row per weight row instead of seeking once per weight, which is the 4x fewer reads its own TODO asked for. I also built the wasm port on the host and confirmed it agrees with pytorch on greedy tokens, which as far as I can tell nobody had verified before.
Two things to be suspicious of. This is a breaking format change: the version goes to 2 and every loader rejects v1 with a message pointing at the exporter, so any existing .bin has to be re-exported. And bit_linear_flash in the esp32 port is still dead code that nothing calls, so my rewrite of it is reasoned rather than run; while I was in there I noticed its x_q and normed buffers are sized DIM but would be indexed to HIDDEN_DIM if it were ever called, which I left alone. The remaining size lever is the float32 embedding table, now the larger half of a v2 export.