Skip to content

feat: petr: add fully-dynamic MX6/MX9 quantization recipes - #39

Open
Jiarwang77 wants to merge 1 commit into
jiarwang/mx9-fake-quantfrom
jiarwang/mx6_9-allmodels
Open

feat: petr: add fully-dynamic MX6/MX9 quantization recipes#39
Jiarwang77 wants to merge 1 commit into
jiarwang/mx9-fake-quantfrom
jiarwang/mx6_9-allmodels

Conversation

@Jiarwang77

Copy link
Copy Markdown
Collaborator

Summary

What this adds

alto.models.petr.quantize.apply_mx_quantization(model, mx_format, ignore=())
— one entry point that quantizes every nn.Conv2d and nn.Linear with a
fully-dynamic MX6 or MX9 recipe. Scales come from the tensor being quantized on
each forward pass, so no calibration data, observers or persisted scales are
involved.
Scope is named by pattern rather than by an attribute name a particular
architecture happens to use, so the same call works on a detector, an LLM, or
anything else built out of those two module types. Mixed precision is expressed
by omission: a model already runs in some dtype, and what ignore skips keeps
it.
Two things needed care:

  • A bare ignore name is widened to its subtree. compressed_tensors
    matches a bare pattern against the module name exactly, so "pts_bbox_head"
    would select only the container — which holds no quantizable leaf and
    therefore excludes nothing, quietly quantizing the part it was meant to
    protect.
  • Two silent failure modes now raise: a recipe that would reach no module,
    and an ignore that would exclude none. Both are checked with
    match_named_modules — the matcher apply_quantization_config itself uses —
    before the model is touched, so a rejected call leaves it unmodified.
    block_axis is injected into QuantizationArgs alongside format and
    forwarded into the packed MX codec. Blocks have to follow the GEMM's reduction
    dimension: Linear reduces along the last axis, but a Conv2d reduces along
    its input channels, and blocking a 3x3 weight along the kernel width instead
    leaves three elements in a sixteen-element block. The field defaults to the
    last axis, so existing recipes are unaffected.

How it was validated

End to end, PETRv2 (VoVNet-99, 800x320) on nuScenes v1.0-mini val, 81
frames, --fuse-conv-bn:

mAP
fp32 0.3907
uniform MX6 0.3092

The PETR-side change that calls this from tools/test.py is a separate PR
against AMD-AGI/PETR and depends on this one landing first.

Test plan

  • pytest tests/unittest/mx9_mx6/test_mx_recipe_scope.py — 13 passed,
    1 xfailed. The xfail is test_finalized_modules_are_frozen:
    freeze_module_quantization only promotes COMPRESSED to FROZEN, so a
    fully-dynamic recipe stays at CALIBRATION after finalize.
  • pytest tests/unittest/mx9_mx6/ — the pre-existing MX6/MX9 kernel and
    dispatch tests still pass.

- Add `alto.models.petr.quantize.apply_mx_quantization`, a model-agnostic
  entry point that quantizes every `nn.Conv2d` and `nn.Linear` with a
  fully-dynamic MX6 or MX9 recipe.
- Per-module-type MX block axis: Conv2d blocks along input channels (axis 1),
  Linear along the last axis; the chosen axis is stored in the recipe and
  forwarded into the packed MX codec.
- `ignore` accepts module names or `re:` patterns, widens bare names to cover
  the subtree, and is validated before the model is touched so a no-op scope
  fails fast.
@hann-wang

Copy link
Copy Markdown
Collaborator

better make models.petr a model-independent API.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants