Repository navigation
Commit ddaabb9
authored
Add QAT support and pass hooks to QuantizationRecipe. (pytorch#21935)
### Summary
Extend `QuantizationRecipe` with QAT support and four ordered pass
hooks, then wire them into `QuantizeStage`.
`QuantizationRecipe` gains `is_qat`, `train_fn`,
`calibration_inputs_fn`, and four pass-list fields
(`pre_prepare_passes`, `post_prepare_passes`, `pre_convert_passes`,
`post_convert_passes`). When `is_qat=True`, `QuantizeStage` calls
`prepare_qat_pt2e` instead of `prepare_pt2e`, invokes `train_fn` on the
prepared model, and skips calibration. When `is_qat=False` (default),
PTQ proceeds as before, but callers may now supply a
`calibration_inputs_fn` to supply their own calibration data. If
omitted, the existing example-inputs fallback is used. The four pass
hooks are applied at the appropriate points in both flows and compose
cleanly with `_combine_recipes`.
`_combine_recipes` is overhauled to handle all previously dropped
fields: the four pass lists are concatenated, `is_qat` must agree across
all combined recipes, at most one `train_fn` is allowed, and multiple
`calibration_inputs_fn` values are chained into a single factory. The
`strict`, `mode`, `pipeline_stages`, and `source_transform_in_place`
scalar fields are now also validated for agreement and propagated to the
combined recipe. `edge_manager_transform_passes` from `LoweringRecipe`
is similarly merged.
### Test Plan
```
pytest -q export/tests/test_export_recipe.py export/tests/test_export_stages.py
```1 parent 8f50849 commit ddaabb9
4 files changed
Lines changed: 1423 additions & 104 deletions
0 commit comments