[BE] Convert quant_primitives methods private #2350
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/2350
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New FailureAs of commit b04a5f6 with merge base ab66083 ( NEW FAILURE - The following job has failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
jainapurva
changed the title
|
@jainapurva has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
jainapurva
added
topic: not user facing
There was a problem hiding this comment.
Pull Request Overview
This PR converts several quantization primitive methods from public to private by renaming them with an underscore prefix. Key changes include updating internal module calls, re-exporting updated functions in the public API, and modifying tests and documentation to reflect the new function names.
Reviewed Changes
Copilot reviewed 23 out of 23 changed files in this pull request and generated no comments.
| File | Description |
|---|---|
| torchao/quantization/qat/utils.py | Updated fake quantization method calls from public to private. |
| torchao/quantization/qat/affine_fake_quantized_tensor.py | Replaced public quant methods with their private counterparts. |
| torchao/quantization/init.py | Updated the public export list to include the new private names. |
| Other files (in prototype, dtypes, tests, docs) | Consistent renaming of quantization functions to enforce internal usage. |
Comments suppressed due to low confidence (5)
torchao/quantization/init.py:172
- Confirm that re-exporting functions with a leading underscore in the public API is intentional; if these functions are meant to be internal only, consider not exposing them here to prevent external usage.
- "choose_qparams_affine_tinygemm",
test/quantization/test_quant_primitives.py:755
- [nitpick] The test method name now uses a leading double underscore; consider renaming it (e.g., test_fake_quantize_affine_internal) for clarity and to avoid potential confusion with Python's name mangling.
def test__fake_quantize_affine(self):
docs/source/api_ref_quantization.rst:66
- Update the API reference to clearly indicate that functions prefixed with an underscore are internal and not part of the public API to prevent misuse.
_choose_qparams_affine_floatx
torchao/dtypes/affine_quantized_tensor.py:293
- Ensure that switching to the private function for quantization parameters is reflected throughout downstream processing, and update any associated type annotations or documentation if necessary.
scale, zero_point = _choose_qparams_affine_tinygemm(input_float, mapping_type, block_size)
torchao/prototype/quantization/gguf/gguf_quantized_tensor.py:201
- The renaming to _choose_qparams_gguf is consistent; please verify that all related modules now reference this private version to avoid any potential discrepancies.
) = _choose_qparams_gguf(input_float, block_size, target_dtype)
|
@jainapurva has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
| @@ -63,14 +63,14 @@ Quantization Primitives | |||
|
|
|||
| choose_qparams_affine | |||
| choose_qparams_affine_with_min_max | |||
| choose_qparams_affine_floatx | |||
| _choose_qparams_affine_floatx | |||
There was a problem hiding this comment.
we can remove non-public methods from doc I think
| @@ -172,17 +172,17 @@ | |||
| "AffineQuantizedObserverBase", | |||
| # quant primitive ops | |||
| "choose_qparams_affine", | |||
| "choose_qparams_affine_tinygemm", | |||
| "choose_qparams_affine_dont_preserve_zero", | |||
| "_choose_qparams_affine_tinygemm", | |||
There was a problem hiding this comment.
no need to add these to torchao.quantization either I think
| @@ -24,32 +24,32 @@ | |||
|
|
|||
| __all__ = [ | |||
| "choose_qparams_affine", | |||
There was a problem hiding this comment.
nit: reorder these to put the public ones first
| @@ -961,12 +961,12 @@ def fake_quantize_affine_cachemask( | |||
| This is equivalent to calling `quantize_affine` + `dequantize_affine` | |||
| but without the dtype casts. | |||
|
|
|||
| Note: Compared to :func:`~torchao.quantization.quant_primitives.fake_quantize_affine`, | |||
| Note: Compared to :func:`~torchao.quantization.quant_primitives._fake_quantize_affine`, | |||
There was a problem hiding this comment.
can we refer to docs for non-public functions?
| @@ -993,7 +993,7 @@ def fake_quantize_affine_cachemask( | |||
| return (dq, mask) | |||
|
|
|||
|
|
|||
| def _do_fake_quantize_affine( | |||
| def _do__fake_quantize_affine( | |||
There was a problem hiding this comment.
nit: probably don't need to add an extra _ for fake_quantize_affine
Following methods have now been converted to private methods
_choose_qparams_affine_tinygemm
_choose_qparams_affine_dont_preserve_zero
_choose_qparams_affine_floatx
_quantize_affine_no_zero_point
_quantize_affine_tinygemm
_dequantize_affine_no_zero_point
_dequantize_affine_tinygemm
_quantize_affine_floatx
_dequantize_affine_floatx
_fake_quantize_affine
_fake_quantize_affine_cachemask
_choose_qparams_and_quantize_affine_hqq
_choose_qparams_and_quantize_affine_qqq
_dequantize_affine_qqq
_choose_qparams_affine_float8
_quantize_affine_float8
_dequantize_affine_float8
_choose_qparams_gguf
_quantize_gguf
_dequantize_gguf