diff --git a/README_WEIGHT_SPARSITY.md b/README_WEIGHT_SPARSITY.md new file mode 100644 index 00000000..ecf6a005 --- /dev/null +++ b/README_WEIGHT_SPARSITY.md @@ -0,0 +1,280 @@ +# TICO post-convert weight sparsity + +This patch adds a `weight_sparsity` recipe stage for measuring affine-quantized +weights immediately after the PTQ `convert()` stage. It does not load or inspect +Circle artifacts. + +The implementation targets: + +- `meta-llama/Llama-3.2-3B-Instruct` +- `Qwen/Qwen3-VL-4B-Instruct` + +Two reports are produced from a single scan of the quantized weights. + +The summary report contains exactly three columns: + +```text +Scope | Qdtype | Sparsity (%) +``` + +The layer-level report adds the normalized layer path: + +```text +Layer | Scope | Qdtype | Sparsity (%) +``` + +`Qscheme` is used internally to interpret per-tensor and per-channel qparams, but +it is intentionally not emitted in either report. + +## Base revision + +The patch was prepared against Samsung/TICO `main` at: + +```text +ef42e2d2323187fc41406e1787b1d840aa316f9a +``` + +## Added functionality + +The analyzer: + +- visits direct TICO weight observers on leaf quantization wrappers; +- verifies that each measured wrapper is in post-convert `QUANT` mode; +- skips `IdentityObserver` weights because they remain floating point; +- reconstructs affine integer codes in bounded chunks; +- counts semantic zeros using `qcode == zero_point`; +- aggregates by total zero count and total element count, not by an unweighted + average of tensor-level percentages; +- deduplicates shared weights only when storage, qdtype, and finalized qparams + are identical; +- emits projection-level scopes for Llama and Qwen3-VL; +- emits an overall row and projection rows for each normalized model layer; +- produces summary and layer reports from the same collected tensor statistics; +- writes CSV and Markdown reports without a Qscheme column. + +## Apply the patch + +From a TICO checkout at the base revision: + +```bash +git checkout ef42e2d2323187fc41406e1787b1d840aa316f9a +git apply /path/to/tico_weight_sparsity.patch +``` + +Alternatively, copy the `tico/` and `test/` directories from this archive over a +TICO checkout. + +## Run the tests + +```bash +python -m unittest \ + test.quantization.analysis.test_weight_sparsity \ + test.quantization.recipes.stages.test_weight_sparsity +``` + +## Run Llama 3.2 3B + +The Meta model is gated, so configure a Hugging Face token through the normal +TICO/Hugging Face authentication path before running. + +```bash +python -m tico.quantization.examples.quantize \ + --config tico/quantization/examples/configs/llama32_3b_weight_sparsity.yaml +``` + +Generated files: + +```text +./out/llama32_3b_weight_sparsity/weight_sparsity.csv +./out/llama32_3b_weight_sparsity/weight_sparsity.md +./out/llama32_3b_weight_sparsity/weight_sparsity_by_layer.csv +./out/llama32_3b_weight_sparsity/weight_sparsity_by_layer.md +``` + +Llama summary scopes are ordered as follows when present: + +```text +All quantized weights +Token embedding +Attention / q_proj +Attention / k_proj +Attention / v_proj +Attention / o_proj +MLP / gate_proj +MLP / up_proj +MLP / down_proj +LM head +Norm weights +SpinQuant rotation weights +Unclassified quantized weights +``` + +The layer report uses normalized paths such as: + +```text +model.embed_tokens +model.rotate_embedding +model.layers.0 +model.layers.1 +... +model.norm +rotate_lm_head +lm_head +``` + +For each decoder layer, the report contains an `All quantized weights` row and +projection-level rows for Q/K/V/O and gate/up/down projections. Norm weights are +reported in the same decoder-layer group. + +## Run Qwen3-VL 4B + +```bash +python -m tico.quantization.examples.quantize \ + --config tico/quantization/examples/configs/qwen3_vl_4b_weight_sparsity.yaml +``` + +Generated files: + +```text +./out/qwen3_vl_4b_weight_sparsity/weight_sparsity.csv +./out/qwen3_vl_4b_weight_sparsity/weight_sparsity.md +./out/qwen3_vl_4b_weight_sparsity/weight_sparsity_by_layer.csv +./out/qwen3_vl_4b_weight_sparsity/weight_sparsity_by_layer.md +``` + +Qwen3-VL summary scopes are ordered as follows when present: + +```text +All quantized weights +Vision / patch_embed.proj +Vision attention / qkv +Vision attention / proj +Vision MLP / linear_fc1 +Vision MLP / linear_fc2 +Vision merger / linear_fc1 +Vision merger / linear_fc2 +Deepstack merger / linear_fc1 +Deepstack merger / linear_fc2 +Text / token embedding +Text attention / q_proj +Text attention / k_proj +Text attention / v_proj +Text attention / o_proj +Text MLP / gate_proj +Text MLP / up_proj +Text MLP / down_proj +LM head +Norm weights +SpinQuant rotation weights +Unclassified quantized weights +``` + +The Qwen3-VL layer report separates vision and text layers with normalized paths: + +```text +model.visual.patch_embed +model.visual.blocks.0 +model.visual.blocks.1 +... +model.visual.merger +model.visual.deepstack_merger_list.0 +... +model.language_model.embed_tokens +model.language_model.rotate_embedding +model.language_model.layers.0 +model.language_model.layers.1 +... +model.language_model.norm +rotate_lm_head +lm_head +``` + +## Pipeline placement + +The stage must immediately follow `ptq`: + +```yaml +pipeline: + - name: gptq + enabled: true + # ... + + - name: ptq + enabled: true + # ... + + - name: weight_sparsity + enabled: true + output_dir: ./out/model_weight_sparsity + filename_stem: weight_sparsity + layer_filename_stem: weight_sparsity_by_layer + formats: [csv, markdown] + precision: 6 + max_chunk_elements: 4194304 + deduplicate_shared_weights: true + include_empty_scopes: false + include_layer_report: true + include_layer_totals: true + print_layer_report: false +``` + +Configuration fields for the layer report: + +- `include_layer_report`: Generate the layer CSV and Markdown files. Default: `true`. +- `include_layer_totals`: Add `All quantized weights` within every layer. Default: `true`. +- `layer_filename_stem`: Set the layer report filename stem. The default is + `_by_layer`. +- `print_layer_report`: Print the full layer table to stdout. When an output + directory exists, the default is `false` to avoid large console output. + +If the stage runs before PTQ conversion, it raises an error instead of reporting +raw floating-point zero counts. + +## Programmatic use + +```python +from tico.quantization import convert +from tico.quantization.analysis import ( + format_layer_weight_sparsity_table, + format_weight_sparsity_table, + measure_weight_sparsity_report, +) + +quantized_model = convert(prepared_model, ptq_config) +report = measure_weight_sparsity_report(quantized_model, family="llama") + +print(format_weight_sparsity_table(report.summary_rows)) +print(format_layer_weight_sparsity_table(report.layer_rows)) +``` + +Use `family="qwen3_vl"` for Qwen3-VL. + +## Sparsity definition + +For an affine-quantized weight element, the integer code is reconstructed as: + +```text +qcode = clamp(round(weight / scale) + zero_point, qmin, qmax) +``` + +The element is counted as zero when: + +```text +qcode == zero_point +``` + +A scope percentage is calculated as: + +```text +sum(scope semantic-zero counts) / sum(scope element counts) * 100 +``` + +The same weighted calculation is used for each layer total and each +layer-projection row. + +Shared weights are deduplicated inside each individual row. Because the same +shared weight may be intentionally visible under more than one logical layer, +layer rows should not be summed to reconstruct the global total. Use the summary +`All quantized weights` row for the globally deduplicated value. + +This is the effective quantized-value sparsity at Torch level after conversion. diff --git a/pyproject.toml b/pyproject.toml index fee54713..a87ed174 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,6 +20,7 @@ dependencies = [ [project.scripts] pt2-to-circle = "tico.pt2_to_circle:main" +circle-weight-sparsity = "tico.quantization.circle_weight_sparsity:main" [tool.setuptools] # Automatic package discovery (replaces the old find_packages role) diff --git a/test/quantization/analysis/test_weight_sparsity.py b/test/quantization/analysis/test_weight_sparsity.py new file mode 100644 index 00000000..9b17fec1 --- /dev/null +++ b/test/quantization/analysis/test_weight_sparsity.py @@ -0,0 +1,451 @@ +# 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. + +import tempfile +import unittest +from pathlib import Path +from types import SimpleNamespace + +import torch + +from tico.quantization.analysis.weight_sparsity import ( + format_layer_weight_sparsity_table, + format_weight_sparsity_table, + measure_layer_weight_sparsity, + measure_weight_sparsity, + measure_weight_sparsity_report, + WeightSparsityError, + write_layer_weight_sparsity_csv, + write_weight_sparsity_csv, +) + + +class _FakeDType: + def __init__(self, bits: int, signed: bool = False): + self.bits = bits + self.signed = signed + + @property + def qmin(self) -> int: + if self.signed: + return -(1 << (self.bits - 1)) + return 0 + + @property + def qmax(self) -> int: + if self.signed: + return (1 << (self.bits - 1)) - 1 + return (1 << self.bits) - 1 + + def __str__(self) -> str: + prefix = "int" if self.signed else "uint" + return f"{prefix}{self.bits}" + + +class _FakeObserver(torch.nn.Module): + def __init__( + self, + *, + scale, + zero_point, + bits: int = 4, + signed: bool = False, + channel_axis: int | None = 0, + ): + super().__init__() + self.dtype = _FakeDType(bits, signed=signed) + self.channel_axis = channel_axis + self.register_buffer( + "_cached_scale", torch.as_tensor(scale, dtype=torch.float32) + ) + self.register_buffer( + "_cached_zp", torch.as_tensor(zero_point, dtype=torch.int32) + ) + + +class IdentityObserver(_FakeObserver): + pass + + +class _WeightOwner(torch.nn.Module): + def __init__(self, weight: torch.Tensor): + super().__init__() + self.weight = torch.nn.Parameter(weight.clone()) + + +class _QuantLeaf(torch.nn.Module): + def __init__( + self, + fp_name: str, + weight: torch.Tensor, + observer: torch.nn.Module, + *, + shared_parameter: torch.nn.Parameter | None = None, + mode: str = "QUANT", + ): + super().__init__() + self.fp_name = fp_name + self._mode = SimpleNamespace(name=mode) + self.module = _WeightOwner(weight) + if shared_parameter is not None: + self.module.weight = shared_parameter + self.observer = observer + + def get_observer(self, name: str, *, recurse: bool = True): + if name == "weight" and not recurse: + return self.observer + return None + + +class _Model(torch.nn.Module): + def __init__(self, *leaves: _QuantLeaf): + super().__init__() + self.leaves = torch.nn.ModuleList(leaves) + + +class TestWeightSparsity(unittest.TestCase): + def test_semantic_zero_uses_zero_point(self): + leaf = _QuantLeaf( + "model.layers.0.self_attn.q_proj", + torch.tensor([[-0.49, 0.49, 0.51, -0.51]]), + _FakeObserver(scale=[1.0], zero_point=[7], channel_axis=0), + ) + + rows = measure_weight_sparsity(_Model(leaf), "llama") + q_proj = next(row for row in rows if row.scope == "Attention / q_proj") + + self.assertEqual(q_proj.qdtype, "uint4") + self.assertEqual(q_proj.numel, 4) + self.assertEqual(q_proj.zero_count, 2) + self.assertAlmostEqual(q_proj.sparsity_percent, 50.0) + + def test_rounding_matches_torch_fake_quantization(self): + leaf = _QuantLeaf( + "model.layers.0.self_attn.k_proj", + torch.tensor([[0.5, -0.5, 1.5, -1.5]]), + _FakeObserver(scale=[1.0], zero_point=[7], channel_axis=0), + ) + + rows = measure_weight_sparsity(_Model(leaf), "llama") + k_proj = next(row for row in rows if row.scope == "Attention / k_proj") + + reference = torch.fake_quantize_per_channel_affine( + leaf.module.weight.detach(), + scale=torch.tensor([1.0]), + zero_point=torch.tensor([7], dtype=torch.int32), + axis=0, + quant_min=0, + quant_max=15, + ) + self.assertEqual(k_proj.zero_count, int(torch.count_nonzero(reference == 0))) + self.assertEqual(k_proj.zero_count, 2) + + def test_per_channel_zero_points_are_broadcast(self): + leaf = _QuantLeaf( + "model.layers.0.mlp.up_proj", + torch.tensor([[0.4, 0.6], [0.9, 1.1]]), + _FakeObserver( + scale=[1.0, 2.0], + zero_point=[5, 9], + channel_axis=0, + ), + ) + + rows = measure_weight_sparsity(_Model(leaf), "llama", max_chunk_elements=2) + up_proj = next(row for row in rows if row.scope == "MLP / up_proj") + + self.assertEqual(up_proj.zero_count, 2) + self.assertEqual(up_proj.numel, 4) + self.assertAlmostEqual(up_proj.sparsity_percent, 50.0) + + def test_scope_aggregation_is_weighted_by_numel(self): + dense_zero = _QuantLeaf( + "model.layers.0.self_attn.q_proj", + torch.tensor([[0.0, 0.0]]), + _FakeObserver(scale=[1.0], zero_point=[7], channel_axis=0), + ) + no_zero = _QuantLeaf( + "model.layers.1.self_attn.q_proj", + torch.ones(1, 8), + _FakeObserver(scale=[1.0], zero_point=[7], channel_axis=0), + ) + + rows = measure_weight_sparsity(_Model(dense_zero, no_zero), "llama") + q_proj = next(row for row in rows if row.scope == "Attention / q_proj") + + self.assertEqual(q_proj.zero_count, 2) + self.assertEqual(q_proj.numel, 10) + self.assertAlmostEqual(q_proj.sparsity_percent, 20.0) + + def test_identity_weight_observer_is_skipped(self): + quantized = _QuantLeaf( + "model.layers.0.mlp.down_proj", + torch.zeros(1, 2), + _FakeObserver(scale=[1.0], zero_point=[7], channel_axis=0), + ) + floating_point = _QuantLeaf( + "lm_head", + torch.zeros(1, 4), + IdentityObserver(scale=[1.0], zero_point=[0], channel_axis=0), + ) + + rows = measure_weight_sparsity(_Model(quantized, floating_point), "llama") + scopes = {row.scope for row in rows} + + self.assertIn("MLP / down_proj", scopes) + self.assertNotIn("LM head", scopes) + + def test_shared_weight_with_identical_qparams_is_deduplicated(self): + shared = torch.nn.Parameter(torch.zeros(1, 4)) + embedding = _QuantLeaf( + "model.embed_tokens", + torch.empty(1, 4), + _FakeObserver(scale=[1.0], zero_point=[7], channel_axis=0), + shared_parameter=shared, + ) + lm_head = _QuantLeaf( + "lm_head", + torch.empty(1, 4), + _FakeObserver(scale=[1.0], zero_point=[7], channel_axis=0), + shared_parameter=shared, + ) + + rows = measure_weight_sparsity(_Model(embedding, lm_head), "llama") + all_weights = rows[0] + + self.assertEqual(all_weights.numel, 4) + self.assertEqual(all_weights.zero_count, 4) + self.assertIn("Token embedding", {row.scope for row in rows}) + self.assertIn("LM head", {row.scope for row in rows}) + + def test_qwen_projection_scopes_are_reported(self): + paths = [ + "visual.patch_embed.proj", + "model.visual.blocks.0.attn.qkv", + "model.visual.blocks.0.attn.proj", + "model.visual.blocks.0.mlp.linear_fc1", + "model.visual.blocks.0.mlp.linear_fc2", + "model.visual.merger.linear_fc1", + "model.visual.merger.linear_fc2", + "model.visual.deepstack_merger_list.0.linear_fc1", + "model.visual.deepstack_merger_list.0.linear_fc2", + "model.language_model.embed_tokens", + "model.language_model.layers.0.self_attn.q_proj", + "model.language_model.layers.0.self_attn.k_proj", + "model.language_model.layers.0.self_attn.v_proj", + "model.language_model.layers.0.self_attn.o_proj", + "model.language_model.layers.0.mlp.gate_proj", + "model.language_model.layers.0.mlp.up_proj", + "model.language_model.layers.0.mlp.down_proj", + "lm_head", + ] + leaves = [ + _QuantLeaf( + path, + torch.zeros(1, 1), + _FakeObserver(scale=[1.0], zero_point=[7], channel_axis=0), + ) + for path in paths + ] + + rows = measure_weight_sparsity(_Model(*leaves), "qwen3_vl") + scopes = {row.scope for row in rows} + + expected = { + "Vision / patch_embed.proj", + "Vision attention / qkv", + "Vision attention / proj", + "Vision MLP / linear_fc1", + "Vision MLP / linear_fc2", + "Vision merger / linear_fc1", + "Vision merger / linear_fc2", + "Deepstack merger / linear_fc1", + "Deepstack merger / linear_fc2", + "Text / token embedding", + "Text attention / q_proj", + "Text attention / k_proj", + "Text attention / v_proj", + "Text attention / o_proj", + "Text MLP / gate_proj", + "Text MLP / up_proj", + "Text MLP / down_proj", + "LM head", + } + self.assertTrue(expected.issubset(scopes)) + + def test_llama_layer_report_contains_totals_and_projection_rows(self): + layer_zero = _QuantLeaf( + "model.layers.0.self_attn.q_proj", + torch.zeros(1, 2), + _FakeObserver(scale=[1.0], zero_point=[7], channel_axis=0), + ) + layer_nonzero = _QuantLeaf( + "model.layers.1.self_attn.q_proj", + torch.ones(1, 4), + _FakeObserver(scale=[1.0], zero_point=[7], channel_axis=0), + ) + layer_mlp = _QuantLeaf( + "model.layers.0.mlp.up_proj", + torch.tensor([[0.0, 1.0]]), + _FakeObserver(scale=[1.0], zero_point=[7], channel_axis=0), + ) + + report = measure_weight_sparsity_report( + _Model(layer_zero, layer_nonzero, layer_mlp), + "llama", + ) + rows = list(report.layer_rows) + + layer_names = list(dict.fromkeys(row.layer for row in rows)) + self.assertEqual(layer_names, ["model.layers.0", "model.layers.1"]) + + layer0_total = next( + row + for row in rows + if row.layer == "model.layers.0" and row.scope == "All quantized weights" + ) + layer0_q = next( + row + for row in rows + if row.layer == "model.layers.0" and row.scope == "Attention / q_proj" + ) + layer0_up = next( + row + for row in rows + if row.layer == "model.layers.0" and row.scope == "MLP / up_proj" + ) + layer1_total = next( + row + for row in rows + if row.layer == "model.layers.1" and row.scope == "All quantized weights" + ) + + self.assertEqual((layer0_total.zero_count, layer0_total.numel), (3, 4)) + self.assertAlmostEqual(layer0_total.sparsity_percent, 75.0) + self.assertEqual((layer0_q.zero_count, layer0_q.numel), (2, 2)) + self.assertEqual((layer0_up.zero_count, layer0_up.numel), (1, 2)) + self.assertEqual((layer1_total.zero_count, layer1_total.numel), (0, 4)) + + def test_qwen_layer_report_uses_architecture_order(self): + paths = [ + "model.language_model.layers.10.self_attn.q_proj", + "model.visual.deepstack_merger_list.1.linear_fc1", + "model.visual.blocks.2.attn.qkv", + "model.language_model.layers.2.self_attn.q_proj", + "model.visual.blocks.0.attn.qkv", + "model.visual.patch_embed.proj", + "model.visual.merger.linear_fc1", + "model.language_model.embed_tokens", + "model.language_model.norm", + "lm_head", + ] + leaves = [ + _QuantLeaf( + path, + torch.zeros(1, 1), + _FakeObserver(scale=[1.0], zero_point=[7], channel_axis=0), + ) + for path in paths + ] + + rows = measure_layer_weight_sparsity(_Model(*leaves), "qwen3_vl") + layers = list(dict.fromkeys(row.layer for row in rows)) + + self.assertEqual( + layers, + [ + "model.visual.patch_embed", + "model.visual.blocks.0", + "model.visual.blocks.2", + "model.visual.merger", + "model.visual.deepstack_merger_list.1", + "model.language_model.embed_tokens", + "model.language_model.layers.2", + "model.language_model.layers.10", + "model.language_model.norm", + "lm_head", + ], + ) + + def test_layer_totals_can_be_disabled(self): + leaf = _QuantLeaf( + "model.layers.0.self_attn.o_proj", + torch.zeros(1, 2), + _FakeObserver(scale=[1.0], zero_point=[7], channel_axis=0), + ) + + rows = measure_layer_weight_sparsity( + _Model(leaf), + "llama", + include_layer_totals=False, + ) + + self.assertEqual(len(rows), 1) + self.assertEqual(rows[0].layer, "model.layers.0") + self.assertEqual(rows[0].scope, "Attention / o_proj") + + def test_layer_output_contains_requested_columns(self): + leaf = _QuantLeaf( + "model.layers.0.mlp.gate_proj", + torch.zeros(1, 2), + _FakeObserver(scale=[1.0], zero_point=[7], channel_axis=0), + ) + rows = measure_layer_weight_sparsity(_Model(leaf), "llama") + + markdown = format_layer_weight_sparsity_table(rows, precision=2) + self.assertIn( + "| Layer | Scope | Qdtype | Sparsity (%) |", + markdown, + ) + self.assertNotIn("Qscheme", markdown) + + with tempfile.TemporaryDirectory() as tmpdir: + path = write_layer_weight_sparsity_csv( + rows, + Path(tmpdir) / "report_by_layer.csv", + ) + header = path.read_text(encoding="utf-8").splitlines()[0] + self.assertEqual(header, "layer,scope,qdtype,sparsity_percent") + + def test_non_quant_mode_is_rejected(self): + leaf = _QuantLeaf( + "model.layers.0.self_attn.q_proj", + torch.zeros(1, 1), + _FakeObserver(scale=[1.0], zero_point=[7], channel_axis=0), + mode="CALIB", + ) + + with self.assertRaises(WeightSparsityError): + measure_weight_sparsity(_Model(leaf), "llama") + + def test_output_contains_only_requested_columns(self): + leaf = _QuantLeaf( + "model.layers.0.mlp.gate_proj", + torch.zeros(1, 2), + _FakeObserver(scale=[1.0], zero_point=[7], channel_axis=0), + ) + rows = measure_weight_sparsity(_Model(leaf), "llama") + + markdown = format_weight_sparsity_table(rows, precision=2) + self.assertIn("| Scope | Qdtype | Sparsity (%) |", markdown) + self.assertNotIn("Qscheme", markdown) + + with tempfile.TemporaryDirectory() as tmpdir: + path = write_weight_sparsity_csv(rows, Path(tmpdir) / "report.csv") + header = path.read_text(encoding="utf-8").splitlines()[0] + self.assertEqual(header, "scope,qdtype,sparsity_percent") + + +if __name__ == "__main__": + unittest.main() diff --git a/test/quantization/recipes/stages/test_weight_sparsity.py b/test/quantization/recipes/stages/test_weight_sparsity.py new file mode 100644 index 00000000..cb9f4704 --- /dev/null +++ b/test/quantization/recipes/stages/test_weight_sparsity.py @@ -0,0 +1,133 @@ +# 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. + +import tempfile +import unittest +from pathlib import Path +from types import SimpleNamespace + +import torch + +from tico.quantization.recipes.context import RecipeContext +from tico.quantization.recipes.stages.weight_sparsity import WeightSparsityStage + + +class _FakeDType: + qmin = 0 + qmax = 15 + + def __str__(self) -> str: + return "uint4" + + +class _FakeObserver(torch.nn.Module): + def __init__(self): + super().__init__() + self.dtype = _FakeDType() + self.channel_axis = 0 + self.register_buffer("_cached_scale", torch.tensor([1.0])) + self.register_buffer("_cached_zp", torch.tensor([7], dtype=torch.int32)) + + +class _Leaf(torch.nn.Module): + def __init__(self): + super().__init__() + self.fp_name = "model.layers.0.self_attn.q_proj" + self._mode = SimpleNamespace(name="QUANT") + self.module = torch.nn.Linear(2, 1, bias=False) + self.module.weight.data.zero_() + self.observer = _FakeObserver() + + def get_observer(self, name: str, *, recurse: bool = True): + if name == "weight" and not recurse: + return self.observer + return None + + +class _Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.leaf = _Leaf() + + +class TestWeightSparsityStage(unittest.TestCase): + def test_stage_writes_csv_and_markdown_artifacts(self): + with tempfile.TemporaryDirectory() as tmpdir: + ctx = RecipeContext( + cfg={}, + adapter=SimpleNamespace(family="llama"), + model=_Model(), + ) + stage = WeightSparsityStage() + returned = stage.run( + ctx, + { + "output_dir": tmpdir, + "formats": ["csv", "markdown"], + "precision": 3, + }, + ) + + self.assertIs(returned, ctx) + csv_path = Path(ctx.artifacts["weight_sparsity_csv"]) + markdown_path = Path(ctx.artifacts["weight_sparsity_markdown"]) + layer_csv_path = Path(ctx.artifacts["weight_sparsity_by_layer_csv"]) + layer_markdown_path = Path( + ctx.artifacts["weight_sparsity_by_layer_markdown"] + ) + self.assertTrue(csv_path.is_file()) + self.assertTrue(markdown_path.is_file()) + self.assertTrue(layer_csv_path.is_file()) + self.assertTrue(layer_markdown_path.is_file()) + self.assertEqual( + csv_path.read_text(encoding="utf-8").splitlines()[0], + "scope,qdtype,sparsity_percent", + ) + self.assertEqual( + layer_csv_path.read_text(encoding="utf-8").splitlines()[0], + "layer,scope,qdtype,sparsity_percent", + ) + self.assertIn( + "| Scope | Qdtype | Sparsity (%) |", + markdown_path.read_text(encoding="utf-8"), + ) + self.assertIn( + "| Layer | Scope | Qdtype | Sparsity (%) |", + layer_markdown_path.read_text(encoding="utf-8"), + ) + + def test_stage_can_disable_layer_report(self): + with tempfile.TemporaryDirectory() as tmpdir: + ctx = RecipeContext( + cfg={}, + adapter=SimpleNamespace(family="llama"), + model=_Model(), + ) + stage = WeightSparsityStage() + stage.run( + ctx, + { + "output_dir": tmpdir, + "formats": ["csv"], + "include_layer_report": False, + }, + ) + + self.assertIn("weight_sparsity_csv", ctx.artifacts) + self.assertNotIn("weight_sparsity_by_layer_csv", ctx.artifacts) + self.assertFalse((Path(tmpdir) / "weight_sparsity_by_layer.csv").exists()) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/quantization/test_circle_weight_sparsity.py b/test/quantization/test_circle_weight_sparsity.py new file mode 100644 index 00000000..6dcb303a --- /dev/null +++ b/test/quantization/test_circle_weight_sparsity.py @@ -0,0 +1,776 @@ +# 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. + +from __future__ import annotations + +from dataclasses import replace +from pathlib import Path +from typing import Any + +import numpy as np +import pytest + +import tico.quantization.circle_weight_sparsity as circle_sparsity +from tico.quantization.circle_weight_sparsity import ( + _count_tensor_semantic_zeros, + _deduplicate_across_files, + aggregate_circle_weight_stats, + analyze_circle_binary, + CircleWeightSparsityReport, + CircleWeightSparsityRow, + CircleWeightTensorStats, + discover_circle_files, + QuantizationInfo, + render_csv, + render_json, + render_markdown, +) + + +class _TensorTypeValues: + """Provide the tensor type values used by fake Circle models.""" + + FLOAT32 = 0 + INT32 = 2 + UINT8 = 3 + INT64 = 4 + BOOL = 6 + INT16 = 7 + INT8 = 9 + FLOAT16 = 10 + UINT4 = 17 + INT4 = 18 + MXFP4 = 30 + + +class _BuiltinOperatorValues: + """Provide the builtin operator values used by fake Circle models.""" + + ADD = 0 + CONV_2D = 3 + FULLY_CONNECTED = 9 + GATHER = 36 + DEQUANTIZE = 60 + RMS_NORM = 120 + + +class _FakeQuantization: + """Mimic generated Circle quantization accessors.""" + + def __init__( + self, + scales: list[float] | None = None, + zero_points: list[int] | None = None, + axis: int = 0, + ) -> None: + self._scales = np.asarray(scales or [], dtype=np.float32) + self._zero_points = np.asarray(zero_points or [], dtype=np.int64) + self._axis = axis + + def ScaleAsNumpy(self) -> np.ndarray: + """Return scale values.""" + + return self._scales + + def ScaleLength(self) -> int: + """Return the number of scale values.""" + + return int(self._scales.size) + + def Scale(self, index: int) -> float: + """Return one scale value.""" + + return float(self._scales[index]) + + def ZeroPointAsNumpy(self) -> np.ndarray: + """Return zero-point values.""" + + return self._zero_points + + def ZeroPointLength(self) -> int: + """Return the number of zero-point values.""" + + return int(self._zero_points.size) + + def ZeroPoint(self, index: int) -> int: + """Return one zero-point value.""" + + return int(self._zero_points[index]) + + def QuantizedDimension(self) -> int: + """Return the per-channel quantized dimension.""" + + return self._axis + + +class _FakeTensor: + """Mimic a generated Circle tensor object.""" + + def __init__( + self, + name: str, + shape: tuple[int, ...], + tensor_type: int, + buffer_index: int, + quantization: _FakeQuantization | None = None, + *, + is_variable: bool = False, + ) -> None: + self._name = name + self._shape = np.asarray(shape, dtype=np.int32) + self._type = tensor_type + self._buffer_index = buffer_index + self._quantization = quantization + self._is_variable = is_variable + + def Name(self) -> bytes: + """Return the encoded tensor name.""" + + return self._name.encode("utf-8") + + def ShapeAsNumpy(self) -> np.ndarray: + """Return the tensor shape.""" + + return self._shape + + def ShapeLength(self) -> int: + """Return the tensor rank.""" + + return int(self._shape.size) + + def Shape(self, index: int) -> int: + """Return one tensor dimension.""" + + return int(self._shape[index]) + + def Type(self) -> int: + """Return the tensor type.""" + + return self._type + + def Buffer(self) -> int: + """Return the backing buffer index.""" + + return self._buffer_index + + def Quantization(self) -> _FakeQuantization | None: + """Return optional quantization metadata.""" + + return self._quantization + + def IsVariable(self) -> bool: + """Return whether the tensor is mutable.""" + + return self._is_variable + + +class _FakeBuffer: + """Mimic a generated Circle buffer object.""" + + def __init__(self, data: bytes = b"") -> None: + self._data = np.frombuffer(data, dtype=np.uint8) + + def DataAsNumpy(self) -> np.ndarray: + """Return buffer bytes.""" + + return self._data + + def DataLength(self) -> int: + """Return the buffer byte length.""" + + return int(self._data.size) + + def Data(self, index: int) -> int: + """Return one buffer byte.""" + + return int(self._data[index]) + + +class _FakeOperator: + """Mimic a generated Circle operator object.""" + + def __init__( + self, + opcode_index: int, + inputs: tuple[int, ...], + outputs: tuple[int, ...], + ) -> None: + self._opcode_index = opcode_index + self._inputs = np.asarray(inputs, dtype=np.int32) + self._outputs = np.asarray(outputs, dtype=np.int32) + + def OpcodeIndex(self) -> int: + """Return the opcode table index.""" + + return self._opcode_index + + def InputsAsNumpy(self) -> np.ndarray: + """Return input tensor IDs.""" + + return self._inputs + + def InputsLength(self) -> int: + """Return the number of inputs.""" + + return int(self._inputs.size) + + def Inputs(self, index: int) -> int: + """Return one input tensor ID.""" + + return int(self._inputs[index]) + + def OutputsAsNumpy(self) -> np.ndarray: + """Return output tensor IDs.""" + + return self._outputs + + def OutputsLength(self) -> int: + """Return the number of outputs.""" + + return int(self._outputs.size) + + def Outputs(self, index: int) -> int: + """Return one output tensor ID.""" + + return int(self._outputs[index]) + + +class _FakeOperatorCode: + """Mimic a generated Circle operator code object.""" + + def __init__(self, builtin_code: int) -> None: + self._builtin_code = builtin_code + + def BuiltinCode(self) -> int: + """Return the current builtin code.""" + + return self._builtin_code + + def DeprecatedBuiltinCode(self) -> int: + """Return the legacy builtin code.""" + + return min(self._builtin_code, 127) + + +class _FakeSubgraph: + """Mimic a generated Circle subgraph object.""" + + def __init__( + self, + tensors: list[_FakeTensor], + operators: list[_FakeOperator], + inputs: tuple[int, ...], + ) -> None: + self._tensors = tensors + self._operators = operators + self._inputs = np.asarray(inputs, dtype=np.int32) + + def TensorsLength(self) -> int: + """Return the number of tensors.""" + + return len(self._tensors) + + def Tensors(self, index: int) -> _FakeTensor: + """Return one tensor.""" + + return self._tensors[index] + + def OperatorsLength(self) -> int: + """Return the number of operators.""" + + return len(self._operators) + + def Operators(self, index: int) -> _FakeOperator: + """Return one operator.""" + + return self._operators[index] + + def InputsAsNumpy(self) -> np.ndarray: + """Return runtime input tensor IDs.""" + + return self._inputs + + def InputsLength(self) -> int: + """Return the number of runtime inputs.""" + + return int(self._inputs.size) + + def Inputs(self, index: int) -> int: + """Return one runtime input tensor ID.""" + + return int(self._inputs[index]) + + +class _FakeModelRoot: + """Mimic a generated Circle model object.""" + + def __init__( + self, + subgraphs: list[_FakeSubgraph], + buffers: list[_FakeBuffer], + operator_codes: list[_FakeOperatorCode], + ) -> None: + self._subgraphs = subgraphs + self._buffers = buffers + self._operator_codes = operator_codes + + def SubgraphsLength(self) -> int: + """Return the number of subgraphs.""" + + return len(self._subgraphs) + + def Subgraphs(self, index: int) -> _FakeSubgraph: + """Return one subgraph.""" + + return self._subgraphs[index] + + def BuffersLength(self) -> int: + """Return the number of buffers.""" + + return len(self._buffers) + + def Buffers(self, index: int) -> _FakeBuffer: + """Return one buffer.""" + + return self._buffers[index] + + def OperatorCodesLength(self) -> int: + """Return the number of operator codes.""" + + return len(self._operator_codes) + + def OperatorCodes(self, index: int) -> _FakeOperatorCode: + """Return one operator code.""" + + return self._operator_codes[index] + + +class _FakeCircle: + """Provide the generated Circle module hierarchy used by the analyzer.""" + + class TensorType: + """Hold the fake tensor type enum class.""" + + TensorType = _TensorTypeValues + + class BuiltinOperator: + """Hold the fake builtin operator enum class.""" + + BuiltinOperator = _BuiltinOperatorValues + + class Model: + """Hold a fake root accessor class.""" + + class Model: + """Return the model registered by a test.""" + + root: _FakeModelRoot | None = None + + @classmethod + def GetRootAsModel(cls, data: Any, offset: int) -> _FakeModelRoot: + """Return the test model regardless of the binary payload.""" + + assert offset == 0 + assert cls.root is not None + return cls.root + + +def _pack_uint4(values: list[int]) -> bytes: + """Pack low-nibble-first UINT4 values for fake Circle buffers.""" + + array = np.asarray(values, dtype=np.uint8) + packed = np.zeros((array.size + 1) // 2, dtype=np.uint8) + packed[:] = array[0::2] + packed[: array.size // 2] |= array[1::2] << np.uint8(4) + return packed.tobytes() + + +def _make_model_with_linear_embedding_and_norm() -> _FakeModelRoot: + """Build a fake mixed-dtype model with three selected weights.""" + + buffers = [ + _FakeBuffer(), + _FakeBuffer(_pack_uint4([2, 2, 3, 4, 5, 1, 5, 0])), + _FakeBuffer(np.asarray([0, 0], dtype=" CircleWeightSparsityReport: + """Build a compact report for serialization tests.""" + + return CircleWeightSparsityReport( + row=CircleWeightSparsityRow( + scope="All model weights", + qdtype="mixed (uint4, uint8)", + sparsity_pct=sparsity_pct, + ), + source_count=1, + tensor_count=2, + zero_count=1, + numel=8, + duplicate_tensor_count=0, + skipped_tensor_count=0, + skipped_messages=(), + ) + + +def test_auto_selection_counts_model_weights_and_excludes_biases() -> None: + _FakeCircle.Model.Model.root = _make_model_with_linear_embedding_and_norm() + + stats, duplicates, skipped = analyze_circle_binary( + b"fake-circle", + circle_module=_FakeCircle, + selection="auto", + chunk_numel=4, + ) + + assert duplicates == 0 + assert skipped == [] + assert {item.tensor_name for item in stats} == { + "tico::p_model_layers_0_q_proj_weight", + "tico::p_model_embed_tokens_weight", + "tico::p_model_norm_weight", + } + assert sum(item.zero_count for item in stats) == 8 + assert sum(item.numel for item in stats) == 16 + assert {item.qdtype for item in stats} == {"uint4", "uint8", "int16"} + + +def test_aggregate_report_is_element_weighted_and_has_stable_dtype_order() -> None: + _FakeCircle.Model.Model.root = _make_model_with_linear_embedding_and_norm() + stats, duplicates, skipped = analyze_circle_binary( + b"fake-circle", + circle_module=_FakeCircle, + selection="auto", + ) + + report = aggregate_circle_weight_stats( + stats, + source_count=1, + duplicate_tensor_count=duplicates, + skipped_messages=skipped, + ) + + assert report.row.scope == "All model weights" + assert report.row.qdtype == "mixed (uint4, uint8, int16)" + assert report.row.sparsity_pct == pytest.approx(50.0) + assert report.zero_count == 8 + assert report.numel == 16 + + +def test_quantized_constant_mode_includes_unconsumed_quantized_constants() -> None: + _FakeCircle.Model.Model.root = _make_model_with_linear_embedding_and_norm() + + stats, _, _ = analyze_circle_binary( + b"fake-circle", + circle_module=_FakeCircle, + selection="quantized-constants", + ) + + names = {item.tensor_name for item in stats} + assert "const_quantized_table" in names + assert "tico::p_model_layers_0_q_proj_bias" in names + + +def test_dequantize_passthrough_resolves_the_stored_weight() -> None: + buffers = [ + _FakeBuffer(), + _FakeBuffer(np.asarray([0, 1, 0, 2], dtype=np.uint8).tobytes()), + ] + tensors = [ + _FakeTensor("input", (1, 4), _TensorTypeValues.FLOAT32, 0), + _FakeTensor( + "tico::p_weight", + (2, 2), + _TensorTypeValues.UINT8, + 1, + _FakeQuantization([1.0], [0]), + ), + _FakeTensor("dequantized_weight", (2, 2), _TensorTypeValues.FLOAT32, 0), + _FakeTensor("output", (1, 2), _TensorTypeValues.FLOAT32, 0), + ] + operator_codes = [ + _FakeOperatorCode(_BuiltinOperatorValues.DEQUANTIZE), + _FakeOperatorCode(_BuiltinOperatorValues.FULLY_CONNECTED), + ] + operators = [ + _FakeOperator(0, (1,), (2,)), + _FakeOperator(1, (0, 2, -1), (3,)), + ] + _FakeCircle.Model.Model.root = _FakeModelRoot( + [_FakeSubgraph(tensors, operators, inputs=(0,))], + buffers, + operator_codes, + ) + + stats, _, _ = analyze_circle_binary( + b"fake-circle", + circle_module=_FakeCircle, + selection="operator-inputs", + ) + + assert len(stats) == 1 + assert stats[0].tensor_name == "tico::p_weight" + assert stats[0].zero_count == 2 + assert stats[0].numel == 4 + + +def test_uint4_unpacking_trims_unused_high_nibble() -> None: + raw = np.frombuffer(_pack_uint4([2, 3, 2, 4, 2]), dtype=np.uint8) + + zero_count, numel = _count_tensor_semantic_zeros( + raw, + "uint4", + (5,), + QuantizationInfo((2,), 0, True), + chunk_numel=4, + ) + + assert zero_count == 3 + assert numel == 5 + + +def test_float_weights_count_positive_and_negative_zero() -> None: + values = np.asarray([0.0, -0.0, 1.0, -2.0], dtype=" None: + raw = np.frombuffer(_pack_uint4([0, 8, 15, 1]), dtype=np.uint8) + + zero_count, numel = _count_tensor_semantic_zeros( + raw, + "int4", + (4,), + QuantizationInfo((0,), 0, True), + chunk_numel=2, + ) + + assert zero_count == 1 + assert numel == 4 + + +def test_per_channel_zero_points_support_nonzero_quantized_axis() -> None: + values = np.asarray( + [ + [[1, 0], [2, 2], [3, 0]], + [[0, 1], [2, 4], [5, 3]], + ], + dtype=np.uint8, + ) + + zero_count, numel = _count_tensor_semantic_zeros( + values.reshape(-1), + "uint8", + (2, 3, 2), + QuantizationInfo((1, 2, 3), 1, True), + chunk_numel=5, + ) + + assert zero_count == 7 + assert numel == 12 + + +def test_model_buffer_reuse_is_counted_once() -> None: + model = _make_model_with_linear_embedding_and_norm() + shared = model.Subgraphs(0).Tensors(5) + model.Subgraphs(0)._tensors.append( + _FakeTensor( + "tico::p_tied_lm_head_weight", + (2, 2), + _TensorTypeValues.UINT8, + shared.Buffer(), + shared.Quantization(), + ) + ) + _FakeCircle.Model.Model.root = model + + stats, duplicates, _ = analyze_circle_binary( + b"fake-circle", + circle_module=_FakeCircle, + selection="auto", + ) + + assert len(stats) == 3 + assert duplicates == 1 + + +def test_cross_file_deduplication_requires_name_and_payload_match() -> None: + base = CircleWeightTensorStats( + source="a.circle", + subgraph_index=0, + tensor_index=1, + tensor_name="tico::p_weight", + buffer_index=1, + qdtype="uint4", + shape=(2, 2), + zero_count=2, + numel=4, + roles=("FULLY_CONNECTED:input1",), + fingerprint="abc", + ) + duplicate = replace(base, source="b.circle", tensor_index=2) + different_name = replace( + base, + source="c.circle", + tensor_name="tico::p_other_weight", + ) + + unique, duplicate_count = _deduplicate_across_files( + [base, duplicate, different_name] + ) + + assert unique == [base, different_name] + assert duplicate_count == 1 + + +def test_renderers_expose_exactly_three_columns() -> None: + report = _report() + + markdown = render_markdown(report, precision=3) + csv_text = render_csv(report, precision=3) + json_payload = __import__("json").loads(render_json(report, precision=3)) + + assert markdown.splitlines()[0] == "| Scope | Qdtype | Sparsity (%) |" + assert csv_text.splitlines()[0] == "Scope,Qdtype,Sparsity (%)" + assert list(json_payload[0]) == ["Scope", "Qdtype", "Sparsity (%)"] + assert json_payload[0]["Sparsity (%)"] == 12.5 + + +def test_cli_writes_one_model_level_csv_row( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _FakeCircle.Model.Model.root = _make_model_with_linear_embedding_and_norm() + model_path = tmp_path / "model.circle" + output_path = tmp_path / "result.csv" + model_path.write_bytes(b"fake-circle") + monkeypatch.setattr(circle_sparsity, "_load_circle_schema", lambda: _FakeCircle) + + status = circle_sparsity.main( + [ + str(model_path), + "--format", + "csv", + "--output", + str(output_path), + "--precision", + "3", + ] + ) + + assert status == 0 + assert output_path.read_text(encoding="utf-8").splitlines() == [ + "Scope,Qdtype,Sparsity (%)", + 'All model weights,"mixed (uint4, uint8, int16)",50.000', + ] + + +def test_discover_circle_files_supports_recursive_directories(tmp_path: Path) -> None: + direct = tmp_path / "a.circle" + nested_dir = tmp_path / "nested" + nested_dir.mkdir() + nested = nested_dir / "b.circle" + ignored = nested_dir / "notes.txt" + direct.write_bytes(b"a") + nested.write_bytes(b"b") + ignored.write_text("ignored", encoding="utf-8") + + assert discover_circle_files([tmp_path], recursive=False) == [direct.resolve()] + assert discover_circle_files([tmp_path], recursive=True) == sorted( + [direct.resolve(), nested.resolve()] + ) + + +def test_unsupported_mx_dtype_fails_clearly() -> None: + raw = np.asarray([0], dtype=np.uint8) + + with pytest.raises(Exception, match="MX formats are not yet decoded"): + _count_tensor_semantic_zeros( + raw, + "mxfp4", + (2,), + QuantizationInfo((), None, False), + chunk_numel=2, + ) diff --git a/tico/quantization/analysis/__init__.py b/tico/quantization/analysis/__init__.py new file mode 100644 index 00000000..481329f7 --- /dev/null +++ b/tico/quantization/analysis/__init__.py @@ -0,0 +1,53 @@ +# 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. + +from tico.quantization.analysis.weight_sparsity import ( + aggregate_layer_weight_sparsity, + aggregate_weight_sparsity, + collect_weight_tensor_sparsity, + format_layer_weight_sparsity_table, + format_weight_sparsity_table, + LayerSparsityRow, + measure_layer_weight_sparsity, + measure_weight_sparsity, + measure_weight_sparsity_report, + SparsityRow, + WeightSparsityError, + WeightSparsityReport, + WeightTensorSparsity, + write_layer_weight_sparsity_csv, + write_layer_weight_sparsity_markdown, + write_weight_sparsity_csv, + write_weight_sparsity_markdown, +) + +__all__ = [ + "LayerSparsityRow", + "SparsityRow", + "WeightSparsityError", + "WeightSparsityReport", + "WeightTensorSparsity", + "aggregate_layer_weight_sparsity", + "aggregate_weight_sparsity", + "collect_weight_tensor_sparsity", + "format_layer_weight_sparsity_table", + "format_weight_sparsity_table", + "measure_layer_weight_sparsity", + "measure_weight_sparsity", + "measure_weight_sparsity_report", + "write_layer_weight_sparsity_csv", + "write_layer_weight_sparsity_markdown", + "write_weight_sparsity_csv", + "write_weight_sparsity_markdown", +] diff --git a/tico/quantization/analysis/weight_sparsity.py b/tico/quantization/analysis/weight_sparsity.py new file mode 100644 index 00000000..e94e2fba --- /dev/null +++ b/tico/quantization/analysis/weight_sparsity.py @@ -0,0 +1,1113 @@ +# 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. + +from __future__ import annotations + +import csv +import hashlib +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable, Iterable, Sequence + +import torch + + +_DEFAULT_MAX_CHUNK_ELEMENTS = 4 * 1024 * 1024 +_ALL_SCOPE = "All quantized weights" +_UNCLASSIFIED_SCOPE = "Unclassified quantized weights" + + +@dataclass(frozen=True) +class WeightTensorSparsity: + """Store semantic-zero statistics for one quantized weight use.""" + + module_name: str + weight_name: str + qdtype: str + zero_count: int + numel: int + dedup_key: str + + @property + def sparsity_percent(self) -> float: + """Return the semantic-zero percentage for this tensor.""" + + if self.numel == 0: + return 0.0 + return 100.0 * self.zero_count / self.numel + + +@dataclass(frozen=True) +class SparsityRow: + """Store one aggregated summary-table row.""" + + scope: str + qdtype: str + zero_count: int + numel: int + + @property + def sparsity_percent(self) -> float | None: + """Return the weighted semantic-zero percentage for this scope.""" + + if self.numel == 0: + return None + return 100.0 * self.zero_count / self.numel + + +@dataclass(frozen=True) +class LayerSparsityRow: + """Store one layer-and-scope sparsity-table row.""" + + layer: str + scope: str + qdtype: str + zero_count: int + numel: int + + @property + def sparsity_percent(self) -> float | None: + """Return the weighted semantic-zero percentage for this layer scope.""" + + if self.numel == 0: + return None + return 100.0 * self.zero_count / self.numel + + +@dataclass(frozen=True) +class WeightSparsityReport: + """Store summary and layer-level rows produced from one tensor scan.""" + + summary_rows: tuple[SparsityRow, ...] + layer_rows: tuple[LayerSparsityRow, ...] + + +class WeightSparsityError(RuntimeError): + """Report an invalid or unsupported post-convert sparsity state.""" + + +def _is_identity_observer(observer: Any) -> bool: + """Return whether an observer intentionally leaves a weight in floating point.""" + + return observer.__class__.__name__ == "IdentityObserver" + + +def _extract_direct_weight(module: torch.nn.Module) -> torch.Tensor | None: + """Return the weight owned by a leaf quantization wrapper.""" + + wrapped_module = getattr(module, "module", None) + weight = getattr(wrapped_module, "weight", None) + if isinstance(weight, torch.Tensor): + return weight + + weight = getattr(module, "weight", None) + if isinstance(weight, torch.Tensor): + return weight + return None + + +def _get_direct_weight_observer(module: torch.nn.Module) -> Any | None: + """Return the direct weight observer without traversing child wrappers.""" + + get_observer = getattr(module, "get_observer", None) + if not callable(get_observer): + return None + + try: + return get_observer("weight", recurse=False) + except TypeError as exc: + raise WeightSparsityError( + f"{type(module).__name__}.get_observer() must accept recurse=False." + ) from exc + + +def _require_post_convert_mode(module: torch.nn.Module, weight_name: str) -> None: + """Verify that a quantization wrapper is in its final QUANT mode.""" + + mode = getattr(module, "_mode", None) + mode_name = getattr(mode, "name", None) + if mode is not None and mode_name != "QUANT": + raise WeightSparsityError( + f"Weight {weight_name!r} is not in post-convert QUANT mode: {mode!r}." + ) + + +def _require_affine_qparams( + observer: Any, weight_name: str +) -> tuple[torch.Tensor, torch.Tensor, int, int, int | None, str]: + """Read finalized affine quantization parameters from a TICO observer.""" + + scale = getattr(observer, "_cached_scale", None) + zero_point = getattr(observer, "_cached_zp", None) + dtype = getattr(observer, "dtype", None) + channel_axis = getattr(observer, "channel_axis", None) + + if not isinstance(scale, torch.Tensor) or scale.numel() == 0: + raise WeightSparsityError( + f"Weight {weight_name!r} does not have a finalized affine scale." + ) + if not isinstance(zero_point, torch.Tensor) or zero_point.numel() == 0: + raise WeightSparsityError( + f"Weight {weight_name!r} does not have a finalized affine zero-point." + ) + if dtype is None or not hasattr(dtype, "qmin") or not hasattr(dtype, "qmax"): + raise WeightSparsityError( + f"Weight {weight_name!r} uses an unsupported non-affine observer " + f"{type(observer).__name__}." + ) + + qmin = int(dtype.qmin) + qmax = int(dtype.qmax) + qdtype = str(dtype) + if channel_axis is not None: + channel_axis = int(channel_axis) + + return scale, zero_point, qmin, qmax, channel_axis, qdtype + + +def _count_block_zeros( + block: torch.Tensor, + scale: torch.Tensor, + zero_point: torch.Tensor, + qmin: int, + qmax: int, +) -> int: + """Count integer codes equal to the broadcast zero-point in one block.""" + + values = block.detach().to(dtype=torch.float32) + scale_f = scale.to(device=values.device, dtype=torch.float32) + zero_f = zero_point.to(device=values.device, dtype=torch.float32) + + quantized = torch.round(values / scale_f) + zero_f + quantized.clamp_(qmin, qmax) + return int(torch.count_nonzero(quantized == zero_f).item()) + + +def _count_per_tensor_zeros( + weight: torch.Tensor, + scale: torch.Tensor, + zero_point: torch.Tensor, + qmin: int, + qmax: int, + max_chunk_elements: int, +) -> int: + """Count semantic zeros for a per-tensor affine weight.""" + + if scale.numel() != 1 or zero_point.numel() != 1: + raise WeightSparsityError( + "Per-tensor quantization requires scalar scale and zero-point values." + ) + + zero_count = 0 + flattened = weight.detach().reshape(-1) + for start in range(0, flattened.numel(), max_chunk_elements): + block = flattened[start : start + max_chunk_elements] + zero_count += _count_block_zeros( + block, + scale.reshape(()), + zero_point.reshape(()), + qmin, + qmax, + ) + return zero_count + + +def _count_per_channel_zeros( + weight: torch.Tensor, + scale: torch.Tensor, + zero_point: torch.Tensor, + qmin: int, + qmax: int, + channel_axis: int, + max_chunk_elements: int, +) -> int: + """Count semantic zeros for a per-channel affine weight.""" + + axis = channel_axis if channel_axis >= 0 else weight.ndim + channel_axis + if axis < 0 or axis >= weight.ndim: + raise WeightSparsityError( + f"Invalid channel axis {channel_axis} for weight shape " + f"{tuple(weight.shape)}." + ) + + channels = int(weight.shape[axis]) + scale_flat = scale.detach().reshape(-1) + zero_flat = zero_point.detach().reshape(-1) + if scale_flat.numel() != channels or zero_flat.numel() != channels: + raise WeightSparsityError( + "Per-channel scale and zero-point lengths must match the quantized axis: " + f"channels={channels}, scale={scale_flat.numel()}, " + f"zero_point={zero_flat.numel()}." + ) + + channel_major = weight.detach().movedim(axis, 0).reshape(channels, -1) + inner_size = int(channel_major.shape[1]) + columns_per_chunk = min(inner_size, max_chunk_elements) + rows_per_chunk = max(1, max_chunk_elements // columns_per_chunk) + + zero_count = 0 + for row_start in range(0, channels, rows_per_chunk): + row_end = min(channels, row_start + rows_per_chunk) + row_scale = scale_flat[row_start:row_end].reshape(-1, 1) + row_zero = zero_flat[row_start:row_end].reshape(-1, 1) + + for column_start in range(0, inner_size, columns_per_chunk): + column_end = min(inner_size, column_start + columns_per_chunk) + block = channel_major[row_start:row_end, column_start:column_end] + zero_count += _count_block_zeros( + block, + row_scale, + row_zero, + qmin, + qmax, + ) + return zero_count + + +def _count_semantic_zeros( + weight: torch.Tensor, + scale: torch.Tensor, + zero_point: torch.Tensor, + qmin: int, + qmax: int, + channel_axis: int | None, + max_chunk_elements: int, +) -> int: + """Count weight values whose integer code equals their affine zero-point.""" + + if weight.device.type == "meta": + raise WeightSparsityError("Meta-device weights cannot be measured.") + if weight.numel() == 0: + return 0 + + with torch.no_grad(): + if channel_axis is None: + return _count_per_tensor_zeros( + weight, + scale, + zero_point, + qmin, + qmax, + max_chunk_elements, + ) + return _count_per_channel_zeros( + weight, + scale, + zero_point, + qmin, + qmax, + channel_axis, + max_chunk_elements, + ) + + +def _tensor_digest(tensor: torch.Tensor) -> str: + """Return a stable digest for a small quantization-parameter tensor.""" + + data = tensor.detach().cpu().contiguous().numpy().tobytes() + return hashlib.sha1(data).hexdigest() + + +def _weight_storage_key(weight: torch.Tensor) -> str: + """Return a storage-and-view identity for a weight tensor.""" + + try: + storage_pointer = weight.untyped_storage().data_ptr() + except (AttributeError, RuntimeError): + storage_pointer = id(weight) + + device_index = -1 if weight.device.index is None else int(weight.device.index) + return ":".join( + [ + weight.device.type, + str(device_index), + str(storage_pointer), + str(weight.storage_offset()), + repr(tuple(weight.shape)), + repr(tuple(weight.stride())), + ] + ) + + +def _make_dedup_key( + weight: torch.Tensor, + qdtype: str, + channel_axis: int | None, + scale: torch.Tensor, + zero_point: torch.Tensor, + module_name: str, + deduplicate_shared_weights: bool, +) -> str: + """Build a key that deduplicates shared weights with identical qparams.""" + + if not deduplicate_shared_weights: + return f"module:{module_name}" + + return ":".join( + [ + _weight_storage_key(weight), + qdtype, + str(channel_axis), + _tensor_digest(scale), + _tensor_digest(zero_point), + ] + ) + + +def collect_weight_tensor_sparsity( + model: torch.nn.Module, + *, + max_chunk_elements: int = _DEFAULT_MAX_CHUNK_ELEMENTS, + deduplicate_shared_weights: bool = True, +) -> list[WeightTensorSparsity]: + """Collect post-convert semantic-zero statistics from TICO weight observers. + + The function visits only direct ``weight`` observers on leaf quantization + wrappers. Identity observers are skipped because they intentionally keep the + associated weight in floating point. Affine integer codes are reconstructed + in bounded chunks, and a semantic zero is counted when ``qcode == zero_point``. + + Args: + model: A TICO model immediately after PTQ ``convert()``. + max_chunk_elements: Maximum number of weight elements processed at once. + deduplicate_shared_weights: Deduplicate shared storage only when the + quantization dtype and finalized qparams are also identical. + + Returns: + One statistics record per quantized weight use. + + Raises: + ValueError: If ``max_chunk_elements`` is not positive. + WeightSparsityError: If the model is not in a measurable post-convert + affine quantization state. + """ + + if max_chunk_elements <= 0: + raise ValueError("max_chunk_elements must be positive.") + + records: list[WeightTensorSparsity] = [] + for module_name, module in model.named_modules(): + observer = _get_direct_weight_observer(module) + if observer is None or _is_identity_observer(observer): + continue + + weight_name = str(getattr(module, "fp_name", None) or module_name) + _require_post_convert_mode(module, weight_name) + + weight = _extract_direct_weight(module) + if weight is None: + raise WeightSparsityError( + f"Quantized weight observer for {weight_name!r} has no direct " + "weight tensor." + ) + + scale, zero_point, qmin, qmax, channel_axis, qdtype = _require_affine_qparams( + observer, weight_name + ) + zero_count = _count_semantic_zeros( + weight, + scale, + zero_point, + qmin, + qmax, + channel_axis, + max_chunk_elements, + ) + dedup_key = _make_dedup_key( + weight, + qdtype, + channel_axis, + scale, + zero_point, + module_name, + deduplicate_shared_weights, + ) + records.append( + WeightTensorSparsity( + module_name=module_name, + weight_name=weight_name, + qdtype=qdtype, + zero_count=zero_count, + numel=weight.numel(), + dedup_key=dedup_key, + ) + ) + + if not records: + raise WeightSparsityError( + "No finalized affine weight observers were found. Run the analysis " + "immediately after the PTQ convert stage." + ) + return records + + +def _path_has_suffix(path: str, suffix: str) -> bool: + """Return whether a module path matches a complete dotted suffix.""" + + return path == suffix or path.endswith(f".{suffix}") + + +def _contains_segment(path: str, *segments: str) -> bool: + """Return whether any complete path segment is present.""" + + path_segments = set(path.split(".")) + return any(segment in path_segments for segment in segments) + + +def _contains_path(path: str, fragment: str) -> bool: + """Return whether a complete dotted subpath is present.""" + + return ( + path == fragment + or path.startswith(f"{fragment}.") + or path.endswith(f".{fragment}") + or f".{fragment}." in path + ) + + +def _classify_llama(weight_name: str) -> str: + """Return the projection-level Llama scope for a weight path.""" + + if _path_has_suffix(weight_name, "model.embed_tokens") or _path_has_suffix( + weight_name, "embed_tokens" + ): + return "Token embedding" + if _path_has_suffix(weight_name, "lm_head"): + return "LM head" + if _contains_segment(weight_name, "rotate_embedding", "rotate_lm_head"): + return "SpinQuant rotation weights" + + projection_scopes = { + "self_attn.q_proj": "Attention / q_proj", + "self_attn.k_proj": "Attention / k_proj", + "self_attn.v_proj": "Attention / v_proj", + "self_attn.o_proj": "Attention / o_proj", + "mlp.gate_proj": "MLP / gate_proj", + "mlp.up_proj": "MLP / up_proj", + "mlp.down_proj": "MLP / down_proj", + } + for suffix, scope in projection_scopes.items(): + if _path_has_suffix(weight_name, suffix): + return scope + + if _contains_segment( + weight_name, + "input_layernorm", + "post_attention_layernorm", + "norm", + ): + return "Norm weights" + return _UNCLASSIFIED_SCOPE + + +def _classify_qwen3_vl(weight_name: str) -> str: + """Return the projection-level Qwen3-VL scope for a weight path.""" + + if _contains_path(weight_name, "visual.patch_embed") and _path_has_suffix( + weight_name, "patch_embed.proj" + ): + return "Vision / patch_embed.proj" + + vision_projection_scopes = { + "attn.qkv": "Vision attention / qkv", + "attn.proj": "Vision attention / proj", + "mlp.linear_fc1": "Vision MLP / linear_fc1", + "mlp.linear_fc2": "Vision MLP / linear_fc2", + } + if _contains_path(weight_name, "visual.blocks"): + for suffix, scope in vision_projection_scopes.items(): + if _path_has_suffix(weight_name, suffix): + return scope + + merger_scopes = { + "linear_fc1": "Vision merger / linear_fc1", + "linear_fc2": "Vision merger / linear_fc2", + } + if _contains_path(weight_name, "visual.merger"): + for suffix, scope in merger_scopes.items(): + if _path_has_suffix(weight_name, suffix): + return scope + if _contains_path(weight_name, "visual.deepstack_merger_list"): + for suffix, scope in merger_scopes.items(): + if _path_has_suffix(weight_name, suffix): + return scope.replace("Vision merger", "Deepstack merger") + + if _contains_path(weight_name, "language_model"): + if _path_has_suffix(weight_name, "embed_tokens"): + return "Text / token embedding" + + text_projection_scopes = { + "self_attn.q_proj": "Text attention / q_proj", + "self_attn.k_proj": "Text attention / k_proj", + "self_attn.v_proj": "Text attention / v_proj", + "self_attn.o_proj": "Text attention / o_proj", + "mlp.gate_proj": "Text MLP / gate_proj", + "mlp.up_proj": "Text MLP / up_proj", + "mlp.down_proj": "Text MLP / down_proj", + } + for suffix, scope in text_projection_scopes.items(): + if _path_has_suffix(weight_name, suffix): + return scope + + if _path_has_suffix(weight_name, "lm_head"): + return "LM head" + if _contains_segment(weight_name, "rotate_embedding", "rotate_lm_head"): + return "SpinQuant rotation weights" + if _contains_segment( + weight_name, + "norm", + "norm1", + "norm2", + "input_layernorm", + "post_attention_layernorm", + "q_norm", + "k_norm", + ): + return "Norm weights" + return _UNCLASSIFIED_SCOPE + + +_LLAMA_SCOPE_ORDER = ( + "Token embedding", + "Attention / q_proj", + "Attention / k_proj", + "Attention / v_proj", + "Attention / o_proj", + "MLP / gate_proj", + "MLP / up_proj", + "MLP / down_proj", + "LM head", + "Norm weights", + "SpinQuant rotation weights", +) + +_QWEN3_VL_SCOPE_ORDER = ( + "Vision / patch_embed.proj", + "Vision attention / qkv", + "Vision attention / proj", + "Vision MLP / linear_fc1", + "Vision MLP / linear_fc2", + "Vision merger / linear_fc1", + "Vision merger / linear_fc2", + "Deepstack merger / linear_fc1", + "Deepstack merger / linear_fc2", + "Text / token embedding", + "Text attention / q_proj", + "Text attention / k_proj", + "Text attention / v_proj", + "Text attention / o_proj", + "Text MLP / gate_proj", + "Text MLP / up_proj", + "Text MLP / down_proj", + "LM head", + "Norm weights", + "SpinQuant rotation weights", +) + +_LLAMA_DECODER_LAYER_RE = re.compile(r"(?:^|\.)(?:model\.)?layers\.(\d+)(?:\.|$)") +_QWEN_VISION_BLOCK_RE = re.compile(r"(?:^|\.)(?:model\.)?visual\.blocks\.(\d+)(?:\.|$)") +_QWEN_DEEPSTACK_MERGER_RE = re.compile( + r"(?:^|\.)(?:model\.)?visual\.deepstack_merger_list\.(\d+)(?:\.|$)" +) +_QWEN_TEXT_LAYER_RE = re.compile( + r"(?:^|\.)(?:model\.)?language_model\.layers\.(\d+)(?:\.|$)" +) + + +def _matched_index(pattern: re.Pattern[str], path: str) -> int | None: + """Return an integer index captured from a module path.""" + + match = pattern.search(path) + if match is None: + return None + return int(match.group(1)) + + +def _classify_llama_layer(weight_name: str) -> str: + """Return a normalized Llama layer name for a weight path.""" + + layer_idx = _matched_index(_LLAMA_DECODER_LAYER_RE, weight_name) + if layer_idx is not None: + return f"model.layers.{layer_idx}" + if _path_has_suffix(weight_name, "model.embed_tokens") or _path_has_suffix( + weight_name, "embed_tokens" + ): + return "model.embed_tokens" + if _contains_segment(weight_name, "rotate_embedding"): + return "model.rotate_embedding" + if _path_has_suffix(weight_name, "model.norm") or weight_name == "norm": + return "model.norm" + if _contains_segment(weight_name, "rotate_lm_head"): + return "rotate_lm_head" + if _path_has_suffix(weight_name, "lm_head"): + return "lm_head" + return weight_name or "" + + +def _classify_qwen3_vl_layer(weight_name: str) -> str: + """Return a normalized Qwen3-VL layer name for a weight path.""" + + vision_idx = _matched_index(_QWEN_VISION_BLOCK_RE, weight_name) + if vision_idx is not None: + return f"model.visual.blocks.{vision_idx}" + + deepstack_idx = _matched_index(_QWEN_DEEPSTACK_MERGER_RE, weight_name) + if deepstack_idx is not None: + return f"model.visual.deepstack_merger_list.{deepstack_idx}" + + text_idx = _matched_index(_QWEN_TEXT_LAYER_RE, weight_name) + if text_idx is not None: + return f"model.language_model.layers.{text_idx}" + + if _contains_path(weight_name, "visual.patch_embed"): + return "model.visual.patch_embed" + if _contains_path(weight_name, "visual.merger"): + return "model.visual.merger" + if _contains_path(weight_name, "language_model.embed_tokens"): + return "model.language_model.embed_tokens" + if _contains_path(weight_name, "language_model.rotate_embedding"): + return "model.language_model.rotate_embedding" + if _path_has_suffix(weight_name, "language_model.norm"): + return "model.language_model.norm" + if _contains_segment(weight_name, "rotate_lm_head"): + return "rotate_lm_head" + if _path_has_suffix(weight_name, "lm_head"): + return "lm_head" + return weight_name or "" + + +def _llama_layer_sort_key(layer: str) -> tuple[int, int, str]: + """Return a stable architecture-aware sort key for Llama layer names.""" + + if layer == "model.embed_tokens": + return (0, 0, layer) + if layer == "model.rotate_embedding": + return (1, 0, layer) + match = re.fullmatch(r"model\.layers\.(\d+)", layer) + if match is not None: + return (2, int(match.group(1)), layer) + if layer == "model.norm": + return (3, 0, layer) + if layer == "rotate_lm_head": + return (4, 0, layer) + if layer == "lm_head": + return (5, 0, layer) + return (99, 0, layer) + + +def _qwen3_vl_layer_sort_key(layer: str) -> tuple[int, int, str]: + """Return a stable architecture-aware sort key for Qwen3-VL layer names.""" + + if layer == "model.visual.patch_embed": + return (0, 0, layer) + match = re.fullmatch(r"model\.visual\.blocks\.(\d+)", layer) + if match is not None: + return (1, int(match.group(1)), layer) + if layer == "model.visual.merger": + return (2, 0, layer) + match = re.fullmatch(r"model\.visual\.deepstack_merger_list\.(\d+)", layer) + if match is not None: + return (3, int(match.group(1)), layer) + if layer == "model.language_model.embed_tokens": + return (4, 0, layer) + if layer == "model.language_model.rotate_embedding": + return (5, 0, layer) + match = re.fullmatch(r"model\.language_model\.layers\.(\d+)", layer) + if match is not None: + return (6, int(match.group(1)), layer) + if layer == "model.language_model.norm": + return (7, 0, layer) + if layer == "rotate_lm_head": + return (8, 0, layer) + if layer == "lm_head": + return (9, 0, layer) + return (99, 0, layer) + + +def _family_definition( + family: str, +) -> tuple[ + Callable[[str], str], + tuple[str, ...], + Callable[[str], str], + Callable[[str], tuple[int, int, str]], +]: + """Return scope and layer classifiers for a supported model family.""" + + normalized = family.strip().lower().replace("-", "_") + if normalized == "llama": + return ( + _classify_llama, + _LLAMA_SCOPE_ORDER, + _classify_llama_layer, + _llama_layer_sort_key, + ) + if normalized in {"qwen3_vl", "qwen3vl"}: + return ( + _classify_qwen3_vl, + _QWEN3_VL_SCOPE_ORDER, + _classify_qwen3_vl_layer, + _qwen3_vl_layer_sort_key, + ) + raise ValueError( + f"Unsupported model family {family!r}. Supported families: llama, qwen3_vl." + ) + + +def _aggregate_records( + scope: str, records: Iterable[WeightTensorSparsity] +) -> SparsityRow: + """Aggregate records by element count while deduplicating equivalent weights.""" + + seen: set[str] = set() + qdtypes: set[str] = set() + zero_count = 0 + numel = 0 + + for record in records: + if record.dedup_key in seen: + continue + seen.add(record.dedup_key) + qdtypes.add(record.qdtype) + zero_count += record.zero_count + numel += record.numel + + if not qdtypes: + qdtype = "-" + elif len(qdtypes) == 1: + qdtype = next(iter(qdtypes)) + else: + qdtype = f"mixed ({', '.join(sorted(qdtypes))})" + + return SparsityRow( + scope=scope, + qdtype=qdtype, + zero_count=zero_count, + numel=numel, + ) + + +def aggregate_weight_sparsity( + records: Sequence[WeightTensorSparsity], + family: str, + *, + include_empty_scopes: bool = False, +) -> list[SparsityRow]: + """Aggregate collected tensor statistics into the three-column summary report.""" + + classifier, scope_order, _, _ = _family_definition(family) + grouped: dict[str, list[WeightTensorSparsity]] = { + scope: [] for scope in scope_order + } + grouped[_UNCLASSIFIED_SCOPE] = [] + for record in records: + grouped.setdefault(classifier(record.weight_name), []).append(record) + + rows = [_aggregate_records(_ALL_SCOPE, records)] + for scope in scope_order: + row = _aggregate_records(scope, grouped.get(scope, ())) + if include_empty_scopes or row.numel > 0: + rows.append(row) + + unclassified = _aggregate_records(_UNCLASSIFIED_SCOPE, grouped[_UNCLASSIFIED_SCOPE]) + if unclassified.numel > 0: + rows.append(unclassified) + return rows + + +def _to_layer_row(layer: str, row: SparsityRow) -> LayerSparsityRow: + """Attach a normalized layer name to an aggregated scope row.""" + + return LayerSparsityRow( + layer=layer, + scope=row.scope, + qdtype=row.qdtype, + zero_count=row.zero_count, + numel=row.numel, + ) + + +def aggregate_layer_weight_sparsity( + records: Sequence[WeightTensorSparsity], + family: str, + *, + include_layer_totals: bool = True, +) -> list[LayerSparsityRow]: + """Aggregate tensor statistics by normalized layer and projection scope.""" + + classifier, scope_order, layer_classifier, layer_sort_key = _family_definition( + family + ) + grouped: dict[str, dict[str, list[WeightTensorSparsity]]] = {} + layer_records: dict[str, list[WeightTensorSparsity]] = {} + + for record in records: + layer = layer_classifier(record.weight_name) + scope = classifier(record.weight_name) + layer_records.setdefault(layer, []).append(record) + grouped.setdefault(layer, {}).setdefault(scope, []).append(record) + + rows: list[LayerSparsityRow] = [] + for layer in sorted(layer_records, key=layer_sort_key): + if include_layer_totals: + rows.append( + _to_layer_row( + layer, + _aggregate_records(_ALL_SCOPE, layer_records[layer]), + ) + ) + + by_scope = grouped[layer] + for scope in scope_order: + row = _aggregate_records(scope, by_scope.get(scope, ())) + if row.numel > 0: + rows.append(_to_layer_row(layer, row)) + + unclassified = _aggregate_records( + _UNCLASSIFIED_SCOPE, + by_scope.get(_UNCLASSIFIED_SCOPE, ()), + ) + if unclassified.numel > 0: + rows.append(_to_layer_row(layer, unclassified)) + return rows + + +def measure_weight_sparsity( + model: torch.nn.Module, + family: str, + *, + max_chunk_elements: int = _DEFAULT_MAX_CHUNK_ELEMENTS, + deduplicate_shared_weights: bool = True, + include_empty_scopes: bool = False, +) -> list[SparsityRow]: + """Measure post-convert weight sparsity and build a three-column summary report. + + Scope-level sparsity is computed from the total semantic-zero count divided by + the total number of elements in that scope. It is not an unweighted average of + per-tensor sparsity values. + + Args: + model: A TICO model immediately after PTQ ``convert()``. + family: ``llama`` or ``qwen3_vl``. + max_chunk_elements: Maximum number of elements processed in one block. + deduplicate_shared_weights: Deduplicate shared weights with identical + finalized affine qparams inside each report row. + include_empty_scopes: Include predefined scopes that have no matching weight. + + Returns: + Rows ordered as ``Scope``, ``Qdtype``, and ``Sparsity (%)``. + """ + + records = collect_weight_tensor_sparsity( + model, + max_chunk_elements=max_chunk_elements, + deduplicate_shared_weights=deduplicate_shared_weights, + ) + return aggregate_weight_sparsity( + records, + family, + include_empty_scopes=include_empty_scopes, + ) + + +def measure_layer_weight_sparsity( + model: torch.nn.Module, + family: str, + *, + max_chunk_elements: int = _DEFAULT_MAX_CHUNK_ELEMENTS, + deduplicate_shared_weights: bool = True, + include_layer_totals: bool = True, +) -> list[LayerSparsityRow]: + """Measure post-convert sparsity by layer and projection scope.""" + + records = collect_weight_tensor_sparsity( + model, + max_chunk_elements=max_chunk_elements, + deduplicate_shared_weights=deduplicate_shared_weights, + ) + return aggregate_layer_weight_sparsity( + records, + family, + include_layer_totals=include_layer_totals, + ) + + +def measure_weight_sparsity_report( + model: torch.nn.Module, + family: str, + *, + max_chunk_elements: int = _DEFAULT_MAX_CHUNK_ELEMENTS, + deduplicate_shared_weights: bool = True, + include_empty_scopes: bool = False, + include_layer_totals: bool = True, +) -> WeightSparsityReport: + """Measure summary and layer reports from a single post-convert tensor scan.""" + + records = collect_weight_tensor_sparsity( + model, + max_chunk_elements=max_chunk_elements, + deduplicate_shared_weights=deduplicate_shared_weights, + ) + return WeightSparsityReport( + summary_rows=tuple( + aggregate_weight_sparsity( + records, + family, + include_empty_scopes=include_empty_scopes, + ) + ), + layer_rows=tuple( + aggregate_layer_weight_sparsity( + records, + family, + include_layer_totals=include_layer_totals, + ) + ), + ) + + +def _format_sparsity_percent(value: float | None, precision: int) -> str: + """Format an optional sparsity percentage for text output.""" + + if value is None: + return "-" + return f"{value:.{precision}f}" + + +def format_weight_sparsity_table( + rows: Sequence[SparsityRow], *, precision: int = 6 +) -> str: + """Format the scope-level Markdown sparsity table.""" + + if precision < 0: + raise ValueError("precision must be non-negative.") + + lines = [ + "| Scope | Qdtype | Sparsity (%) |", + "|---|---|---:|", + ] + for row in rows: + sparsity = _format_sparsity_percent(row.sparsity_percent, precision) + scope = row.scope.replace("|", "\\|") + qdtype = row.qdtype.replace("|", "\\|") + lines.append(f"| {scope} | {qdtype} | {sparsity} |") + return "\n".join(lines) + + +def format_layer_weight_sparsity_table( + rows: Sequence[LayerSparsityRow], *, precision: int = 6 +) -> str: + """Format the layer-and-scope Markdown sparsity table.""" + + if precision < 0: + raise ValueError("precision must be non-negative.") + + lines = [ + "| Layer | Scope | Qdtype | Sparsity (%) |", + "|---|---|---|---:|", + ] + for row in rows: + sparsity = _format_sparsity_percent(row.sparsity_percent, precision) + layer = row.layer.replace("|", "\\|") + scope = row.scope.replace("|", "\\|") + qdtype = row.qdtype.replace("|", "\\|") + lines.append(f"| {layer} | {scope} | {qdtype} | {sparsity} |") + return "\n".join(lines) + + +def write_weight_sparsity_csv( + rows: Sequence[SparsityRow], + path: str | Path, + *, + precision: int = 6, +) -> Path: + """Write the three-column summary report as CSV.""" + + if precision < 0: + raise ValueError("precision must be non-negative.") + + output_path = Path(path) + output_path.parent.mkdir(parents=True, exist_ok=True) + with output_path.open("w", encoding="utf-8", newline="") as output_file: + writer = csv.writer(output_file) + writer.writerow(["scope", "qdtype", "sparsity_percent"]) + for row in rows: + value = ( + "" + if row.sparsity_percent is None + else f"{row.sparsity_percent:.{precision}f}" + ) + writer.writerow([row.scope, row.qdtype, value]) + return output_path + + +def write_layer_weight_sparsity_csv( + rows: Sequence[LayerSparsityRow], + path: str | Path, + *, + precision: int = 6, +) -> Path: + """Write the layer-and-scope sparsity report as CSV.""" + + if precision < 0: + raise ValueError("precision must be non-negative.") + + output_path = Path(path) + output_path.parent.mkdir(parents=True, exist_ok=True) + with output_path.open("w", encoding="utf-8", newline="") as output_file: + writer = csv.writer(output_file) + writer.writerow(["layer", "scope", "qdtype", "sparsity_percent"]) + for row in rows: + value = ( + "" + if row.sparsity_percent is None + else f"{row.sparsity_percent:.{precision}f}" + ) + writer.writerow([row.layer, row.scope, row.qdtype, value]) + return output_path + + +def write_weight_sparsity_markdown( + rows: Sequence[SparsityRow], + path: str | Path, + *, + precision: int = 6, +) -> Path: + """Write the three-column summary report as Markdown.""" + + output_path = Path(path) + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text( + format_weight_sparsity_table(rows, precision=precision) + "\n", + encoding="utf-8", + ) + return output_path + + +def write_layer_weight_sparsity_markdown( + rows: Sequence[LayerSparsityRow], + path: str | Path, + *, + precision: int = 6, +) -> Path: + """Write the layer-and-scope sparsity report as Markdown.""" + + output_path = Path(path) + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text( + format_layer_weight_sparsity_table(rows, precision=precision) + "\n", + encoding="utf-8", + ) + return output_path diff --git a/tico/quantization/circle_weight_sparsity.py b/tico/quantization/circle_weight_sparsity.py new file mode 100644 index 00000000..35dca0d0 --- /dev/null +++ b/tico/quantization/circle_weight_sparsity.py @@ -0,0 +1,1220 @@ +# 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. + +from __future__ import annotations + +import argparse +import csv +import hashlib +import io +import json +import math +import mmap +import sys +from collections import defaultdict +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Iterable, Mapping, Sequence + +import numpy as np + + +_DEFAULT_CHUNK_NUMEL = 1024 * 1024 +_SCOPE_NAME = "All model weights" +_SELECTION_MODES = ("auto", "operator-inputs", "quantized-constants") + +# Input positions follow the Circle operator signatures emitted by TICO. +# Dynamic inputs are ignored because only tensors backed by constant buffers +# are accepted as weights. +_WEIGHT_INPUT_INDICES: Mapping[str, tuple[int, ...]] = { + "FULLY_CONNECTED": (1,), + "CONV_2D": (1,), + "DEPTHWISE_CONV_2D": (1,), + "TRANSPOSE_CONV": (1,), + "GATHER": (0,), + "RMS_NORM": (1,), + "PRELU": (1,), + "INSTANCE_NORM": (1, 2), + "BATCH_MATMUL": (0, 1), + "SVDF": (1, 2), +} + +# These operators preserve the number and location of zero values in a +# constant weight tensor. Following input zero lets the analyzer resolve a +# stored weight even when a serializer leaves a lightweight transform between +# the buffer and the weight-bearing operator. +_PASSTHROUGH_DATA_INPUT: Mapping[str, int] = { + "DEQUANTIZE": 0, + "QUANTIZE": 0, + "RESHAPE": 0, + "TRANSPOSE": 0, + "CAST": 0, + "SQUEEZE": 0, + "EXPAND_DIMS": 0, +} + +_WEIGHT_NAME_MARKERS = ( + "weight", + "kernel", + "embedding", + "embed_tokens", + "lm_head", + "rotation", + "hadamard", +) +_NAME_EXCLUSION_MARKERS = ( + "bias", + "zero_point", + "zeropoint", +) + +_DTYPE_ORDER = { + "int4": 0, + "uint4": 1, + "int8": 2, + "uint8": 3, + "int16": 4, + "uint16": 5, + "int32": 6, + "uint32": 7, + "int64": 8, + "uint64": 9, + "float16": 10, + "bfloat16": 11, + "float32": 12, + "float64": 13, + "bool": 14, +} + +_NUMPY_DTYPES: Mapping[str, np.dtype[Any]] = { + "int8": np.dtype("i1"), + "uint8": np.dtype("u1"), + "int16": np.dtype(" Any: + """Import the Circle schema lazily so pure helper tests need no FlatBuffers.""" + + try: + from circle_schema import circle + except ImportError as exc: + raise CircleWeightSparsityError( + "circle-schema is required. Install TICO or run " + "`pip install circle-schema`." + ) from exc + return circle + + +def _call_int(obj: Any, method_name: str, default: int = 0) -> int: + """Call an integer-valued FlatBuffer accessor when it exists.""" + + method = getattr(obj, method_name, None) + if method is None or not callable(method): + return default + return int(method()) + + +def _vector_as_numpy(obj: Any, prefix: str, dtype: np.dtype[Any]) -> np.ndarray: + """Read a generated FlatBuffer vector through NumPy or scalar accessors.""" + + numpy_method = getattr(obj, f"{prefix}AsNumpy", None) + if numpy_method is not None and callable(numpy_method): + values = numpy_method() + if isinstance(values, np.ndarray): + return np.asarray(values, dtype=dtype).reshape(-1) + if isinstance(values, memoryview): + return np.frombuffer(values, dtype=dtype) + + length = _call_int(obj, f"{prefix}Length") + scalar_method = getattr(obj, prefix, None) + if length == 0 or scalar_method is None or not callable(scalar_method): + return np.empty(0, dtype=dtype) + return np.asarray([scalar_method(index) for index in range(length)], dtype=dtype) + + +def _decode_name(value: Any, default: str) -> str: + """Decode a FlatBuffer string accessor into a stable Python string.""" + + if value is None: + return default + if isinstance(value, bytes): + return value.decode("utf-8", errors="replace") + return str(value) + + +def _enum_name(enum_class: Any, value: int, prefix: str) -> str: + """Return the symbolic name of a generated FlatBuffer enum value.""" + + for name, enum_value in vars(enum_class).items(): + if name.startswith("_") or not isinstance(enum_value, int): + continue + if int(enum_value) == int(value): + return name + return f"{prefix}_{value}" + + +def _tensor_dtype_name(circle: Any, tensor_type: int) -> str: + """Return a lowercase Circle tensor dtype name.""" + + enum_class = circle.TensorType.TensorType + return _enum_name(enum_class, tensor_type, "TENSOR_TYPE").lower() + + +def _operator_builtin_code(operator_code: Any) -> int: + """Resolve an operator code across old and new Circle schema fields.""" + + builtin_method = getattr(operator_code, "BuiltinCode", None) + deprecated_method = getattr(operator_code, "DeprecatedBuiltinCode", None) + builtin = int(builtin_method()) if callable(builtin_method) else None + deprecated = int(deprecated_method()) if callable(deprecated_method) else None + + if builtin is None: + if deprecated is None: + raise CircleWeightSparsityError("OperatorCode has no builtin code field.") + return deprecated + if builtin == 0 and deprecated not in (None, 0): + return deprecated + return builtin + + +def _operator_name(circle: Any, model: Any, operator: Any) -> str: + """Return a symbolic Circle builtin operator name.""" + + opcode_index = _call_int(operator, "OpcodeIndex", -1) + if opcode_index < 0 or opcode_index >= _call_int(model, "OperatorCodesLength"): + raise CircleWeightSparsityError(f"Invalid operator code index {opcode_index}.") + operator_code = model.OperatorCodes(opcode_index) + builtin_code = _operator_builtin_code(operator_code) + return _enum_name( + circle.BuiltinOperator.BuiltinOperator, + builtin_code, + "BUILTIN_OPERATOR", + ) + + +def _tensor_shape(tensor: Any) -> tuple[int, ...]: + """Read a static Circle tensor shape.""" + + values = _vector_as_numpy(tensor, "Shape", np.dtype(np.int64)) + shape = tuple(int(value) for value in values) + if any(dimension < 0 for dimension in shape): + raise CircleWeightSparsityError( + f"Negative tensor dimension is unsupported: {shape}." + ) + return shape + + +def _tensor_name(tensor: Any, tensor_index: int) -> str: + """Return a readable Circle tensor name.""" + + name_method = getattr(tensor, "Name", None) + value = name_method() if callable(name_method) else None + return _decode_name(value, f"tensor_{tensor_index}") + + +def _buffer_data(model: Any, buffer_index: int) -> np.ndarray: + """Return a zero-copy uint8 view of one Circle buffer when possible.""" + + buffer_count = _call_int(model, "BuffersLength") + if buffer_index < 0 or buffer_index >= buffer_count: + raise CircleWeightSparsityError( + f"Invalid buffer index {buffer_index}; model has {buffer_count} buffers." + ) + buffer = model.Buffers(buffer_index) + data = _vector_as_numpy(buffer, "Data", np.dtype(np.uint8)) + if data.dtype != np.uint8: + data = data.astype(np.uint8, copy=False) + return np.ascontiguousarray(data).reshape(-1) + + +def _buffer_length(model: Any, buffer_index: int) -> int: + """Return a Circle buffer byte length without decoding its payload.""" + + buffer_count = _call_int(model, "BuffersLength") + if buffer_index < 0 or buffer_index >= buffer_count: + return 0 + buffer = model.Buffers(buffer_index) + return _call_int(buffer, "DataLength") + + +def _is_constant_tensor(model: Any, tensor: Any) -> bool: + """Return whether a Circle tensor owns a non-empty immutable buffer.""" + + is_variable = bool(_call_int(tensor, "IsVariable", 0)) + if is_variable: + return False + buffer_index = _call_int(tensor, "Buffer", -1) + return _buffer_length(model, buffer_index) > 0 + + +def _has_quantization_metadata(tensor: Any) -> bool: + """Return whether a Circle tensor contains affine quantization metadata.""" + + quantization_method = getattr(tensor, "Quantization", None) + if quantization_method is None or not callable(quantization_method): + return False + quantization = quantization_method() + if quantization is None: + return False + return ( + _call_int(quantization, "ScaleLength") > 0 + or _call_int(quantization, "ZeroPointLength") > 0 + ) + + +def _quantization_info(tensor: Any) -> QuantizationInfo: + """Extract zero points and the quantized dimension from a Circle tensor.""" + + quantization_method = getattr(tensor, "Quantization", None) + quantization = ( + quantization_method() + if quantization_method is not None and callable(quantization_method) + else None + ) + if quantization is None: + return QuantizationInfo((), None, False) + + scales = _vector_as_numpy(quantization, "Scale", np.dtype(np.float64)) + zero_points = _vector_as_numpy(quantization, "ZeroPoint", np.dtype(np.int64)) + if scales.size == 0 and zero_points.size == 0: + return QuantizationInfo((), None, False) + + if zero_points.size == 0: + zero_points = np.asarray([0], dtype=np.int64) + axis = _call_int(quantization, "QuantizedDimension", 0) + return QuantizationInfo( + tuple(int(value) for value in zero_points), + axis, + True, + tuple(float(value) for value in scales), + ) + + +def _operator_inputs(operator: Any) -> tuple[int, ...]: + """Return all tensor IDs consumed by one Circle operator.""" + + values = _vector_as_numpy(operator, "Inputs", np.dtype(np.int64)) + return tuple(int(value) for value in values) + + +def _operator_outputs(operator: Any) -> tuple[int, ...]: + """Return all tensor IDs produced by one Circle operator.""" + + values = _vector_as_numpy(operator, "Outputs", np.dtype(np.int64)) + return tuple(int(value) for value in values) + + +def _subgraph_inputs(subgraph: Any) -> set[int]: + """Return the tensor IDs declared as runtime subgraph inputs.""" + + values = _vector_as_numpy(subgraph, "Inputs", np.dtype(np.int64)) + return {int(value) for value in values if int(value) >= 0} + + +def _build_producer_map( + circle: Any, + model: Any, + subgraph: Any, +) -> dict[int, tuple[str, Any]]: + """Map every produced tensor ID to its operator and symbolic name.""" + + producers: dict[int, tuple[str, Any]] = {} + for operator_index in range(_call_int(subgraph, "OperatorsLength")): + operator = subgraph.Operators(operator_index) + name = _operator_name(circle, model, operator) + for tensor_id in _operator_outputs(operator): + if tensor_id >= 0: + producers[tensor_id] = (name, operator) + return producers + + +def _resolve_constant_source( + model: Any, + subgraph: Any, + tensor_id: int, + producers: Mapping[int, tuple[str, Any]], +) -> int | None: + """Follow zero-preserving transforms until a stored constant is found.""" + + tensor_count = _call_int(subgraph, "TensorsLength") + current = tensor_id + visited: set[int] = set() + + while current not in visited: + visited.add(current) + if current < 0 or current >= tensor_count: + return None + tensor = subgraph.Tensors(current) + if _is_constant_tensor(model, tensor): + return current + + producer = producers.get(current) + if producer is None: + return None + operator_name, operator = producer + data_input_index = _PASSTHROUGH_DATA_INPUT.get(operator_name) + if data_input_index is None: + return None + inputs = _operator_inputs(operator) + if data_input_index >= len(inputs): + return None + current = inputs[data_input_index] + + return None + + +def _looks_like_weight_name(name: str) -> bool: + """Return whether a tensor name strongly suggests a model weight.""" + + normalized = name.lower() + if any(marker in normalized for marker in _NAME_EXCLUSION_MARKERS): + return False + return any(marker in normalized for marker in _WEIGHT_NAME_MARKERS) + + +def _select_weight_tensor_ids( + circle: Any, + model: Any, + subgraph: Any, + selection: str, +) -> dict[int, set[str]]: + """Select stored weight tensors and record how each tensor was recognized.""" + + if selection not in _SELECTION_MODES: + raise ValueError( + f"Unsupported selection mode {selection!r}; " + f"choose one of {_SELECTION_MODES}." + ) + + tensor_count = _call_int(subgraph, "TensorsLength") + graph_inputs = _subgraph_inputs(subgraph) + producers = _build_producer_map(circle, model, subgraph) + selected: dict[int, set[str]] = defaultdict(set) + + if selection in ("auto", "operator-inputs"): + for operator_index in range(_call_int(subgraph, "OperatorsLength")): + operator = subgraph.Operators(operator_index) + operator_name = _operator_name(circle, model, operator) + input_indices = _WEIGHT_INPUT_INDICES.get(operator_name, ()) + inputs = _operator_inputs(operator) + for input_index in input_indices: + if input_index >= len(inputs): + continue + source_id = _resolve_constant_source( + model, + subgraph, + inputs[input_index], + producers, + ) + if source_id is None or source_id in graph_inputs: + continue + selected[source_id].add(f"{operator_name}:input{input_index}") + + if selection == "auto": + for tensor_index in range(tensor_count): + if tensor_index in graph_inputs: + continue + tensor = subgraph.Tensors(tensor_index) + if not _is_constant_tensor(model, tensor): + continue + name = _tensor_name(tensor, tensor_index) + if _looks_like_weight_name(name): + selected[tensor_index].add("name-fallback") + + if selection == "quantized-constants": + for tensor_index in range(tensor_count): + if tensor_index in graph_inputs: + continue + tensor = subgraph.Tensors(tensor_index) + if not _is_constant_tensor(model, tensor): + continue + if _has_quantization_metadata(tensor): + selected[tensor_index].add("quantized-constant") + + return selected + + +def _validate_shape_and_numel(shape: tuple[int, ...]) -> int: + """Return the number of logical elements in a static tensor shape.""" + + if not shape: + return 1 + return math.prod(shape) + + +def _required_buffer_bytes(qdtype: str, numel: int) -> int: + """Return the minimum number of bytes required by a logical tensor.""" + + if qdtype in ("uint4", "int4"): + return (numel + 1) // 2 + if qdtype == "bfloat16": + return numel * 2 + numpy_dtype = _NUMPY_DTYPES.get(qdtype) + if numpy_dtype is None: + raise CircleWeightSparsityError( + f"Unsupported Circle weight dtype {qdtype!r}. " + "Affine integer, standard floating-point, and boolean weights " + "are supported; MX formats are not yet decoded." + ) + return numel * numpy_dtype.itemsize + + +def _semantic_zero_points( + qdtype: str, + quantization: QuantizationInfo, +) -> tuple[int, ...]: + """Return integer codes that represent real-valued zero.""" + + if qdtype.startswith("float") or qdtype == "bfloat16": + return () + if quantization.is_affine and quantization.zero_points: + return quantization.zero_points + return (0,) + + +def _count_integer_chunk_zeros( + values: np.ndarray, + start_index: int, + shape: tuple[int, ...], + zero_points: tuple[int, ...], + quantized_dimension: int | None, +) -> int: + """Count integer semantic zeros in one flattened logical tensor chunk.""" + + if len(zero_points) == 1: + return int(np.count_nonzero(values == zero_points[0])) + + if not shape: + raise CircleWeightSparsityError( + "A scalar tensor cannot use multiple zero points." + ) + axis = int(quantized_dimension or 0) % len(shape) + channel_count = shape[axis] + if len(zero_points) != channel_count: + raise CircleWeightSparsityError( + "Zero-point count does not match the quantized dimension: " + f"shape={shape}, axis={axis}, zero_points={len(zero_points)}." + ) + + elements_after_axis = math.prod(shape[axis + 1 :]) + indices = np.arange( + start_index, + start_index + values.size, + dtype=np.int64, + ) + channel_indices = (indices // elements_after_axis) % channel_count + targets = np.asarray(zero_points, dtype=np.int64)[channel_indices] + return int(np.count_nonzero(values == targets)) + + +def _count_packed_4bit_zeros( + raw: np.ndarray, + qdtype: str, + shape: tuple[int, ...], + numel: int, + zero_points: tuple[int, ...], + quantized_dimension: int | None, + chunk_numel: int, +) -> int: + """Count semantic zeros in low-nibble-first packed 4-bit data.""" + + logical_chunk = max(2, chunk_numel) + logical_chunk -= logical_chunk % 2 + zero_count = 0 + + for start in range(0, numel, logical_chunk): + end = min(start + logical_chunk, numel) + byte_start = start // 2 + byte_end = (end + 1) // 2 + packed = raw[byte_start:byte_end] + decoded = np.empty(packed.size * 2, dtype=np.uint8) + decoded[0::2] = packed & np.uint8(0x0F) + decoded[1::2] = packed >> np.uint8(4) + values = decoded[: end - start] + if qdtype == "int4": + signed = values.astype(np.int8) + signed[signed >= 8] -= 16 + values = signed + zero_count += _count_integer_chunk_zeros( + values, + start, + shape, + zero_points, + quantized_dimension, + ) + return zero_count + + +def _count_dense_zeros( + raw: np.ndarray, + qdtype: str, + shape: tuple[int, ...], + numel: int, + zero_points: tuple[int, ...], + quantized_dimension: int | None, + chunk_numel: int, +) -> int: + """Count semantic zeros in a byte-aligned Circle tensor buffer.""" + + if qdtype == "bfloat16": + values = np.frombuffer(raw, dtype=np.dtype(" tuple[int, int]: + """Count semantic zeros in one stored Circle weight tensor.""" + + if chunk_numel <= 0: + raise ValueError(f"chunk_numel must be positive, got {chunk_numel}.") + numel = _validate_shape_and_numel(shape) + if numel == 0: + return 0, 0 + + required_bytes = _required_buffer_bytes(qdtype, numel) + if raw.size < required_bytes: + raise CircleWeightSparsityError( + f"Buffer is too small for {qdtype} tensor {shape}: " + f"required={required_bytes}, available={raw.size}." + ) + raw = raw[:required_bytes] + zero_points = _semantic_zero_points(qdtype, quantization) + + if qdtype in ("uint4", "int4"): + zero_count = _count_packed_4bit_zeros( + raw, + qdtype, + shape, + numel, + zero_points, + quantization.quantized_dimension, + chunk_numel, + ) + else: + zero_count = _count_dense_zeros( + raw, + qdtype, + shape, + numel, + zero_points, + quantization.quantized_dimension, + chunk_numel, + ) + return zero_count, numel + + +def _tensor_fingerprint( + raw: np.ndarray, + qdtype: str, + shape: tuple[int, ...], + quantization: QuantizationInfo, +) -> str: + """Hash a stored tensor and its zero-relevant metadata.""" + + digest = hashlib.sha256() + digest.update(qdtype.encode("utf-8")) + digest.update(repr(shape).encode("utf-8")) + digest.update(repr(quantization.scales).encode("utf-8")) + digest.update(repr(quantization.zero_points).encode("utf-8")) + digest.update(repr(quantization.quantized_dimension).encode("utf-8")) + numel = _validate_shape_and_numel(shape) + required_bytes = _required_buffer_bytes(qdtype, numel) + digest.update(memoryview(raw[:required_bytes])) + return digest.hexdigest() + + +def _analyze_root_model( + model: Any, + circle: Any, + source: str, + *, + selection: str, + chunk_numel: int, + strict: bool, + fingerprint: bool, +) -> tuple[list[CircleWeightTensorStats], int, list[str]]: + """Analyze all subgraphs in one parsed Circle root model.""" + + stats: list[CircleWeightTensorStats] = [] + duplicate_count = 0 + skipped_messages: list[str] = [] + seen_model_buffers: set[int] = set() + + for subgraph_index in range(_call_int(model, "SubgraphsLength")): + subgraph = model.Subgraphs(subgraph_index) + selected = _select_weight_tensor_ids( + circle, + model, + subgraph, + selection, + ) + for tensor_index in sorted(selected): + tensor = subgraph.Tensors(tensor_index) + tensor_name = _tensor_name(tensor, tensor_index) + try: + shape = _tensor_shape(tensor) + tensor_type = _call_int(tensor, "Type", -1) + qdtype = _tensor_dtype_name(circle, tensor_type) + buffer_index = _call_int(tensor, "Buffer", -1) + quantization = _quantization_info(tensor) + if buffer_index in seen_model_buffers: + duplicate_count += 1 + continue + + raw = _buffer_data(model, buffer_index) + zero_count, numel = _count_tensor_semantic_zeros( + raw, + qdtype, + shape, + quantization, + chunk_numel, + ) + tensor_digest = ( + _tensor_fingerprint(raw, qdtype, shape, quantization) + if fingerprint + else None + ) + stats.append( + CircleWeightTensorStats( + source=source, + subgraph_index=subgraph_index, + tensor_index=tensor_index, + tensor_name=tensor_name, + buffer_index=buffer_index, + qdtype=qdtype, + shape=shape, + zero_count=zero_count, + numel=numel, + roles=tuple(sorted(selected[tensor_index])), + fingerprint=tensor_digest, + ) + ) + seen_model_buffers.add(buffer_index) + except (CircleWeightSparsityError, ValueError, TypeError) as exc: + message = ( + f"{source}:subgraph={subgraph_index}:tensor={tensor_index}" + f"({tensor_name}): {exc}" + ) + if strict: + raise CircleWeightSparsityError(message) from exc + skipped_messages.append(message) + + return stats, duplicate_count, skipped_messages + + +def analyze_circle_binary( + circle_binary: Any, + *, + source: str = "", + selection: str = "auto", + chunk_numel: int = _DEFAULT_CHUNK_NUMEL, + strict: bool = True, + fingerprint: bool = False, + circle_module: Any | None = None, +) -> tuple[list[CircleWeightTensorStats], int, list[str]]: + """Analyze weight tensors from an in-memory Circle FlatBuffer.""" + + circle = circle_module or _load_circle_schema() + model = circle.Model.Model.GetRootAsModel(circle_binary, 0) + return _analyze_root_model( + model, + circle, + source, + selection=selection, + chunk_numel=chunk_numel, + strict=strict, + fingerprint=fingerprint, + ) + + +def analyze_circle_file( + path: str | Path, + *, + selection: str = "auto", + chunk_numel: int = _DEFAULT_CHUNK_NUMEL, + strict: bool = True, + fingerprint: bool = False, +) -> tuple[list[CircleWeightTensorStats], int, list[str]]: + """Memory-map and analyze one Circle model file.""" + + circle_path = Path(path) + if not circle_path.is_file(): + raise CircleWeightSparsityError(f"Circle file does not exist: {circle_path}") + if circle_path.stat().st_size == 0: + raise CircleWeightSparsityError(f"Circle file is empty: {circle_path}") + + circle = _load_circle_schema() + with circle_path.open("rb") as file: + with mmap.mmap(file.fileno(), 0, access=mmap.ACCESS_READ) as mapped: + model = circle.Model.Model.GetRootAsModel(mapped, 0) + result = _analyze_root_model( + model, + circle, + str(circle_path), + selection=selection, + chunk_numel=chunk_numel, + strict=strict, + fingerprint=fingerprint, + ) + del model + return result + + +def discover_circle_files( + inputs: Iterable[str | Path], + *, + recursive: bool = False, +) -> list[Path]: + """Resolve file and directory inputs into a sorted unique Circle file list.""" + + discovered: list[Path] = [] + for raw_input in inputs: + path = Path(raw_input).expanduser() + if path.is_file(): + discovered.append(path.resolve()) + continue + if path.is_dir(): + iterator = path.rglob("*.circle") if recursive else path.glob("*.circle") + discovered.extend(candidate.resolve() for candidate in iterator) + continue + raise CircleWeightSparsityError(f"Input path does not exist: {path}") + + unique = sorted(dict.fromkeys(discovered)) + if not unique: + raise CircleWeightSparsityError("No Circle files were found.") + return unique + + +def _format_qdtype(stats: Sequence[CircleWeightTensorStats]) -> str: + """Return a single dtype or a stable compact mixed-dtype label.""" + + qdtypes = sorted( + {item.qdtype for item in stats}, + key=lambda value: (_DTYPE_ORDER.get(value, 10_000), value), + ) + if len(qdtypes) == 1: + return qdtypes[0] + return f"mixed ({', '.join(qdtypes)})" + + +def _deduplicate_across_files( + stats: Sequence[CircleWeightTensorStats], +) -> tuple[list[CircleWeightTensorStats], int]: + """Deduplicate matching tensor names and payloads across Circle files.""" + + unique: list[CircleWeightTensorStats] = [] + seen: set[tuple[str, str, tuple[int, ...], str]] = set() + duplicates = 0 + for item in stats: + if item.fingerprint is None: + raise CircleWeightSparsityError( + "Cross-file deduplication requires tensor fingerprints." + ) + key = ( + item.tensor_name, + item.qdtype, + item.shape, + item.fingerprint, + ) + if key in seen: + duplicates += 1 + continue + seen.add(key) + unique.append(item) + return unique, duplicates + + +def aggregate_circle_weight_stats( + stats: Sequence[CircleWeightTensorStats], + *, + source_count: int, + duplicate_tensor_count: int = 0, + skipped_messages: Sequence[str] = (), +) -> CircleWeightSparsityReport: + """Aggregate tensor statistics into the single model-level report row.""" + + if not stats: + details = "" + if skipped_messages: + details = "\n" + "\n".join(f" - {item}" for item in skipped_messages) + raise CircleWeightSparsityError( + "No supported Circle weight tensors were selected." + details + ) + + zero_count = sum(item.zero_count for item in stats) + numel = sum(item.numel for item in stats) + if numel == 0: + raise CircleWeightSparsityError("Selected Circle weights contain no elements.") + row = CircleWeightSparsityRow( + scope=_SCOPE_NAME, + qdtype=_format_qdtype(stats), + sparsity_pct=100.0 * zero_count / numel, + ) + return CircleWeightSparsityReport( + row=row, + source_count=source_count, + tensor_count=len(stats), + zero_count=zero_count, + numel=numel, + duplicate_tensor_count=duplicate_tensor_count, + skipped_tensor_count=len(skipped_messages), + skipped_messages=tuple(skipped_messages), + ) + + +def measure_circle_weight_sparsity( + inputs: Iterable[str | Path], + *, + recursive: bool = False, + selection: str = "auto", + chunk_numel: int = _DEFAULT_CHUNK_NUMEL, + strict: bool = True, + deduplicate_across_files: bool = False, +) -> CircleWeightSparsityReport: + """Measure model-level Circle weight sparsity across one or more files.""" + + files = discover_circle_files(inputs, recursive=recursive) + all_stats: list[CircleWeightTensorStats] = [] + duplicate_count = 0 + skipped_messages: list[str] = [] + + for path in files: + file_stats, file_duplicates, file_skipped = analyze_circle_file( + path, + selection=selection, + chunk_numel=chunk_numel, + strict=strict, + fingerprint=deduplicate_across_files, + ) + all_stats.extend(file_stats) + duplicate_count += file_duplicates + skipped_messages.extend(file_skipped) + + if deduplicate_across_files: + all_stats, cross_file_duplicates = _deduplicate_across_files(all_stats) + duplicate_count += cross_file_duplicates + + return aggregate_circle_weight_stats( + all_stats, + source_count=len(files), + duplicate_tensor_count=duplicate_count, + skipped_messages=skipped_messages, + ) + + +def render_markdown( + report: CircleWeightSparsityReport, + *, + precision: int = 6, +) -> str: + """Render the model-level result as a three-column Markdown table.""" + + if precision < 0: + raise ValueError(f"precision must be non-negative, got {precision}.") + row = report.row + return "\n".join( + ( + "| Scope | Qdtype | Sparsity (%) |", + "|---|---|---:|", + f"| {row.scope} | {row.qdtype} | {row.sparsity_pct:.{precision}f} |", + ) + ) + + +def render_csv( + report: CircleWeightSparsityReport, + *, + precision: int = 6, +) -> str: + """Render the model-level result as a three-column CSV document.""" + + if precision < 0: + raise ValueError(f"precision must be non-negative, got {precision}.") + output = io.StringIO(newline="") + writer = csv.writer(output) + writer.writerow(("Scope", "Qdtype", "Sparsity (%)")) + writer.writerow( + ( + report.row.scope, + report.row.qdtype, + f"{report.row.sparsity_pct:.{precision}f}", + ) + ) + return output.getvalue() + + +def render_json( + report: CircleWeightSparsityReport, + *, + precision: int = 6, +) -> str: + """Render the model-level result as a one-row JSON array.""" + + if precision < 0: + raise ValueError(f"precision must be non-negative, got {precision}.") + payload = [ + { + "Scope": report.row.scope, + "Qdtype": report.row.qdtype, + "Sparsity (%)": round(report.row.sparsity_pct, precision), + } + ] + return json.dumps(payload, indent=2, ensure_ascii=False) + "\n" + + +def render_report( + report: CircleWeightSparsityReport, + output_format: str, + *, + precision: int = 6, +) -> str: + """Render a report in Markdown, CSV, or JSON format.""" + + normalized = output_format.lower() + if normalized in ("markdown", "md"): + return render_markdown(report, precision=precision) + "\n" + if normalized == "csv": + return render_csv(report, precision=precision) + if normalized == "json": + return render_json(report, precision=precision) + raise ValueError(f"Unsupported output format {output_format!r}.") + + +def _build_argument_parser() -> argparse.ArgumentParser: + """Build the command-line parser for Circle weight sparsity analysis.""" + + parser = argparse.ArgumentParser( + description=( + "Measure model-level semantic weight sparsity from Circle files. " + "Affine integer zero is detected by qcode == zero_point." + ) + ) + parser.add_argument( + "inputs", + nargs="+", + help="Circle file paths or directories containing .circle files.", + ) + parser.add_argument( + "--recursive", + action="store_true", + help="Search input directories recursively for .circle files.", + ) + parser.add_argument( + "--selection", + choices=_SELECTION_MODES, + default="auto", + help=( + "Weight selection policy. 'auto' uses known operator input roles and " + "a conservative tensor-name fallback." + ), + ) + parser.add_argument( + "--chunk-numel", + type=int, + default=_DEFAULT_CHUNK_NUMEL, + help="Maximum logical elements decoded per temporary chunk.", + ) + parser.add_argument( + "--strict", + action=argparse.BooleanOptionalAction, + default=True, + help="Fail instead of skipping malformed or unsupported selected tensors.", + ) + parser.add_argument( + "--deduplicate-across-files", + action=argparse.BooleanOptionalAction, + default=False, + help=( + "Deduplicate tensors with the same name, dtype, shape, and payload " + "across multiple input files." + ), + ) + parser.add_argument( + "--format", + choices=("markdown", "csv", "json"), + default="markdown", + dest="output_format", + help="Output format for the single model-level result row.", + ) + parser.add_argument( + "--output", + type=Path, + default=None, + help="Write the report to this file instead of standard output.", + ) + parser.add_argument( + "--precision", + type=int, + default=6, + help="Number of decimal places in the sparsity percentage.", + ) + parser.add_argument( + "--verbose", + action="store_true", + help="Print internal file, tensor, and element counters to standard error.", + ) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + """Run the Circle weight sparsity command-line tool.""" + + parser = _build_argument_parser() + args = parser.parse_args(argv) + + try: + report = measure_circle_weight_sparsity( + args.inputs, + recursive=args.recursive, + selection=args.selection, + chunk_numel=args.chunk_numel, + strict=args.strict, + deduplicate_across_files=args.deduplicate_across_files, + ) + rendered = render_report( + report, + args.output_format, + precision=args.precision, + ) + except (CircleWeightSparsityError, ValueError, OSError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 2 + + if args.output is None: + sys.stdout.write(rendered) + else: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(rendered, encoding="utf-8") + + if report.skipped_tensor_count: + print( + f"warning: skipped {report.skipped_tensor_count} selected tensor(s)", + file=sys.stderr, + ) + if args.verbose: + for message in report.skipped_messages: + print(f" - {message}", file=sys.stderr) + + if args.verbose: + print( + "files={files} tensors={tensors} elements={elements} zeros={zeros} " + "duplicates={duplicates} skipped={skipped}".format( + files=report.source_count, + tensors=report.tensor_count, + elements=report.numel, + zeros=report.zero_count, + duplicates=report.duplicate_tensor_count, + skipped=report.skipped_tensor_count, + ), + file=sys.stderr, + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tico/quantization/examples/configs/llama32_3b_weight_sparsity.yaml b/tico/quantization/examples/configs/llama32_3b_weight_sparsity.yaml new file mode 100644 index 00000000..9d645557 --- /dev/null +++ b/tico/quantization/examples/configs/llama32_3b_weight_sparsity.yaml @@ -0,0 +1,92 @@ +model: + family: llama + name_or_path: meta-llama/Llama-3.2-3B-Instruct + trust_remote_code: false + hf_token: null + cache_dir: null + +runtime: + device: cuda + dtype: float32 + seed: 42 + show_progress: true + +model_args: + profile: npu_export + +calibration: + dataset: Salesforce/wikitext + dataset_config: wikitext-2-raw-v1 + split: train + n_samples: 128 + seq_len: 2048 + decode_steps: 8 + +pipeline: + - name: spinquant + enabled: true + + - name: cle + enabled: false + pairs: + - model.layers.*.mlp.up_proj:model.layers.*.mlp.down_proj + method: absmax + max_iter: 1 + + - name: gptq + enabled: true + weight_bits: 4 + weight_bits_overrides: {} + perchannel: true + symmetric: false + mse: mse + sensitivity: + mode: compute + path: null + percdamp: 0.01 + groupsize: -1 + actorder: true + static_groups: false + quantize_lm_head: false + use_orig_model_inference: false + verbose: false + show_progress: true + + - name: ptq + enabled: true + profile: npu_export + activation: int16 + linear_weight: uint4 + embedding_weight: uint8 + lm_head_weight: uint8 + spin_rotation_weight: int16 + norm_weight: int16 + strict_wrap: true + decode_calibration_steps: 8 + + - name: weight_sparsity + enabled: true + output_dir: ./out/llama32_3b_weight_sparsity + filename_stem: weight_sparsity + layer_filename_stem: weight_sparsity_by_layer + formats: [csv, markdown] + precision: 6 + include_layer_report: true + include_layer_totals: true + print_layer_report: false + max_chunk_elements: 4194304 + deduplicate_shared_weights: true + include_empty_scopes: false + +evaluation: + enabled: false + perplexity: null + lm_eval_tasks: null + max_seq_len: 2048 + +export: + enabled: false + output_dir: ./out/llama32_3b_weight_sparsity + max_seq_len: 2048 + prefill_decode: false + artifacts: [] diff --git a/tico/quantization/examples/configs/qwen3_vl_4b_weight_sparsity.yaml b/tico/quantization/examples/configs/qwen3_vl_4b_weight_sparsity.yaml new file mode 100644 index 00000000..1094d566 --- /dev/null +++ b/tico/quantization/examples/configs/qwen3_vl_4b_weight_sparsity.yaml @@ -0,0 +1,117 @@ +model: + family: qwen3_vl + name_or_path: Qwen/Qwen3-VL-4B-Instruct + trust_remote_code: true + hf_token: null + cache_dir: null + +runtime: + device: cuda + dtype: float32 + seed: 42 + show_progress: true + +calibration: + datasets: + - dataset: vqav2 + split: testdev + n_samples: 128 + - dataset: wikitext2 + split: train + n_samples: 128 + seq_len: 2048 + +model_args: + vision: + grid_thw: [8, 24, 24] + visual_start_idx: 0 + spatial_merge_size: 2 + +pipeline: + - name: spinquant + enabled: true + init_method: random + r1_path: null + r2_map_path: null + enable_r1: true + enable_r2: true + fuse_deepstack_visual_outputs: true + vision_init_method: random + vision_r1_path: null + vision_r2_map_path: null + enable_vision_r1: true + enable_vision_r2: true + fuse_vision_layer_norms: true + require_vision_r1_layernorm_compatible: true + vision_rotation_tolerance: 1.0e-4 + + - name: smoothquant + enabled: false + alpha: 0.5 + components: both + custom_alpha_map: null + + - name: gptq + enabled: true + weight_bits: 4 + weight_bits_overrides: {} + perchannel: true + symmetric: false + mse: null + sensitivity: + mode: compute + path: null + percdamp: 0.01 + groupsize: -1 + actorder: true + static_groups: false + verbose: false + show_progress: true + quantize_vision: true + quantize_text: true + quantize_lm_head: false + quantize_vision_patch_embed: true + quantize_vision_blocks: true + quantize_vision_merger: true + quantize_vision_deepstack_mergers: true + quantize_text_layers: true + move_cache_to_cpu: false + + - name: ptq + enabled: true + activation: int16 + linear_weight: uint4 + vision_patch_embed_weight: uint8 + embedding_weight: uint8 + lm_head_weight: uint8 + spin_rotation_weight: int16 + norm: int16 + norm_weight: int16 + strict_wrap: true + + - name: weight_sparsity + enabled: true + output_dir: ./out/qwen3_vl_4b_weight_sparsity + filename_stem: weight_sparsity + layer_filename_stem: weight_sparsity_by_layer + formats: [csv, markdown] + precision: 6 + include_layer_report: true + include_layer_totals: true + print_layer_report: false + max_chunk_elements: 4194304 + deduplicate_shared_weights: true + include_empty_scopes: false + +evaluation: + enabled: false + vlm_tasks: [] + coco: false + llava_bench: false + n_samples: 50 + max_seq_len: 2048 + +export: + enabled: false + output_dir: ./out/qwen3_vl_4b_weight_sparsity + artifacts: [] diff --git a/tico/quantization/recipes/stages/__init__.py b/tico/quantization/recipes/stages/__init__.py index 6bf16961..c937181f 100644 --- a/tico/quantization/recipes/stages/__init__.py +++ b/tico/quantization/recipes/stages/__init__.py @@ -4,6 +4,7 @@ from tico.quantization.recipes.stages.ptq import PTQStage from tico.quantization.recipes.stages.smoothquant import SmoothQuantStage from tico.quantization.recipes.stages.spinquant import SpinQuantStage +from tico.quantization.recipes.stages.weight_sparsity import WeightSparsityStage _STAGE_REGISTRY = { "gptq": GPTQStage(), @@ -11,6 +12,7 @@ "spinquant": SpinQuantStage(), "cle": CLEStage(), "smoothquant": SmoothQuantStage(), + "weight_sparsity": WeightSparsityStage(), } diff --git a/tico/quantization/recipes/stages/weight_sparsity.py b/tico/quantization/recipes/stages/weight_sparsity.py new file mode 100644 index 00000000..f89748e2 --- /dev/null +++ b/tico/quantization/recipes/stages/weight_sparsity.py @@ -0,0 +1,183 @@ +# 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. + +from collections.abc import Sequence +from pathlib import Path +from typing import Any, Mapping + +from tico.quantization.analysis.weight_sparsity import ( + format_layer_weight_sparsity_table, + format_weight_sparsity_table, + measure_weight_sparsity, + measure_weight_sparsity_report, + write_layer_weight_sparsity_csv, + write_layer_weight_sparsity_markdown, + write_weight_sparsity_csv, + write_weight_sparsity_markdown, +) +from tico.quantization.recipes.context import RecipeContext +from tico.quantization.recipes.stages.base import Stage + + +_DEFAULT_FORMATS = ("csv", "markdown") + + +def _normalize_formats(value: Any) -> tuple[str, ...]: + """Normalize configured output formats to canonical names.""" + + if value is None: + return _DEFAULT_FORMATS + if isinstance(value, str): + values = [value] + elif isinstance(value, Sequence): + values = list(value) + else: + raise TypeError("weight_sparsity.formats must be a string or sequence.") + + normalized: list[str] = [] + for item in values: + name = str(item).strip().lower() + if name == "md": + name = "markdown" + if name not in {"csv", "markdown"}: + raise ValueError( + "Unsupported weight sparsity format " + f"{item!r}. Supported formats: csv, markdown." + ) + if name not in normalized: + normalized.append(name) + return tuple(normalized) + + +def _resolve_output_dir( + ctx: RecipeContext, stage_cfg: Mapping[str, Any] +) -> Path | None: + """Resolve the report directory from the stage or export configuration.""" + + configured = stage_cfg.get("output_dir") + if configured is None: + export_cfg = ctx.cfg.get("export", {}) + if isinstance(export_cfg, Mapping): + configured = export_cfg.get("output_dir") + if configured is None: + return None + return Path(str(configured)) + + +class WeightSparsityStage(Stage): + """Measure summary and layer-level weight sparsity after PTQ conversion.""" + + name = "weight_sparsity" + + def run(self, ctx: RecipeContext, stage_cfg: Mapping[str, Any]) -> RecipeContext: + """Measure, print, and optionally save post-convert sparsity reports.""" + + precision = int(stage_cfg.get("precision", 6)) + max_chunk_elements = int(stage_cfg.get("max_chunk_elements", 4 * 1024 * 1024)) + deduplicate_shared_weights = bool( + stage_cfg.get("deduplicate_shared_weights", True) + ) + include_empty_scopes = bool(stage_cfg.get("include_empty_scopes", False)) + include_layer_report = bool(stage_cfg.get("include_layer_report", True)) + include_layer_totals = bool(stage_cfg.get("include_layer_totals", True)) + output_dir = _resolve_output_dir(ctx, stage_cfg) + print_layer_report = bool( + stage_cfg.get("print_layer_report", output_dir is None) + ) + + if include_layer_report: + report = measure_weight_sparsity_report( + ctx.require_model(), + family=ctx.adapter.family, + max_chunk_elements=max_chunk_elements, + deduplicate_shared_weights=deduplicate_shared_weights, + include_empty_scopes=include_empty_scopes, + include_layer_totals=include_layer_totals, + ) + rows = report.summary_rows + layer_rows = report.layer_rows + else: + rows = tuple( + measure_weight_sparsity( + ctx.require_model(), + family=ctx.adapter.family, + max_chunk_elements=max_chunk_elements, + deduplicate_shared_weights=deduplicate_shared_weights, + include_empty_scopes=include_empty_scopes, + ) + ) + layer_rows = () + + print("=== Post-convert weight sparsity ===") + print(format_weight_sparsity_table(rows, precision=precision)) + print() + + if include_layer_report and print_layer_report: + print("=== Post-convert weight sparsity by layer ===") + print(format_layer_weight_sparsity_table(layer_rows, precision=precision)) + print() + + if output_dir is None: + return ctx + + filename_stem = str(stage_cfg.get("filename_stem", "weight_sparsity")) + layer_filename_stem = str( + stage_cfg.get("layer_filename_stem", f"{filename_stem}_by_layer") + ) + formats = _normalize_formats(stage_cfg.get("formats")) + output_dir.mkdir(parents=True, exist_ok=True) + + if "csv" in formats: + csv_path = write_weight_sparsity_csv( + rows, + output_dir / f"{filename_stem}.csv", + precision=precision, + ) + ctx.artifacts["weight_sparsity_csv"] = csv_path + print(f"Saved weight sparsity CSV to {csv_path.resolve()}") + + if include_layer_report: + layer_csv_path = write_layer_weight_sparsity_csv( + layer_rows, + output_dir / f"{layer_filename_stem}.csv", + precision=precision, + ) + ctx.artifacts["weight_sparsity_by_layer_csv"] = layer_csv_path + print( + "Saved layer-level weight sparsity CSV to " + f"{layer_csv_path.resolve()}" + ) + + if "markdown" in formats: + markdown_path = write_weight_sparsity_markdown( + rows, + output_dir / f"{filename_stem}.md", + precision=precision, + ) + ctx.artifacts["weight_sparsity_markdown"] = markdown_path + print(f"Saved weight sparsity Markdown to {markdown_path.resolve()}") + + if include_layer_report: + layer_markdown_path = write_layer_weight_sparsity_markdown( + layer_rows, + output_dir / f"{layer_filename_stem}.md", + precision=precision, + ) + ctx.artifacts["weight_sparsity_by_layer_markdown"] = layer_markdown_path + print( + "Saved layer-level weight sparsity Markdown to " + f"{layer_markdown_path.resolve()}" + ) + + return ctx