Isolate Gemma4 per-layer embedding quantization

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>

Author: justinchuby

src/mobius/_builder.py

@@ -592,29 +592,17 @@ def build(
         # ``embedding_bits`` is only for Gemma4's per-layer embedding table.  Do not
         # attach a QuantizationConfig to ordinary text models, because that changes
         # their regular token embedding/Linear modules.
-        if getattr(config, "hidden_size_per_layer_input", 0) and getattr(
-            config, "vocab_size_per_layer_input", 0
+        if (
+            hasattr(config, "per_layer_embedding_bits")
+            and getattr(config, "hidden_size_per_layer_input", 0)
+            and getattr(config, "vocab_size_per_layer_input", 0)
         ):
-            from mobius._configs import QuantizationConfig
-
-            if config.quantization is None:
-                config = dataclasses.replace(
-                    config,
-                    quantization=QuantizationConfig(
-                        bits=embedding_bits,
-                        group_size=32,
-                        quant_method="none",
-                        sym=False,
-                        quantize_embeddings=True,
-                    ),
-                )
-            else:
-                qc = dataclasses.replace(
-                    config.quantization,
-                    bits=embedding_bits,
-                    quantize_embeddings=True,
-                )
-                config = dataclasses.replace(config, quantization=qc)
+            config = dataclasses.replace(
+                config,
+                per_layer_embedding_bits=embedding_bits,
+                per_layer_embedding_group_size=32,
+                per_layer_embedding_sym=False,
+            )
 
     if task is None:
         task = _default_task_for_model(model_type)

src/mobius/_configs/_base.py

@@ -324,7 +324,9 @@ def _extract_vision_config(config, parent_config, model_type: str) -> dict:
     hooks live under :mod:`mobius._configs.per_model` and are
     registered with :mod:`mobius._configs._extractors` at import time.
     """
-    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._extractors import extract_vision_config as _dispatch
 
     return _dispatch(config, parent_config, model_type)
@@ -337,7 +339,9 @@ def _extract_audio_config(config, parent_config, model_type: str) -> dict:
     hooks live under :mod:`mobius._configs.per_model` and are
     registered with :mod:`mobius._configs._extractors` at import time.
     """
-    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._extractors import extract_audio_config as _dispatch
 
     return _dispatch(config, parent_config, model_type)
@@ -1593,6 +1597,12 @@ class Gemma4Config(VisionLanguageConfig):
     # is too small for the fused [V, L*D] per-layer embedding table, to split
     # it into L separate [V, D] tables that each fit within the EP's buffer limit.
     split_per_layer_embedding: bool = False
+    # Dedicated quantization settings for the split per-layer embedding table.
+    # Kept separate from ``quantization`` so embedding_bits does not affect text
+    # token embeddings, LM head, or decoder Linear projection quantization.
+    per_layer_embedding_bits: int | None = None
+    per_layer_embedding_group_size: int = 32
+    per_layer_embedding_sym: bool = False
 
     @classmethod
     def from_transformers(cls, config, parent_config=None) -> Gemma4Config:

src/mobius/models/gemma4.py

@@ -230,6 +230,26 @@ def _per_layer_embedding_block_size(embedding_dim: int, group_size: int) -> int:
     return group_size
 
 
+def _per_layer_embedding_quantization(config: Gemma4Config) -> tuple[int, int, bool] | None:
+    """Return split per-layer embedding quantization as (bits, group_size, symmetric)."""
+    bits = getattr(config, "per_layer_embedding_bits", None)
+    if bits is None:
+        quantization_config = getattr(config, "quantization", None)
+        if quantization_config is None or not getattr(
+            quantization_config, "quantize_embeddings", False
+        ):
+            return None
+        bits = quantization_config.bits
+        group_size = quantization_config.group_size
+        symmetric = quantization_config.sym
+    else:
+        group_size = getattr(config, "per_layer_embedding_group_size", 32)
+        symmetric = getattr(config, "per_layer_embedding_sym", False)
+    if bits not in (4, 8):
+        raise ValueError(f"quantize_embeddings requires bits=4 or bits=8, got {bits}")
+    return bits, group_size, symmetric
+
+
 def _dtype_safe_compress(
     op: OpBuilder, data: ir.Value, condition: ir.Value, *, axis: int
 ) -> ir.Value:
@@ -1830,25 +1850,20 @@ def __init__(self, config: Gemma4Config):
             # 256 MiB limit; ~128 MiB each vs ~4.7 GB fused).
             # Only the table actually called in forward() is realized as an
             # ONNX initializer, so the unused one adds no graph weight.
-            qc = getattr(config, "quantization", None)
-            if (
-                qc is not None
-                and getattr(qc, "quantize_embeddings", False)
-                and config.split_per_layer_embedding
-            ):
-                block_size = _per_layer_embedding_block_size(
-                    self._per_layer_dim, qc.group_size
-                )
+            per_layer_quant = _per_layer_embedding_quantization(config)
+            if per_layer_quant is not None and config.split_per_layer_embedding:
+                bits, group_size, symmetric = per_layer_quant
+                block_size = _per_layer_embedding_block_size(self._per_layer_dim, group_size)
                 self.embed_tokens_per_layer_split = nn.ModuleList(
                     [
                         QuantizedScaledWordEmbedding(
                             vocab_per_layer,
                             self._per_layer_dim,
                             config.pad_token_id,
                             embed_scale=float(self._per_layer_dim**0.5),
-                            bits=qc.bits,
+                            bits=bits,
                             block_size=block_size,
-                            has_zero_point=not qc.sym,
+                            has_zero_point=not symmetric,
                         )
                         for _ in range(self._num_layers)
                     ]
@@ -2402,24 +2417,18 @@ def preprocess_weights(
                     f"got {fused.shape[1]}"
                 )
                 chunks = fused.chunk(num_layers, dim=1)
-                qc = getattr(self.config, "quantization", None)
-                quantize_per_layer = qc is not None and getattr(
-                    qc, "quantize_embeddings", False
-                )
-                if quantize_per_layer:
-                    if qc.bits not in (4, 8):
-                        raise ValueError(
-                            f"quantize_embeddings requires bits=4 or bits=8, got {qc.bits}"
-                        )
+                per_layer_quant = _per_layer_embedding_quantization(self.config)
+                if per_layer_quant is not None:
+                    bits, group_size, symmetric = per_layer_quant
                     from mobius._weight_utils import quantize_embedding_rtn
 
-                    block_size = _per_layer_embedding_block_size(per_layer_dim, qc.group_size)
+                    block_size = _per_layer_embedding_block_size(per_layer_dim, group_size)
                     for i, chunk in enumerate(chunks):
                         qweight, scales, zero_points = quantize_embedding_rtn(
                             chunk.contiguous(),
-                            bits=qc.bits,
+                            bits=bits,
                             block_size=block_size,
-                            symmetric=qc.sym,
+                            symmetric=symmetric,
                         )
                         state_dict[f"model.embed_tokens_per_layer_split.{i}.qweight"] = qweight
                         state_dict[f"model.embed_tokens_per_layer_split.{i}.scales"] = scales
@@ -3229,24 +3238,18 @@ def preprocess_weights(
                     f"got {fused.shape[1]}"
                 )
                 chunks = fused.chunk(num_layers, dim=1)
-                qc = getattr(self.config, "quantization", None)
-                quantize_per_layer = qc is not None and getattr(
-                    qc, "quantize_embeddings", False
-                )
-                if quantize_per_layer:
-                    if qc.bits not in (4, 8):
-                        raise ValueError(
-                            f"quantize_embeddings requires bits=4 or bits=8, got {qc.bits}"
-                        )
+                per_layer_quant = _per_layer_embedding_quantization(self.config)
+                if per_layer_quant is not None:
+                    bits, group_size, symmetric = per_layer_quant
                     from mobius._weight_utils import quantize_embedding_rtn
 
-                    block_size = _per_layer_embedding_block_size(per_layer_dim, qc.group_size)
+                    block_size = _per_layer_embedding_block_size(per_layer_dim, group_size)
                     for i, chunk in enumerate(chunks):
                         qweight, scales, zero_points = quantize_embedding_rtn(
                             chunk.contiguous(),
-                            bits=qc.bits,
+                            bits=bits,
                             block_size=block_size,
-                            symmetric=qc.sym,
+                            symmetric=symmetric,
                         )
                         renamed[f"decoder.model.embed_tokens_per_layer_split.{i}.qweight"] = (
                             qweight

src/mobius/models/gemma4_test.py

@@ -8,7 +8,7 @@
 import onnx_ir as ir
 import torch
 
-from mobius._configs import AudioConfig, Gemma4Config
+from mobius._configs import AudioConfig, Gemma4Config, QuantizationConfig
 from mobius.models.gemma4 import Gemma4CausalLMModel, Gemma4EmbeddingModel, Gemma4Model
 
 
@@ -120,6 +120,49 @@ def test_per_expert_scale_not_folded(self):
         assert torch.allclose(result[key], torch.ones(4))
 
 
+class TestGemma4PerLayerEmbeddingQuantization:
+    """Per-layer embedding quantization stays independent from text quantization."""
+
+    def test_split_defaults_do_not_mutate_existing_quantization_config(self, monkeypatch):
+        import dataclasses
+
+        import mobius.tasks._gemma4 as gemma4_task_module
+        from mobius.tasks._gemma4 import Gemma4Task
+
+        class _TinyCaps:
+            max_buffer_size = 1
+
+        monkeypatch.setattr(gemma4_task_module, "ep_capabilities", lambda: _TinyCaps())
+
+        config = _tiny_gemma4_config(
+            enable_moe_block=False,
+            hidden_size_per_layer_input=32,
+            vocab_size_per_layer_input=64,
+        )
+        config = dataclasses.replace(
+            config,
+            quantization=QuantizationConfig(
+                bits=8,
+                group_size=64,
+                quant_method="none",
+                sym=True,
+                quantize_embeddings=False,
+            ),
+        )
+
+        pkg = Gemma4Task().build(Gemma4Model(config), config)
+
+        assert config.quantization is not None
+        assert config.quantization.bits == 8
+        assert config.quantization.group_size == 64
+        assert config.quantization.sym is True
+        assert config.quantization.quantize_embeddings is False
+        assert config.per_layer_embedding_bits == 4
+        assert config.per_layer_embedding_group_size == 32
+        assert config.per_layer_embedding_sym is False
+        assert any(node.op_type == "GatherBlockQuantized" for node in pkg["decoder"].graph)
+
+
 class TestGemma4EmbeddingModel:
     def test_reuses_token_id_fields_for_masking(self):
         config = _tiny_gemma4_config(

src/mobius/tasks/_gemma4.py

@@ -25,7 +25,7 @@
 from onnxscript import GraphBuilder, nn
 
 from mobius._build_context import ep_capabilities
-from mobius._configs import Gemma4Config, QuantizationConfig
+from mobius._configs import Gemma4Config
 from mobius._model_package import ModelPackage
 from mobius._pipeline_contract import (
     declare_component_presence,
@@ -444,28 +444,24 @@ def build(
         # in mobius.build() / the MobiusBuilder Olive pass config.
         if config.split_per_layer_embedding:
             qc = config.quantization
-            if qc is not None and qc.quantize_embeddings:
-                # embedding_bits was set by the caller — use its bits value
+            if getattr(config, "per_layer_embedding_bits", None) is not None:
+                if config.per_layer_embedding_bits not in (4, 8):
+                    raise ValueError(
+                        "quantize_embeddings requires bits=4 or bits=8, "
+                        f"got {config.per_layer_embedding_bits}"
+                    )
+            elif qc is not None and qc.quantize_embeddings:
                 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
+                config.per_layer_embedding_bits = qc.bits
+                config.per_layer_embedding_group_size = qc.group_size
+                config.per_layer_embedding_sym = qc.sym
             else:
-                config.quantization = QuantizationConfig(
-                    bits=4,
-                    group_size=32,
-                    quant_method="mobius",
-                    sym=False,
-                    quantize_embeddings=True,
-                )
+                config.per_layer_embedding_bits = 4
+                config.per_layer_embedding_group_size = 32
+                config.per_layer_embedding_sym = False
             # The module was already constructed before build() runs, so the
             # per-layer embeddings are plain Embedding (Gather).  Replace them
             # with QuantizedScaledWordEmbedding (GatherBlockQuantized).
@@ -474,23 +470,25 @@ def build(
                 _per_layer_embedding_block_size,
             )
 
-            qc = config.quantization
+            bits = config.per_layer_embedding_bits
+            group_size = config.per_layer_embedding_group_size
+            symmetric = config.per_layer_embedding_sym
             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))
-            block_size = _per_layer_embedding_block_size(per_layer_dim, qc.group_size)
+            block_size = _per_layer_embedding_block_size(per_layer_dim, group_size)
             module.decoder.model.embed_tokens_per_layer_split = nn.ModuleList(
                 [
                     QuantizedScaledWordEmbedding(
                         vocab_per_layer,
                         per_layer_dim,
                         config.pad_token_id,
                         embed_scale=embed_scale,
-                        bits=qc.bits,
+                        bits=bits,
                         block_size=block_size,
-                        has_zero_point=not qc.sym,
+                        has_zero_point=not symmetric,
                     )
                     for _ in range(config.num_hidden_layers)
                 ]