feat: petr: add fully-dynamic MX6/MX9 quantization recipes - #39
Open
Jiarwang77 wants to merge 1 commit into
Open
feat: petr: add fully-dynamic MX6/MX9 quantization recipes#39Jiarwang77 wants to merge 1 commit into
Jiarwang77 wants to merge 1 commit into
Conversation
- 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.
Collaborator
|
better make |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
What this adds
alto.models.petr.quantize.apply_mx_quantization(model, mx_format, ignore=())— one entry point that quantizes every
nn.Conv2dandnn.Linearwith afully-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
ignoreskips keepsit.
Two things needed care:
ignorename is widened to its subtree.compressed_tensorsmatches 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.
and an
ignorethat would exclude none. Both are checked withmatch_named_modules— the matcherapply_quantization_configitself uses —before the model is touched, so a rejected call leaves it unmodified.
block_axisis injected intoQuantizationArgsalongsideformatandforwarded into the packed MX codec. Blocks have to follow the GEMM's reduction
dimension:
Linearreduces along the last axis, but aConv2dreduces alongits 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:The PETR-side change that calls this from
tools/test.pyis a separate PRagainst
AMD-AGI/PETRand 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_quantizationonly promotesCOMPRESSEDtoFROZEN, so afully-dynamic recipe stays at
CALIBRATIONafterfinalize.pytest tests/unittest/mx9_mx6/— the pre-existing MX6/MX9 kernel anddispatch tests still pass.