diff --git a/src/mobius/_builder.py b/src/mobius/_builder.py index 4638ecd4..7e8dc044 100644 --- a/src/mobius/_builder.py +++ b/src/mobius/_builder.py @@ -355,6 +355,7 @@ def build( execution_provider: str = "default", trace_optimization: bool = False, text_only: bool = False, + embedding_bits: int | None = None, ) -> ModelPackage: """Build an ONNX :class:`ModelPackage` from a HuggingFace model ID. @@ -418,6 +419,14 @@ def build( Raises :class:`ValueError` if the resolved ``model_type`` has no text-only sibling. Currently supported for ``gemma4_unified`` (``google/gemma-4-12B``). + embedding_bits: Quantization bit-width for per-layer embedding tables + (4 or 8). When set, large embedding tables that exceed the EP's + buffer limit are block-quantized with ``GatherBlockQuantized`` + instead of stored at full precision. Only affects models that + split per-layer embeddings (e.g. Gemma4 on WebGPU). When ``None`` + (default), the model task decides — currently defaults to INT4 + when splitting is required. Ignored by models without per-layer + embedding tables. Returns: A :class:`ModelPackage` containing the built model(s). @@ -577,6 +586,24 @@ def build( # states. See ``ArchitectureConfig.output_layer_indices``. config = dataclasses.replace(config, output_layer_indices=list(output_layer_indices)) + if embedding_bits is not None: + if embedding_bits not in (4, 8): + raise ValueError(f"embedding_bits must be 4 or 8, got {embedding_bits}") + # ``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 ( + hasattr(config, "per_layer_embedding_bits") + and getattr(config, "hidden_size_per_layer_input", 0) + and getattr(config, "vocab_size_per_layer_input", 0) + ): + 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) diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index 7ebbf390..d7f225d7 100644 --- a/src/mobius/_configs/_base.py +++ b/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: diff --git a/src/mobius/_exporter_test.py b/src/mobius/_exporter_test.py index 557560b8..b662602c 100644 --- a/src/mobius/_exporter_test.py +++ b/src/mobius/_exporter_test.py @@ -375,3 +375,32 @@ def test_build_sets_graph_name(self, monkeypatch): ) pkg = build("fake/model-id", load_weights=False) assert pkg["model"].graph.name == "fake/model-id/model" + + def test_embedding_bits_ignored_for_regular_text_models(self, monkeypatch): + """embedding_bits must not opt ordinary text models into quantized modules.""" + import transformers + + fake_config = self._mock_hf_config() + monkeypatch.setattr( + transformers.AutoConfig, + "from_pretrained", + lambda *a, **kw: fake_config, + ) + pkg = build("fake/model-id", embedding_bits=4, load_weights=False) + op_types = {node.op_type for node in pkg["model"].graph} + assert "MatMulNBits" not in op_types + assert "GatherBlockQuantized" not in op_types + assert pkg.config.quantization is None + + def test_embedding_bits_validates_bit_width(self, monkeypatch): + """embedding_bits only accepts the bit-widths supported by GatherBlockQuantized.""" + import transformers + + fake_config = self._mock_hf_config() + monkeypatch.setattr( + transformers.AutoConfig, + "from_pretrained", + lambda *a, **kw: fake_config, + ) + with pytest.raises(ValueError, match="embedding_bits must be 4 or 8"): + build("fake/model-id", embedding_bits=3, load_weights=False) diff --git a/src/mobius/_weight_utils.py b/src/mobius/_weight_utils.py index f7b198b9..ce0ae9bd 100644 --- a/src/mobius/_weight_utils.py +++ b/src/mobius/_weight_utils.py @@ -706,6 +706,111 @@ def preprocess_gptq_weights( return result +def quantize_embedding_rtn( + weight: torch.Tensor, + bits: int = 4, + block_size: int = 32, + symmetric: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Block-wise RTN quantization of a 2D embedding table. + + Produces the packed uint8 format expected by GatherBlockQuantized / + QuantizedEmbedding: quantize along axis=1 (embedding dim), gather along + axis=0 (vocabulary). + + Args: + weight: [num_embeddings, embedding_dim] float tensor. + bits: Quantization bit-width (4 or 8). + block_size: Number of elements per quantization group. + symmetric: If True, use symmetric quantization (no zero_points). + + Returns: + (qweight, scales, zero_points) where: + - qweight: [num_embeddings, embedding_dim * bits // 8] uint8 + - scales: [num_embeddings, n_blocks] same dtype as input + - zero_points: [num_embeddings, ceil(n_blocks * bits / 8)] uint8, + or None if symmetric. + """ + if bits not in (4, 8): + raise ValueError(f"bits must be 4 or 8, got {bits}") + if weight.ndim != 2: + raise ValueError(f"weight must be 2D, got {weight.ndim}D") + 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})" + ) + + n_blocks = embedding_dim // block_size + w = weight.detach().cpu().float().numpy() + + import numpy as np + + # Reshape to [num_embeddings, n_blocks, block_size] for per-block quantization + blocks = w.reshape(num_embeddings, n_blocks, block_size) + + if symmetric: + qmax = (1 << (bits - 1)) - 1 # 7 for 4-bit + qmin = -(1 << (bits - 1)) # -8 for 4-bit + abs_max = np.maximum(np.abs(blocks.max(axis=2, keepdims=True)), + np.abs(blocks.min(axis=2, keepdims=True))) + scales_np = abs_max / float(qmax) + scales_np = np.where(scales_np == 0, 1.0, scales_np) + quantized = np.clip(np.round(blocks / scales_np), qmin, qmax).astype(np.int8) + # Convert signed to unsigned for packing: [-8,7] -> [0,15] + quantized_unsigned = (quantized.astype(np.int16) + (1 << (bits - 1))).astype(np.uint8) + scales_np = scales_np.squeeze(2) # [num_embeddings, n_blocks] + zero_points_np = None + else: + qmax = (1 << bits) - 1 # 15 for 4-bit + block_min = blocks.min(axis=2, keepdims=True) + block_max = blocks.max(axis=2, keepdims=True) + # Ensure range includes zero + block_min = np.minimum(block_min, 0.0) + block_max = np.maximum(block_max, 0.0) + scales_np = (block_max - block_min) / float(qmax) + scales_np = np.where(scales_np == 0, 1.0, scales_np) + zp = np.clip(np.round(-block_min / scales_np), 0, qmax).astype(np.uint8) + quantized_unsigned = np.clip( + np.round(blocks / scales_np + zp), 0, qmax + ).astype(np.uint8) + scales_np = scales_np.squeeze(2) # [num_embeddings, n_blocks] + zero_points_np = zp.squeeze(2) # [num_embeddings, n_blocks] + + # Pack into uint8 + if bits == 4: + # Flatten quantized to [num_embeddings, embedding_dim] + flat = quantized_unsigned.reshape(num_embeddings, embedding_dim) + # Pack pairs of 4-bit values into uint8 (low nibble first) + packed = (flat[:, 0::2] & 0x0F) | ((flat[:, 1::2] & 0x0F) << 4) + qweight = packed.astype(np.uint8) # [num_embeddings, embedding_dim // 2] + + if zero_points_np is not None: + # Pack zero_points the same way + if n_blocks % 2 == 0: + zp_packed = (zero_points_np[:, 0::2] & 0x0F) | ((zero_points_np[:, 1::2] & 0x0F) << 4) + else: + # Pad to even count + padded = np.pad(zero_points_np, ((0, 0), (0, 1)), constant_values=0) + zp_packed = (padded[:, 0::2] & 0x0F) | ((padded[:, 1::2] & 0x0F) << 4) + zero_points_out = torch.from_numpy(zp_packed.astype(np.uint8)) + else: + zero_points_out = None + else: + # 8-bit: no packing + qweight = quantized_unsigned.reshape(num_embeddings, embedding_dim) + if zero_points_np is not None: + zero_points_out = torch.from_numpy(zero_points_np.astype(np.uint8)) + else: + zero_points_out = None + + return ( + torch.from_numpy(qweight), + torch.from_numpy(scales_np).to(weight.dtype), + zero_points_out, + ) + + def preprocess_olive_weights( state_dict: dict[str, torch.Tensor], bits: int = 4, diff --git a/src/mobius/_weight_utils_test.py b/src/mobius/_weight_utils_test.py index a4e4b9e9..9b6c7ad3 100644 --- a/src/mobius/_weight_utils_test.py +++ b/src/mobius/_weight_utils_test.py @@ -1085,3 +1085,92 @@ def test_indivisible_mp_num_raises(self): """mp_num that doesn't divide hidden raises ValueError.""" with pytest.raises(ValueError, match="divisible"): split_codegen_qkv(torch.zeros(96, 32), num_heads=4, head_dim=8, mp_num=3) + + +class TestQuantizeEmbeddingRTN: + """Tests for quantize_embedding_rtn.""" + + def setup_method(self): + from mobius._weight_utils import quantize_embedding_rtn + + self.quantize = quantize_embedding_rtn + + def test_int4_asymmetric_shapes(self): + """INT4 asymmetric: qweight is packed to half, scales/zp have correct shapes.""" + weight = torch.randn(128, 64, dtype=torch.float16) + qweight, scales, zp = self.quantize(weight, bits=4, block_size=32, symmetric=False) + assert qweight.shape == (128, 32) # 64 * 4 / 8 = 32 + assert qweight.dtype == torch.uint8 + assert scales.shape == (128, 2) # 64 / 32 = 2 blocks + assert scales.dtype == torch.float16 + assert zp is not None + assert zp.shape == (128, 1) # 2 blocks packed: ceil(2*4/8) = 1 + assert zp.dtype == torch.uint8 + + def test_int8_asymmetric_shapes(self): + """INT8 asymmetric: qweight is same shape (no packing), zp unpacked.""" + weight = torch.randn(128, 64, dtype=torch.float16) + qweight, scales, zp = self.quantize(weight, bits=8, block_size=32, symmetric=False) + assert qweight.shape == (128, 64) + assert qweight.dtype == torch.uint8 + assert scales.shape == (128, 2) + assert zp is not None + assert zp.shape == (128, 2) + assert zp.dtype == torch.uint8 + + def test_int4_symmetric_no_zero_points(self): + """Symmetric mode returns None for zero_points.""" + weight = torch.randn(64, 32, dtype=torch.float16) + qweight, scales, zp = self.quantize(weight, bits=4, block_size=32, symmetric=True) + assert zp is None + assert qweight.shape == (64, 16) + assert scales.shape == (64, 1) + + def test_int8_symmetric_no_zero_points(self): + """INT8 symmetric returns None for zero_points.""" + weight = torch.randn(64, 32, dtype=torch.float16) + qweight, scales, zp = self.quantize(weight, bits=8, block_size=32, symmetric=True) + assert zp is None + assert qweight.shape == (64, 32) + assert scales.shape == (64, 1) + + def test_invalid_bits_raises(self): + """bits not in (4, 8) raises ValueError.""" + weight = torch.randn(32, 32, dtype=torch.float16) + with pytest.raises(ValueError, match="bits must be 4 or 8"): + self.quantize(weight, bits=3) + + def test_non_2d_raises(self): + """Non-2D input raises ValueError.""" + weight = torch.randn(4, 8, 16, dtype=torch.float16) + with pytest.raises(ValueError, match="2D"): + self.quantize(weight, bits=4) + + def test_indivisible_dim_raises(self): + """embedding_dim not divisible by block_size raises ValueError.""" + weight = torch.randn(32, 30, dtype=torch.float16) + with pytest.raises(ValueError, match="divisible by block_size"): + self.quantize(weight, bits=4, block_size=32) + + def test_roundtrip_accuracy_int8(self): + """INT8 quantize-dequantize error is small.""" + torch.manual_seed(42) + weight = torch.randn(64, 64, dtype=torch.float16) + qweight, scales, zp = self.quantize(weight, bits=8, block_size=32, symmetric=False) + # Dequantize manually + import numpy as np + + q = qweight.numpy().astype(np.float32).reshape(64, 2, 32) + s = scales.numpy().astype(np.float32).reshape(64, 2, 1) + z = zp.numpy().astype(np.float32).reshape(64, 2, 1) + deq = (q - z) * s + deq = deq.reshape(64, 64) + error = np.abs(deq - weight.float().numpy()).mean() + assert error < 0.02, f"Mean abs error too large: {error}" + + def test_cpu_tensor_from_cuda_mock(self): + """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) + assert qweight.shape == (32, 16) diff --git a/src/mobius/models/gemma4.py b/src/mobius/models/gemma4.py index 265f8b0f..66899b36 100644 --- a/src/mobius/models/gemma4.py +++ b/src/mobius/models/gemma4.py @@ -200,6 +200,56 @@ def forward(self, op: OpBuilder, input_ids: ir.Value) -> ir.Value: return op.Mul(embeddings, self.embed_scale) +class QuantizedScaledWordEmbedding(QuantizedEmbedding): + """Quantized embedding with scaling — INT4 GatherBlockQuantized + scale multiply.""" + + def __init__( + self, + num_embeddings: int, + embedding_dim: int, + padding_idx: int, + embed_scale: float = 1.0, + bits: int = 4, + block_size: int = 32, + has_zero_point: bool = True, + ): + super().__init__( + num_embeddings, embedding_dim, bits, block_size, has_zero_point, padding_idx + ) + self.embed_scale = embed_scale + + def forward(self, op: OpBuilder, input_ids: ir.Value): + embeddings = super().forward(op, input_ids) + return op.Mul(embeddings, self.embed_scale) + + +def _per_layer_embedding_block_size(embedding_dim: int, group_size: int) -> int: + """Use one valid block for small test embeddings whose dim is below the real group.""" + if group_size > embedding_dim >= 16 and embedding_dim & (embedding_dim - 1) == 0: + return embedding_dim + 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: @@ -1800,17 +1850,36 @@ 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. - self.embed_tokens_per_layer_split = nn.ModuleList( - [ - Gemma3TextScaledWordEmbedding( - vocab_per_layer, - self._per_layer_dim, - config.pad_token_id, - embed_scale=float(self._per_layer_dim**0.5), - ) - for _ in range(self._num_layers) - ] - ) + 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=bits, + block_size=block_size, + has_zero_point=not symmetric, + ) + for _ in range(self._num_layers) + ] + ) + else: + self.embed_tokens_per_layer_split = nn.ModuleList( + [ + Gemma3TextScaledWordEmbedding( + vocab_per_layer, + self._per_layer_dim, + config.pad_token_id, + embed_scale=float(self._per_layer_dim**0.5), + ) + for _ in range(self._num_layers) + ] + ) self.per_layer_model_projection = Linear( config.hidden_size, config.num_hidden_layers * self._per_layer_dim, @@ -2348,8 +2417,28 @@ def preprocess_weights( f"got {fused.shape[1]}" ) chunks = fused.chunk(num_layers, dim=1) - for i, chunk in enumerate(chunks): - state_dict[f"model.embed_tokens_per_layer_split.{i}.weight"] = chunk + 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, group_size) + for i, chunk in enumerate(chunks): + qweight, scales, zero_points = quantize_embedding_rtn( + chunk.contiguous(), + bits=bits, + block_size=block_size, + 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 + if zero_points is not None: + state_dict[ + f"model.embed_tokens_per_layer_split.{i}.zero_points" + ] = zero_points + else: + for i, chunk in enumerate(chunks): + state_dict[f"model.embed_tokens_per_layer_split.{i}.weight"] = chunk return state_dict @@ -3149,8 +3238,34 @@ def preprocess_weights( f"got {fused.shape[1]}" ) chunks = fused.chunk(num_layers, dim=1) - for i, chunk in enumerate(chunks): - renamed[f"decoder.model.embed_tokens_per_layer_split.{i}.weight"] = chunk + 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, group_size) + for i, chunk in enumerate(chunks): + qweight, scales, zero_points = quantize_embedding_rtn( + chunk.contiguous(), + bits=bits, + block_size=block_size, + symmetric=symmetric, + ) + renamed[f"decoder.model.embed_tokens_per_layer_split.{i}.qweight"] = ( + qweight + ) + renamed[f"decoder.model.embed_tokens_per_layer_split.{i}.scales"] = ( + scales + ) + if zero_points is not None: + renamed[ + f"decoder.model.embed_tokens_per_layer_split.{i}.zero_points" + ] = zero_points + else: + for i, chunk in enumerate(chunks): + renamed[f"decoder.model.embed_tokens_per_layer_split.{i}.weight"] = ( + chunk + ) # Re-route the projection weights from embedding.* → decoder.model.* for k in list(renamed.keys()): if k.startswith( diff --git a/src/mobius/models/gemma4_test.py b/src/mobius/models/gemma4_test.py index 66e24e19..165fd19a 100644 --- a/src/mobius/models/gemma4_test.py +++ b/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( diff --git a/src/mobius/tasks/_gemma4.py b/src/mobius/tasks/_gemma4.py index a8233a79..a54c8502 100644 --- a/src/mobius/tasks/_gemma4.py +++ b/src/mobius/tasks/_gemma4.py @@ -439,6 +439,64 @@ def build( config.split_per_layer_embedding = fused_bytes > caps.max_buffer_size else: config.split_per_layer_embedding = False + # When splitting per-layer embeddings, quantize them to reduce size. + # Bits default to 4 (INT4) but can be overridden via embedding_bits + # in mobius.build() / the MobiusBuilder Olive pass config. + if config.split_per_layer_embedding: + qc = config.quantization + 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}" + ) + 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.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). + from mobius.models.gemma4 import ( + QuantizedScaledWordEmbedding, + _per_layer_embedding_block_size, + ) + + 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, 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=bits, + block_size=block_size, + has_zero_point=not symmetric, + ) + for _ in range(config.num_hidden_layers) + ] + ) + # Cast scales to model dtype (modules created after _cast_module_dtype) + from mobius._builder import _cast_module_dtype + + _cast_module_dtype(module.decoder.model.embed_tokens_per_layer_split, config.dtype) models: dict[str, ir.Model] = {} models["decoder"] = self._build_decoder(module.decoder, config) models["vision_encoder"] = self._build_vision(module.vision_encoder, config)