Skip to content

Commit ddaabb9

Browse files
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

File tree

0 commit comments

Comments
 (0)