From 42b8c8667358b85e6304bdf4ca863158cbc683bc Mon Sep 17 00:00:00 2001 From: seongwoo Date: Thu, 6 Aug 2026 12:39:29 +0900 Subject: [PATCH 1/2] [quantization] Export gemma4 This commit supports gemma4 export. TICO-DCO-1.0-Signed-off-by: seongwoo --- .../recipes/test_export_gemma4.py | 309 +++++++++ tico/quantization/examples/README.md | 20 +- .../examples/configs/gemma4_export.yaml | 8 +- tico/quantization/recipes/adapters/gemma4.py | 21 + tico/quantization/recipes/export/gemma4.py | 603 ++++++++++++++++++ .../wrappers/gemma4/quant_vision_model.py | 26 +- .../gemma4/quant_vision_patch_embedder.py | 6 +- .../wrappers/gemma4/quant_vision_pooler.py | 34 +- 8 files changed, 997 insertions(+), 30 deletions(-) create mode 100644 test/quantization/recipes/test_export_gemma4.py create mode 100644 tico/quantization/recipes/export/gemma4.py diff --git a/test/quantization/recipes/test_export_gemma4.py b/test/quantization/recipes/test_export_gemma4.py new file mode 100644 index 000000000..404b46e30 --- /dev/null +++ b/test/quantization/recipes/test_export_gemma4.py @@ -0,0 +1,309 @@ +# Copyright (c) 2026 Samsung Electronics Co., Ltd. All Rights Reserved +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +try: + from quantization.recipes.optional_dependency_stubs import ( + install_optional_dependency_stubs, + ) +except ModuleNotFoundError: + from optional_dependency_stubs import install_optional_dependency_stubs + +install_optional_dependency_stubs() + +import tempfile +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import tico.quantization.recipes.adapters.gemma4 as gemma_adapter_mod +import tico.quantization.recipes.export.gemma4 as gemma_export + +import torch +from tico.quantization.recipes.adapters.gemma4 import Gemma4Adapter +from tico.quantization.recipes.context import RecipeContext + + +class FakePTQWrapper(torch.nn.Module): + """Minimal PTQWrapper-like container.""" + + def __init__(self, wrapped): + super().__init__() + self.wrapped = wrapped + + def forward(self, *args, **kwargs): + """Forward to the wrapped module.""" + return self.wrapped(*args, **kwargs) + + +class FakeVisionExport(torch.nn.Module): + """Return a model-output-like object for the vision stage.""" + + def forward(self, pixel_values, pixel_position_ids): + """Return four tiny visual tokens.""" + del pixel_values, pixel_position_ids + return SimpleNamespace(last_hidden_state=torch.zeros(4, 8)) + + +class FakeVision(torch.nn.Module): + """Expose the fixed Gemma4 vision geometry.""" + + def __init__(self): + super().__init__() + self.config = SimpleNamespace( + patch_size=2, + pooling_kernel_size=2, + default_output_length=4, + ) + + def as_export_module(self, mode, *, pixel_position_ids): + """Return the trivial static vision module.""" + del mode, pixel_position_ids + return FakeVisionExport() + + +class FakeAttention(torch.nn.Module): + """Expose one Gemma4 text-attention export contract.""" + + def __init__(self, *, layer_idx, is_sliding, is_shared): + super().__init__() + self.layer_idx = layer_idx + self.config = SimpleNamespace(num_attention_heads=2) + self.num_key_value_groups = 2 + self.head_dim = 4 + self.sliding_window = 2 if is_sliding else None + self.is_sliding = is_sliding + self.is_kv_shared_layer = is_shared + self.max_seq = 4 + + +class FakeDecoderLayer(torch.nn.Module): + """Expose prefill and decode export modules.""" + + def __init__(self, *, layer_idx, is_sliding, is_shared): + super().__init__() + self.self_attn = FakePTQWrapper( + FakeAttention( + layer_idx=layer_idx, + is_sliding=is_sliding, + is_shared=is_shared, + ) + ) + + def as_export_module(self, mode, *, return_kv=True): + """Return a placeholder module for the requested export mode.""" + del mode, return_kv + return torch.nn.Identity() + + +class FakeText(torch.nn.Module): + """Minimal Gemma4 text wrapper hierarchy.""" + + def __init__(self): + super().__init__() + self.config = SimpleNamespace( + hidden_size=8, + hidden_size_per_layer_input=2, + max_position_embeddings=4, + num_hidden_layers=2, + vocab_size=32, + enable_moe_block=False, + ) + self.embed_tokens = torch.nn.Embedding(32, 8) + self.layers = torch.nn.ModuleList( + [ + FakePTQWrapper( + FakeDecoderLayer( + layer_idx=0, + is_sliding=True, + is_shared=False, + ) + ), + FakePTQWrapper( + FakeDecoderLayer( + layer_idx=1, + is_sliding=False, + is_shared=True, + ) + ), + ] + ) + self.norm = torch.nn.Identity() + + +class FakeGemmaModel(torch.nn.Module): + """Minimal multimodal Gemma4 wrapper hierarchy.""" + + def __init__(self): + super().__init__() + self.vision_tower = FakePTQWrapper(FakeVision()) + self.language_model = FakePTQWrapper(FakeText()) + self.embed_vision = torch.nn.Identity() + self.visual_start_idx = 0 + self.num_visual_tokens = 4 + + +class FakeTopLevelGemma(torch.nn.Module): + """Minimal conditional-generation wrapper hierarchy.""" + + def __init__(self): + super().__init__() + self.model = FakePTQWrapper(FakeGemmaModel()) + self.lm_head = torch.nn.Linear(8, 32, bias=False) + + +class FakeExportModel(torch.nn.Module): + """Outer PTQWrapper-like model returned by prepare/convert.""" + + def __init__(self): + super().__init__() + self.wrapped = FakeTopLevelGemma() + + +def _model_args(): + """Return the fixed tiny vision contract used by exporter tests.""" + return { + "vision": { + "visual_start_idx": 0, + "num_visual_tokens": 4, + "max_soft_tokens": 4, + } + } + + +class TestGemma4PerLayerExport(unittest.TestCase): + def test_exports_all_static_runtime_stages(self): + """Gemma4 staged export should emit vision, prefill, and decode graphs.""" + calls = [] + export_model = FakeExportModel() + dynamic_shapes = {"input_ids": {1: "S"}} + + def fake_convert_and_save(module, example_inputs, save_path, **kwargs): + del module, example_inputs + calls.append((save_path.name, kwargs.get("dynamic_shapes"))) + + with tempfile.TemporaryDirectory() as tmpdir, patch.object( + gemma_export, + "_prepare_gemma4_export_model", + return_value=(export_model, "q"), + ), patch.object( + gemma_export, + "make_token_embedding_dynamic_shapes", + return_value=dynamic_shapes, + ), patch.object( + gemma_export, + "_convert_and_save", + fake_convert_and_save, + ): + gemma_export.export_gemma4_per_layer( + q_model=torch.nn.Identity(), + max_seq_len=4, + output_dir=tmpdir, + model_args=_model_args(), + prefill_decode=True, + ) + + self.assertEqual( + [name for name, _ in calls], + [ + "vision_prefill.q.circle", + "token_embedding.q.circle", + "multimodal_fusion_prefill.q.circle", + "decoder_layer_prefill_0.q.circle", + "decoder_layer_decode_0.q.circle", + "decoder_layer_prefill_1.q.circle", + "decoder_layer_decode_1.q.circle", + "lm_head.q.circle", + ], + ) + token_embedding_shapes = [ + shapes for name, shapes in calls if name == "token_embedding.q.circle" + ] + self.assertEqual(token_embedding_shapes, [dynamic_shapes]) + + def test_prefill_only_export_uses_unsuffixed_stage_names(self): + """Disabling decode export should omit all decode artifacts.""" + names = [] + export_model = FakeExportModel() + + def fake_convert_and_save(module, example_inputs, save_path, **kwargs): + del module, example_inputs, kwargs + names.append(save_path.name) + + with tempfile.TemporaryDirectory() as tmpdir, patch.object( + gemma_export, + "_prepare_gemma4_export_model", + return_value=(export_model, "f32"), + ), patch.object( + gemma_export, + "_convert_and_save", + fake_convert_and_save, + ): + gemma_export.export_gemma4_per_layer( + q_model=torch.nn.Identity(), + max_seq_len=4, + output_dir=tmpdir, + model_args=_model_args(), + prefill_decode=False, + ) + + self.assertEqual( + names, + [ + "vision_prefill.f32.circle", + "token_embedding.f32.circle", + "multimodal_fusion.f32.circle", + "decoder_layer_0.f32.circle", + "decoder_layer_1.f32.circle", + "lm_head.f32.circle", + ], + ) + + def test_adapter_routes_circle_per_layer_artifact(self): + """The Gemma4 adapter should dispatch the generic Circle artifact key.""" + model = torch.nn.Identity() + ctx = RecipeContext( + cfg={ + "calibration": {"seq_len": 2048}, + "model_args": _model_args(), + "export": { + "enabled": True, + "output_dir": "./out/gemma4", + "max_seq_len": 1024, + "prefill_decode": True, + "strict": True, + "artifacts": ["circle_per_layer"], + }, + }, + adapter=Gemma4Adapter(), + model=model, + ) + + with patch.object( + gemma_adapter_mod, + "export_gemma4_per_layer", + ) as export_per_layer: + Gemma4Adapter().export(ctx) + + export_per_layer.assert_called_once_with( + q_model=model, + max_seq_len=1024, + output_dir=gemma_adapter_mod.Path("./out/gemma4"), + model_args=_model_args(), + prefill_decode=True, + strict=True, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tico/quantization/examples/README.md b/tico/quantization/examples/README.md index 43f6cfb69..d4fec7175 100644 --- a/tico/quantization/examples/README.md +++ b/tico/quantization/examples/README.md @@ -99,9 +99,8 @@ Use this three-step loop for the common LLaMA post-quantization benchmark flow. ### Gemma4 E2B flow -Gemma4 E2B uses the same config-driven three-step interface. The current export -preset writes a serialized model checkpoint; static Circle submodule export is -not yet exposed through the Gemma4 recipe adapter. +Gemma4 E2B uses the same config-driven three-step interface. The export preset +writes static per-stage Circle files for the vision and text runtime. 1. Quantize Gemma4 E2B and save the checkpoint. @@ -119,7 +118,7 @@ not yet exposed through the Gemma4 recipe adapter. --checkpoint ./out/gemma4_quantized/quantized_model.pt ``` -3. Export the configured checkpoint artifact. +3. Export quantized per-stage Circle artifacts. ```bash python -m tico.quantization.examples.export \ @@ -127,6 +126,19 @@ not yet exposed through the Gemma4 recipe adapter. --checkpoint ./out/gemma4_quantized/quantized_model.pt ``` + Export the floating-point model with the same static runtime contract: + + ```bash + python -m tico.quantization.examples.export \ + --config tico/quantization/examples/configs/gemma4_export.yaml \ + --source model \ + --device cpu \ + --output-dir ./out/gemma4_float/circle_layers + ``` + + Floating-point files use the `.f32.circle` suffix. Quantized checkpoint + files use `.q.circle`. + ### Quantize `quantize.py` is the command that executes the recipe pipeline. It runs the diff --git a/tico/quantization/examples/configs/gemma4_export.yaml b/tico/quantization/examples/configs/gemma4_export.yaml index 64cbe06b8..bfcf9a1a1 100644 --- a/tico/quantization/examples/configs/gemma4_export.yaml +++ b/tico/quantization/examples/configs/gemma4_export.yaml @@ -19,6 +19,7 @@ model_args: vision: visual_start_idx: 0 num_visual_tokens: 256 + max_soft_tokens: 280 image_height: 896 image_width: 896 validate_static_layout: false @@ -31,6 +32,9 @@ evaluation: export: enabled: true - output_dir: ./out/gemma4_export + output_dir: ./out/gemma4_quantized/circle_layers + max_seq_len: 2048 + prefill_decode: true + strict: false artifacts: - - ptq_checkpoint + - circle_per_layer diff --git a/tico/quantization/recipes/adapters/gemma4.py b/tico/quantization/recipes/adapters/gemma4.py index 1e867ed3f..3b7b97315 100644 --- a/tico/quantization/recipes/adapters/gemma4.py +++ b/tico/quantization/recipes/adapters/gemma4.py @@ -42,6 +42,7 @@ print_vqa_results, ) from tico.quantization.recipes.export.checkpoint import save_checkpoint +from tico.quantization.recipes.export.gemma4 import export_gemma4_per_layer from tico.quantization.recipes.utils import ( move_to_device, quant_spec_from_config, @@ -476,3 +477,23 @@ def export(self, ctx: RecipeContext) -> None: artifacts = set(export_cfg.get("artifacts", [])) if "ptq_checkpoint" in artifacts or "checkpoint" in artifacts: save_checkpoint(ctx.require_model(), output_dir) + + if "circle_per_layer" in artifacts: + calibration_cfg = ctx.cfg.get("calibration", {}) + max_seq_len = int( + export_cfg.get( + "max_seq_len", + calibration_cfg.get("seq_len", 2048), + ) + ) + model_args = ctx.cfg.get("model_args", {}) + if not isinstance(model_args, Mapping): + raise TypeError("model_args must be a mapping for Gemma4 export.") + export_gemma4_per_layer( + q_model=ctx.require_model(), + max_seq_len=max_seq_len, + output_dir=output_dir, + model_args=model_args, + prefill_decode=bool(export_cfg.get("prefill_decode", True)), + strict=bool(export_cfg.get("strict", False)), + ) diff --git a/tico/quantization/recipes/export/gemma4.py b/tico/quantization/recipes/export/gemma4.py new file mode 100644 index 000000000..d9e692ad6 --- /dev/null +++ b/tico/quantization/recipes/export/gemma4.py @@ -0,0 +1,603 @@ +# Copyright (c) 2026 Samsung Electronics Co., Ltd. All Rights Reserved +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Static per-stage Circle export for Gemma4 E2B.""" + +from copy import deepcopy +from math import isqrt +from pathlib import Path +from typing import Any, Mapping + +import torch + +import tico +from tico.quantization.config.ptq import PTQConfig +from tico.quantization.wrapq.wrap_helper import PTQWrapHelper +from tico.quantization.wrapq.wrappers.gemma4.export_adapters import ( + Gemma4LMHeadExportAdapter, + Gemma4MMFusionExportAdapter, + Gemma4TokenEmbeddingExportAdapter, +) +from tico.quantization.wrapq.wrappers.gemma4.utils import assert_gemma4_e2b_no_moe +from tico.quantization.wrapq.wrappers.llama.export_adapters import ( + make_token_embedding_dynamic_shapes, +) +from tico.utils.utils import SuppressWarning + + +class _Gemma4VisionPrefillStage(torch.nn.Module): + """Run the static vision model and project soft tokens to text width.""" + + def __init__( + self, + vision_model: torch.nn.Module, + vision_projection: torch.nn.Module, + ) -> None: + super().__init__() + self.vision_model = vision_model + self.vision_projection = vision_projection + + def forward( + self, + pixel_values: torch.Tensor, + pixel_position_ids: torch.Tensor, + ) -> torch.Tensor: + """Return projected visual soft tokens for one static image profile.""" + vision_outputs = self.vision_model(pixel_values, pixel_position_ids) + return self.vision_projection(vision_outputs.last_hidden_state) + + +def _convert_and_save( + module: torch.nn.Module, + example_inputs: tuple[Any, ...], + save_path: Path, + *, + kwargs: dict[str, Any] | None = None, + dynamic_shapes: Any | None = None, + strict: bool = False, +) -> None: + """Convert one Gemma4 export stage to Circle and save it.""" + print(f"Saving {save_path.name} to {save_path.resolve()}") + with torch.no_grad(), SuppressWarning(UserWarning, ".*"): + circle_model = tico.convert( + module.eval(), + example_inputs, + kwargs=kwargs, + dynamic_shapes=dynamic_shapes, + strict=strict, + ) + circle_model.save(save_path) + + +def _is_wrapped_export_model(model: torch.nn.Module) -> bool: + """Return whether a model already exposes the PTQ wrapper export layout.""" + wrapped = getattr(model, "wrapped", None) + return ( + wrapped is not None + and hasattr(wrapped, "model") + and hasattr(wrapped, "lm_head") + ) + + +def _float_artifact_tag(model: torch.nn.Module) -> str: + """Validate a floating-point export model and return its precision tag.""" + try: + dtype = next(model.parameters()).dtype + except StopIteration: + dtype = torch.float32 + + if dtype is not torch.float32: + raise TypeError( + "Floating-point Gemma4 export currently supports float32 only. " + f"Got parameter dtype {dtype}." + ) + return "f32" + + +def _normalize_model_args( + model_args: Mapping[str, Any] | None, + *, + max_seq_len: int, +) -> dict[str, Any]: + """Normalize the fixed Gemma4 runtime contract used by export wrappers.""" + normalized = deepcopy(dict(model_args or {})) + + vision = normalized.setdefault("vision", {}) + if not isinstance(vision, dict): + raise TypeError("model_args.vision must be a mapping.") + + if "visual_start_idx" not in vision: + raise ValueError( + "Gemma4 Circle export requires model_args.vision.visual_start_idx." + ) + vision["visual_start_idx"] = int(vision["visual_start_idx"]) + if vision["visual_start_idx"] < 0: + raise ValueError("model_args.vision.visual_start_idx must be non-negative.") + + if "num_visual_tokens" not in vision: + raise ValueError( + "Gemma4 Circle export requires model_args.vision.num_visual_tokens." + ) + vision["num_visual_tokens"] = int(vision["num_visual_tokens"]) + if vision["num_visual_tokens"] <= 0: + raise ValueError("model_args.vision.num_visual_tokens must be positive.") + + if "max_soft_tokens" in vision: + vision["max_soft_tokens"] = int(vision["max_soft_tokens"]) + if vision["max_soft_tokens"] <= 0: + raise ValueError("model_args.vision.max_soft_tokens must be positive.") + + text = normalized.setdefault("text", {}) + if not isinstance(text, dict): + raise TypeError("model_args.text must be a mapping.") + configured_max_seq = text.get("max_seq") + if configured_max_seq is not None and int(configured_max_seq) != max_seq_len: + raise ValueError( + "model_args.text.max_seq must match export.max_seq_len: " + f"text.max_seq={int(configured_max_seq)}, " + f"export.max_seq_len={max_seq_len}." + ) + text["max_seq"] = int(max_seq_len) + + return normalized + + +def _prepare_gemma4_export_model( + model: torch.nn.Module, + model_args: Mapping[str, Any], +) -> tuple[torch.nn.Module, str]: + """Normalize a checkpoint or FP model for staged Gemma4 export. + + Floating-point models are structurally wrapped in ``NO_QUANT`` mode. No + calibration or fake quantization is introduced. Converted checkpoints keep + their existing quantization state. + """ + model = model.eval().cpu() + if _is_wrapped_export_model(model): + return model, "q" + + artifact_tag = _float_artifact_tag(model) + wrapper_config = PTQConfig( + model_args=deepcopy(dict(model_args)), + strict_wrap=True, + ) + export_model = PTQWrapHelper(strict_wrap=True).wrap_supported( + model, + wrapper_config, + ) + if not _is_wrapped_export_model(export_model): + raise TypeError( + "Gemma4 staged export requires a top-level PTQ wrapper exposing " + "the wrapped model and LM head." + ) + return export_model, artifact_tag + + +def _circle_name(stem: str, artifact_tag: str) -> str: + """Build a Circle artifact name with an explicit precision tag.""" + return f"{stem}.{artifact_tag}.circle" + + +def _unwrap_gemma4_components( + export_model: torch.nn.Module, +) -> tuple[torch.nn.Module, torch.nn.Module, torch.nn.Module, torch.nn.Module]: + """Return top-level, multimodal, vision, and text Gemma4 wrappers.""" + qmodel = export_model.wrapped + gemma_model = qmodel.model.wrapped + if getattr(gemma_model, "vision_tower", None) is None: + raise ValueError("Gemma4 Circle export requires an image vision tower.") + qvision = gemma_model.vision_tower.wrapped + qtext = gemma_model.language_model.wrapped + return qmodel, gemma_model, qvision, qtext + + +def _resolve_vision_contract( + *, + gemma_model: torch.nn.Module, + qvision: torch.nn.Module, + model_args: Mapping[str, Any], + max_seq_len: int, +) -> tuple[int, int, int, int]: + """Validate the fixed padded-patch and visual-token layout.""" + vision_args = model_args["vision"] + visual_start_idx = int(vision_args["visual_start_idx"]) + num_visual_tokens = int(vision_args["num_visual_tokens"]) + + wrapped_start_idx = int(getattr(gemma_model, "visual_start_idx")) + if visual_start_idx != wrapped_start_idx: + raise ValueError( + "Configured visual_start_idx does not match the wrapped checkpoint: " + f"configured={visual_start_idx}, wrapped={wrapped_start_idx}." + ) + + wrapped_visual_tokens = int(getattr(gemma_model, "num_visual_tokens")) + if num_visual_tokens != wrapped_visual_tokens: + raise ValueError( + "Configured num_visual_tokens does not match the wrapped checkpoint: " + f"configured={num_visual_tokens}, wrapped={wrapped_visual_tokens}." + ) + + if visual_start_idx + num_visual_tokens > max_seq_len: + raise ValueError( + "The fixed visual-token span exceeds max_seq_len: " + f"start={visual_start_idx}, visual_tokens={num_visual_tokens}, " + f"max_seq_len={max_seq_len}." + ) + + vision_config = qvision.config + pooling_kernel_size = int(vision_config.pooling_kernel_size) + if pooling_kernel_size <= 0: + raise ValueError( + "Gemma4 vision pooling_kernel_size must be positive, got " + f"{pooling_kernel_size}." + ) + + max_soft_tokens = int( + vision_args.get( + "max_soft_tokens", + getattr(vision_config, "default_output_length", num_visual_tokens), + ) + ) + if max_soft_tokens < num_visual_tokens: + raise ValueError( + "model_args.vision.max_soft_tokens must be greater than or equal " + f"to num_visual_tokens, got {max_soft_tokens} < {num_visual_tokens}." + ) + + visual_side = isqrt(num_visual_tokens) + if visual_side * visual_side != num_visual_tokens: + raise ValueError( + "Gemma4 num_visual_tokens must form a square visual grid, got " + f"{num_visual_tokens}." + ) + + num_valid_patches = num_visual_tokens * pooling_kernel_size**2 + num_patches = max_soft_tokens * pooling_kernel_size**2 + return visual_start_idx, num_visual_tokens, num_valid_patches, num_patches + + +def _make_vision_inputs( + *, + qvision: torch.nn.Module, + num_visual_tokens: int, + num_valid_patches: int, + num_patches: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Create processor-style padded patch values and 2-D position IDs.""" + pooling_kernel_size = int(qvision.config.pooling_kernel_size) + visual_side = isqrt(num_visual_tokens) + patch_grid_side = visual_side * pooling_kernel_size + + coords = torch.arange(num_valid_patches, dtype=torch.long) + valid_positions = torch.stack( + (coords % patch_grid_side, coords // patch_grid_side), + dim=-1, + ) + padding_positions = torch.full( + (num_patches - num_valid_patches, 2), + -1, + dtype=torch.long, + ) + pixel_position_ids = torch.cat( + (valid_positions, padding_positions), + dim=0, + ).unsqueeze(0) + + patch_size = int(qvision.config.patch_size) + patch_vector_size = 3 * patch_size * patch_size + pixel_values = torch.rand(1, num_patches, patch_vector_size, device="cpu") + if num_patches > num_valid_patches: + pixel_values[:, num_valid_patches:, :] = 0.0 + return pixel_values, pixel_position_ids + + +def _make_prefill_attention_mask( + *, + max_seq_len: int, + sliding_window: int | None, +) -> torch.Tensor: + """Create a full or sliding additive causal mask for tracing.""" + query_positions = torch.arange(max_seq_len).unsqueeze(1) + key_positions = torch.arange(max_seq_len).unsqueeze(0) + allowed = key_positions <= query_positions + if sliding_window is not None: + allowed = allowed & (key_positions > query_positions - sliding_window) + + mask = torch.full((max_seq_len, max_seq_len), -120.0) + mask.masked_fill_(allowed, 0.0) + return mask.unsqueeze(0).unsqueeze(0) + + +def _make_decode_attention_mask( + *, + max_seq_len: int, + sliding_window: int | None, +) -> torch.Tensor: + """Create an additive single-token decode mask at maximum cache length.""" + mask = torch.full((1, 1, max_seq_len), -120.0) + start = 0 if sliding_window is None else max(0, max_seq_len - sliding_window) + mask[..., start:max_seq_len] = 0.0 + return mask + + +def _make_position_embeddings( + *, + seq_len: int, + head_dim: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Create example Gemma4 RoPE cosine and sine tensors.""" + return ( + torch.randn(1, seq_len, head_dim, device="cpu"), + torch.randn(1, seq_len, head_dim, device="cpu"), + ) + + +def _attention_contract( + layer: torch.nn.Module, + *, + max_seq_len: int, +) -> tuple[torch.nn.Module, int, int, int | None, bool]: + """Return attention dimensions and sharing mode for one text layer.""" + attention = layer.wrapped.self_attn.wrapped + attention_capacity = int(getattr(attention, "max_seq", max_seq_len)) + if max_seq_len > attention_capacity: + raise ValueError( + "max_seq_len exceeds the wrapped Gemma4 attention capacity: " + f"layer={int(getattr(attention, 'layer_idx', -1))}, " + f"max_seq_len={max_seq_len}, capacity={attention_capacity}." + ) + + num_heads = int(attention.config.num_attention_heads) + num_kv_groups = int(attention.num_key_value_groups) + if num_kv_groups <= 0 or num_heads % num_kv_groups: + raise ValueError( + "Invalid Gemma4 grouped-query attention contract: " + f"num_heads={num_heads}, num_key_value_groups={num_kv_groups}." + ) + num_kv_heads = num_heads // num_kv_groups + head_dim = int(attention.head_dim) + sliding_window = ( + int(attention.sliding_window) + if bool(getattr(attention, "is_sliding", False)) + else None + ) + is_shared = bool(getattr(attention, "is_kv_shared_layer", False)) + return attention, num_kv_heads, head_dim, sliding_window, is_shared + + +def export_gemma4_per_layer( + *, + q_model: torch.nn.Module, + max_seq_len: int, + output_dir: str | Path, + model_args: Mapping[str, Any], + prefill_decode: bool = True, + strict: bool = False, +) -> None: + """Export a floating-point or PTQ-wrapped Gemma4 E2B by runtime stage. + + The generated Circle set contains the image vision prefill stage, dynamic + token embedding, fixed-slot multimodal fusion, every text decoder layer, + and final norm/LM head. With ``prefill_decode=True``, each text layer is + emitted once for full prefill and once for single-token decode. + + PLE token lookup and the packed context projection intentionally remain CPU + runtime responsibilities. Each decoder Circle receives only its sliced + ``per_layer_input`` tensor, matching ``StaticGemma4Runtime``. + """ + if max_seq_len < 1: + raise ValueError(f"max_seq_len must be positive, got {max_seq_len}.") + + normalized_model_args = _normalize_model_args( + model_args, + max_seq_len=max_seq_len, + ) + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + export_model, artifact_tag = _prepare_gemma4_export_model( + q_model, + normalized_model_args, + ) + qmodel, gemma_model, qvision, qtext = _unwrap_gemma4_components(export_model) + assert_gemma4_e2b_no_moe(qtext.config) + + text_capacity = int(qtext.config.max_position_embeddings) + if max_seq_len > text_capacity: + raise ValueError( + "max_seq_len exceeds the wrapped Gemma4 text capacity: " + f"max_seq_len={max_seq_len}, capacity={text_capacity}." + ) + + ( + visual_start_idx, + num_visual_tokens, + num_valid_patches, + num_patches, + ) = _resolve_vision_contract( + gemma_model=gemma_model, + qvision=qvision, + model_args=normalized_model_args, + max_seq_len=max_seq_len, + ) + + config = qtext.config + hidden_size = int(config.hidden_size) + vocab_size = int(config.vocab_size) + ple_dim = int(getattr(config, "hidden_size_per_layer_input", 0) or 0) + + pixel_values, pixel_position_ids = _make_vision_inputs( + qvision=qvision, + num_visual_tokens=num_visual_tokens, + num_valid_patches=num_valid_patches, + num_patches=num_patches, + ) + vision_model = qvision.as_export_module( + mode="prefill", + pixel_position_ids=pixel_position_ids, + ) + _convert_and_save( + _Gemma4VisionPrefillStage(vision_model, gemma_model.embed_vision), + (pixel_values, pixel_position_ids), + output_dir / _circle_name("vision_prefill", artifact_tag), + strict=strict, + ) + + token_input_ids = torch.randint( + low=0, + high=vocab_size, + size=(1, max_seq_len), + dtype=torch.long, + device="cpu", + ) + _convert_and_save( + Gemma4TokenEmbeddingExportAdapter(qtext), + (token_input_ids,), + output_dir / _circle_name("token_embedding", artifact_tag), + dynamic_shapes=make_token_embedding_dynamic_shapes(max_seq_len), + strict=strict, + ) + + text_embeds = torch.randn(1, max_seq_len, hidden_size, device="cpu") + visual_embeds = torch.randn(num_visual_tokens, hidden_size, device="cpu") + fusion_name = "multimodal_fusion_prefill" if prefill_decode else "multimodal_fusion" + _convert_and_save( + Gemma4MMFusionExportAdapter( + visual_start_idx=visual_start_idx, + num_visual_tokens=num_visual_tokens, + ), + (text_embeds, visual_embeds), + output_dir / _circle_name(fusion_name, artifact_tag), + strict=strict, + ) + + prefill_hidden = torch.randn(1, max_seq_len, hidden_size, device="cpu") + decode_hidden = torch.randn(1, 1, hidden_size, device="cpu") + + for layer_idx, layer in enumerate(qtext.layers): + ( + _attention, + num_kv_heads, + head_dim, + sliding_window, + is_shared, + ) = _attention_contract(layer, max_seq_len=max_seq_len) + + prefill_kwargs: dict[str, Any] = { + "attention_mask": _make_prefill_attention_mask( + max_seq_len=max_seq_len, + sliding_window=sliding_window, + ), + "position_embeddings": _make_position_embeddings( + seq_len=max_seq_len, + head_dim=head_dim, + ), + } + if ple_dim: + prefill_kwargs["per_layer_input"] = torch.randn( + 1, + max_seq_len, + ple_dim, + device="cpu", + ) + if is_shared: + shared_key = torch.randn( + 1, + num_kv_heads, + max_seq_len, + head_dim, + device="cpu", + ) + prefill_kwargs["shared_key_value"] = ( + shared_key, + torch.randn_like(shared_key), + ) + + prefill_stem = ( + f"decoder_layer_prefill_{layer_idx}" + if prefill_decode + else f"decoder_layer_{layer_idx}" + ) + _convert_and_save( + layer.wrapped.as_export_module( + "prefill", + return_kv=prefill_decode, + ), + (prefill_hidden,), + output_dir / _circle_name(prefill_stem, artifact_tag), + kwargs=prefill_kwargs, + strict=strict, + ) + + if not prefill_decode: + continue + + decode_kwargs: dict[str, Any] = { + "attention_mask": _make_decode_attention_mask( + max_seq_len=max_seq_len, + sliding_window=sliding_window, + ), + "position_embeddings": _make_position_embeddings( + seq_len=1, + head_dim=head_dim, + ), + } + if ple_dim: + decode_kwargs["per_layer_input"] = torch.randn( + 1, + 1, + ple_dim, + device="cpu", + ) + if is_shared: + shared_key = torch.randn( + 1, + num_kv_heads, + max_seq_len, + head_dim, + device="cpu", + ) + decode_kwargs["shared_key_value"] = ( + shared_key, + torch.randn_like(shared_key), + ) + else: + past_key = torch.randn( + 1, + num_kv_heads, + max_seq_len - 1, + head_dim, + device="cpu", + ) + decode_kwargs["past_key_value"] = ( + past_key, + torch.randn_like(past_key), + ) + + _convert_and_save( + layer.wrapped.as_export_module("decode", return_kv=True), + (decode_hidden,), + output_dir + / _circle_name(f"decoder_layer_decode_{layer_idx}", artifact_tag), + kwargs=decode_kwargs, + strict=strict, + ) + + lm_head_hidden = torch.randn(1, 1, hidden_size, device="cpu") + _convert_and_save( + Gemma4LMHeadExportAdapter(qmodel), + (lm_head_hidden,), + output_dir / _circle_name("lm_head", artifact_tag), + strict=strict, + ) diff --git a/tico/quantization/wrapq/wrappers/gemma4/quant_vision_model.py b/tico/quantization/wrapq/wrappers/gemma4/quant_vision_model.py index 891952102..1149b0a0f 100644 --- a/tico/quantization/wrapq/wrappers/gemma4/quant_vision_model.py +++ b/tico/quantization/wrapq/wrappers/gemma4/quant_vision_model.py @@ -269,8 +269,8 @@ def forward_export( if self.config.standardize: assert self.obs_std_bias is not None assert self.obs_std_scale is not None - std_bias = self.obs_std_bias.fake_quant(self.std_bias) - std_scale = self.obs_std_scale.fake_quant(self.std_scale) + std_bias = self._fq(self.std_bias, self.obs_std_bias) + std_scale = self._fq(self.std_scale, self.obs_std_scale) hidden_states = hidden_states - std_bias.float() hidden_states = self._fq(hidden_states, self.obs_minus_bias) hidden_states = hidden_states * std_scale.float() @@ -293,7 +293,7 @@ def as_export_module( """Prepare the model for torch.export by precomputing static tensors. This method: - 1. Asserts that the model is in QUANT mode + 1. Requires either NO_QUANT or QUANT mode 2. Recursively converts submodules to their export adapters 3. Registers output_length and padding tensors for static export @@ -312,12 +312,19 @@ def as_export_module( Returns: Gemma4VisionModelPrefillExportAdapter wrapping this module. """ - # Assert QUANT mode - assert self._mode is Mode.QUANT, "Must be in QUANT mode for export" + if mode != "prefill": + raise ValueError(f"Unsupported Gemma4 VisionModel export mode: {mode!r}") - # Make sure that all observers are calibrated - for obs in self._all_observers(): - assert obs.has_qparams, f"Observer {obs.name} has not been calibrated" + if self._mode not in (Mode.NO_QUANT, Mode.QUANT): + raise RuntimeError( + "Gemma4 VisionModel export requires NO_QUANT or QUANT mode, " + f"got {self._mode}." + ) + + if self._mode is Mode.QUANT: + # Make sure that all observers are calibrated. + for obs in self._all_observers(): + assert obs.has_qparams, f"Observer {obs.name} has not been calibrated" # Store output_length for use in forward_export pooling_kernel_size = self.config.pooling_kernel_size @@ -351,7 +358,8 @@ def as_export_module( padding_positions = (pixel_position_ids == -1).all(dim=-1) self.register_buffer("padding_positions", padding_positions) - register_fake_quant_meta_kernels_for_dynamic_export() + if self._mode is Mode.QUANT: + register_fake_quant_meta_kernels_for_dynamic_export() from tico.quantization.wrapq.wrappers.gemma4.export_adapters import ( Gemma4VisionModelPrefillExportAdapter, diff --git a/tico/quantization/wrapq/wrappers/gemma4/quant_vision_patch_embedder.py b/tico/quantization/wrapq/wrappers/gemma4/quant_vision_patch_embedder.py index 685b4a578..b9ce58c1e 100644 --- a/tico/quantization/wrapq/wrappers/gemma4/quant_vision_patch_embedder.py +++ b/tico/quantization/wrapq/wrappers/gemma4/quant_vision_patch_embedder.py @@ -183,7 +183,11 @@ def _all_observers(self) -> Iterable: def as_export_module(self, mode: str = "prefill", **kwargs) -> nn.Module: """Return self for export (this wrapper is already exportable).""" - assert self._mode is Mode.QUANT, "Must be in QUANT mode for export" + if self._mode not in (Mode.NO_QUANT, Mode.QUANT): + raise RuntimeError( + "Gemma4 VisionPatchEmbedder export requires NO_QUANT or " + f"QUANT mode, got {self._mode}." + ) if mode != "prefill": raise ValueError( f"Unsupported Gemma4 VisionPatchEmbedder export mode: {mode!r}" diff --git a/tico/quantization/wrapq/wrappers/gemma4/quant_vision_pooler.py b/tico/quantization/wrapq/wrappers/gemma4/quant_vision_pooler.py index 3b89f4645..6a72e2de7 100644 --- a/tico/quantization/wrapq/wrappers/gemma4/quant_vision_pooler.py +++ b/tico/quantization/wrapq/wrappers/gemma4/quant_vision_pooler.py @@ -230,7 +230,7 @@ def forward_export( hidden_states = self._fq(hidden_states, self.obs_pool_in) # Step 3: Fake-quantize pool_weights (weight quantization for matmul). - pool_weights_q = self.obs_pool_weight.fake_quant(self.pool_weights) + pool_weights_q = self._fq(self.pool_weights, self.obs_pool_weight) # Step 4: Spatial pooling via precomputed weight matrix. # pool_weights_q: (1, V, S), hidden_states.float(): (1, S, D) @@ -240,7 +240,7 @@ def forward_export( pooled = self._fq(pooled, self.obs_pool_matmul_out) # Step 6: Scale by sqrt(hidden_size) in float32. - root_hidden_size = self.obs_root_hidden_size.fake_quant(self.root_hidden_size) + root_hidden_size = self._fq(self.root_hidden_size, self.obs_root_hidden_size) pooled = pooled * root_hidden_size # Step 7: Fake-quantize final output (collects stats in CALIB, applies Q-DQ in QUANT). @@ -265,12 +265,17 @@ def as_export_module( Gemma4VisionPoolerPrefillExportAdapter, ) - assert self._mode is Mode.QUANT + if self._mode not in (Mode.NO_QUANT, Mode.QUANT): + raise RuntimeError( + "Gemma4 VisionPooler export requires NO_QUANT or QUANT mode, " + f"got {self._mode}." + ) - # Make sure that all observers are calibrated - for obs in self._all_observers(): - if isinstance(obs, AffineObserverBase): - assert obs.has_qparams + if self._mode is Mode.QUANT: + # Make sure that all observers are calibrated. + for obs in self._all_observers(): + if isinstance(obs, AffineObserverBase): + assert obs.has_qparams # Precompute static tensors weights, mask = self._build_pool_weights( @@ -281,12 +286,13 @@ def as_export_module( self.register_buffer("pool_weights", weights) self.register_buffer("pool_mask", mask) - # Collect statistics about pool_weights and compute qparams - obs_pool_weight_enabled: bool = self.obs_pool_weight.enabled - self.obs_pool_weight.enabled = True - self.obs_pool_weight.reset() - self.obs_pool_weight.collect(self.pool_weights) - self.obs_pool_weight.compute_qparams() - self.obs_pool_weight.enabled = obs_pool_weight_enabled + if self._mode is Mode.QUANT: + # Collect statistics about generated pool weights and compute qparams. + obs_pool_weight_enabled: bool = self.obs_pool_weight.enabled + self.obs_pool_weight.enabled = True + self.obs_pool_weight.reset() + self.obs_pool_weight.collect(self.pool_weights) + self.obs_pool_weight.compute_qparams() + self.obs_pool_weight.enabled = obs_pool_weight_enabled return Gemma4VisionPoolerPrefillExportAdapter(self) From 68ca2703cfe8585de22ec09ac0a355ade4951eb3 Mon Sep 17 00:00:00 2001 From: seongwoo Date: Thu, 6 Aug 2026 13:09:51 +0900 Subject: [PATCH 2/2] test fix. --- .../recipes/test_gemma4_example_configs.py | 8 +++-- .../gemma4/test_quant_vision_model.py | 34 ++++++++++++++++--- .../test_quant_vision_patch_embedder.py | 19 ++++++++--- .../gemma4/test_quant_vision_pooler.py | 26 +++++++++++--- 4 files changed, 71 insertions(+), 16 deletions(-) diff --git a/test/quantization/recipes/test_gemma4_example_configs.py b/test/quantization/recipes/test_gemma4_example_configs.py index ed1603222..33c10a29d 100644 --- a/test/quantization/recipes/test_gemma4_example_configs.py +++ b/test/quantization/recipes/test_gemma4_example_configs.py @@ -86,14 +86,16 @@ def test_eval_config_does_not_run_quantization_stages(self): ) self.assertEqual(cfg["evaluation"]["n_samples"], 1000) - def test_export_config_only_requests_the_checkpoint_artifact(self): - """The export preset should avoid calibration and pipeline execution.""" + def test_export_config_requests_circle_per_layer_artifact(self): + """The export preset should directly request static Circle artifacts.""" cfg = self._load("gemma4_export.yaml") self.assertEqual(cfg["pipeline"], []) self.assertFalse(cfg["evaluation"]["enabled"]) self.assertTrue(cfg["export"]["enabled"]) - self.assertEqual(cfg["export"]["artifacts"], ["ptq_checkpoint"]) + self.assertEqual(cfg["export"]["artifacts"], ["circle_per_layer"]) + self.assertEqual(cfg["export"]["max_seq_len"], 2048) + self.assertTrue(cfg["export"]["prefill_decode"]) if __name__ == "__main__": diff --git a/test/quantization/wrapq/wrappers/gemma4/test_quant_vision_model.py b/test/quantization/wrapq/wrappers/gemma4/test_quant_vision_model.py index c016f0953..cbc64f1a7 100644 --- a/test/quantization/wrapq/wrappers/gemma4/test_quant_vision_model.py +++ b/test/quantization/wrapq/wrappers/gemma4/test_quant_vision_model.py @@ -249,18 +249,42 @@ def test_standardize_false_no_buffers(self): self.assertFalse(hasattr(q_model, "std_bias")) self.assertFalse(hasattr(q_model, "std_scale")) - def test_as_export_module_requires_quant_mode(self): - """as_export_module should assert that mode is QUANT.""" + def test_as_export_module_supports_no_quant_mode(self): + """Floating-point export should be available in NO_QUANT mode.""" + from tico.quantization.wrapq.wrappers.gemma4.export_adapters import ( + Gemma4VisionModelPrefillExportAdapter, + ) from tico.quantization.wrapq.wrappers.gemma4.quant_vision_model import ( QuantGemma4VisionModel, ) fp_model = self._make_vision_model() q_model = QuantGemma4VisionModel(fp_model).eval() + sample = self._sample_inputs() + + export_module = q_model.as_export_module( + mode="prefill", + pixel_position_ids=sample["pixel_position_ids"], + ) + + self.assertIsInstance(export_module, Gemma4VisionModelPrefillExportAdapter) + + def test_as_export_module_rejects_calibration_mode(self): + """Export should reject CALIB mode because its qparams are incomplete.""" + from tico.quantization.wrapq.wrappers.gemma4.quant_vision_model import ( + QuantGemma4VisionModel, + ) + + fp_model = self._make_vision_model() + q_model = QuantGemma4VisionModel(fp_model).eval() + q_model.enable_calibration() + sample = self._sample_inputs() - # Should fail in NO_QUANT mode - with self.assertRaises(AssertionError): - q_model.as_export_module(mode="prefill", pixel_position_ids=None) + with self.assertRaisesRegex(RuntimeError, "NO_QUANT or QUANT"): + q_model.as_export_module( + mode="prefill", + pixel_position_ids=sample["pixel_position_ids"], + ) def test_as_export_module_requires_standardize(self): """as_export_module should assert that config.standardize is True.""" diff --git a/test/quantization/wrapq/wrappers/gemma4/test_quant_vision_patch_embedder.py b/test/quantization/wrapq/wrappers/gemma4/test_quant_vision_patch_embedder.py index da122262d..94cd7637a 100644 --- a/test/quantization/wrapq/wrappers/gemma4/test_quant_vision_patch_embedder.py +++ b/test/quantization/wrapq/wrappers/gemma4/test_quant_vision_patch_embedder.py @@ -281,8 +281,8 @@ def test_config_attributes_are_stored(self): # as_export_module # ------------------------------------------------------------------ - def test_as_export_module_requires_quant_mode(self): - """as_export_module should assert that mode is QUANT.""" + def test_as_export_module_supports_no_quant_mode(self): + """Floating-point export should be available in NO_QUANT mode.""" fp_module = _make_patch_embedder( hidden_size=self.hidden_size, patch_size=self.patch_size, @@ -290,8 +290,19 @@ def test_as_export_module_requires_quant_mode(self): ) q_module = QuantGemma4VisionPatchEmbedder(fp_module).eval() - # Should fail in NO_QUANT mode - with self.assertRaises(AssertionError): + self.assertIs(q_module.as_export_module(mode="prefill"), q_module) + + def test_as_export_module_rejects_calibration_mode(self): + """Export should reject CALIB mode because its qparams are incomplete.""" + fp_module = _make_patch_embedder( + hidden_size=self.hidden_size, + patch_size=self.patch_size, + position_embedding_size=self.position_embedding_size, + ) + q_module = QuantGemma4VisionPatchEmbedder(fp_module).eval() + q_module.enable_calibration() + + with self.assertRaisesRegex(RuntimeError, "NO_QUANT or QUANT"): q_module.as_export_module(mode="prefill") def test_as_export_module_returns_self(self): diff --git a/test/quantization/wrapq/wrappers/gemma4/test_quant_vision_pooler.py b/test/quantization/wrapq/wrappers/gemma4/test_quant_vision_pooler.py index ef0d1c365..3a8e45728 100644 --- a/test/quantization/wrapq/wrappers/gemma4/test_quant_vision_pooler.py +++ b/test/quantization/wrapq/wrappers/gemma4/test_quant_vision_pooler.py @@ -354,15 +354,33 @@ def test_forward_export_matches_forward(self): # as_export_module # ------------------------------------------------------------------ - def test_as_export_module_requires_quant_mode(self): - """as_export_module should assert that mode is QUANT.""" + def test_as_export_module_supports_no_quant_mode(self): + """Floating-point export should be available in NO_QUANT mode.""" + from tico.quantization.wrapq.wrappers.gemma4.export_adapters import ( + Gemma4VisionPoolerPrefillExportAdapter, + ) + fp_pooler = _make_pooler() q_pooler = QuantGemma4VisionPooler(fp_pooler).eval() + pixel_pos_ids = _pixel_position_ids(self.batch_size, self.seq_len) + export_module = q_pooler.as_export_module( + output_length=self.output_length, + pixel_position_ids=pixel_pos_ids, + ) + + self.assertIsInstance(export_module, Gemma4VisionPoolerPrefillExportAdapter) + self.assertTrue(hasattr(q_pooler, "pool_weights")) + self.assertTrue(hasattr(q_pooler, "pool_mask")) + + def test_as_export_module_rejects_calibration_mode(self): + """Export should reject CALIB mode because its qparams are incomplete.""" + fp_pooler = _make_pooler() + q_pooler = QuantGemma4VisionPooler(fp_pooler).eval() + q_pooler.enable_calibration() pixel_pos_ids = _pixel_position_ids(self.batch_size, self.seq_len) - # Should fail in NO_QUANT mode - with self.assertRaises(AssertionError): + with self.assertRaisesRegex(RuntimeError, "NO_QUANT or QUANT"): q_pooler.as_export_module( output_length=self.output_length, pixel_position_ids=pixel_pos_ids,