mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-11-03 08:14:29 +08:00
Extends ``suggest_exclusion`` with two new keyword arguments so operators can
turn per-node sensitivity scores into a block-level exclusion set without
reimplementing the coverage / threshold / near-tie logic themselves. Motivated
by empirical validation on ViT-tiny where per-node picking hits a ~60% top-1
ceiling due to intra-block precision fragmentation; block-level exclusion
recovers ~75% top-1 (within 1pp of native ``trtexec --int8 --fp16``).
New API surface:
* ``blocks: Mapping[str, Sequence[str | re.Pattern]] | None = None`` -- maps
group name to a list of regex patterns matching node paths. Each node is
assigned to at most one group (first-match wins across the ``blocks``
dict). Nodes matching no pattern become their own singleton group named
after themselves, so architecturally-important standalone nodes (final
``LayerNormalization`` before the head, patch-embed ``Conv``, etc.) compete
for exclusion on equal footing with multi-node blocks. When ``blocks`` is
``None`` (default) the picker behaves exactly as before -- fully
backward-compatible.
* ``block_agg: Literal["sum", "max", "mean"] = "sum"`` -- aggregation used to
compute a group's score from its members. Kept as ``Literal`` for
IDE/mypy support, plus a runtime ``ValueError`` in ``suggest_exclusion``
for defensive validation.
Semantics:
* When ``blocks`` is set, the picker computes per-group aggregated scores
and applies coverage / threshold / near-tie / ``max_nodes`` /
``min_score_floor`` semantics identically to the per-node path. The
returned exclusion list is the union of member node names across the
selected groups, ready to pass as ``nodes_to_exclude=`` to
``modelopt.onnx.quantization.quantize``.
* Natural pairings between ``block_agg`` and picker mode -- documented in
the docstring and RST guide:
- ``block_agg="sum"`` with ``coverage`` (recommended default): identical
"fraction of total KL mass" semantic as per-node coverage. Portable
across per-node and per-block grouping on the same model.
- ``block_agg="max"`` with ``threshold``: same units as per-node
threshold (excludes any group whose peak-node score exceeds the
cutoff). Preserves operator intuition when transferring per-node
threshold guidance to the block level.
- Other combinations remain valid but change what ``coverage`` and
``threshold`` mean in units; the docstring calls this out explicitly.
Implementation notes:
* The existing per-node core (coverage / threshold / near-tie logic) is
extracted into a private ``_pick_from_scores`` helper. Both the per-node
and per-block paths call it, so both share identical semantics for every
future behavior change. Zero duplication.
* Two additional private helpers: ``_assign_groups`` (regex-based
first-match-wins assignment with singleton fallback) and
``_aggregate_group_scores`` (dispatches on ``block_agg``).
* Backward compatibility: every existing ``suggest_exclusion`` call site
behaves exactly as before because ``blocks`` defaults to ``None`` and the
``block_agg`` value is only inspected when ``blocks`` is set.
Documentation:
* ``docs/source/guides/_onnx_quantization.rst`` gains a new subsection at
the end -- "Grouping per-node scores into architectural blocks" -- with:
- A ``vit_tiny_patch16_224`` (timm) worked example at ``coverage=0.95``
with ``block_agg="sum"`` that selects blocks 8, 10, 9, 11, 7 (~100
nodes across 5 whole transformer blocks), which recovers ~75% top-1
on ImageNet-1k versus ~60% for the best per-node picking.
- A depth-2 example showing how to split each transformer block into
``blocks.N.attn`` and ``blocks.N.mlp`` sub-groups.
- Guidance on the ``block_agg`` / picker-mode pairings.
- A "when block-level grouping doesn't help" note calling out
Conv-heavy architectures (MobileNet, ResNet families) where diffuse
per-node sensitivity means the per-node picker still wins.
Signed-off-by: gcunhase <4861122+gcunhase@users.noreply.github.com>
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>