Quantize Gemma4 per-layer embedding tables for WebGPU (INT4/INT8 configurable) - #403
Draft
feich-ms wants to merge 6 commits into
Draft
Quantize Gemma4 per-layer embedding tables for WebGPU (INT4/INT8 configurable)#403feich-ms wants to merge 6 commits into
feich-ms wants to merge 6 commits into
Conversation
feich-ms
changed the title
feich-ms
changed the title
Performance Comparison
|
|
The author of this PR, feich-ms, is not an activated member of this organization on Codecov. |
The 35 per-layer embedding tables ([262144, 256] FP16 each) take 4.48 GB unquantized — 77% of the decoder. Quantize them to INT4 using GatherBlockQuantized, reducing the decoder from 5.8 GB to 2.7 GB and total model from 6.9 GB to 3.8 GB. Co-Authored-By: Claude <noreply@anthropic.com>
Add embedding_bits parameter to mobius.build() that threads through to the Gemma4 task's per-layer embedding quantization. When set, the specified bit-width (4 or 8) is used for GatherBlockQuantized embedding tables instead of the previous hardcoded INT4. When omitted, the default INT4 behavior is preserved. Co-Authored-By: Claude <noreply@anthropic.com>
feich-ms
force-pushed
the
user/feich/quantize_per_layer_embedding_table
branch
from
July 14, 2026 02:26
6b105f6 to
ded37e9
Compare
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds a WebGPU-focused path to reduce Gemma4 decoder size by block-quantizing the split per-layer embedding tables (INT4 by default, INT8 optional), and introduces a new embedding_bits knob plus an RTN embedding-quantization utility.
Changes:
- Add
embedding_bitsparameter tomobius.build()and plumb it intoQuantizationConfig. - Introduce
QuantizedScaledWordEmbedding(GatherBlockQuantized + Gemma scaling) and use it for Gemma4 per-layer split embeddings when enabled. - Add
quantize_embedding_rtn()to generate packedqweight/scales/zero_pointstensors for GatherBlockQuantized.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 7 comments.
| File | Description |
|---|---|
| src/mobius/tasks/_gemma4.py | Enables embedding quantization when per-layer embeddings must be split (WebGPU max buffer constraint) and swaps in quantized embedding modules. |
| src/mobius/models/gemma4.py | Adds a quantized scaled embedding module and updates Gemma4 weight preprocessing to emit qweight/scales/zero_points for per-layer embeddings. |
| src/mobius/_weight_utils.py | Adds RTN embedding-table quantization utility producing GatherBlockQuantized-compatible packed outputs. |
| src/mobius/_builder.py | Adds public API flag embedding_bits and maps it onto config quantization settings. |
- Add embedding_bits to build() docstring with description and valid values - Validate embedding_bits is 4 or 8 in build(), raise ValueError otherwise - Replace assert-based validation in quantize_embedding_rtn with ValueError - Add .detach().cpu() before .numpy() to handle non-CPU tensors - Add bits validation (4/8) in both gemma4.py preprocess_weights paths - Fix _gemma4.py to explicitly set bits/group_size/sym when enabling quantize_embeddings on an existing QuantizationConfig - Add 9 unit tests for quantize_embedding_rtn covering shapes, dtypes, error cases, and roundtrip accuracy Co-Authored-By: Claude <noreply@anthropic.com>
Restrict embedding_bits to Gemma4 per-layer embedding tables so ordinary text models are not opted into quantized embeddings or linears. Defer per-layer quantized table construction until splitting is active and use a valid embedding block size for smaller configs. Add regression coverage for embedding_bits validation and no-op behavior on regular text models.\n\nCo-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby <justinchuby@users.noreply.github.com>
|
|
|
You are seeing this message because GitHub Code Scanning has recently been set up for this repository, or this pull request contains the workflow file for the Code Scanning tool. What Enabling Code Scanning Means:
For more information about GitHub Code Scanning, check out the documentation. |
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| assert scales.shape == (64, 1) | ||
|
|
||
| def test_invalid_bits_raises(self): | ||
| """bits not in (4, 8) raises ValueError.""" |
| """Handles non-CPU tensors gracefully (detach().cpu() path).""" | ||
| weight = torch.randn(32, 32, dtype=torch.float16) | ||
| # Just verify it doesn't crash — the .detach().cpu() handles it | ||
| qweight, scales, zp = self.quantize(weight, bits=4, block_size=32) |
| """Handles non-CPU tensors gracefully (detach().cpu() path).""" | ||
| weight = torch.randn(32, 32, dtype=torch.float16) | ||
| # Just verify it doesn't crash — the .detach().cpu() handles it | ||
| qweight, scales, zp = self.quantize(weight, bits=4, block_size=32) |
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.
Comments suppressed due to low confidence (3)
src/mobius/_builder.py:617
- When config.quantization already exists (e.g., from a quantized HF/GGUF/Olive config), this code unconditionally overwrites qc.bits to embedding_bits. That can make the build configuration inconsistent with the checkpoint’s actual quantized weights and break weight loading. Consider rejecting overrides when an existing weight-quantization config is present, or only updating bits when quant_method == "none".
else:
qc = dataclasses.replace(
config.quantization,
bits=embedding_bits,
quantize_embeddings=True,
)
config = dataclasses.replace(config, quantization=qc)
src/mobius/tasks/_gemma4.py:482
- embed_scale is computed via float(np.float16(...)), which rounds the scale to FP16 and then back to Python float. That changes the constant compared to the existing embedding scaling logic (which uses full-precision float), and can introduce avoidable numerical differences.
import numpy as np
embed_scale = float(np.float16(per_layer_dim**0.5))
src/mobius/models/gemma4.py:221
- QuantizedScaledWordEmbedding.forward is missing an explicit return type annotation (other embedding forwards in this file return ir.Value). This can trip strict type-checking and is inconsistent with surrounding code.
def forward(self, op: OpBuilder, input_ids: ir.Value):
Comment on lines
+446
to
+468
| qc = config.quantization | ||
| if qc is not None and qc.quantize_embeddings: | ||
| # embedding_bits was set by the caller — use its bits value | ||
| if qc.bits not in (4, 8): | ||
| raise ValueError( | ||
| f"quantize_embeddings requires bits=4 or bits=8, got {qc.bits}" | ||
| ) | ||
| elif qc is not None: | ||
| # Existing quantization config without embedding quantization — | ||
| # enable it with INT4 (the existing bits may be for MatMul | ||
| # quantization which is unrelated, so override for embeddings) | ||
| config.quantization.quantize_embeddings = True | ||
| config.quantization.bits = 4 | ||
| config.quantization.group_size = 32 | ||
| config.quantization.sym = False | ||
| else: | ||
| config.quantization = QuantizationConfig( | ||
| bits=4, | ||
| group_size=32, | ||
| quant_method="mobius", | ||
| sym=False, | ||
| quantize_embeddings=True, | ||
| ) |
Keep embedding_bits and split per-layer embedding quantization separate from the shared text QuantizationConfig so decoder Linear, LM head, and token embedding quantization are not changed incidentally. Add dedicated Gemma4 config fields for per-layer embedding quantization and a regression test that split defaults do not mutate an existing quantization config.\n\nCo-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby <justinchuby@users.noreply.github.com>
| """ | ||
| from mobius._configs import per_model # noqa: F401 - imported for registration side effect | ||
| from mobius._configs import ( | ||
| per_model, # ruff: ignore[unused-import] - imported for registration side effect |
| """ | ||
| from mobius._configs import per_model # noqa: F401 - imported for registration side effect | ||
| from mobius._configs import ( | ||
| per_model, # ruff: ignore[unused-import] - imported for registration side effect |
| """ | ||
| from mobius._configs import per_model # noqa: F401 - imported for registration side effect | ||
| from mobius._configs import ( | ||
| per_model, # ruff: ignore[unused-import] - imported for registration side effect |
| """ | ||
| from mobius._configs import per_model # noqa: F401 - imported for registration side effect | ||
| from mobius._configs import ( | ||
| per_model, # ruff: ignore[unused-import] - imported for registration side effect |
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
| def test_split_defaults_do_not_mutate_existing_quantization_config(self, monkeypatch): | ||
| import dataclasses | ||
|
|
||
| import mobius.tasks._gemma4 as gemma4_task_module |
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 8 out of 8 changed files in this pull request and generated 1 comment.
Comments suppressed due to low confidence (5)
src/mobius/_configs/_base.py:344
- Same as above:
# ruff: ignore[unused-import]is not the suppression style used in this repo and is likely to fail linting. Prefer the existing# noqa: F401pattern for side-effect imports.
from mobius._configs import (
per_model, # ruff: ignore[unused-import] - imported for registration side effect
)
src/mobius/tasks/_gemma4.py:481
embed_scaleis computed withfloat(np.float16(...))here, but Gemma4’s constructor usesfloat(self._per_layer_dim**0.5)when it creates the same embedding module. That means the scale constant depends on whether split/quantization is decided before vs during task.build(), which can change numerics. Also, NumPy is only imported for this cast.
per_layer_dim = getattr(config, "hidden_size_per_layer_input", 0)
vocab_per_layer = getattr(config, "vocab_size_per_layer_input", 0)
import numpy as np
embed_scale = float(np.float16(per_layer_dim**0.5))
src/mobius/models/gemma4.py:205
- This docstring hard-codes “INT4”, but the class is used with both 4-bit and 8-bit embedding quantization (bits=4/8). Updating the docstring avoids misleading API/docs readers.
class QuantizedScaledWordEmbedding(QuantizedEmbedding):
"""Quantized embedding with scaling — INT4 GatherBlockQuantized + scale multiply."""
src/mobius/models/gemma4.py:221
QuantizedScaledWordEmbedding.forwardis missing a return type annotation (-> ir.Value), unlike the surrounding modules. This can trip mypy/typing checks and reduces API clarity.
def forward(self, op: OpBuilder, input_ids: ir.Value):
src/mobius/_weight_utils.py:742
- For
bits==4, packing assumesembedding_dimis even (pairs of nibbles). Ifembedding_dimis odd, this will throw a NumPy broadcasting error at pack time. Add an explicit validation so callers get a clearValueError.
num_embeddings, embedding_dim = weight.shape
if embedding_dim % block_size != 0:
raise ValueError(
f"embedding_dim ({embedding_dim}) must be divisible by block_size ({block_size})"
)
Comment on lines
+327
to
+329
| from mobius._configs import ( | ||
| per_model, # ruff: ignore[unused-import] - imported for registration side effect | ||
| ) |
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.
Summary
embed_tokens_per_layer_split) from FP16 to INT4 usingGatherBlockQuantized, reducing decoder size from 5.8 GB to 2.7 GB for WebGPUembedding_bitsparameter tomobius.build()so callers can choose INT4 (default) or INT8 without editing source codequantize_embedding_rtn()utility for block-wise RTN quantization of embedding tables (supports 4-bit and 8-bit, symmetric and asymmetric)QuantizedScaledWordEmbeddingclass that wrapsGatherBlockQuantizedwith the Gemma4 embed scalingDetails
The Gemma4 E2B multimodal model splits a fused [262144, 35×256] embedding table into 35 separate [262144, 256] tables for WebGPU's 256 MiB buffer limit. At FP16, these tables alone take 4.48 GB (77% of the decoder). This PR quantizes them to INT4 (560 MB) or INT8 (1.12 GB) using the
GatherBlockQuantizedcustom op.The
embedding_bitsparameter can be set via:mobius.build(..., embedding_bits=4)(Python API)"embedding_bits": 4in the MobiusBuilder Olive pass config JSONWhen omitted, defaults to INT4 for backward compatibility.
Test plan
embedding_bitspreserves existing INT4 default behavior