Compress the embedding table to SFP at load time - #993
Conversation
`c_embedding` is capped at 16 bits by `.min_size = Type::kBF16` in tensor_info.cc, so it stays BF16 even in an SFP model where every other large tensor is 8-bit. For 2.0-2b-pt-sfp.sbs that single 256000x2304 tensor is 1.18 GB, 37% of the file. Add `--sfp_embedding` (default off) to compress it to SFP while loading, as suggested by google#164. This halves its footprint and, for models with tied input/output embeddings, the weight bandwidth of the per-token logits MatMul. Models with a separate output head receive only the memory benefit. The read paths already handle SFP: `EmbedToken` seeks by `Stride()`, and `CallMatMul` dispatches to matmul_static_sfp. `MakeBatches` reads directly into destination rows, so flagged tensors bypass it and are read via `ReadAllToSFP`. That reads 4 MiB row chunks rather than staging the whole tensor, which would raise peak RSS by more than the conversion saves. The file is read twice because SFP encodes a limited range of magnitudes, hence the per-tensor scale must be known before encoding; the second read comes from the OS cache. Conversion requires owned memory, so an explicit SFP request disables automatic mapping. An explicit `--map=1` still takes precedence and warns that SFP embedding conversion is ignored. Validate the source blob size before chunked reads so inconsistent metadata cannot cross into an adjacent blob. Add focused tests for F32 and BF16 conversion, scale propagation, final partial chunks, malformed blob sizes, and mapping-mode selection. Measured on an M2 (4 workers), 2.0-2b-pt-sfp.sbs, Mode::kRead: Gen.EmbeddingMatmul 654 -> 468 us/token (-28.5%) decode 11.21 -> 12.42 tok/s (+10.8%, median of 3) max RSS 3482 -> 2928 MiB (-554 MiB) cross entropy/byte 1.372889 -> 1.368894 (no degradation) Startup does not regress: ReadBatches drops more than ReadAllToSFP adds, because MakeBatches issues row-wise I/O for this padded 256000-row tensor.
jan-wassenberg
left a comment
jan-wassenberg
left a comment
There was a problem hiding this comment.
Nice work, thanks! Two minor suggestions.
| buffers.Decompress(tensor, reader, begin, end); | ||
|
|
||
| float maxabs = 0.0f; | ||
| for (size_t i = 0; i < (end - begin) * cols; ++i) { |
There was a problem hiding this comment.
Is this worth vectorizing? It's nice to minimize startup time. Recommend two separate vectors and then ReduceMax(d, Max(v1, v2)).
There was a problem hiding this comment.
Done in a075d4e. MaxAbs now uses two independent Highway vector accumulators and combines them with ReduceMax(df, Max(max0, max1)); the remaining vector and scalar tail are handled explicitly. On a Gemma 3 4B-sized BF16 chunk (819 x 2560 floats), the loop benchmark improved from 2.110 ms to 0.270 ms (~7.8x), with a bit-identical maximum. Thanks!
| } | ||
|
|
||
| template <> | ||
| float SourceValue(const BF16 value) { |
There was a problem hiding this comment.
Both conversions can just be hwy::ConvertScalarTo.
There was a problem hiding this comment.
Done in 4b49a8f. Both SourceValue and MakeSourceValue now use hwy::ConvertScalarTo, so the explicit BF16 specializations are gone. Thanks!
jan-wassenberg
left a comment
jan-wassenberg
left a comment
There was a problem hiding this comment.
Nice, congrats on the speedup. Interesting the compiler was not able to autovectorize this effectively.
Note that weights.cc does not yet use dynamic dispatch, but that's probably fine, there's only light compute happening here.
One more minor nit for readability:
| Chunk<T> buffers(cols, rows_per_chunk); | ||
| const float* raw = buffers.Decompress(tensor, reader, begin, end); | ||
|
|
||
| using DF = hwy::HWY_NAMESPACE::ScalableTag<float>; |
There was a problem hiding this comment.
Let's add namespace hn = hwy::HWY_NAMESPACE; (typically at namespace scope), then write hn::Zero to shorten the code :)
There was a problem hiding this comment.
Done in f3407af. Added namespace hn = hwy::HWY_NAMESPACE; at namespace scope and used the hn:: shorthand for the Highway operations, including hn::Zero. Also ran clang-format and verified all four weights_test cases pass. Thanks for giving me feedback!
jan-wassenberg
left a comment
jan-wassenberg
left a comment
There was a problem hiding this comment.
LGTM, thanks for updating!
jan-wassenberg
added
the
copybara-import
2 participants
Summary
--sfp_embeddingto convert the input embedding from F32/BF16 to SFP while loading--map 1still takes precedence and warnsThis implements the embedding-compression approach discussed in #164. Existing SFP read and MatMul paths require no changes after loading.
Results
Measured on an M2 with four workers using
2.0-2b-pt-sfp.sbsin read mode:Gen.EmbeddingMatmulTesting
cmake --build build --target weights_test -j4ctest --test-dir build -R '^WeightsTest\.' --output-on-failurecmake --build build --target gemma -j4--map 1