tico 0.1.0.dev250831__py3-none-any.whl → 0.1.0.dev250901__py3-none-any.whl

tico/__init__.py

@@ -29,7 +29,7 @@ __all__ = [
 ]
 
 # THIS LINE IS AUTOMATICALLY GENERATED BY setup.py
-__version__ = "0.1.0.dev250831"
+__version__ = "0.1.0.dev250901"
 
 MINIMUM_SUPPORTED_VERSION = "2.5.0"
 SECURE_TORCH_VERSION = "2.6.0"

tico/experimental/quantization/ptq/examples/quantize_llama_decoder_layer.py

@@ -0,0 +1,124 @@
+# Copyright (c) 2025 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.
+
+# =============================================================================
+# POST-TRAINING QUANTIZATION EXAMPLE — Llama Decoder Layer (Self-Attn + MLP)
+# -----------------------------------------------------------------------------
+# This demo shows how to:
+#   1. Replace a single FP32 `LlamaDecoderLayer` with `QuantLlamaDecoderLayer`.
+#   2. Collect activation statistics in one calibration sweep.
+#   3. Freeze scales / zero-points and switch to INT-simulation mode.
+#   4. Compare INT-8 vs FP32 outputs with a quick mean-absolute-diff check.
+#   5. Export the calibrated, quantized block to a Circle model.
+# -----------------------------------------------------------------------------
+# Style / layout is kept identical to the `quantize_llama_attn.py` and
+# `quantize_llama_mlp.py` examples for easy side-by-side reading.
+# =============================================================================
+
+import pathlib
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from tico.experimental.quantization.evaluation.metric import compute_peir
+from tico.experimental.quantization.evaluation.utils import plot_two_outputs
+from tico.experimental.quantization.ptq.mode import Mode
+from tico.experimental.quantization.ptq.wrappers.llama.quant_decoder_layer import (
+    QuantLlamaDecoderLayer,
+)
+from tico.utils.utils import SuppressWarning
+
+MODEL_NAME = "Maykeye/TinyLLama-v0"
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+
+model.eval()  # disable dropout, etc.
+rotary = model.model.rotary_emb  # RoPE helper
+
+# -------------------------------------------------------------------------
+# 1. Swap in the quant wrapper
+# -------------------------------------------------------------------------
+fp32_layer = model.model.layers[0]  # keep a reference for diff check
+model.model.layers[0] = QuantLlamaDecoderLayer(
+    fp32_layer
+)  # PTQWrapper(fp32_layer) is also fine
+model.eval()
+
+qlayer = model.model.layers[0]  # alias for brevity
+
+# -------------------------------------------------------------------------
+# 2. Single-pass calibration (gather activation ranges)
+# -------------------------------------------------------------------------
+PROMPTS = [
+    "The quick brown fox jumps over the lazy dog.",
+    "In 2025, AI systems accelerated hardware-software co-design at scale.",
+    "양자화는 왜 어려울까? 분포, 길이, 마스크가 관건이다.",
+    "今日はいい天気ですね。ところでRoPE角度は長さに依存します。",
+    "def quicksort(arr):\n    if len(arr) <= 1: return arr\n    ...",
+    "Prices rose 3.14% — see Figure 2; emails: foo@bar.com!",
+]
+
+with torch.no_grad():
+    qlayer.enable_calibration()
+    for prompt in PROMPTS:
+        ids = tokenizer(prompt, return_tensors="pt")
+        hidden = model.model.embed_tokens(ids["input_ids"])
+        pos = rotary(hidden, ids["input_ids"])  # (cos, sin) tuple
+        S = pos[0].shape[1]
+        attn_mask = torch.zeros(1, 1, S, S)  # causal-mask placeholder
+        _ = qlayer(hidden, attention_mask=attn_mask, position_embeddings=pos)
+    qlayer.freeze_qparams()
+
+assert qlayer._mode is Mode.QUANT, "Quantization mode should be active now."
+
+# -------------------------------------------------------------------------
+# 3. Quick INT-sim vs FP32 sanity check
+# -------------------------------------------------------------------------
+ids = tokenizer("check", return_tensors="pt")
+hidden = model.model.embed_tokens(ids["input_ids"])
+pos = rotary(hidden, ids["input_ids"])
+S = pos[0].shape[1]
+attn_mask = torch.zeros(1, 1, S, S)
+
+with torch.no_grad():
+    int8_out = qlayer(hidden, attention_mask=attn_mask, position_embeddings=pos)
+    int8 = int8_out[0] if isinstance(int8_out, tuple) else int8_out
+    fp32_out = fp32_layer(hidden, attention_mask=attn_mask, position_embeddings=pos)
+    fp32 = fp32_out[0] if isinstance(fp32_out, tuple) else fp32_out
+
+print("┌───────────── Quantization Error Summary ─────────────")
+print(f"│ Mean |diff|: {(int8 - fp32).abs().mean().item():.6f}")
+print(f"│ PEIR       : {compute_peir(fp32, int8) * 100:.6f} %")
+print("└──────────────────────────────────────────────────────")
+print(plot_two_outputs(fp32, int8))
+
+# -------------------------------------------------------------------------
+# 4. Export the calibrated layer to Circle
+# -------------------------------------------------------------------------
+import tico
+
+save_path = pathlib.Path("decoder_layer.q.circle")
+B, S, D = 1, 4, model.config.hidden_size
+example_hidden = torch.randn(B, S, D)
+example_pos = rotary(example_hidden, torch.arange(S)[None, :])
+attn_mask = torch.zeros(1, 1, S, S)
+
+with SuppressWarning(UserWarning, ".*"):
+    cm = tico.convert(
+        qlayer, (example_hidden, attn_mask), {"position_embeddings": example_pos}
+    )
+# Note that the model is not fully quantized.
+cm.save(save_path)
+
+print(f"Quantized Circle model saved to {save_path.resolve()}")

tico/experimental/quantization/ptq/wrappers/llama/__init__.py

@@ -1,6 +1,9 @@
 from tico.experimental.quantization.ptq.wrappers.llama.quant_attn import (
     QuantLlamaAttention,
 )
+from tico.experimental.quantization.ptq.wrappers.llama.quant_decoder_layer import (
+    QuantLlamaDecoderLayer,
+)
 from tico.experimental.quantization.ptq.wrappers.llama.quant_mlp import QuantLlamaMLP
 
-__all__ = ["QuantLlamaAttention", "QuantLlamaMLP"]
+__all__ = ["QuantLlamaAttention", "QuantLlamaDecoderLayer", "QuantLlamaMLP"]

tico/experimental/quantization/ptq/wrappers/llama/quant_decoder_layer.py

@@ -0,0 +1,168 @@
+# Copyright (c) 2025 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.
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+
+from tico.experimental.quantization.ptq.quant_config import QuantConfig
+from tico.experimental.quantization.ptq.wrappers.llama.quant_attn import (
+    QuantLlamaAttention,
+)
+from tico.experimental.quantization.ptq.wrappers.llama.quant_mlp import QuantLlamaMLP
+from tico.experimental.quantization.ptq.wrappers.ptq_wrapper import PTQWrapper
+from tico.experimental.quantization.ptq.wrappers.quant_module_base import (
+    QuantModuleBase,
+)
+from tico.experimental.quantization.ptq.wrappers.registry import try_register
+
+
+@try_register("transformers.models.llama.modeling_llama.LlamaDecoderLayer")
+class QuantLlamaDecoderLayer(QuantModuleBase):
+    """
+    Quant-aware drop-in replacement for HF `LlamaDecoderLayer`.
+    Signature and return-value are identical to the original.
+
+    ▸ Attention & MLP blocks are replaced by their quantized counterparts
+    ▸ LayerNorms remain FP32 (no fake-quant)
+    ▸ A "static" causal mask is pre-built in `__init__` to avoid
+      dynamic boolean-to-float casts inside `forward`.
+
+    Notes on the causal mask
+    ------------------------
+    Building a boolean mask "inside" `forward` would introduce
+    non-deterministic dynamic ops that an integer-only accelerator cannot
+    fuse easily.  Therefore we:
+
+    1. Pre-compute a full upper-triangular mask of size
+       `[1, 1, max_seq, max_seq]` in `__init__`.
+    2. In `forward`, if the caller passes `attention_mask=None`, we
+       slice the pre-computed template to the current sequence length.
+    """
+
+    def __init__(
+        self,
+        fp_layer: nn.Module,
+        *,
+        qcfg: Optional[QuantConfig] = None,
+        fp_name: Optional[str] = None,
+        return_type: Optional[str] = None,
+    ):
+        """
+        Q) Why do we need `return_type`?
+        A) Different versions of `transformers` wrap the decoder output in
+            different containers: a plain Tensor or a tuple.
+        """
+        self.return_type = return_type
+        if self.return_type is None:
+            import transformers
+
+            v = tuple(map(int, transformers.__version__.split(".")[:2]))
+            self.return_type = "tensor" if v >= (4, 54) else "tuple"
+        assert self.return_type is not None
+        super().__init__(qcfg, fp_name=fp_name)
+
+        # Child QuantConfigs -------------------------------------------------
+        attn_cfg = qcfg.child("self_attn") if qcfg else None
+        mlp_cfg = qcfg.child("mlp") if qcfg else None
+
+        # Quantized sub-modules ---------------------------------------------
+        assert hasattr(fp_layer, "self_attn") and isinstance(
+            fp_layer.self_attn, torch.nn.Module
+        )
+        assert hasattr(fp_layer, "mlp") and isinstance(fp_layer.mlp, torch.nn.Module)
+        self.self_attn = PTQWrapper(
+            fp_layer.self_attn, qcfg=attn_cfg, fp_name=f"{fp_name}.self_attn"
+        )
+        self.mlp = PTQWrapper(fp_layer.mlp, qcfg=mlp_cfg, fp_name=f"{fp_name}.mlp")
+
+        # LayerNorms remain FP (copied from fp_layer to keep weights)
+        assert hasattr(fp_layer, "input_layernorm") and isinstance(
+            fp_layer.input_layernorm, torch.nn.Module
+        )
+        assert hasattr(fp_layer, "post_attention_layernorm") and isinstance(
+            fp_layer.post_attention_layernorm, torch.nn.Module
+        )
+        self.input_layernorm = fp_layer.input_layernorm
+        self.post_attention_layernorm = fp_layer.post_attention_layernorm
+
+        # Static causal mask template ---------------------------------------
+        assert hasattr(fp_layer.self_attn, "config") and hasattr(
+            fp_layer.self_attn.config, "max_position_embeddings"
+        )
+        assert isinstance(fp_layer.self_attn.config.max_position_embeddings, int)
+        max_seq = fp_layer.self_attn.config.max_position_embeddings
+        mask = torch.full((1, 1, max_seq, max_seq), float("-120"))
+        mask.triu_(1)
+        self.register_buffer("causal_mask_template", mask, persistent=False)
+
+    def _slice_causal(self, seq_len: int, device: torch.device) -> torch.Tensor:
+        """Return `[1,1,L,L]` causal mask slice on *device*."""
+        assert isinstance(self.causal_mask_template, torch.Tensor)
+        return self.causal_mask_template[..., :seq_len, :seq_len].to(device)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional["Cache"] = None,  # type: ignore[name-defined]
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor] | torch.Tensor:
+        if output_attentions:
+            raise NotImplementedError(
+                "QuantLlamaDecoderLayer does not support output attention yet."
+            )
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+
+        if attention_mask is None or attention_mask.dtype == torch.bool:
+            L = hidden_states.size(1)
+            attention_mask = self._slice_causal(L, hidden_states.device)
+
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+
+        # ─── MLP block ─────────────────────────────────────────────────
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        if self.return_type == "tuple":
+            return (hidden_states,)
+        elif self.return_type == "tensor":
+            return hidden_states
+        else:
+            raise RuntimeError("Invalid return type.")
+
+    # No local observers; just recurse into children
+    def _all_observers(self):
+        yield from self.self_attn._all_observers()
+        yield from self.mlp._all_observers()

tico/experimental/quantization/ptq/wrappers/registry.py

@@ -29,6 +29,7 @@ _CORE_MODULES = (
     "tico.experimental.quantization.ptq.wrappers.nn.quant_linear",
     "tico.experimental.quantization.ptq.wrappers.nn.quant_silu",
     "tico.experimental.quantization.ptq.wrappers.llama.quant_attn",
+    "tico.experimental.quantization.ptq.wrappers.llama.quant_decoder_layer",
     "tico.experimental.quantization.ptq.wrappers.llama.quant_mlp",
     # add future core wrappers here
 )

{tico-0.1.0.dev250831.dist-info → tico-0.1.0.dev250901.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tico
-Version: 0.1.0.dev250831
+Version: 0.1.0.dev250901
 Summary: Convert exported Torch module to circle
 Home-page: UNKNOWN
 License: UNKNOWN

{tico-0.1.0.dev250831.dist-info → tico-0.1.0.dev250901.dist-info}/RECORD

@@ -1,4 +1,4 @@
-tico/__init__.py,sha256=AGSSd8oWeJ_KCyfANsnPaH1Sf2i33iYNqoeT153KsC0,1883
+tico/__init__.py,sha256=MgvCVXWMpNL2dxPn54C8fdQaTJPdtHivhuNHH4qN5R8,1883
 tico/pt2_to_circle.py,sha256=gu3MD4Iqc0zMZcCZ2IT8oGbyj21CTSbT3Rgd9s2B_9A,2767
 tico/config/__init__.py,sha256=xZzCXjZ84qE-CsBi-dfaL05bqpQ3stKKfTXhnrJRyVs,142
 tico/config/base.py,sha256=q5xMqGxTUZs4mFqt5c7i_y9U00fYgdMGl9nUqIVMlCo,1248
@@ -64,6 +64,7 @@ tico/experimental/quantization/ptq/quant_config.py,sha256=nm7570Y1X2mOT_8s27ilWi
 tico/experimental/quantization/ptq/examples/__init__.py,sha256=IO6FP_xYbGy0dW0HL26GXD3ouxARaxCK7bz9dn4blPQ,26
 tico/experimental/quantization/ptq/examples/quantize_linear.py,sha256=8zq-ZJDYgam0xQ-PbC6Xb1I7W1mv0Wi-b--IP2wwXtw,4539
 tico/experimental/quantization/ptq/examples/quantize_llama_attn.py,sha256=cVWUSSzaZWFp5QZkNkrlpHU3kXyP84QtnZbahVml_yQ,4329
+tico/experimental/quantization/ptq/examples/quantize_llama_decoder_layer.py,sha256=mBWrjkyEovYQsPC4Rrsri6Pm1rlFmDb3NiP0DQQhFyM,5751
 tico/experimental/quantization/ptq/examples/quantize_llama_mlp.py,sha256=N1qZQgt1S-xZrdv-PW7OfXEcv0gsO2q9faOF4aD-zKo,4147
 tico/experimental/quantization/ptq/observers/__init__.py,sha256=WF2MvL9M_jl-B1FqcY9zic34NOCRp17HkRYv-TMxMr4,613
 tico/experimental/quantization/ptq/observers/affine_base.py,sha256=e2Eba64nrxKQyE4F_WJ7WTSsk3xe6bkdGUKaoLFWGFw,4638
@@ -78,9 +79,10 @@ tico/experimental/quantization/ptq/wrappers/__init__.py,sha256=47DEQpj8HBSa-_TIm
 tico/experimental/quantization/ptq/wrappers/ptq_wrapper.py,sha256=F9sK_DiRaXiGNHULcwIbs5EUtHz6ZJ7N4r5CWTTfhsM,2442
 tico/experimental/quantization/ptq/wrappers/quant_elementwise.py,sha256=LhEoobfvto6zKrBOKL4gmxfFFc31jHzyQV_zfps-iQM,3604
 tico/experimental/quantization/ptq/wrappers/quant_module_base.py,sha256=vkcDos_knGSS29rIZuEIWkAJLHrENbGz8nCH2-iara8,5969
-tico/experimental/quantization/ptq/wrappers/registry.py,sha256=TwH-MD-qkTkG6M-f1VqFLmSNcXLNYsh21yjyzCcojJc,4706
-tico/experimental/quantization/ptq/wrappers/llama/__init__.py,sha256=b360gkQ0RxExmiV-ZaaxwJdMPJ53g6uCRlR2-_dOby0,240
+tico/experimental/quantization/ptq/wrappers/registry.py,sha256=M1D_foC0PR-Ii4G0lbOO3_pmhvHlMF28NolK_q2DZtw,4783
+tico/experimental/quantization/ptq/wrappers/llama/__init__.py,sha256=4xuAYnJcohMTtBzrH4cxq8WKG2GQo8nbhektVg8w7F0,380
 tico/experimental/quantization/ptq/wrappers/llama/quant_attn.py,sha256=WIUI6EFMTvvruvqu8pBxWy6qJeDyjkaYbJk1R3pAmwE,8578
+tico/experimental/quantization/ptq/wrappers/llama/quant_decoder_layer.py,sha256=2XsIf5rcabDXXkahqriSxfo2curFq0Y5bnRPcYkJPg8,7187
 tico/experimental/quantization/ptq/wrappers/llama/quant_mlp.py,sha256=uZMnrX66oZwxhKhcNbLXXeri-WxxRBiZnr15aBXJMm0,3562
 tico/experimental/quantization/ptq/wrappers/nn/__init__.py,sha256=I9uTt5HfcRoMEDYHpAeATMv2TbCQiX0ZbfUFMzSJ4Qw,336
 tico/experimental/quantization/ptq/wrappers/nn/quant_layernorm.py,sha256=G5Sgt-tXnzh0Rxyk-2honmZIfEQOZlRfOsoDBdSGmA4,6887
@@ -240,9 +242,9 @@ tico/utils/mx/__init__.py,sha256=IO6FP_xYbGy0dW0HL26GXD3ouxARaxCK7bz9dn4blPQ,26
 tico/utils/mx/elemwise_ops.py,sha256=V6glyAHsVR1joqpsgnNytatCD_ew92xNWZ19UFDoMTA,10281
 tico/utils/mx/formats.py,sha256=uzNWyu-1onUlwQfX5cZ6fZSUfHMRqorper7_T1k3jfk,3404
 tico/utils/mx/mx_ops.py,sha256=RcfUTYVi-wilGB2sC35OeARdwDqnixv7dG5iyZ-fQT8,8555
-tico-0.1.0.dev250831.dist-info/LICENSE,sha256=kp4JLII7bzRhPb0CPD5XTDZMh22BQ7h3k3B7t8TiSbw,12644
-tico-0.1.0.dev250831.dist-info/METADATA,sha256=1tkuUVzG_pExk9oBSJO3s5opMoRAVMkWns4biwGPDu0,8450
-tico-0.1.0.dev250831.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
-tico-0.1.0.dev250831.dist-info/entry_points.txt,sha256=kBKYSS_IYrSXmUYevmmepqIVPScq5vF8ulQRu3I_Zf0,59
-tico-0.1.0.dev250831.dist-info/top_level.txt,sha256=oqs7UPoNSKZEwqsX8B-KAWdQwfAa7i60pbxW_Jk7P3w,5
-tico-0.1.0.dev250831.dist-info/RECORD,,
+tico-0.1.0.dev250901.dist-info/LICENSE,sha256=kp4JLII7bzRhPb0CPD5XTDZMh22BQ7h3k3B7t8TiSbw,12644
+tico-0.1.0.dev250901.dist-info/METADATA,sha256=LMgoYoHYFT8cJU9VNYiiX89tMSxEX30x17x_6eWAr4o,8450
+tico-0.1.0.dev250901.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
+tico-0.1.0.dev250901.dist-info/entry_points.txt,sha256=kBKYSS_IYrSXmUYevmmepqIVPScq5vF8ulQRu3I_Zf0,59
+tico-0.1.0.dev250901.dist-info/top_level.txt,sha256=oqs7UPoNSKZEwqsX8B-KAWdQwfAa7i60pbxW_Jk7P3w,5
+tico-0.1.0.dev250901.dist-info/RECORD,,