Quantized token embedding via GatherBlockQuantized (272MB→76MB) - #400
Merged
Conversation
Performance Comparison
|
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR extends the GGUF “keep quantized” path to keep the token embedding table block-quantized on disk and in the ONNX graph, switching embeddings from a float Gather to com.microsoft.GatherBlockQuantized to reduce size and bandwidth (notably for WebGPU’s 256 MiB buffer constraint).
Changes:
- Emit
QuantizedEmbedding(GatherBlockQuantized) when GGUF import can preserve embedding quantization, and plumbquantize_embeddings/quantize_lm_head/tie_word_embeddingsinto the GGUF quantization config. - Adjust
CausalLMModel.preprocess_weights()to avoid applying float tie logic when the model uses a tied quantized embedding/head table. - Add GGUF tests covering quantized embeddings and tied quantized embedding+head behavior; update CLI/help text and quantization config docs accordingly.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
src/mobius/models/base.py |
Avoids applying float tie logic in preprocess_weights() when embed+lm_head are tied via the quantized-table path. |
src/mobius/integrations/gguf/_builder.py |
Detects whether embeddings can stay quantized; repacks embedding tensors for GatherBlockQuantized; updates GGUF quantization config wiring and logging. |
src/mobius/integrations/gguf/_builder_test.py |
Adds coverage verifying GatherBlockQuantized is present and that tied quantized embedding+head doesn’t create duplicate lm_head initializers. |
src/mobius/_configs/_quantization.py |
Updates docstrings to reflect both Olive RTN and GGUF using quantized embeddings / tied quantized head. |
src/mobius/__main__.py |
Updates GGUF CLI messaging/help to mention GatherBlockQuantized alongside MatMulNBits. |
Comments suppressed due to low confidence (1)
src/mobius/integrations/gguf/_builder.py:599
- _load_quantized_state_dict always emits a *.zero_points tensor when repacked.zero_points is present. For symmetric GGUF types (config.quantization.sym=True), the QuantizedLinear/QuantizedEmbedding modules don’t have a zero_points initializer, so these entries will be skipped with warning spam and unnecessary memory. Gate zero_points emission on the config’s sym flag.
state_dict[f"{stem}.scales"] = s
if repacked.zero_points is not None:
zp = torch.from_numpy(repacked.zero_points)
if _needs_qk_permute(hf_name, num_heads, num_kv_heads, model_type):
zp = _reverse_permute(zp, n_head)
state_dict[f"{stem}.zero_points"] = zp
Emit GatherBlockQuantized for repackable GGUF token embeddings and preserve tied quantized LM heads. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
justinchuby
force-pushed
the
feat/gguf-quantized-embedding
branch
from
July 13, 2026 05:15
9f5a9f7 to
71e65bc
Compare
justinchuby
added a commit
that referenced
this pull request
Jul 13, 2026
Follow-up to #400 (quantized embedding). Profiling showed the CPU model shipped an untied `lm_head` as a plain fp32 MatMul (~544 MB) run every token. This emits the output head as Q4 `MatMulNBits` (untied) / shares the packed embedding table (tied). - Result: 169 MatMulNBits, 1 GatherBlockQuantized, **0 plain fp32 MatMul**; no 544 MB fp32 tables. Model ~1.2 GB → ~399 MB. - Coherent output verified ('Paris'). - Note: a quick 6-thread decode check showed ~38 vs ~40 tok/s (no clear speedup) — a rigorous benchmark is pending; the win is model size + correct all-quantized graph shape (matches llama.cpp). Stacks on #400's embedding commit. lintrunner clean; gguf pytest 157 passed. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
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.
The token-embedding table was emitted as fp16
[151936,896]≈ 272 MB (a plain Gather) even in Q4 models — larger than WebGPU's 256 MiB buffer limit and a big bandwidth cost. This keeps it quantized:com.microsoft.GatherBlockQuantizedwith Q4 packed data[151936,448]+ fp16 scales.WebGpuExecutionProvider(on-device) with coherent output.Note: a runtime-side WebGPU stability issue (stale KV validation) is tracked separately in onnx-genai. lintrunner clean; gguf pytest 153 passed.