tico 0.2.0.dev260708__tar.gz → 0.2.0.dev260709__tar.gz

{tico-0.2.0.dev260708 → tico-0.2.0.dev260709}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tico
-Version: 0.2.0.dev260708
+Version: 0.2.0.dev260709
 Summary: Convert Exported Torch Module To Circle
 License: This file provides full text of licenses used in this project
         

tico-0.2.0.dev260709/tico/_version.py

@@ -0,0 +1 @@
+__version__ = "0.2.0.dev260709"

{tico-0.2.0.dev260708 → tico-0.2.0.dev260709}/tico/quantization/recipes/adapters/gemma4.py

@@ -12,7 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Mapping, Sequence
+from contextlib import contextmanager
+from typing import Any, Iterator, Mapping, Sequence
 
 import torch
 import tqdm
@@ -111,19 +112,95 @@ class Gemma4Adapter(ModelAdapter):
         if text_config is not None and hasattr(text_config, "use_cache"):
             text_config.use_cache = False
 
+    @staticmethod
+    def _static_calibration_image_size(
+        cfg: Mapping[str, Any],
+    ) -> tuple[int, int] | None:
+        """Return the configured static image size for calibration samples."""
+        model_args = cfg.get("model_args", {})
+        if not isinstance(model_args, Mapping):
+            return None
+
+        vision_cfg = model_args.get("vision", {})
+        if not isinstance(vision_cfg, Mapping):
+            return None
+
+        height = vision_cfg.get("image_height")
+        width = vision_cfg.get("image_width")
+        if height is None and width is None:
+            return None
+        if height is None or width is None:
+            raise ValueError(
+                "Both model_args.vision.image_height and "
+                "model_args.vision.image_width must be set for static Gemma4 "
+                "calibration resizing."
+            )
+
+        parsed_height = int(height)
+        parsed_width = int(width)
+        if parsed_height <= 0 or parsed_width <= 0:
+            raise ValueError(
+                "Gemma4 static calibration image dimensions must be positive: "
+                f"image_height={parsed_height}, image_width={parsed_width}."
+            )
+        return parsed_height, parsed_width
+
+    @staticmethod
+    @contextmanager
+    def _fixed_image_processor_size(
+        processor: Any,
+        image_size: tuple[int, int] | None,
+    ) -> Iterator[None]:
+        """Temporarily force processor image resizing for static calibration."""
+        if image_size is None:
+            yield
+            return
+
+        image_processor = getattr(processor, "image_processor", None)
+        if image_processor is None:
+            yield
+            return
+
+        height, width = image_size
+        updates: dict[str, Any] = {
+            "do_resize": True,
+            "size": {"height": height, "width": width},
+        }
+        if hasattr(image_processor, "crop_size"):
+            updates["crop_size"] = {"height": height, "width": width}
+
+        originals: dict[str, tuple[bool, Any]] = {}
+        for attr, value in updates.items():
+            originals[attr] = (
+                hasattr(image_processor, attr),
+                getattr(image_processor, attr, None),
+            )
+            setattr(image_processor, attr, value)
+
+        try:
+            yield
+        finally:
+            for attr, (had_attr, original) in originals.items():
+                if had_attr:
+                    setattr(image_processor, attr, original)
+                elif hasattr(image_processor, attr):
+                    delattr(image_processor, attr)
+
     def build_calibration_inputs(self, ctx: RecipeContext) -> list[dict]:
         """Build VLM calibration inputs for fixed image-text PTQ."""
         calib = ctx.cfg.get("calibration", {})
         runtime = ctx.cfg.get("runtime", {})
-        return build_vlm_calibration_inputs(
-            processor=ctx.processor,
-            dataset=calib.get("dataset", "vqav2"),
-            datasets=calib.get("datasets"),
-            n_samples=int(calib.get("n_samples", 128)),
-            split=calib.get("split", "testdev"),
-            max_seq_len=calib.get("seq_len"),
-            seed=int(runtime.get("seed", 42)),
-        )
+        image_size = self._static_calibration_image_size(ctx.cfg)
+        with self._fixed_image_processor_size(ctx.processor, image_size):
+            return build_vlm_calibration_inputs(
+                processor=ctx.processor,
+                dataset=calib.get("dataset", "vqav2"),
+                datasets=calib.get("datasets"),
+                n_samples=int(calib.get("n_samples", 128)),
+                split=calib.get("split", "testdev"),
+                max_seq_len=calib.get("seq_len"),
+                seed=int(runtime.get("seed", 42)),
+            )
 
     def forward_calibration(
         self,

{tico-0.2.0.dev260708 → tico-0.2.0.dev260709}/tico/quantization/recipes/debug/wrapper_smoke/cases/qwen3_vl.py

@@ -464,14 +464,14 @@ class QwenTextMLPCase(QwenBaseCase):
         return ForwardInput((torch.randn(2, 8, self.text_cfg.hidden_size),))
 
 
-class QwenTextDecoderLayerCase(QwenBaseCase):
-    """Smoke case for qwen/quantize_text_decoder_layer.py."""
+class QwenTextDecoderLayerBaseCase(QwenBaseCase):
+    """Base class for Qwen3-VL text decoder-layer smoke cases."""
 
-    name = "qwen3_vl_text_decoder_layer"
-    description = "Quantize one tiny Qwen3-VL text decoder layer."
-    tags = ("qwen3_vl", "text", "decoder_layer")
+    tags: tuple[str, ...] = ("qwen3_vl", "text", "decoder_layer")
     max_mean_abs_diff = 3.0
     inplace_convert = True
+    seq_len = 8
+    export_mode = "prefill"
 
     def build(self, cfg: Mapping[str, Any]) -> tuple[torch.nn.Module, torch.nn.Module]:
         """Build a tiny text decoder layer and reference copy."""
@@ -484,16 +484,92 @@ class QwenTextDecoderLayerCase(QwenBaseCase):
         module = Qwen3VLTextDecoderLayer(self.text_cfg, layer_idx=0).eval()
         return module, clone_module(module)
 
+    def export_module(
+        self, quantized: torch.nn.Module, cfg: Mapping[str, Any]
+    ) -> torch.nn.Module:
+        """Export the wrapped text decoder layer in the case-specific mode."""
+        wrapped = getattr(quantized, "wrapped", quantized)
+        return (
+            wrapped.as_export_module(self.export_mode).eval()
+            if hasattr(wrapped, "as_export_module")
+            else quantized
+        )
+
+
+class QwenTextDecoderLayerPrefillCase(QwenTextDecoderLayerBaseCase):
+    """Smoke case for the Qwen3-VL text decoder-layer prefill wrapper path."""
+
+    name = "qwen3_vl_text_decoder_layer_prefill"
+    description = "Quantize one tiny Qwen3-VL text decoder layer in prefill mode."
+    tags = ("qwen3_vl", "text", "decoder_layer", "prefill")
+    export_mode = "prefill"
+
     def _sample(self) -> ForwardInput:
-        """Create one synthetic text decoder-layer input."""
-        seq_len = 8
-        hidden = torch.randn(1, seq_len, self.text_cfg.hidden_size)
-        pos = (
-            torch.randn(1, seq_len, self.text_cfg.head_dim),
-            torch.randn(1, seq_len, self.text_cfg.head_dim),
+        """Create one synthetic text decoder-layer prefill input."""
+        hidden = torch.randn(1, self.seq_len, self.text_cfg.hidden_size)
+        pos = _text_rope(1, self.seq_len, self.text_cfg.head_dim)
+        mask = _causal_mask(self.seq_len)
+        position_ids = torch.arange(self.seq_len).unsqueeze(0)
+        return ForwardInput(
+            (),
+            {
+                "hidden_states": hidden,
+                "position_embeddings": pos,
+                "attention_mask": mask,
+                "position_ids": position_ids,
+            },
         )
-        mask = torch.zeros(1, 1, seq_len, seq_len)
-        position_ids = torch.arange(seq_len).unsqueeze(0)
+
+    def calibration_inputs(
+        self, prepared: torch.nn.Module, cfg: Mapping[str, Any]
+    ) -> list[ForwardInput]:
+        """Create text decoder-layer prefill calibration samples."""
+        return [self._sample() for _ in range(5)]
+
+    def eval_input(
+        self, prepared: torch.nn.Module, cfg: Mapping[str, Any]
+    ) -> ForwardInput:
+        """Create the text decoder-layer prefill evaluation sample."""
+        return self._sample()
+
+    def export_input(
+        self, eval_sample: ForwardInput, cfg: Mapping[str, Any]
+    ) -> ForwardInput:
+        """Create positional prefill inputs expected by the export adapter."""
+        hidden = eval_sample.kwargs["hidden_states"]
+        mask = eval_sample.kwargs["attention_mask"]
+        pos = eval_sample.kwargs["position_embeddings"]
+        return ForwardInput((hidden, mask, pos))
+
+
+class QwenTextDecoderLayerDecodeCase(QwenTextDecoderLayerBaseCase):
+    """Smoke case for the Qwen3-VL text decoder-layer decode wrapper path."""
+
+    name = "qwen3_vl_text_decoder_layer_decode"
+    description = "Quantize one tiny Qwen3-VL text decoder layer in static decode mode."
+    tags = ("qwen3_vl", "text", "decoder_layer", "decode")
+    compare_reference_source = "prepared"
+    export_mode = "decode"
+
+    def after_prepare(self, prepared: torch.nn.Module, cfg: Mapping[str, Any]) -> None:
+        """Force tuple return so hidden states and cache tensors are available."""
+        wrapped = getattr(prepared, "wrapped", prepared)
+        if hasattr(wrapped, "return_type"):
+            wrapped.return_type = "tuple"
+
+    def _sample(self) -> ForwardInput:
+        """Create one synthetic static decode text decoder-layer input."""
+        hidden = torch.randn(1, 1, self.text_cfg.hidden_size)
+        pos = _text_rope(1, 1, self.text_cfg.head_dim)
+        mask = torch.zeros(1, 1, 1, self.seq_len)
+        position_ids = torch.full((1, 1), self.seq_len - 1, dtype=torch.long)
+        past_k = torch.randn(
+            1,
+            self.text_cfg.num_key_value_heads,
+            self.seq_len - 1,
+            self.text_cfg.head_dim,
+        )
+        past_v = torch.randn_like(past_k)
         return ForwardInput(
             (),
             {
@@ -501,21 +577,33 @@ class QwenTextDecoderLayerCase(QwenBaseCase):
                 "position_embeddings": pos,
                 "attention_mask": mask,
                 "position_ids": position_ids,
+                "past_key_values": (past_k, past_v),
+                "use_cache": True,
             },
         )
 
     def calibration_inputs(
         self, prepared: torch.nn.Module, cfg: Mapping[str, Any]
     ) -> list[ForwardInput]:
-        """Create text decoder-layer calibration samples."""
+        """Create text decoder-layer decode calibration samples."""
         return [self._sample() for _ in range(3)]
 
     def eval_input(
         self, prepared: torch.nn.Module, cfg: Mapping[str, Any]
     ) -> ForwardInput:
-        """Create the text decoder-layer evaluation sample."""
+        """Create the text decoder-layer decode evaluation sample."""
         return self._sample()
 
+    def export_input(
+        self, eval_sample: ForwardInput, cfg: Mapping[str, Any]
+    ) -> ForwardInput:
+        """Create positional static decode inputs expected by the export adapter."""
+        hidden = eval_sample.kwargs["hidden_states"]
+        mask = eval_sample.kwargs["attention_mask"]
+        pos = eval_sample.kwargs["position_embeddings"]
+        past = eval_sample.kwargs["past_key_values"]
+        return ForwardInput((hidden, mask, pos, past))
+
 
 class QwenTextModelCase(QwenBaseCase):
     """Smoke case for qwen/quantize_text_model.py."""
@@ -871,7 +959,8 @@ QWEN3_VL_CASES: tuple[WrapperSmokeCase, ...] = (
     QwenTextAttentionPrefillCase(),
     QwenTextAttentionDecodeCase(),
     QwenTextMLPCase(),
-    QwenTextDecoderLayerCase(),
+    QwenTextDecoderLayerPrefillCase(),
+    QwenTextDecoderLayerDecodeCase(),
     QwenTextModelCase(),
     QwenVisionAttentionCase(),
     QwenVisionMLPCase(),